pangolin/server/lib/aiInferenceResource.ts

import { and, eq, inArray } from "drizzle-orm";
import {
    aiModels,
    aiProviders,
    db,
    resourceAiModels,
    resourceAiProviders,
    siteResourceAiModels,
    siteResourceAiProviders,
    type Transaction
} from "@server/db";
import { z } from "zod";
import { modelKeysConflict } from "@server/lib/aiModelKeyMatch";

type DbOrTrx = Transaction | typeof db;

export const modelListTypeSchema = z.enum(["allow", "block"]);

export type ModelListType = z.infer<typeof modelListTypeSchema>;

export const accessModeSchema = z.enum(["inherit", "select"]);

export type AccessMode = z.infer<typeof accessModeSchema>;

export const resourceAiProviderAttachmentSchema = z.strictObject({
    providerId: z.number().int().positive(),
    accessMode: accessModeSchema.optional().default("inherit")
});

export type ResourceAiProviderInput = z.infer<
    typeof resourceAiProviderAttachmentSchema
>;

export type ResourceAiProviderAttachment = {
    providerId: number;
    accessMode: AccessMode;
};

export const resourceAiModelEntrySchema = z.strictObject({
    modelId: z.number().int().positive(),
    listType: modelListTypeSchema
});

export type ResourceAiModelEntry = z.infer<typeof resourceAiModelEntrySchema>;

export type InferenceFieldsError = {
    error: string;
};

export function isInferenceFieldsError(
    value: { error: string } | object
): value is InferenceFieldsError {
    return "error" in value;
}

/**
 * Resolve which allow/block patterns apply for an attachment.
 * inherit → provider lists; select → resource-selected lists (replace).
 */
export function resolveEffectiveLists(input: {
    accessMode: AccessMode;
    providerAllows: string[];
    providerBlocks: string[];
    resourceAllows: string[];
    resourceBlocks: string[];
}): { allows: string[]; blocks: string[] } {
    if (input.accessMode === "select") {
        return {
            allows: input.resourceAllows,
            blocks: input.resourceBlocks
        };
    }
    return {
        allows: input.providerAllows,
        blocks: input.providerBlocks
    };
}

function normalizeAttachments(
    inputs: ResourceAiProviderInput[]
): ResourceAiProviderAttachment[] {
    const byProviderId = new Map<number, AccessMode>();
    for (const input of inputs) {
        byProviderId.set(input.providerId, input.accessMode ?? "inherit");
    }
    return [...byProviderId.entries()].map(([providerId, accessMode]) => ({
        providerId,
        accessMode
    }));
}

type EffectiveAllowRow = {
    providerId: number;
    modelKey: string;
};

/**
 * Ensure effective allow modelKeys do not conflict across attached providers.
 * inherit uses provider allows; select uses resource-selected allows (or the
 * optional override map). Block patterns are ignored for overlap checks.
 */
