Merge branch 'aig' of https://github.com/fosrl/pangolin into aig

server/aiGatewayServer.ts

@@ -0,0 +1,39 @@
+import express from "express";
+import helmet from "helmet";
+import cors from "cors";
+import config from "@server/lib/config";
+import logger from "@server/logger";
+import {
+    errorHandlerMiddleware,
+    notFoundMiddleware
+} from "@server/middlewares";
+import * as aiGateway from "@server/routers/aiGateway";
+
+const aiGatewayPort = config.getRawConfig().server.ai_gateway_port;
+
+export function createAiGatewayServer() {
+    const aiGatewayServer = express();
+
+    const trustProxy = config.getRawConfig().server.trust_proxy;
+    if (trustProxy) {
+        aiGatewayServer.set("trust proxy", trustProxy);
+    }
+
+    aiGatewayServer.use(helmet());
+    aiGatewayServer.use(cors());
+    aiGatewayServer.use(express.json());
+
+    aiGatewayServer.post("/chat/completions", aiGateway.chatCompletions);
+
+    aiGatewayServer.use(notFoundMiddleware);
+    aiGatewayServer.use(errorHandlerMiddleware);
+
+    aiGatewayServer.listen(aiGatewayPort, (err?: any) => {
+        if (err) throw err;
+        logger.info(
+            `AI gateway server is running on http://localhost:${aiGatewayPort}`
+        );
+    });
+
+    return aiGatewayServer;
+}

server/index.ts

@@ -5,6 +5,7 @@ import { runSetupFunctions } from "./setup";
 import { createApiServer } from "./apiServer";
 import { createNextServer } from "./nextServer";
 import { createInternalServer } from "./internalServer";
