@nvae/llmswitch 0.2.0 → 0.5.0

package/dist/bridge/server.js

@@ -1,17 +1,88 @@
 import { createServer, } from "node:http";
-import { buildProxyEnv } from "../utils/proxy.js";
-import { emptyProxy } from "../types.js";
-import { readBridgeState, readBridgeUpstreams } from "./state.js";
+import { parseBridgeRuntimeLimits } from "./runtime.js";
+import { requestWithNodeTransport, } from "./transport.js";
+import { constantTimeTokenEqual, readBridgeState, readBridgeUpstreams, } from "./state.js";
 import { anthropicToChatRequest } from "./anthropic-translate-request.js";
 import { chatChunkToAnthropicEvents, chatCompletionToAnthropicMessage, createAnthropicStreamState, forceCompleteAnthropicStream, parseChatSseLine, } from "./anthropic-translate-response.js";
 import { collectCustomToolNames, responsesToChatRequest, responsesToCompletionsRequest, } from "./translate-request.js";
 import { chatChunkToResponsesEvents, chatCompletionToResponse, createStreamState, forceCompleteStream, parseChatSseLine as parseChatSseLineResponses, } from "./translate-response.js";
-function readBody(req) {
+function headerValue(value) {
+    return Array.isArray(value) ? value[0] : value;
+}
+function controlToken(req) {
+    return headerValue(req.headers["x-llm-switch-control"]);
+}
+class RequestBodyTooLargeError extends Error {
+    maxBytes;
+    constructor(maxBytes) {
+        super(`Request body exceeded ${maxBytes} bytes`);
+        this.maxBytes = maxBytes;
+        this.name = "RequestBodyTooLargeError";
+    }
+}
+function bearerToken(req) {
+    const authorization = headerValue(req.headers.authorization);
+    const match = authorization?.match(/^Bearer\s+([^\s]+)$/i);
+    return match?.[1];
+}
+function authenticateDataRequest(req, upstream, tool) {
+    if (!upstream?.clientToken || upstream.migrationRequired)
+        return false;
+    const bearer = bearerToken(req);
+    if (tool === "codex") {
+        return constantTimeTokenEqual(upstream.clientToken, bearer);
+    }
+    const apiKey = headerValue(req.headers["x-api-key"]);
+    if (bearer && apiKey && !constantTimeTokenEqual(bearer, apiKey)) {
+        return false;
+    }
+    return constantTimeTokenEqual(upstream.clientToken, bearer || apiKey);
+}
+function authenticateModelsRequest(req, upstreams) {
+    return (authenticateDataRequest(req, upstreams.codex, "codex") ||
+        authenticateDataRequest(req, upstreams.claude, "claude"));
+}
+function readBody(req, maxBytes = parseBridgeRuntimeLimits().maxBodyBytes) {
     return new Promise((resolve, reject) => {
         const chunks = [];
-        req.on("data", (c) => chunks.push(Buffer.isBuffer(c) ? c : Buffer.from(c)));
-        req.on("end", () => resolve(Buffer.concat(chunks)));
-        req.on("error", reject);
+        let bytes = 0;
+        let settled = false;
+        const onData = (value) => {
+            if (settled)
+                return;
+            const chunk = Buffer.isBuffer(value) ? value : Buffer.from(value);
+            bytes += chunk.length;
+            if (bytes > maxBytes) {
+                settled = true;
+                cleanup();
+                req.resume();
+                reject(new RequestBodyTooLargeError(maxBytes));
+                return;
+            }
+            chunks.push(chunk);
+        };
+        const onEnd = () => {
+            if (settled)
+                return;
+            settled = true;
+            cleanup();
+            resolve(Buffer.concat(chunks));
+        };
+        const onError = (error) => {
+            if (settled)
+                return;
+            settled = true;
+            cleanup();
+            reject(error);
+        };
+        const cleanup = () => {
+            req.off("data", onData);
+            req.off("end", onEnd);
+            req.off("error", onError);
+        };
+        req.on("data", onData);
+        req.on("end", onEnd);
+        req.on("error", onError);
     });
 }
 function sendJson(res, status, body) {
@@ -33,56 +104,45 @@ function joinUrl(baseUrl, path) {
     }
     return `${base}${p}`;
 }
-function upstreamHeaders(upstream, incoming) {
+function upstreamHeaders(upstream) {
     const headers = {
         Accept: "application/json",
         "Content-Type": "application/json",
     };
-    const incomingAuth = incoming.headers.authorization;
-    const incomingKey = incoming.headers["x-api-key"];
-    if (incomingAuth) {
-        headers.Authorization = Array.isArray(incomingAuth)
-            ? incomingAuth[0]
-            : incomingAuth;
-    }
-    else if (incomingKey) {
-        const key = Array.isArray(incomingKey) ? incomingKey[0] : incomingKey;
-        headers.Authorization = `Bearer ${key}`;
-    }
-    else if (upstream.apiKey) {
-        headers.Authorization = `Bearer ${upstream.apiKey}`;
-    }
+    let hasAuthorization = false;
     if (upstream.headers) {
-        for (const [k, v] of Object.entries(upstream.headers)) {
-            headers[k] = v;
+        for (const [name, value] of Object.entries(upstream.headers)) {
+            if (/^(connection|keep-alive|proxy-authenticate|proxy-authorization|te|trailer|transfer-encoding|upgrade)$/i.test(name)) {
+                continue;
+            }
+            headers[name] = value;
+            if (name.toLowerCase() === "authorization")
+                hasAuthorization = true;
         }
     }
+    if (upstream.apiKey && !hasAuthorization) {
+        headers.Authorization = `Bearer ${upstream.apiKey}`;
+    }
     return headers;
 }
-function withProxyEnv(upstream, fn) {
-    if (emptyProxy(upstream.proxy))
-        return fn();
-    const next = buildProxyEnv(upstream.proxy);
-    const backup = new Map();
-    for (const [k, v] of Object.entries(next)) {
-        backup.set(k, process.env[k]);
-        process.env[k] = v;
-    }
-    return fn().finally(() => {
-        for (const [k, v] of backup) {
-            if (v === undefined)
-                delete process.env[k];
-            else
-                process.env[k] = v;
-        }
+function requestUpstream(upstream, url, method, body, signal) {
+    const limits = parseBridgeRuntimeLimits();
+    return requestWithNodeTransport({
+        url,
+        method,
+        headers: upstreamHeaders(upstream),
+        body,
+        proxy: upstream.proxy,
+        signal,
+        connectTimeoutMs: limits.connectTimeoutMs,
+        idleTimeoutMs: limits.idleTimeoutMs,
+        totalTimeoutMs: limits.totalTimeoutMs,
+        maxResponseBytes: limits.maxResponseBytes,
     });
 }
-async function fetchModelsJson(upstream, req) {
+async function fetchModelsJson(upstream) {
     const url = joinUrl(upstream.baseUrl, "/models");
-    const response = await withProxyEnv(upstream, () => fetch(url, {
-        method: "GET",
-        headers: upstreamHeaders(upstream, req),
-    }));
+    const response = await requestUpstream(upstream, url, "GET");
     if (!response.ok) {
         return { ok: false, status: response.status, data: [] };
     }
@@ -95,7 +155,7 @@ async function fetchModelsJson(upstream, req) {
         return { ok: false, status: 502, data: [] };
     }
 }
-async function proxyModelsMerged(req, res, upstreams) {
+async function proxyModelsMerged(_req, res, upstreams) {
     const sides = [upstreams.codex, upstreams.claude].filter((u) => Boolean(u?.baseUrl));
     if (!sides.length) {
         sendJson(res, 503, {
@@ -103,7 +163,7 @@ async function proxyModelsMerged(req, res, upstreams) {
         });
         return;
     }
-    const results = await Promise.all(sides.map((u) => fetchModelsJson(u, req).catch(() => ({
+    const results = await Promise.all(sides.map((u) => fetchModelsJson(u).catch(() => ({
         ok: false,
         status: 502,
         data: [],
@@ -148,7 +208,7 @@ async function handleResponses(req, res, upstream, bodyBuf) {
     }
     await forwardChatResponses(req, res, upstream, body, wantStream);
 }
-async function handleMessages(req, res, upstream, bodyBuf) {
+async function handleMessages(_req, res, upstream, bodyBuf) {
     let body;
     try {
         body = JSON.parse(bodyBuf.toString("utf8"));
@@ -165,11 +225,7 @@ async function handleMessages(req, res, upstream, bodyBuf) {
     const url = joinUrl(upstream.baseUrl, "/chat/completions");
     let response;
     try {
-        response = await withProxyEnv(upstream, () => fetch(url, {
-            method: "POST",
-            headers: upstreamHeaders(upstream, req),
-            body: JSON.stringify(chatReq),
-        }));
+        response = await requestUpstream(upstream, url, "POST", JSON.stringify(chatReq));
     }
     catch (err) {
         sendJson(res, 502, {
@@ -208,11 +264,7 @@ async function forwardChatResponses(req, res, upstream, body, wantStream) {
     const url = joinUrl(upstream.baseUrl, "/chat/completions");
     let response;
     try {
-        response = await withProxyEnv(upstream, () => fetch(url, {
-            method: "POST",
-            headers: upstreamHeaders(upstream, req),
-            body: JSON.stringify(chatReq),
-        }));
+        response = await requestUpstream(upstream, url, "POST", JSON.stringify(chatReq));
     }
     catch (err) {
         sendJson(res, 502, {
@@ -236,22 +288,18 @@ async function forwardChatResponses(req, res, upstream, body, wantStream) {
     }
     if (!wantStream) {
         const json = (await response.json());
-        sendJson(res, 200, chatCompletionToResponse(json, String(body.model || ""), customTools));
+        sendJson(res, 200, chatCompletionToResponse(json, String(body.model || ""), customTools, true));
         return;
     }
-    await pipeChatStreamToResponses(response, res, String(body.model || ""), customTools);
+    await pipeChatStreamToResponses(response, res, String(body.model || ""), customTools, true);
 }
-async function forwardCompletions(req, res, upstream, body, wantStream) {
+async function forwardCompletions(_req, res, upstream, body, wantStream) {
     const completionReq = responsesToCompletionsRequest(body);
     const customTools = collectCustomToolNames(body.tools);
     const url = joinUrl(upstream.baseUrl, "/completions");
     let response;
     try {
-        response = await withProxyEnv(upstream, () => fetch(url, {
-            method: "POST",
-            headers: upstreamHeaders(upstream, req),
-            body: JSON.stringify(completionReq),
-        }));
+        response = await requestUpstream(upstream, url, "POST", JSON.stringify(completionReq));
     }
     catch (err) {
         sendJson(res, 502, {
@@ -271,19 +319,19 @@ async function forwardCompletions(req, res, upstream, body, wantStream) {
     }
     if (!wantStream) {
         const json = (await response.json());
-        sendJson(res, 200, chatCompletionToResponse(json, String(body.model || ""), customTools));
+        sendJson(res, 200, chatCompletionToResponse(json, String(body.model || ""), customTools, true));
         return;
     }
-    await pipeChatStreamToResponses(response, res, String(body.model || ""), customTools);
+    await pipeChatStreamToResponses(response, res, String(body.model || ""), customTools, true);
 }
-async function pipeChatStreamToResponses(upstream, res, model, customTools) {
+async function pipeChatStreamToResponses(upstream, res, model, customTools, webSearchEnabled = false) {
     res.writeHead(200, {
         "Content-Type": "text/event-stream; charset=utf-8",
         "Cache-Control": "no-cache, no-transform",
         Connection: "keep-alive",
         "X-Accel-Buffering": "no",
     });
-    const state = createStreamState(model, undefined, customTools);
+    const state = createStreamState(model, undefined, customTools, webSearchEnabled);
     const reader = upstream.body?.getReader();
     if (!reader) {
         for (const frame of forceCompleteStream(state))
@@ -396,40 +444,94 @@ async function pipeChatStreamToAnthropic(upstream, res, model) {
         res.end();
     }
 }
-export function createBridgeServer() {
+export function createBridgeServer(options = {}) {
     return createServer(async (req, res) => {
         try {
-            const upstreams = readBridgeUpstreams();
             const state = readBridgeState();
-            const merged = {
-                codex: upstreams.codex || state.upstreams.codex,
-                claude: upstreams.claude || state.upstreams.claude,
-            };
+            const expectedControlToken = options.controlToken ?? state.instance?.controlToken;
+            const expectedInstanceId = options.instanceId ?? state.instance?.id;
+            const upstreams = readBridgeUpstreams();
+            const merged = upstreams;
             const url = new URL(req.url || "/", `http://${req.headers.host || "127.0.0.1"}`);
             const path = url.pathname.replace(/\/+$/, "") || "/";
             if (req.method === "GET" && (path === "/health" || path === "/v1/health")) {
+                const suppliedControl = controlToken(req);
+                if (!suppliedControl) {
+                    sendJson(res, 200, { ok: true, service: "llm-switch-bridge" });
+                    return;
+                }
+                if (!expectedControlToken ||
+                    !constantTimeTokenEqual(expectedControlToken, suppliedControl)) {
+                    sendJson(res, 401, {
+                        ok: false,
+                        error: { code: "invalid_control_token", message: "Unauthorized" },
+                    });
+                    return;
+                }
                 sendJson(res, 200, {
                     ok: true,
+                    service: "llm-switch-bridge",
+                    instanceId: expectedInstanceId,
                     upstreams: {
                         codex: merged.codex
                             ? {
-                                baseUrl: merged.codex.baseUrl,
                                 mode: merged.codex.mode,
                                 profile: merged.codex.profileName || null,
+                                migrationRequired: merged.codex.migrationRequired === true,
                             }
                             : null,
                         claude: merged.claude
                             ? {
-                                baseUrl: merged.claude.baseUrl,
                                 mode: merged.claude.mode,
                                 profile: merged.claude.profileName || null,
+                                migrationRequired: merged.claude.migrationRequired === true,
                             }
                             : null,
                     },
                 });
                 return;
             }
+            if (req.method === "POST" && path === "/_control/shutdown") {
+                const suppliedControl = controlToken(req);
+                if (!expectedControlToken ||
+                    !constantTimeTokenEqual(expectedControlToken, suppliedControl)) {
+                    sendJson(res, 401, {
+                        ok: false,
+                        error: { code: "invalid_control_token", message: "Unauthorized" },
+                    });
+                    return;
+                }
+                let instanceId = "";
+                try {
+                    const body = JSON.parse((await readBody(req)).toString("utf8"));
+                    instanceId = typeof body.instanceId === "string" ? body.instanceId : "";
+                }
+                catch {
+                    sendJson(res, 400, { error: { message: "Invalid JSON body" } });
+                    return;
+                }
+                if (!expectedInstanceId || instanceId !== expectedInstanceId) {
+                    sendJson(res, 409, {
+                        error: { code: "instance_mismatch", message: "Bridge instance mismatch" },
+                    });
+                    return;
+                }
+                sendJson(res, 202, { ok: true, instanceId });
+                queueMicrotask(() => {
+                    void options.onShutdown?.(instanceId);
+                });
+                return;
+            }
             if (req.method === "GET" && (path === "/v1/models" || path === "/models")) {
+                if (!authenticateModelsRequest(req, merged)) {
+                    sendJson(res, 401, {
+                        error: {
+                            code: "invalid_bridge_token",
+                            message: "Bridge token 无效；升级后请重新执行 llms <tool> use <profile>",
+                        },
+                    });
+                    return;
+                }
                 await proxyModelsMerged(req, res, merged);
                 return;
             }
@@ -443,6 +545,15 @@ export function createBridgeServer() {
                     });
                     return;
                 }
+                if (!authenticateDataRequest(req, merged.codex, "codex")) {
+                    sendJson(res, 401, {
+                        error: {
+                            code: "invalid_bridge_token",
+                            message: "Bridge token 无效；请重新执行 llms codex use <profile>",
+                        },
+                    });
+                    return;
+                }
                 const body = await readBody(req);
                 await handleResponses(req, res, merged.codex, body);
                 return;
@@ -459,6 +570,16 @@ export function createBridgeServer() {
                     });
                     return;
                 }
+                if (!authenticateDataRequest(req, merged.claude, "claude")) {
+                    sendJson(res, 401, {
+                        type: "error",
+                        error: {
+                            type: "authentication_error",
+                            message: "Bridge token 无效；请重新执行 llms claude use <profile>",
+                        },
+                    });
+                    return;
+                }
                 const body = await readBody(req);
                 await handleMessages(req, res, merged.claude, body);
                 return;
@@ -470,6 +591,12 @@ export function createBridgeServer() {
             });
         }
         catch (err) {
+            if (err instanceof RequestBodyTooLargeError) {
+                sendJson(res, 413, {
+                    error: { code: "request_too_large", message: err.message },
+                });
+                return;
+            }
             sendJson(res, 500, {
                 error: {
                     message: err instanceof Error ? err.message : String(err),
@@ -478,8 +605,8 @@ export function createBridgeServer() {
         }
     });
 }
-export function listenBridge(port, host) {
-    const server = createBridgeServer();
+export function listenBridge(port, host, options = {}) {
+    const server = createBridgeServer(options);
     return new Promise((resolve, reject) => {
         server.once("error", reject);
         server.listen(port, host, () => resolve(server));