pangolin/server/routers/aiGateway/chatCompletions.ts

import { Request, Response } from "express";
import { and, eq, inArray } from "drizzle-orm";
import {
    AiProvider,
    aiModels,
    aiProviders,
    clients,
    db,
    exitNodes,
    resourceAiModels,
    resourceAiProviders,
    resources,
    siteResourceAiModels,
    siteResourceAiProviders,
    siteResources,
    users
} from "@server/db";
import config from "@server/lib/config";
import { decrypt } from "@server/lib/crypto";
import {
    AiProviderAuthType,
    AiProviderRoutingMode,
    AiProviderType,
    resolveAiProviderConfig
} from "@server/lib/aiProviderDefaults";
import {
    SESSION_COOKIE_NAME,
    validateSessionToken
} from "@server/auth/sessions/app";
import { getUserOrgRoles } from "@server/lib/userOrgRoles";
import { isIpInCidr } from "@server/lib/ip";
import { localCache } from "@server/lib/cache";
import logger from "@server/logger";
import HttpCode from "@server/types/HttpCode";
import type { ModelAccessMode } from "@server/lib/aiInferenceResource";

// Short-lived local caches so a burst of requests from the same IP/user
// doesn't hit the database on every single request. None of this is
// security-critical to cache aggressively (identity is re-derived from the
// session cookie or from a client's exit-node-scoped subnet each time), so
// a small TTL is just an efficiency win, not a trust boundary.
const EXIT_NODE_RANGES_CACHE_KEY = "aiGateway:exitNodeRanges";
const EXIT_NODE_RANGES_TTL_SEC = 6000;
const CLIENT_BY_IP_TTL_SEC = 30;
const REQUEST_USER_TTL_SEC = 30;

type CachedClient = { clientId: number; userId: string | null } | null;

// The set of CIDRs an exit node manages; client exitNodeSubnets are always
// /32s carved out of one of these ranges. Checking against this small,
// cacheable list lets us skip the (much more frequent) per-IP client lookup
// entirely for traffic that could never match a client anyway.
async function getExitNodeRanges(): Promise<string[]> {
    const cached = localCache.get<string[]>(EXIT_NODE_RANGES_CACHE_KEY);
    if (cached) {
        return cached;
    }

    const rows = await db
        .select({ address: exitNodes.address })
        .from(exitNodes);
    const ranges = rows.map((r) => r.address);

    localCache.set(
        EXIT_NODE_RANGES_CACHE_KEY,
        ranges,
        EXIT_NODE_RANGES_TTL_SEC
    );
    return ranges;
}

async function findClientByIp(ip: string): Promise<CachedClient> {
    const cacheKey = `aiGateway:clientByIp:${ip}`;
    const cached = localCache.get<CachedClient>(cacheKey);
    if (cached !== undefined) {
        return cached;
    }

    const [client] = await db
        .select({ clientId: clients.clientId, userId: clients.userId })
        .from(clients)
        .where(eq(clients.exitNodeSubnet, `${ip}/32`))
        .limit(1);

    const result: CachedClient = client || null;
    localCache.set(cacheKey, result, CLIENT_BY_IP_TTL_SEC);
    return result;
}

type ProviderAttachment = {
    provider: AiProvider;
    modelAccessMode: ModelAccessMode;
};

type ResolvedTarget = {
    resourceId: number | null;
    siteResourceId: number | null;
    orgId: string | null;
    attachments: ProviderAttachment[];
    // Model IDs on the resource allowlist that belong to allowlist-mode
    // providers. Empty when no attached provider uses allowlist mode.
    allowlistedModelIds: Set<number>;
};

type ProviderSelection =
    | { ok: true; provider: AiProvider }
    | { ok: false; status: number; message: string };

export type RequestUser = {
    userId: string;
    username: string;
    email: string | null;
    name: string | null;
    role: string | null;
};

async function buildRequestUser(
    userId: string,
    orgId: string | null
): Promise<RequestUser | null> {
    const cacheKey = `aiGateway:requestUser:${userId}:${orgId || ""}`;
    const cached = localCache.get<RequestUser | null>(cacheKey);
    if (cached !== undefined) {
        return cached;
    }

    const [user] = await db
        .select()
        .from(users)
        .where(eq(users.userId, userId))
        .limit(1);

    if (!user) {
        localCache.set(cacheKey, null, REQUEST_USER_TTL_SEC);
        return null;
    }

    const orgRoles = orgId ? await getUserOrgRoles(user.userId, orgId) : [];

    const requestUser: RequestUser = {
        userId: user.userId,
        username: user.username,
        email: user.email,
        name: user.name,
        role: orgRoles.map((r) => r.roleName).join(", ") || null
    };

    localCache.set(cacheKey, requestUser, REQUEST_USER_TTL_SEC);
    return requestUser;
}