+import { createAiGatewayServer } from "./aiGatewayServer";
 import { createIntegrationApiServer } from "./integrationApiServer";
 import {
     ApiKey,
@@ -49,6 +50,7 @@ async function startServers() {
     // Start all servers
     const apiServer = createApiServer();
     const internalServer = createInternalServer();
+    const aiGatewayServer = createAiGatewayServer();
 
     const nextServer = await createNextServer();
     if (config.getRawConfig().traefik.file_mode) {
@@ -67,6 +69,7 @@ async function startServers() {
         apiServer,
         nextServer,
         internalServer,
+        aiGatewayServer,
         integrationServer
     };
 }

server/lib/readConfigFile.ts

@@ -79,7 +79,13 @@ export const configSchema = z
                     .default(3001)
                     .transform(stoi)
                     .pipe(portSchema),
+                ai_gateway_port: portSchema
+                    .optional()
+                    .default(3005)
+                    .transform(stoi)
+                    .pipe(portSchema),
                 badger_override: z.string().optional(),
+                ai_gateway_override: z.string().optional(),
                 next_port: portSchema
                     .optional()
                     .default(3002)
@@ -139,6 +145,7 @@ export const configSchema = z
                 integration_port: 3003,
                 external_port: 3000,
                 internal_port: 3001,
+                ai_gateway_port: 3005,
                 next_port: 3002,
                 internal_hostname: "pangolin",
                 session_cookie_name: "p_session_token",

server/lib/traefik/TraefikConfigManager.ts

@@ -516,9 +516,11 @@ export class TraefikConfigManager {
             const maintenanceHost =
                 config.getRawConfig().server.internal_hostname;
             const pangolinUIUrl = `http://${maintenanceHost}:${maintenancePort}`;
-            const aiGatewayUrl = `http://${maintenanceHost}:${
-                config.getRawConfig().server.internal_port
-            }/api/v1/ai-gateway`;
+            const aiGatewayUrl =
+                config.getRawConfig().server.ai_gateway_override ||
+                `http://${maintenanceHost}:${
+                    config.getRawConfig().server.ai_gateway_port
+                }`;
 
             // logger.debug(`Fetching traefik config for exit node: ${currentExitNode}`);
             traefikConfig = await getTraefikConfig(

server/lib/traefik/getTraefikConfig.ts

@@ -785,7 +785,7 @@ export async function getTraefikConfig(
 
             config_output.http.services[serviceName] = {
                 loadBalancer: {
-                    servers: [{ url: `${aiGatewayUrl}/chat/completions` }],
+                    servers: [{ url: aiGatewayUrl }],
                     passHostHeader: true
                 }
             };

server/private/lib/traefik/getTraefikConfig.ts

@@ -433,41 +433,21 @@ export async function getTraefikConfig(
             )
         );
 
-    let siteResourcesInference: {
-        siteResourceId: number;
-        alias: string | null;
-        ssl: boolean | null;
-        enabled: boolean | null;
-    }[] = [];
-    if (build == "enterprise") {
-        siteResourcesInference = await db
-            .selectDistinct({
-                siteResourceId: siteResources.siteResourceId,
-                alias: siteResources.alias,
-                ssl: siteResources.ssl,
-                enabled: siteResources.enabled
-            })
-            .from(siteResources)
-            .innerJoin(
-                siteResourceAiProviders,
-                eq(
-                    siteResources.siteResourceId,
-                    siteResourceAiProviders.siteResourceId
-                )
+    const siteResourcesInference = await db
+        .selectDistinct({
+            siteResourceId: siteResources.siteResourceId,
+            alias: siteResources.alias,
+            ssl: siteResources.ssl,
+            enabled: siteResources.enabled
+        })
+        .from(siteResources)
+        .where(
+            and(
+                eq(siteResources.mode, "inference"),
+                eq(siteResources.enabled, true),
+                isNotNull(siteResources.alias)
             )
-            .innerJoin(
-                aiProviders,
-                eq(siteResourceAiProviders.providerId, aiProviders.providerId)
-            )
-            .where(
-                and(
-                    eq(siteResources.mode, "inference"),
-                    eq(siteResources.enabled, true),
-                    eq(aiProviders.enabled, true),
-                    isNotNull(siteResources.alias)
-                )
-            );
-    }
+        );
 
     let validCerts: CertificateResult[] = [];
     if (privateConfig.getRawPrivateConfig().flags.use_pangolin_dns) {
@@ -1548,6 +1528,20 @@ export async function getTraefikConfig(
     }
 
     if (aiGatewayUrl) {
+        // The AI gateway may live on a different host than the inference
+        // resource itself (e.g. a remote exit node forwarding to the
+        // central dashboard over a tunnel). passHostHeader would forward
+        // the resource's own Host, which that external host won't
+        // recognize, so we pin the Host header to the gateway's own host
+        // and smuggle the original resource host through in "p-host"
+        // instead (same pattern as the maintenance-page routes above).
+        let aiGatewayHost: string | undefined;
+        try {
+            aiGatewayHost = new URL(aiGatewayUrl).host;
+        } catch {
+            aiGatewayHost = undefined;
+        }
+
         // Public inference resources: same TLS/cert-resolver handling as
         // plain http-mode resources, but the service points at the AI
         // gateway instead of any real backend targets.
@@ -1613,10 +1607,21 @@ export async function getTraefikConfig(
                 }
             }
 
+            const irHeadersMiddlewareName = `${irKey}-headers-middleware`;
+            config_output.http.middlewares[irHeadersMiddlewareName] = {
+                headers: {
+                    customRequestHeaders: {
+                        ...(aiGatewayHost ? { Host: aiGatewayHost } : {}),
+                        "p-host": fullDomain
+                    }
+                }
+            };
+
             const additionalMiddlewares =
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
                 badgerMiddlewareName,
+                irHeadersMiddlewareName,
                 ...additionalMiddlewares
             ];
 
@@ -1647,8 +1652,7 @@ export async function getTraefikConfig(
 
             config_output.http.services[serviceName] = {
                 loadBalancer: {
-                    servers: [{ url: `${aiGatewayUrl}/chat/completions` }],
-                    passHostHeader: true
+                    servers: [{ url: aiGatewayUrl }]
                 }
             };
         }
@@ -1701,8 +1705,22 @@ export async function getTraefikConfig(
                 }
             }
 
+            const srHeadersMiddlewareName = `${srKey}-headers-middleware`;
+            config_output.http.middlewares[srHeadersMiddlewareName] = {
+                headers: {
+                    customRequestHeaders: {
+                        ...(aiGatewayHost ? { Host: aiGatewayHost } : {}),
+                        "p-host": alias
+                    }
+                }
+            };
+
             const additionalMiddlewares =
                 config.getRawConfig().traefik.additional_middlewares || [];
+            const routerMiddlewares = [
+                srHeadersMiddlewareName,
+                ...additionalMiddlewares
+            ];
 
             if (sr.ssl) {
                 config_output.http.routers[routerName + "-redirect"] = {
@@ -1722,7 +1740,7 @@ export async function getTraefikConfig(
                         ? config.getRawConfig().traefik.https_entrypoint
                         : config.getRawConfig().traefik.http_entrypoint
                 ],
-                middlewares: additionalMiddlewares,
+                middlewares: routerMiddlewares,
                 service: serviceName,
                 rule,
                 priority: 100,
@@ -1731,8 +1749,7 @@ export async function getTraefikConfig(
 
             config_output.http.services[serviceName] = {
                 loadBalancer: {
-                    servers: [{ url: `${aiGatewayUrl}/chat/completions` }],
-                    passHostHeader: true
+                    servers: [{ url: aiGatewayUrl }]
                 }
             };
         }

server/routers/aiGateway/chatCompletions.ts

@@ -416,13 +416,19 @@ export async function chatCompletions(
     res: Response
 ): Promise<any> {
     try {
-        const host = (req.headers.host || "").split(":")[0];
+        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({

server/routers/internal.ts

@@ -3,7 +3,6 @@ import * as gerbil from "@server/routers/gerbil";
 import * as traefik from "@server/routers/traefik";
 import * as resource from "./resource";
 import * as badger from "./badger";
-import * as aiGateway from "@server/routers/aiGateway";
 import * as auth from "@server/routers/auth";
 import * as supporterKey from "@server/routers/supporterKey";
 import * as idp from "@server/routers/idp";
@@ -65,9 +64,3 @@ badgerRouter.post("/verify-session", badger.verifyResourceSession);
 
 badgerRouter.post("/exchange-session", badger.exchangeSession);
 
-// AI inference gateway - minimal chat-completions proxy for inference-mode
-// resources/siteResources
-internalRouter.post(
-    "/ai-gateway/chat/completions",
-    aiGateway.chatCompletions
-);

server/routers/traefik/traefikConfigProvider.ts

@@ -20,9 +20,11 @@ export async function traefikConfigProvider(
         const maintenancePort = config.getRawConfig().server.next_port;
         const maintenanceHost = config.getRawConfig().server.internal_hostname;
         const pangolinUIUrl = `http://${maintenanceHost}:${maintenancePort}`;
-        const aiGatewayUrl = `http://${maintenanceHost}:${
-            config.getRawConfig().server.internal_port
-        }/api/v1/ai-gateway`;
+        const aiGatewayUrl =
+            config.getRawConfig().server.ai_gateway_override ||
+            `http://${maintenanceHost}:${
+                config.getRawConfig().server.ai_gateway_port
+            }`;
 
         const traefikConfig = await getTraefikConfig(
             currentExitNodeId,