diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index a947f3d86..bdc4b0eb5 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -28,6 +28,7 @@ import { providerHasCapability, type AiCapability } from "@server/lib/aiCapabilities"; +import { proxyAiGatewayToSiteTarget } from "@server/routers/aiGateway/targetRouting"; import { SESSION_COOKIE_NAME, validateSessionToken @@ -486,6 +487,10 @@ export async function handleAiGatewayProxy( const { provider } = selection; + if (provider.type === "custom" && provider.routingMode === "target") { + return await proxyAiGatewayToSiteTarget(req, res, provider); + } + const upstreamUrl = provider.upstreamUrl; const authType = provider.authType as AiProviderAuthType; diff --git a/server/routers/aiGateway/targetRouting.ts b/server/routers/aiGateway/targetRouting.ts new file mode 100644 index 000000000..ad9f08f79 --- /dev/null +++ b/server/routers/aiGateway/targetRouting.ts @@ -0,0 +1,204 @@ +import { Request, Response } from "express"; +import { and, eq } from "drizzle-orm"; +import { AiProvider, db, exitNodes, sites, targets } from "@server/db"; +import { localCache } from "@server/lib/cache"; +import logger from "@server/logger"; +import HttpCode from "@server/types/HttpCode"; + +// Short TTL: long enough to spare the DB on a burst of requests, short +// enough that target/site changes (added, removed, exit node moved) show up +// almost immediately without needing explicit cache invalidation. +const PROVIDER_TARGETS_TTL_SEC = 7; + +// Header gerbil reads to know which host:port (reachable over the +// WireGuard network) to rewrite an incoming /router/* request to. Must +// match gerbil's `pangolinDestHeader` constant. +const PANGOLIN_DEST_HEADER = "p-dest-header"; + +const SKIP_HEADERS = new Set([ + "p-host", + "host", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "content-length", + "accept-encoding" +]); + +type ResolvedProviderTarget = { + targetId: number; + // ":", passed to gerbil as the + // destination to proxy the request to over the WireGuard tunnel. + destination: string; + // The target's site's exit node HTTP API base URL (gerbil's /router/*). + gerbilBaseUrl: string; +}; + +async function fetchProviderTargets( + providerId: number +): Promise { + const rows = await db + .select({ + targetId: targets.targetId, + internalPort: targets.internalPort, + port: targets.port, + exitNodeSubnet: sites.exitNodeSubnet, + reachableAt: exitNodes.reachableAt + }) + .from(targets) + .innerJoin(sites, eq(targets.siteId, sites.siteId)) + .innerJoin(exitNodes, eq(sites.exitNodeId, exitNodes.exitNodeId)) + .where( + and(eq(targets.providerId, providerId), eq(targets.enabled, true)) + ); + + const resolved: ResolvedProviderTarget[] = []; + for (const row of rows) { + // Sites not yet connected to an exit node (no subnet assigned) or + // whose exit node has no known HTTP address can't be routed to. + if (!row.exitNodeSubnet || !row.reachableAt) { + continue; + } + const host = row.exitNodeSubnet.split("/")[0]; + const port = row.internalPort ?? row.port; + resolved.push({ + targetId: row.targetId, + destination: `${host}:${port}`, + gerbilBaseUrl: row.reachableAt + }); + } + + return resolved; +} + +async function getProviderTargets( + providerId: number +): Promise { + const cacheKey = `aiGateway:providerTargets:${providerId}`; + const cached = localCache.get(cacheKey); + if (cached !== undefined) { + return cached; + } + + const resolved = await fetchProviderTargets(providerId); + localCache.set(cacheKey, resolved, PROVIDER_TARGETS_TTL_SEC); + return resolved; +} + +// Round-robin cursor per provider. Process-local and unpersisted - fine +// since it only needs to spread load across targets, not guarantee a +// perfectly even distribution across restarts or multiple server instances. +const roundRobinCursors = new Map(); + +function pickTarget( + providerId: number, + providerTargets: ResolvedProviderTarget[] +): ResolvedProviderTarget { + const cursor = roundRobinCursors.get(providerId) ?? 0; + roundRobinCursors.set(providerId, cursor + 1); + return providerTargets[cursor % providerTargets.length]; +} + +function pathFromRequest(req: Request): string { + const raw = + req.originalUrl?.split("?")[0] || req.url?.split("?")[0] || req.path; + return raw.startsWith("/") ? raw : `/${raw}`; +} + +/** + * Proxies an AI gateway request to one of a "custom" / "target" routing-mode + * provider's site targets, via that site's gerbil sidecar. Gerbil's + * /router/* endpoint forwards the request (untouched body, same path minus + * the /router prefix, and all headers besides PANGOLIN_DEST_HEADER) over the + * WireGuard tunnel to the destination named in that header. Always writes a + * response to `res`, including on failure. + */ +export async function proxyAiGatewayToSiteTarget( + req: Request, + res: Response, + provider: AiProvider +): Promise { + const providerTargets = await getProviderTargets(provider.providerId); + if (providerTargets.length === 0) { + res.status(HttpCode.INTERNAL_SERVER_ERROR).json({ + error: { + message: + "AI provider has no reachable site targets configured" + } + }); + return; + } + + const target = pickTarget(provider.providerId, providerTargets); + const gerbilUrl = `${target.gerbilBaseUrl.replace(/\/+$/, "")}/router${pathFromRequest(req)}`; + + const headers: Record = {}; + for (const [key, value] of Object.entries(req.headers)) { + if (SKIP_HEADERS.has(key.toLowerCase()) || value === undefined) { + continue; + } + headers[key] = Array.isArray(value) ? value.join(", ") : value; + } + headers[PANGOLIN_DEST_HEADER] = target.destination; + + const body = JSON.stringify(req.body); + + logger.debug("AI gateway target-routed request", { + providerId: provider.providerId, + targetId: target.targetId, + destination: target.destination, + url: gerbilUrl + }); + + let upstreamRes: globalThis.Response; + try { + upstreamRes = await fetch(gerbilUrl, { + method: "POST", + headers, + body + }); + } catch (fetchError) { + logger.error({ + message: "AI gateway target proxy request failed", + url: gerbilUrl, + targetId: target.targetId, + error: fetchError, + cause: + fetchError instanceof Error + ? (fetchError as Error & { cause?: unknown }).cause + : undefined + }); + res.status(HttpCode.BAD_GATEWAY).json({ + error: { message: "Failed to reach AI provider target" } + }); + return; + } + + 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); + } + res.end(); + return; + } + + const text = await upstreamRes.text(); + res.send(text); +}