async function resolveRequestUser(
    req: Request,
    resourceId: number | null,
    orgId: string | null
): Promise<RequestUser | null> {
    // Public resources behind badger: badger passes the resource session
    // cookie through to the backend (same mechanism the browser gateway,
    // e.g. the SSH page, relies on), so we can validate it exactly like
    // verifySessionUserMiddleware does for the dashboard.
    const sessionToken = req.cookies?.[SESSION_COOKIE_NAME];
    if (sessionToken) {
        const { session, user } = await validateSessionToken(sessionToken);
        if (session && user) {
            return buildRequestUser(user.userId, orgId);
        }
    }

    // TODO: MAKE SURE THIS CAN NOT BE SPOOFED AND CAN BE TRUSTED AS AN INTERNAL ADDRESS FROM A NODE

    // No session cookie - fall back to identifying the caller by source IP.
    // A client's exitNodeSubnet is a /32 handed out from one of our exit
    // node's address ranges, so an IP that isn't inside any of those ranges
    // can never belong to a client and we can skip the DB entirely.
    const ip = req.ip;
    if (!ip) {
        return null;
    }

    const exitNodeRanges = await getExitNodeRanges();
    const inExitNodeRange = exitNodeRanges.some((range) =>
        isIpInCidr(ip, range)
    );
    if (!inExitNodeRange) {
        return null;
    }

    const client = await findClientByIp(ip);
    if (!client || !client.userId) {
        return null;
    }

    return buildRequestUser(client.userId, orgId);
}

async function resolveTarget(host: string): Promise<ResolvedTarget | null> {
    // 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,
            orgId: resources.orgId
        })
        .from(resources)
        .where(
            and(
                eq(resources.fullDomain, host),
                eq(resources.mode, "inference"),
                eq(resources.enabled, true)
            )
        )
        .limit(1);

    if (resourceRow) {
        const attachmentRows = await db
            .select({
                modelAccessMode: resourceAiProviders.modelAccessMode,
                provider: aiProviders
            })
            .from(resourceAiProviders)
            .innerJoin(
                aiProviders,
                eq(resourceAiProviders.providerId, aiProviders.providerId)
            )
            .where(
                and(
                    eq(resourceAiProviders.resourceId, resourceRow.resourceId),
                    eq(aiProviders.enabled, true)
                )
            );

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

        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<number>();
        if (allowlistProviderIds.length > 0) {
            const restrictions = await db
                .select({ modelId: resourceAiModels.modelId })
                .from(resourceAiModels)
                .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,
            allowlistedModelIds
        };
    }

    const [siteResourceRow] = await db
        .select({
            siteResourceId: siteResources.siteResourceId,
            orgId: siteResources.orgId
        })
        .from(siteResources)
        .where(
            and(
                eq(siteResources.alias, host),
                eq(siteResources.mode, "inference"),
                eq(siteResources.enabled, true)
            )
        )
        .limit(1);

    if (siteResourceRow) {
        const attachmentRows = await db
            .select({
                modelAccessMode: siteResourceAiProviders.modelAccessMode,
                provider: aiProviders
            })
            .from(siteResourceAiProviders)
            .innerJoin(
                aiProviders,
                eq(siteResourceAiProviders.providerId, aiProviders.providerId)
            )
            .where(
                and(
                    eq(
                        siteResourceAiProviders.siteResourceId,
                        siteResourceRow.siteResourceId
                    ),
                    eq(aiProviders.enabled, true)
                )
            );

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

        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<number>();
        if (allowlistProviderIds.length > 0) {
            const restrictions = await db
                .select({ modelId: siteResourceAiModels.modelId })
                .from(siteResourceAiModels)
                .innerJoin(
                    aiModels,
                    eq(siteResourceAiModels.modelId, aiModels.modelId)
                )
                .where(
                    and(
                        eq(
                            siteResourceAiModels.siteResourceId,
                            siteResourceRow.siteResourceId
                        ),
                        inArray(aiModels.providerId, allowlistProviderIds)
                    )
                );
            for (const row of restrictions) {
                allowlistedModelIds.add(row.modelId);
            }
        }

        return {
            resourceId: null,
            siteResourceId: siteResourceRow.siteResourceId,
            orgId: siteResourceRow.orgId,
            attachments,
            allowlistedModelIds
        };
    }

    return null;
}