export async function assertNoOverlappingModelKeys(
    attachments: ResourceAiProviderAttachment[],
    options: {
        trx?: DbOrTrx;
        resourceId?: number;
        siteResourceId?: number;
        selectedAllowsByProvider?: Map<number, string[]>;
    } = {}
): Promise<InferenceFieldsError | null> {
    const trx = options.trx ?? db;

    if (attachments.length < 2) {
        return null;
    }

    const inheritProviderIds = attachments
        .filter((a) => a.accessMode === "inherit")
        .map((a) => a.providerId);
    const selectProviderIds = attachments
        .filter((a) => a.accessMode === "select")
        .map((a) => a.providerId);

    const effectiveAllows: EffectiveAllowRow[] = [];

    if (inheritProviderIds.length > 0) {
        const providerAllows = await trx
            .select({
                providerId: aiModels.providerId,
                modelKey: aiModels.modelKey
            })
            .from(aiModels)
            .where(
                and(
                    inArray(aiModels.providerId, inheritProviderIds),
                    eq(aiModels.enabled, true),
                    eq(aiModels.listType, "allow")
                )
            );
        effectiveAllows.push(...providerAllows);
    }

    if (selectProviderIds.length > 0) {
        if (options.selectedAllowsByProvider) {
            for (const providerId of selectProviderIds) {
                const keys =
                    options.selectedAllowsByProvider.get(providerId) ?? [];
                for (const modelKey of keys) {
                    effectiveAllows.push({ providerId, modelKey });
                }
            }
        } else if (options.resourceId !== undefined) {
            const rows = await trx
                .select({
                    providerId: aiModels.providerId,
                    modelKey: aiModels.modelKey
                })
                .from(resourceAiModels)
                .innerJoin(
                    aiModels,
                    eq(resourceAiModels.modelId, aiModels.modelId)
                )
                .where(
                    and(
                        eq(resourceAiModels.resourceId, options.resourceId),
                        inArray(aiModels.providerId, selectProviderIds),
                        eq(resourceAiModels.listType, "allow"),
                        eq(aiModels.enabled, true)
                    )
                );
            effectiveAllows.push(...rows);
        } else if (options.siteResourceId !== undefined) {
            const rows = await trx
                .select({
                    providerId: aiModels.providerId,
                    modelKey: aiModels.modelKey
                })
                .from(siteResourceAiModels)
                .innerJoin(
                    aiModels,
                    eq(siteResourceAiModels.modelId, aiModels.modelId)
                )
                .where(
                    and(
                        eq(
                            siteResourceAiModels.siteResourceId,
                            options.siteResourceId
                        ),
                        inArray(aiModels.providerId, selectProviderIds),
                        eq(siteResourceAiModels.listType, "allow"),
                        eq(aiModels.enabled, true)
                    )
                );
            effectiveAllows.push(...rows);
        }
    }

    const conflictPairs: string[] = [];
    for (let i = 0; i < effectiveAllows.length; i++) {
        for (let j = i + 1; j < effectiveAllows.length; j++) {
            const left = effectiveAllows[i];
            const right = effectiveAllows[j];
            if (left.providerId === right.providerId) {
                continue;
            }
            if (!modelKeysConflict(left.modelKey, right.modelKey)) {
                continue;
            }
            const pair = [left.modelKey, right.modelKey].sort().join(" vs ");
            if (!conflictPairs.includes(pair)) {
                conflictPairs.push(pair);
            }
        }
    }

    if (conflictPairs.length === 0) {
        return null;
    }

    conflictPairs.sort();
    return {
        error: `Model keys must be unique across providers on a resource. Overlapping keys: ${conflictPairs.join(", ")}`
    };
}

/**
 * Validate provider attachments for an org.
 */
export async function resolveProviderAttachments(input: {
    orgId: string;
    attachments: ResourceAiProviderInput[];
    requireAtLeastOne: boolean;
    resourceId?: number;
    siteResourceId?: number;
}): Promise<ResourceAiProviderAttachment[] | InferenceFieldsError> {
    const attachments = normalizeAttachments(input.attachments);

    if (input.requireAtLeastOne && attachments.length === 0) {
        return {
            error: "At least one AI provider is required for inference-mode resources"
        };
    }

    if (attachments.length === 0) {
        return [];
    }

    const providerIds = attachments.map((a) => a.providerId);
    const providers = await db
        .select({
            providerId: aiProviders.providerId,
            orgId: aiProviders.orgId,
            enabled: aiProviders.enabled
        })
        .from(aiProviders)
        .where(
            and(
                inArray(aiProviders.providerId, providerIds),
                eq(aiProviders.orgId, input.orgId)
            )
        );

    if (providers.length !== providerIds.length) {
        return {
            error: "One or more AI providers were not found in this organization"
        };
    }

    const disabled = providers.find((p) => !p.enabled);
    if (disabled) {
        return {
            error: `AI provider with ID ${disabled.providerId} is disabled`
        };
    }

    const overlapError = await assertNoOverlappingModelKeys(attachments, {
        resourceId: input.resourceId,
        siteResourceId: input.siteResourceId
    });
    if (overlapError) {
        return overlapError;
    }

    return attachments;
}

