diff --git a/src/app/api/v1/relay/chat/completions/route.ts b/src/app/api/v1/relay/chat/completions/route.ts index 5d97fd2b2f..32bae3ef08 100644 --- a/src/app/api/v1/relay/chat/completions/route.ts +++ b/src/app/api/v1/relay/chat/completions/route.ts @@ -18,11 +18,165 @@ import { hashToken, sanitizeForensicHeader, } from "./relaySecurity"; +import { + getBifrostRoutingConfig, + resolveRelayRoutingBackend, + shouldTryBifrost, + type BifrostRoutingConfig, +} from "./routingBackend"; +import type { RelayToken } from "@/lib/db/relayProxies"; const JSON_CORS_HEADERS = { ...CORS_HEADERS, "Content-Type": "application/json" } as const; const injectionGuard = createInjectionGuard(); +type RelayUsageStatus = "success" | "error"; + +function recordUsage( + tokenId: string, + request: Request, + startTime: number, + clientIp: string, + userAgent: string | null, + status: RelayUsageStatus, + statusCode: number +) { + recordRelayUsage(tokenId, { + requestId: request.headers.get("x-request-id") || undefined, + status, + statusCode, + latencyMs: Date.now() - startTime, + clientIp, + userAgent, + }); +} + +function finalizeReadableStream( + body: ReadableStream, + onFinalize: (error?: unknown) => void +): ReadableStream { + const reader = body.getReader(); + let finalized = false; + + const finalizeOnce = (error?: unknown) => { + if (finalized) return; + finalized = true; + onFinalize(error); + }; + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + finalizeOnce(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (error) { + finalizeOnce(error); + controller.error(error); + } + }, + async cancel(reason) { + try { + await reader.cancel(reason); + } finally { + finalizeOnce(reason); + } + }, + }); +} + +async function forwardToBifrost( + request: Request, + body: unknown, + token: RelayToken, + config: BifrostRoutingConfig, + startTime: number, + clientIp: string, + userAgent: string | null +): Promise { + const wantsStream = + Boolean((body as { stream?: boolean } | null)?.stream) && config.streamingEnabled; + const upstreamHeaders: Record = { + "Content-Type": "application/json", + "x-relay-token-id": token.id, + "x-relay-client-ip": clientIp, + }; + if (config.apiKey) { + upstreamHeaders.Authorization = `Bearer ${config.apiKey}`; + } + + const ac = new AbortController(); + let timedOut = false; + const tid = setTimeout(() => { + timedOut = true; + ac.abort(); + }, config.timeoutMs); + + try { + const upstream = await fetch(`${config.baseUrl}/v1/chat/completions`, { + method: "POST", + headers: upstreamHeaders, + body: JSON.stringify(body), + signal: ac.signal, + }); + clearTimeout(tid); + + const headers = new Headers(upstream.headers); + headers.set("X-Routed-By", "bifrost"); + headers.set("X-Routing-Backend", "bifrost"); + headers.set("X-Relay-Token", token.tokenPrefix + "..."); + if (!wantsStream) { + headers.set("Content-Type", upstream.headers.get("Content-Type") ?? "application/json"); + } + + if (wantsStream && upstream.body) { + const stream = finalizeReadableStream(upstream.body, (error) => { + recordUsage( + token.id, + request, + startTime, + clientIp, + userAgent, + error || upstream.status >= 500 ? "error" : "success", + upstream.status + ); + }); + + return new Response(stream, { + status: upstream.status, + headers, + }); + } + + recordUsage( + token.id, + request, + startTime, + clientIp, + userAgent, + upstream.status < 500 ? "success" : "error", + upstream.status + ); + + return new Response(upstream.body, { + status: upstream.status, + headers, + }); + } catch (error) { + clearTimeout(tid); + const isAbort = error instanceof Error && error.name === "AbortError"; + throw new Error( + isAbort + ? `Bifrost sidecar timed out after ${config.timeoutMs}ms` + : `Bifrost sidecar unreachable: ${error instanceof Error ? error.message : String(error)}` + ); + } +} + export async function OPTIONS() { return handleCorsOptions(); } @@ -112,11 +266,13 @@ export async function POST(request: Request) { // 3. Clone request and forward to internal handler const cloned = request.clone(); + let parsedBody: unknown = null; + // Prompt injection guard (same as main endpoint) try { - const body = await cloned.json().catch(() => null); - if (body) { - const { blocked, result } = injectionGuard(body); + parsedBody = await cloned.json().catch(() => null); + if (parsedBody) { + const { blocked, result } = injectionGuard(parsedBody); if (blocked) { recordRelayUsage(token.id, { requestId: request.headers.get("x-request-id") || undefined, @@ -142,7 +298,7 @@ export async function POST(request: Request) { // Check allowed models const allowedModels: string[] = JSON.parse(token.allowedModels); if (allowedModels.length > 0 && !allowedModels.includes("*")) { - const model = (body as { model?: string }).model || ""; + const model = (parsedBody as { model?: string }).model || ""; const allowed = allowedModels.some( (p) => model === p || (p.endsWith("*") && model.startsWith(p.slice(0, -1))) ); @@ -162,6 +318,38 @@ export async function POST(request: Request) { // Continue even if guard fails } + const backend = resolveRelayRoutingBackend(); + const bifrostConfig = getBifrostRoutingConfig(); + if (shouldTryBifrost(backend, bifrostConfig)) { + try { + return await forwardToBifrost( + request, + parsedBody, + token, + bifrostConfig, + startTime, + clientIp, + userAgent + ); + } catch (error) { + if (backend === "bifrost") { + recordUsage(token.id, request, startTime, clientIp, userAgent, "error", 502); + return new Response( + JSON.stringify( + buildErrorBody(502, error instanceof Error ? error.message : String(error)) + ), + { + status: 502, + headers: { + ...JSON_CORS_HEADERS, + "X-Bifrost-Fallback": "/api/v1/relay/chat/completions", + }, + } + ); + } + } + } + // 4. Proxy to internal handler const originalRequest = new Request( request.url.replace("/relay/chat/completions", "/chat/completions"), @@ -183,6 +371,10 @@ export async function POST(request: Request) { // Add relay headers const newHeaders = new Headers(response.headers); newHeaders.set("X-Relay-Token", token.tokenPrefix + "..."); + newHeaders.set("X-Routing-Backend", "ts"); + if (backend === "auto" && bifrostConfig?.enabled) { + newHeaders.set("X-Routing-Fallback", "bifrost"); + } return new Response(response.body, { status: response.status, diff --git a/src/app/api/v1/relay/chat/completions/routingBackend.ts b/src/app/api/v1/relay/chat/completions/routingBackend.ts new file mode 100644 index 0000000000..36d4a6eae0 --- /dev/null +++ b/src/app/api/v1/relay/chat/completions/routingBackend.ts @@ -0,0 +1,41 @@ +export type RelayRoutingBackend = "ts" | "bifrost" | "auto"; + +const VALID_BACKENDS = new Set(["ts", "bifrost", "auto"]); + +export interface BifrostRoutingConfig { + baseUrl: string; + apiKey?: string; + timeoutMs: number; + streamingEnabled: boolean; + enabled: boolean; +} + +export function getBifrostRoutingConfig(env: NodeJS.ProcessEnv = process.env): BifrostRoutingConfig | null { + const baseUrl = env.BIFROST_BASE_URL?.replace(/\/$/, ""); + if (!baseUrl) return null; + const timeoutMs = Number.parseInt(env.BIFROST_TIMEOUT_MS || "", 10); + + return { + baseUrl, + apiKey: env.BIFROST_API_KEY || env.OMNIROUTE_BIFROST_KEY || undefined, + timeoutMs: Number.isFinite(timeoutMs) && timeoutMs > 0 ? timeoutMs : 30000, + streamingEnabled: env.BIFROST_STREAMING_ENABLED !== "0", + enabled: env.BIFROST_ENABLED !== "0", + }; +} + +export function resolveRelayRoutingBackend(env: NodeJS.ProcessEnv = process.env): RelayRoutingBackend { + const configured = env.OMNIROUTE_RELAY_BACKEND || env.RELAY_ROUTING_BACKEND; + if (configured && VALID_BACKENDS.has(configured as RelayRoutingBackend)) { + return configured as RelayRoutingBackend; + } + + return getBifrostRoutingConfig(env)?.enabled ? "auto" : "ts"; +} + +export function shouldTryBifrost( + backend: RelayRoutingBackend, + config: BifrostRoutingConfig | null +): config is BifrostRoutingConfig { + return Boolean(config?.enabled && backend !== "ts"); +} diff --git a/tests/unit/api/v1/relay-routing-backend.test.ts b/tests/unit/api/v1/relay-routing-backend.test.ts new file mode 100644 index 0000000000..1545fbb4f9 --- /dev/null +++ b/tests/unit/api/v1/relay-routing-backend.test.ts @@ -0,0 +1,68 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { + getBifrostRoutingConfig, + resolveRelayRoutingBackend, + shouldTryBifrost, +} from "../../../../src/app/api/v1/relay/chat/completions/routingBackend.ts"; + +test("relay routing backend defaults to TypeScript without bifrost", () => { + const env = {}; + + assert.equal(resolveRelayRoutingBackend(env), "ts"); + assert.equal(getBifrostRoutingConfig(env), null); +}); + +test("relay routing backend auto-enables bifrost when base URL is configured", () => { + const env = { + BIFROST_BASE_URL: "http://127.0.0.1:8080/", + OMNIROUTE_BIFROST_KEY: "sidecar-key", + BIFROST_TIMEOUT_MS: "250", + }; + + const config = getBifrostRoutingConfig(env); + + assert.equal(resolveRelayRoutingBackend(env), "auto"); + assert.equal(config?.baseUrl, "http://127.0.0.1:8080"); + assert.equal(config?.apiKey, "sidecar-key"); + assert.equal(config?.timeoutMs, 250); + assert.equal(config?.streamingEnabled, true); + assert.equal(config?.enabled, true); + assert.equal(shouldTryBifrost("auto", config), true); +}); + +test("relay routing backend honors explicit TS and strict bifrost modes", () => { + const env = { + BIFROST_BASE_URL: "http://127.0.0.1:8080", + OMNIROUTE_RELAY_BACKEND: "ts", + }; + const config = getBifrostRoutingConfig(env); + + assert.equal(resolveRelayRoutingBackend(env), "ts"); + assert.equal(shouldTryBifrost("ts", config), false); + + env.OMNIROUTE_RELAY_BACKEND = "bifrost"; + assert.equal(resolveRelayRoutingBackend(env), "bifrost"); + assert.equal(shouldTryBifrost("bifrost", config), true); +}); + +test("relay routing backend respects bifrost killswitch", () => { + const env = { + BIFROST_BASE_URL: "http://127.0.0.1:8080", + BIFROST_ENABLED: "0", + }; + const config = getBifrostRoutingConfig(env); + + assert.equal(resolveRelayRoutingBackend(env), "ts"); + assert.equal(config?.enabled, false); + assert.equal(shouldTryBifrost("auto", config), false); +}); + +test("relay routing backend falls back on invalid timeout values", () => { + const env = { + BIFROST_BASE_URL: "http://127.0.0.1:8080", + BIFROST_TIMEOUT_MS: "not-a-number", + }; + + assert.equal(getBifrostRoutingConfig(env)?.timeoutMs, 30000); +});