add relay routing backend selection (#5315)

Integrated into release/v3.8.41. Relay routing backend selection (ts/bifrost/auto), env-gated and merge-result-validated.

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<Uint8Array>,
+  onFinalize: (error?: unknown) => void
+): ReadableStream<Uint8Array> {
+  const reader = body.getReader();
+  let finalized = false;
+
+  const finalizeOnce = (error?: unknown) => {
+    if (finalized) return;
+    finalized = true;
+    onFinalize(error);
+  };
+
+  return new ReadableStream<Uint8Array>({
+    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<Response> {
+  const wantsStream =
+    Boolean((body as { stream?: boolean } | null)?.stream) && config.streamingEnabled;
+  const upstreamHeaders: Record<string, string> = {
+    "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,

src/app/api/v1/relay/chat/completions/routingBackend.ts

@@ -0,0 +1,41 @@
+export type RelayRoutingBackend = "ts" | "bifrost" | "auto";
+
+const VALID_BACKENDS = new Set<RelayRoutingBackend>(["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");
+}