export async function assertInferenceModeAllowsProviderFields(input: {
    mode: string;
    hasProviderAttachments: boolean;
}): Promise<InferenceFieldsError | null> {
    if (input.mode === "inference") {
        return null;
    }
    if (input.hasProviderAttachments) {
        return {
            error: "AI providers can only be attached to inference-mode resources"
        };
    }
    return null;
}

/**
 * Attach providers to a resource. Inherit attachments use the provider lists
 * as-is (resource model rows for those providers are pruned). Select
 * attachments keep resource-selected allow/block subsets.
 */
export async function setPublicResourceAiProviders(
    resourceId: number,
    attachments: ResourceAiProviderAttachment[],
    trx: DbOrTrx = db
): Promise<void> {
    await trx
        .delete(resourceAiProviders)
        .where(eq(resourceAiProviders.resourceId, resourceId));

    if (attachments.length > 0) {
        await trx.insert(resourceAiProviders).values(
            attachments.map((a) => ({
                resourceId,
                providerId: a.providerId,
                accessMode: a.accessMode
            }))
        );
    }

    await prunePublicResourceModelsToSelectProviders(
        resourceId,
        attachments,
        trx
    );
}

export async function setSiteResourceAiProviders(
    siteResourceId: number,
    attachments: ResourceAiProviderAttachment[],
    trx: DbOrTrx = db
): Promise<void> {
    await trx
        .delete(siteResourceAiProviders)
        .where(eq(siteResourceAiProviders.siteResourceId, siteResourceId));

    if (attachments.length > 0) {
        await trx.insert(siteResourceAiProviders).values(
            attachments.map((a) => ({
                siteResourceId,
                providerId: a.providerId,
                accessMode: a.accessMode
            }))
        );
    }

    await pruneSiteResourceModelsToSelectProviders(
        siteResourceId,
        attachments,
        trx
    );
}

/**
 * Keep resource model rows only for providers in select mode.
 */
async function prunePublicResourceModelsToSelectProviders(
    resourceId: number,
    attachments: ResourceAiProviderAttachment[],
    trx: DbOrTrx
): Promise<void> {
    const selectProviderIds = attachments
        .filter((a) => a.accessMode === "select")
        .map((a) => a.providerId);

    if (selectProviderIds.length === 0) {
        await trx
            .delete(resourceAiModels)
            .where(eq(resourceAiModels.resourceId, resourceId));
        return;
    }

    const existing = await trx
        .select({
            modelId: resourceAiModels.modelId,
            providerId: aiModels.providerId
        })
        .from(resourceAiModels)
        .innerJoin(aiModels, eq(resourceAiModels.modelId, aiModels.modelId))
        .where(eq(resourceAiModels.resourceId, resourceId));

    const allowed = new Set(selectProviderIds);
    const toRemove = existing
        .filter((row) => !allowed.has(row.providerId))
        .map((row) => row.modelId);

    if (toRemove.length > 0) {
        await trx
            .delete(resourceAiModels)
            .where(
                and(
                    eq(resourceAiModels.resourceId, resourceId),
                    inArray(resourceAiModels.modelId, toRemove)
                )
            );
    }
}

async function pruneSiteResourceModelsToSelectProviders(
    siteResourceId: number,
    attachments: ResourceAiProviderAttachment[],
    trx: DbOrTrx
): Promise<void> {
    const selectProviderIds = attachments
        .filter((a) => a.accessMode === "select")
        .map((a) => a.providerId);

    if (selectProviderIds.length === 0) {
        await trx
            .delete(siteResourceAiModels)
            .where(eq(siteResourceAiModels.siteResourceId, siteResourceId));
        return;
    }

    const existing = await trx
        .select({
            modelId: siteResourceAiModels.modelId,
            providerId: aiModels.providerId
        })
        .from(siteResourceAiModels)
        .innerJoin(aiModels, eq(siteResourceAiModels.modelId, aiModels.modelId))
        .where(eq(siteResourceAiModels.siteResourceId, siteResourceId));

    const allowed = new Set(selectProviderIds);
    const toRemove = existing
        .filter((row) => !allowed.has(row.providerId))
        .map((row) => row.modelId);

    if (toRemove.length > 0) {
        await trx
            .delete(siteResourceAiModels)
            .where(
                and(
                    eq(siteResourceAiModels.siteResourceId, siteResourceId),
                    inArray(siteResourceAiModels.modelId, toRemove)
                )
            );
    }
}