async function selectProvider(
    attachments: ProviderAttachment[],
    allowlistedModelIds: Set<number>,
    requestedModel: string | undefined
): Promise<ProviderSelection> {
    if (!requestedModel) {
        return {
            ok: false,
            status: HttpCode.FORBIDDEN,
            message: "A model must be specified for this resource"
        };
    }

    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)
            )
        );

    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] };
    }

    if (candidates.length > 1) {
        return {
            ok: false,
            status: HttpCode.FORBIDDEN,
            message: `Model "${requestedModel}" is ambiguous across multiple AI providers on this resource`
        };
    }

    return {
        ok: false,
        status: HttpCode.FORBIDDEN,
        message: `Model "${requestedModel}" is not permitted on this resource`
    };
}

export async function chatCompletions(
    req: Request,
    res: Response
): Promise<any> {
    try {
        const host = (
            (req.headers["p-host"] as string | undefined) ||
            req.headers.host ||
            ""
        ).split(":")[0];
        if (!host) {
            return res
                .status(HttpCode.BAD_REQUEST)
                .json({ error: { message: "Missing Host header" } });
        }

        logger.info(`AI gateway request for host: ${host}`);

        const target = await resolveTarget(host);
        if (!target) {
            return res.status(HttpCode.NOT_FOUND).json({
                error: {
                    message: "No inference resource found for this host"
                }
            });
        }

        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
        // restrictions) without another round of plumbing later.
        const requestUser = await resolveRequestUser(req, resourceId, orgId);
        if (requestUser) {
            logger.debug(
                `AI gateway request from user ${requestUser.userId} (${requestUser.username})`
            );
        }

        const requestedModel =
            typeof req.body?.model === "string" ? req.body.model : undefined;

        const selection = await selectProvider(
            attachments,
            allowlistedModelIds,
            requestedModel
        );
        if (!selection.ok) {
            return res.status(selection.status).json({
                error: { message: selection.message }
            });
        }

        const { provider } = selection;

        if (!provider.apiKey) {
            return res.status(HttpCode.INTERNAL_SERVER_ERROR).json({
                error: { message: "AI provider has no API key configured" }
            });
        }

        const secret = config.getRawConfig().server.secret!;
        const apiKey = decrypt(provider.apiKey, secret);

        const { upstreamUrl, authType } = resolveAiProviderConfig({
            type: provider.type as AiProviderType,
            upstreamUrl: provider.upstreamUrl,
            authType: provider.authType as AiProviderAuthType | null,
            routingMode: provider.routingMode as AiProviderRoutingMode | null
        });

        if (!upstreamUrl) {
            return res.status(HttpCode.INTERNAL_SERVER_ERROR).json({
                error: {
                    message: "AI provider has no upstream URL configured"
                }
            });
        }

        const targetUrl = `${upstreamUrl.replace(/\/$/, "")}`;

        // 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<string, string> = {};
        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.
        // Known limitation: this is not safe under concurrent requests mixing
        // skipTlsVerification providers with strict ones.
        const restoreTlsReject = process.env.NODE_TLS_REJECT_UNAUTHORIZED;
        if (provider.skipTlsVerification) {
            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
            });
        } 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) {
                    delete process.env.NODE_TLS_REJECT_UNAUTHORIZED;
                } else {
                    process.env.NODE_TLS_REJECT_UNAUTHORIZED = restoreTlsReject;
                }
            }
        }

        const contentType = upstreamRes.headers.get("content-type") || "";
        const isStream =
            req.body?.stream === true ||
            contentType.includes("text/event-stream");

        res.status(upstreamRes.status);
        res.setHeader("Content-Type", contentType || "application/json");

        if (isStream && upstreamRes.body) {
            res.flushHeaders();
            const reader = upstreamRes.body.getReader();
            while (true) {
                const { done, value } = await reader.read();
                if (done) break;
                res.write(value);
            }
            return res.end();
        }

        const text = await upstreamRes.text();
        return res.send(text);
    } catch (error) {
        logger.error(error);
        return res.status(HttpCode.INTERNAL_SERVER_ERROR).json({
            error: { message: "Failed to proxy inference request" }
        });
    }
}