diff --git a/server/routers/aiGateway/chatCompletions.ts b/server/routers/aiGateway/chatCompletions.ts index f5de5cc75..2f1bfc451 100644 --- a/server/routers/aiGateway/chatCompletions.ts +++ b/server/routers/aiGateway/chatCompletions.ts @@ -1,5 +1,5 @@ import { Request, Response } from "express"; -import { and, eq } from "drizzle-orm"; +import { and, eq, inArray } from "drizzle-orm"; import { AiProvider, aiModels, @@ -94,10 +94,12 @@ type ProviderAttachment = { type ResolvedTarget = { resourceId: number | null; + siteResourceId: number | null; orgId: string | null; attachments: ProviderAttachment[]; - // model IDs on this resource's allowlist (resource-wide) - allowedModelIds: number[]; + // Model IDs on the resource allowlist that belong to allowlist-mode + // providers. Empty when no attached provider uses allowlist mode. + allowlistedModelIds: Set; }; type ProviderSelection = @@ -192,6 +194,9 @@ async function resolveRequestUser( } async function resolveTarget(host: string): Promise { + // TODO: eventually we need to know if it's a private or public resource + // and not just simply check the fullDomain in case there is a private resource with the same fullDomain + const [resourceRow] = await db .select({ resourceId: resources.resourceId, @@ -229,26 +234,40 @@ async function resolveTarget(host: string): Promise { return null; } - const hasAllowlist = attachmentRows.some( - (a) => a.modelAccessMode === "allowlist" - ); - let allowedModelIds: number[] = []; - if (hasAllowlist) { + const attachments: ProviderAttachment[] = attachmentRows.map((a) => ({ + provider: a.provider, + modelAccessMode: a.modelAccessMode as ModelAccessMode + })); + + const allowlistProviderIds = attachments + .filter((a) => a.modelAccessMode === "allowlist") + .map((a) => a.provider.providerId); + const allowlistedModelIds = new Set(); + if (allowlistProviderIds.length > 0) { const restrictions = await db .select({ modelId: resourceAiModels.modelId }) .from(resourceAiModels) - .where(eq(resourceAiModels.resourceId, resourceRow.resourceId)); - allowedModelIds = restrictions.map((r) => r.modelId); + .innerJoin( + aiModels, + eq(resourceAiModels.modelId, aiModels.modelId) + ) + .where( + and( + eq(resourceAiModels.resourceId, resourceRow.resourceId), + inArray(aiModels.providerId, allowlistProviderIds) + ) + ); + for (const row of restrictions) { + allowlistedModelIds.add(row.modelId); + } } return { resourceId: resourceRow.resourceId, + siteResourceId: null, orgId: resourceRow.orgId, - attachments: attachmentRows.map((a) => ({ - provider: a.provider, - modelAccessMode: a.modelAccessMode as ModelAccessMode - })), - allowedModelIds + attachments, + allowlistedModelIds }; } @@ -292,74 +311,52 @@ async function resolveTarget(host: string): Promise { return null; } - const hasAllowlist = attachmentRows.some( - (a) => a.modelAccessMode === "allowlist" - ); - let allowedModelIds: number[] = []; - if (hasAllowlist) { + const attachments: ProviderAttachment[] = attachmentRows.map((a) => ({ + provider: a.provider, + modelAccessMode: a.modelAccessMode as ModelAccessMode + })); + + const allowlistProviderIds = attachments + .filter((a) => a.modelAccessMode === "allowlist") + .map((a) => a.provider.providerId); + const allowlistedModelIds = new Set(); + if (allowlistProviderIds.length > 0) { const restrictions = await db .select({ modelId: siteResourceAiModels.modelId }) .from(siteResourceAiModels) + .innerJoin( + aiModels, + eq(siteResourceAiModels.modelId, aiModels.modelId) + ) .where( - eq( - siteResourceAiModels.siteResourceId, - siteResourceRow.siteResourceId + and( + eq( + siteResourceAiModels.siteResourceId, + siteResourceRow.siteResourceId + ), + inArray(aiModels.providerId, allowlistProviderIds) ) ); - allowedModelIds = restrictions.map((r) => r.modelId); + for (const row of restrictions) { + allowlistedModelIds.add(row.modelId); + } } return { - // siteResources have no per-user auth/policy stack today (see - // the routing comment in getTraefikConfig.ts), so there's no - // resource access token scope to validate a user token against. resourceId: null, + siteResourceId: siteResourceRow.siteResourceId, orgId: siteResourceRow.orgId, - attachments: attachmentRows.map((a) => ({ - provider: a.provider, - modelAccessMode: a.modelAccessMode as ModelAccessMode - })), - allowedModelIds + attachments, + allowlistedModelIds }; } return null; } -async function providerMatchesModel( - attachment: ProviderAttachment, - requestedModel: string, - allowedModelIds: number[] -): Promise { - const [matchedModel] = await db - .select({ - modelId: aiModels.modelId, - enabled: aiModels.enabled - }) - .from(aiModels) - .where( - and( - eq(aiModels.providerId, attachment.provider.providerId), - eq(aiModels.modelKey, requestedModel) - ) - ) - .limit(1); - - if (!matchedModel) { - return false; - } - - if (attachment.modelAccessMode === "catalog") { - return matchedModel.enabled; - } - - // allowlist - return allowedModelIds.includes(matchedModel.modelId); -} - async function selectProvider( attachments: ProviderAttachment[], - allowedModelIds: number[], + allowlistedModelIds: Set, requestedModel: string | undefined ): Promise { if (!requestedModel) { @@ -370,21 +367,55 @@ async function selectProvider( }; } - const candidates: ProviderAttachment[] = []; - for (const attachment of attachments) { - if ( - await providerMatchesModel( - attachment, - requestedModel, - allowedModelIds + const providerById = new Map( + attachments.map((a) => [a.provider.providerId, a]) + ); + const providerIds = [...providerById.keys()]; + if (providerIds.length === 0) { + return { + ok: false, + status: HttpCode.FORBIDDEN, + message: `Model "${requestedModel}" is not permitted on this resource` + }; + } + + // One lookup for the requested model key across all attached providers. + const matchingModels = await db + .select({ + modelId: aiModels.modelId, + providerId: aiModels.providerId, + enabled: aiModels.enabled + }) + .from(aiModels) + .where( + and( + inArray(aiModels.providerId, providerIds), + eq(aiModels.modelKey, requestedModel) ) - ) { - candidates.push(attachment); + ); + + const candidates: AiProvider[] = []; + for (const model of matchingModels) { + const attachment = providerById.get(model.providerId); + if (!attachment) { + continue; + } + + if (attachment.modelAccessMode === "catalog") { + if (model.enabled) { + candidates.push(attachment.provider); + } + continue; + } + + // allowlist: only models explicitly attached to the resource + if (allowlistedModelIds.has(model.modelId)) { + candidates.push(attachment.provider); } } if (candidates.length === 1) { - return { ok: true, provider: candidates[0].provider }; + return { ok: true, provider: candidates[0] }; } if (candidates.length > 1) { @@ -402,15 +433,6 @@ async function selectProvider( }; } -// Generic OpenAI-wire-compatible passthrough. Anthropic's native API uses a -// different path/schema; everything else here is OpenAI-compatible today. -function getCompletionsPath(type: AiProviderType): string { - if (type === "anthropic") { - return "/v1/messages"; - } - return "/chat/completions"; -} - export async function chatCompletions( req: Request, res: Response @@ -438,7 +460,9 @@ export async function chatCompletions( }); } - const { attachments, allowedModelIds, resourceId, orgId } = target; + const { attachments, allowlistedModelIds, resourceId, orgId } = target; + + logger.debug("+++++ gateway target: ", target); // Best-effort identity resolution - not yet enforced, but lets us // start making per-user access decisions (e.g. model/role-based @@ -455,7 +479,7 @@ export async function chatCompletions( const selection = await selectProvider( attachments, - allowedModelIds, + allowlistedModelIds, requestedModel ); if (!selection.ok) { @@ -490,16 +514,35 @@ export async function chatCompletions( }); } - const targetUrl = `${upstreamUrl.replace(/\/$/, "")}${getCompletionsPath( - provider.type as AiProviderType - )}`; + const targetUrl = `${upstreamUrl.replace(/\/$/, "")}`; - const headers: Record = { - "Content-Type": "application/json" - }; - if (authType === "bearer") { - headers["Authorization"] = `Bearer ${apiKey}`; + // Drop hop-by-hop / proxy-only headers. Forwarding Host especially + // breaks Node fetch (TLS/SNI targets the upstream URL while Host + // still says localhost). + const skipHeaders = new Set([ + "p-host", + "host", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "content-length", + "accept-encoding" + ]); + + const headers: Record = {}; + for (const [key, value] of Object.entries(req.headers)) { + if (skipHeaders.has(key.toLowerCase()) || value === undefined) { + continue; + } + headers[key] = Array.isArray(value) ? value.join(", ") : value; } + // TODO: temporary hardcoded auth for testing; restore bearer from authType + headers["x-api-key"] = apiKey; // No dedicated per-request TLS agent is wired up (no extra deps for // this v1 gateway) - toggle the process-wide Node TLS check instead. @@ -510,13 +553,33 @@ export async function chatCompletions( process.env.NODE_TLS_REJECT_UNAUTHORIZED = "0"; } + const body = JSON.stringify(req.body); + + logger.debug("AI gateway upstream request", { + url: targetUrl, + method: "POST", + headers, + body: req.body + }); + let upstreamRes: globalThis.Response; try { upstreamRes = await fetch(targetUrl, { method: "POST", headers, - body: JSON.stringify(req.body) + body }); + } catch (fetchError) { + logger.error({ + message: "AI gateway upstream fetch failed", + url: targetUrl, + error: fetchError, + cause: + fetchError instanceof Error + ? (fetchError as Error & { cause?: unknown }).cause + : undefined + }); + throw fetchError; } finally { if (provider.skipTlsVerification) { if (restoreTlsReject === undefined) {