export async function clearPublicResourceAiConfig(
    resourceId: number,
    trx: DbOrTrx = db
): Promise<void> {
    await trx
        .delete(resourceAiModels)
        .where(eq(resourceAiModels.resourceId, resourceId));
    await trx
        .delete(resourceAiProviders)
        .where(eq(resourceAiProviders.resourceId, resourceId));
}

export async function clearSiteResourceAiConfig(
    siteResourceId: number,
    trx: DbOrTrx = db
): Promise<void> {
    await trx
        .delete(siteResourceAiModels)
        .where(eq(siteResourceAiModels.siteResourceId, siteResourceId));
    await trx
        .delete(siteResourceAiProviders)
        .where(eq(siteResourceAiProviders.siteResourceId, siteResourceId));
}

export async function listPublicResourceAiProviders(resourceId: number) {
    return db
        .select({
            providerId: resourceAiProviders.providerId,
            name: aiProviders.name,
            type: aiProviders.type,
            enabled: aiProviders.enabled,
            accessMode: resourceAiProviders.accessMode
        })
        .from(resourceAiProviders)
        .innerJoin(
            aiProviders,
            eq(resourceAiProviders.providerId, aiProviders.providerId)
        )
        .where(eq(resourceAiProviders.resourceId, resourceId));
}

export async function listSiteResourceAiProviders(siteResourceId: number) {
    return db
        .select({
            providerId: siteResourceAiProviders.providerId,
            name: aiProviders.name,
            type: aiProviders.type,
            enabled: aiProviders.enabled,
            accessMode: siteResourceAiProviders.accessMode
        })
        .from(siteResourceAiProviders)
        .innerJoin(
            aiProviders,
            eq(siteResourceAiProviders.providerId, aiProviders.providerId)
        )
        .where(eq(siteResourceAiProviders.siteResourceId, siteResourceId));
}

/**
 * Model list APIs require an inference resource with at least one select-mode
 * attached provider.
 */
export async function assertPublicModelListApiEligible(resource: {
    resourceId: number;
    mode: string;
}): Promise<string | null> {
    if (resource.mode !== "inference") {
        return "AI model lists are only supported on inference-mode resources";
    }

    const [row] = await db
        .select({ providerId: resourceAiProviders.providerId })
        .from(resourceAiProviders)
        .where(
            and(
                eq(resourceAiProviders.resourceId, resource.resourceId),
                eq(resourceAiProviders.accessMode, "select")
            )
        )
        .limit(1);

    if (!row) {
        return "Set at least one attached AI provider to select mode before managing model lists";
    }
    return null;
}

export async function assertSiteModelListApiEligible(siteResource: {
    siteResourceId: number;
    mode: string;
}): Promise<string | null> {
    if (siteResource.mode !== "inference") {
        return "AI model lists are only supported on inference-mode resources";
    }

    const [row] = await db
        .select({ providerId: siteResourceAiProviders.providerId })
        .from(siteResourceAiProviders)
        .where(
            and(
                eq(
                    siteResourceAiProviders.siteResourceId,
                    siteResource.siteResourceId
                ),
                eq(siteResourceAiProviders.accessMode, "select")
            )
        )
        .limit(1);

    if (!row) {
        return "Set at least one attached AI provider to select mode before managing model lists";
    }
    return null;
}

/**
 * Resource model entries must belong to select-mode attached providers, and
 * listType must match the provider catalog entry (allow→allow, block→block).
 */
export async function assertPublicResourceModelEntriesValid(input: {
    orgId: string;
    resourceId: number;
    models: ResourceAiModelEntry[];
}): Promise<string | null> {
    const uniqueModels = dedupeModelEntries(input.models);
    if (uniqueModels.length === 0) {
        return null;
    }

    const attachments = await db
        .select({
            providerId: resourceAiProviders.providerId,
            accessMode: resourceAiProviders.accessMode
        })
        .from(resourceAiProviders)
        .innerJoin(
            aiProviders,
            eq(resourceAiProviders.providerId, aiProviders.providerId)
        )
        .where(
            and(
                eq(resourceAiProviders.resourceId, input.resourceId),
                eq(aiProviders.orgId, input.orgId)
            )
        );

    return assertModelEntriesValid({
        orgId: input.orgId,
        modelEntries: uniqueModels,
        attachments,
        resourceLabel: "resource"
    });
}

export async function assertSiteResourceModelEntriesValid(input: {
    orgId: string;
    siteResourceId: number;
    models: ResourceAiModelEntry[];
}): Promise<string | null> {
    const uniqueModels = dedupeModelEntries(input.models);
    if (uniqueModels.length === 0) {
        return null;
    }

    const attachments = await db
        .select({
            providerId: siteResourceAiProviders.providerId,
            accessMode: siteResourceAiProviders.accessMode
        })
        .from(siteResourceAiProviders)
        .innerJoin(
            aiProviders,
            eq(siteResourceAiProviders.providerId, aiProviders.providerId)
        )
        .where(
            and(
                eq(
                    siteResourceAiProviders.siteResourceId,
                    input.siteResourceId
                ),
                eq(aiProviders.orgId, input.orgId)
            )
        );

    return assertModelEntriesValid({
        orgId: input.orgId,
        modelEntries: uniqueModels,
        attachments,
        resourceLabel: "site resource"
    });
}

function dedupeModelEntries(
    models: ResourceAiModelEntry[]
): ResourceAiModelEntry[] {
    const byModelId = new Map(
        models.map((m) => [m.modelId, m.listType] as const)
    );
    return [...byModelId.entries()].map(([modelId, listType]) => ({
        modelId,
        listType
    }));
}

async function assertModelEntriesValid(input: {
    orgId: string;
    modelEntries: ResourceAiModelEntry[];
    attachments: ResourceAiProviderAttachment[];
    resourceLabel: string;
}): Promise<string | null> {
    const selectProviderIds = input.attachments
        .filter((a) => a.accessMode === "select")
        .map((a) => a.providerId);

    if (selectProviderIds.length === 0) {
        return "Set at least one attached AI provider to select mode before managing model lists";
    }

    const modelIds = input.modelEntries.map((m) => m.modelId);
    const catalogRows = await db
        .select({
            modelId: aiModels.modelId,
            modelKey: aiModels.modelKey,
            listType: aiModels.listType,
            providerId: aiModels.providerId,
            enabled: aiModels.enabled
        })
        .from(aiModels)
        .innerJoin(aiProviders, eq(aiModels.providerId, aiProviders.providerId))
        .where(
            and(
                inArray(aiModels.modelId, modelIds),
                inArray(aiModels.providerId, selectProviderIds),
                eq(aiProviders.orgId, input.orgId)
            )
        );

    if (catalogRows.length !== modelIds.length) {
        return `One or more model IDs do not exist or do not belong to a select-mode provider on this ${input.resourceLabel}`;
    }

    const catalogById = new Map(catalogRows.map((row) => [row.modelId, row]));
    const selectedAllowsByProvider = new Map<number, string[]>();
    for (const entry of input.modelEntries) {
        const catalog = catalogById.get(entry.modelId);
        if (!catalog) {
            return `One or more model IDs do not exist or do not belong to a select-mode provider on this ${input.resourceLabel}`;
        }
        if (catalog.listType !== entry.listType) {
            return `Model ${entry.modelId} must use listType "${catalog.listType}" to match the provider catalog entry`;
        }
        if (!catalog.enabled) {
            return `Model ${entry.modelId} is disabled on its provider`;
        }
        if (entry.listType === "allow") {
            const keys = selectedAllowsByProvider.get(catalog.providerId) ?? [];
            keys.push(catalog.modelKey);
            selectedAllowsByProvider.set(catalog.providerId, keys);
        }
    }

    const overlapError = await assertNoOverlappingModelKeys(input.attachments, {
        selectedAllowsByProvider
    });
    if (overlapError) {
        return overlapError.error;
    }

    return null;
}