@khanglvm/llm-router 1.3.1 → 2.0.0-beta.0

package/src/runtime/handler/amp-response.js

@@ -0,0 +1,308 @@
+import { withCorsHeaders } from "./http.js";
+
+const AMP_MODEL_FIELD_PATHS = [
+  ["model"],
+  ["message", "model"],
+  ["response", "model"],
+  ["modelVersion"],
+  ["response", "modelVersion"]
+];
+
+function isPlainObject(value) {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function readAmpVisibleModel(requestBody) {
+  return typeof requestBody?.model === "string" ? requestBody.model.trim() : "";
+}
+
+function setPathIfPresent(payload, path, value) {
+  if (!isPlainObject(payload) || !Array.isArray(path) || path.length === 0) return false;
+  let cursor = payload;
+  for (let index = 0; index < path.length - 1; index += 1) {
+    if (!isPlainObject(cursor?.[path[index]])) return false;
+    cursor = cursor[path[index]];
+  }
+
+  const key = path[path.length - 1];
+  if (typeof cursor?.[key] !== "string") return false;
+  if (cursor[key] === value) return false;
+  cursor[key] = value;
+  return true;
+}
+
+function rewriteAmpModelFields(payload, visibleModel) {
+  if (!visibleModel) return false;
+  let changed = false;
+  for (const path of AMP_MODEL_FIELD_PATHS) {
+    changed = setPathIfPresent(payload, path, visibleModel) || changed;
+  }
+  return changed;
+}
+
+function stripClaudeThinkingBlocks(payload) {
+  if (!isPlainObject(payload) || !Array.isArray(payload.content)) return false;
+  const hasToolUse = payload.content.some((block) => block?.type === "tool_use");
+  if (!hasToolUse) return false;
+
+  const filtered = payload.content.filter(
+    (block) => block?.type !== "thinking" && block?.type !== "redacted_thinking"
+  );
+  if (filtered.length === payload.content.length) return false;
+  payload.content = filtered;
+  return true;
+}
+
+function normalizeClaudeToolUseStopReason(payload) {
+  if (!isPlainObject(payload) || !Array.isArray(payload.content)) return false;
+  const hasToolUse = payload.content.some((block) => block?.type === "tool_use");
+  if (!hasToolUse) return false;
+  if (payload.stop_reason === "tool_use") return false;
+  payload.stop_reason = "tool_use";
+  return true;
+}
+
+function rewriteAmpPayload(payload, visibleModel) {
+  if (!isPlainObject(payload)) return false;
+  let changed = false;
+  changed = rewriteAmpModelFields(payload, visibleModel) || changed;
+  changed = normalizeClaudeToolUseStopReason(payload) || changed;
+  changed = stripClaudeThinkingBlocks(payload) || changed;
+  return changed;
+}
+
+function rewriteAmpSseEvent(rawEvent, visibleModel) {
+  const lines = String(rawEvent || "").split(/\r?\n/);
+  let changed = false;
+  const nextLines = lines.map((line) => {
+    if (!line.startsWith("data:")) return line;
+    const value = line.slice(5).trimStart();
+    if (!value || value === "[DONE]") return line;
+
+    try {
+      const payload = JSON.parse(value);
+      if (!rewriteAmpPayload(payload, visibleModel)) return line;
+      changed = true;
+      return `data: ${JSON.stringify(payload)}`;
+    } catch {
+      return line;
+    }
+  });
+
+  return {
+    changed,
+    event: `${nextLines.join("\n")}\n\n`
+  };
+}
+
+const AMP_STREAM_SUPPRESSED_BLOCK_TYPES = new Set([
+  "thinking",
+  "redacted_thinking"
+]);
+
+function parseSseEventBlock(rawEvent) {
+  const lines = String(rawEvent || "").split(/\r?\n/);
+  let eventName = "";
+  const dataLines = [];
+
+  for (const line of lines) {
+    if (!line) continue;
+    if (line.startsWith("event:")) {
+      eventName = line.slice(6).trim();
+      continue;
+    }
+    if (line.startsWith("data:")) {
+      dataLines.push(line.slice(5).trimStart());
+    }
+  }
+
+  return {
+    eventName,
+    dataText: dataLines.join("\n").trim()
+  };
+}
+
+function formatSseEventBlock(eventName, payload) {
+  const lines = [];
+  if (eventName) {
+    lines.push(`event: ${eventName}`);
+  }
+  lines.push(`data: ${JSON.stringify(payload)}`);
+  return `${lines.join("\n")}\n\n`;
+}
+
+function rewriteAmpClaudeStreamPayload(state, payload, visibleModel) {
+  let changed = rewriteAmpPayload(payload, visibleModel);
+  const type = String(payload?.type || "").trim();
+
+  if (type === "content_block_start") {
+    const originalIndex = Number(payload.index);
+    const blockType = String(payload?.content_block?.type || "").trim();
+    if (Number.isFinite(originalIndex) && AMP_STREAM_SUPPRESSED_BLOCK_TYPES.has(blockType)) {
+      state.suppressedIndexes.add(originalIndex);
+      return {
+        changed: false,
+        suppressed: true,
+        payload
+      };
+    }
+
+    if (Number.isFinite(originalIndex)) {
+      const visibleIndex = state.nextVisibleIndex;
+      state.nextVisibleIndex += 1;
+      state.visibleIndexByOriginal.set(originalIndex, visibleIndex);
+      if (visibleIndex !== originalIndex) {
+        payload.index = visibleIndex;
+        changed = true;
+      }
+    }
+
+    return {
+      changed,
+      suppressed: false,
+      payload
+    };
+  }
+
+  if (type === "content_block_delta" || type === "content_block_stop") {
+    const originalIndex = Number(payload.index);
+    if (Number.isFinite(originalIndex) && state.suppressedIndexes.has(originalIndex)) {
+      return {
+        changed: false,
+        suppressed: true,
+        payload
+      };
+    }
+
+    const visibleIndex = state.visibleIndexByOriginal.get(originalIndex);
+    if (Number.isFinite(visibleIndex) && visibleIndex !== originalIndex) {
+      payload.index = visibleIndex;
+      changed = true;
+    }
+
+    return {
+      changed,
+      suppressed: false,
+      payload
+    };
+  }
+
+  return {
+    changed,
+    suppressed: false,
+    payload
+  };
+}
+
+function rewriteAmpStreamResponse(response, visibleModel) {
+  if (!(response instanceof Response) || !response.body || !visibleModel) return response;
+
+  const decoder = new TextDecoder();
+  const encoder = new TextEncoder();
+  let buffer = "";
+  const claudeStreamState = {
+    suppressedIndexes: new Set(),
+    visibleIndexByOriginal: new Map(),
+    nextVisibleIndex: 0
+  };
+
+  const stream = new ReadableStream({
+    async start(controller) {
+      const reader = response.body.getReader();
+
+      const emitEvent = (rawEvent) => {
+        if (!rawEvent) return;
+        const parsedEvent = parseSseEventBlock(rawEvent);
+        if (!parsedEvent.dataText || parsedEvent.dataText === "[DONE]") {
+          const rewritten = rewriteAmpSseEvent(rawEvent, visibleModel);
+          controller.enqueue(encoder.encode(rewritten.event));
+          return;
+        }
+
+        try {
+          const payload = JSON.parse(parsedEvent.dataText);
+          const rewritten = rewriteAmpClaudeStreamPayload(claudeStreamState, payload, visibleModel);
+          if (rewritten.suppressed) return;
+          controller.enqueue(encoder.encode(formatSseEventBlock(parsedEvent.eventName, rewritten.payload)));
+        } catch {
+          const rewritten = rewriteAmpSseEvent(rawEvent, visibleModel);
+          controller.enqueue(encoder.encode(rewritten.event));
+        }
+      };
+
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buffer += decoder.decode(value, { stream: true }).replace(/\r\n/g, "\n");
+
+          while (true) {
+            const separatorIndex = buffer.indexOf("\n\n");
+            if (separatorIndex < 0) break;
+            const rawEvent = buffer.slice(0, separatorIndex);
+            buffer = buffer.slice(separatorIndex + 2);
+            emitEvent(rawEvent);
+          }
+        }
+
+        buffer += decoder.decode();
+        if (buffer.trim()) emitEvent(buffer);
+        controller.close();
+      } catch (error) {
+        controller.error(error);
+      } finally {
+        try {
+          reader.releaseLock();
+        } catch {
+          // Ignore reader cleanup errors.
+        }
+      }
+    }
+  });
+
+  const headers = new Headers(response.headers);
+  headers.delete("content-length");
+
+  return new Response(stream, {
+    status: response.status,
+    statusText: response.statusText,
+    headers
+  });
+}
+
+export async function maybeRewriteAmpClientResponse(response, { clientType, requestBody, stream = false } = {}) {
+  if (clientType !== "amp" || !(response instanceof Response)) return response;
+
+  const visibleModel = readAmpVisibleModel(requestBody);
+  if (!visibleModel) return response;
+
+  const contentType = String(response.headers.get("content-type") || "").toLowerCase();
+  if (stream || contentType.includes("text/event-stream")) {
+    return rewriteAmpStreamResponse(response, visibleModel);
+  }
+
+  if (!contentType.includes("application/json") && !contentType.includes("+json")) {
+    return response;
+  }
+
+  let payload;
+  try {
+    payload = await response.clone().json();
+  } catch {
+    return response;
+  }
+
+  if (!rewriteAmpPayload(payload, visibleModel)) {
+    return response;
+  }
+
+  const headers = new Headers(response.headers);
+  headers.set("content-type", headers.get("content-type") || "application/json");
+  headers.delete("content-length");
+
+  return new Response(JSON.stringify(payload), {
+    status: response.status,
+    statusText: response.statusText,
+    headers: new Headers(withCorsHeaders(Object.fromEntries(headers.entries())))
+  });
+}

package/src/runtime/handler/amp.js

@@ -0,0 +1,290 @@
+import { jsonResponse, passthroughResponseWithCors } from "./http.js";
+
+const AMP_PROXY_SCRUBBED_HEADERS = [
+  "authorization",
+  "x-api-key",
+  "x-goog-api-key",
+  "host",
+  "content-length",
+  "accept-encoding",
+  "proxy-authorization",
+  "proxy-authenticate",
+  "forwarded",
+  "via",
+  "x-forwarded-for",
+  "x-forwarded-host",
+  "x-forwarded-proto",
+  "x-forwarded-port",
+  "cf-connecting-ip",
+  "true-client-ip",
+  "x-real-ip",
+  "x-client-ip",
+  "cf-ray",
+  "cf-ipcountry",
+  "x-client-data",
+  "sec-ch-ua",
+  "sec-ch-ua-mobile",
+  "sec-ch-ua-platform",
+  "sec-fetch-site",
+  "sec-fetch-mode",
+  "sec-fetch-user",
+  "sec-fetch-dest"
+];
+
+function isPlainObject(value) {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function isLoopbackAddress(value) {
+  const ip = String(value || "").trim().toLowerCase();
+  if (!ip) return false;
+  if (["127.0.0.1", "::1", "::ffff:127.0.0.1", "localhost"].includes(ip)) return true;
+  return ip.startsWith("127.");
+}
+
+function readClientIp(request) {
+  const direct = request.headers.get("x-real-ip") || request.headers.get("cf-connecting-ip");
+  if (direct) return String(direct).trim();
+  const forwardedFor = String(request.headers.get("x-forwarded-for") || "").trim();
+  if (!forwardedFor) return "";
+  const firstHop = forwardedFor.split(",")[0]?.trim();
+  return firstHop || "";
+}
+
+export function isAmpProxyEnabled(config) {
+  return Boolean(String(config?.amp?.upstreamUrl || "").trim());
+}
+
+export function isAmpManagementAllowed(request, config) {
+  if (config?.amp?.restrictManagementToLocalhost !== true) return true;
+  return isLoopbackAddress(readClientIp(request));
+}
+
+function readClientAuthToken(request) {
+  const authHeader = request.headers.get("authorization");
+  if (authHeader?.startsWith("Bearer ")) return authHeader.slice(7).trim();
+  if (authHeader) return authHeader.trim();
+
+  for (const headerName of ["x-api-key", "x-goog-api-key"]) {
+    const headerValue = request.headers.get(headerName);
+    if (headerValue) return String(headerValue).trim();
+  }
+
+  try {
+    const requestUrl = new URL(request.url);
+    for (const key of ["key", "api_key", "auth_token"]) {
+      const value = requestUrl.searchParams.get(key);
+      if (value) return value.trim();
+    }
+  } catch {
+    // Ignore malformed request URLs and continue with empty token.
+  }
+
+  return "";
+}
+
+function removeQueryValuesMatching(url, key, match) {
+  if (!(url instanceof URL) || !match) return;
+  const values = url.searchParams.getAll(key);
+  if (values.length === 0) return;
+
+  url.searchParams.delete(key);
+  for (const value of values) {
+    if (value === match) continue;
+    url.searchParams.append(key, value);
+  }
+}
+
+function stripClientCredentialsFromAmpUpstreamUrl(upstreamUrl, request) {
+  const clientToken = readClientAuthToken(request);
+  if (!clientToken) return;
+
+  for (const key of ["key", "api_key", "auth_token"]) {
+    removeQueryValuesMatching(upstreamUrl, key, clientToken);
+  }
+}
+
+function buildAmpProxyHeaders(request, config) {
+  const headers = new Headers(request.headers);
+  const upstreamApiKey = String(config?.amp?.upstreamApiKey || "").trim();
+
+  for (const name of AMP_PROXY_SCRUBBED_HEADERS) {
+    headers.delete(name);
+  }
+
+  if (upstreamApiKey) {
+    headers.set("authorization", `Bearer ${upstreamApiKey}`);
+    headers.set("x-api-key", upstreamApiKey);
+  }
+
+  return headers;
+}
+
+function normalizeAmpUserPayload(payload) {
+  if (!isPlainObject(payload)) return null;
+  return {
+    ...payload,
+    freeTierEligibleIfWorkspaceAllows: true,
+    dailyGrantEnabledIfWorkspaceAllows: false
+  };
+}
+
+function normalizeAmpInternalPayload(requestUrl, payload) {
+  if (!isPlainObject(payload)) return null;
+
+  if (requestUrl.searchParams.has("getUserFreeTierStatus")) {
+    return {
+      ...payload,
+      result: {
+        ...(isPlainObject(payload.result) ? payload.result : {}),
+        canUseAmpFree: true,
+        isDailyGrantEnabled: false
+      }
+    };
+  }
+
+  if (requestUrl.searchParams.has("getUserInfo")) {
+    return {
+      ...payload,
+      result: normalizeAmpUserPayload(payload.result) || payload.result
+    };
+  }
+
+  return null;
+}
+
+function normalizeAmpManagementPayload(requestUrl, payload) {
+  if (requestUrl.pathname === "/api/user") {
+    return normalizeAmpUserPayload(payload);
+  }
+
+  if (requestUrl.pathname === "/api/internal") {
+    return normalizeAmpInternalPayload(requestUrl, payload);
+  }
+
+  return null;
+}
+
+function shouldNormalizeAmpManagementResponse(requestUrl) {
+  if (requestUrl.pathname === "/api/user") return true;
+  if (requestUrl.pathname !== "/api/internal") return false;
+  return requestUrl.searchParams.has("getUserFreeTierStatus") || requestUrl.searchParams.has("getUserInfo");
+}
+
+function isJsonLikeResponse(response) {
+  const contentType = String(response?.headers?.get("content-type") || "").toLowerCase();
+  return contentType.includes("application/json") || contentType.includes("+json");
+}
+
+function isAmpStreamingResponse(response) {
+  const contentType = String(response?.headers?.get("content-type") || "").toLowerCase();
+  return contentType.includes("text/event-stream");
+}
+
+async function gunzipBytes(bytes) {
+  if (typeof DecompressionStream !== "function") return null;
+  try {
+    const gzipStream = new Blob([bytes]).stream().pipeThrough(new DecompressionStream("gzip"));
+    const decompressed = await new Response(gzipStream).arrayBuffer();
+    return new Uint8Array(decompressed);
+  } catch {
+    return null;
+  }
+}
+
+async function maybeDecompressAmpProxyResponse(response) {
+  if (!(response instanceof Response) || !response.ok) return response;
+  if (response.headers.get("content-encoding")) return response;
+  if (isAmpStreamingResponse(response)) return response;
+
+  let bytes;
+  try {
+    bytes = new Uint8Array(await response.clone().arrayBuffer());
+  } catch {
+    return response;
+  }
+
+  if (bytes.length < 2 || bytes[0] !== 0x1f || bytes[1] !== 0x8b) {
+    return response;
+  }
+
+  const decompressed = await gunzipBytes(bytes);
+  if (!(decompressed instanceof Uint8Array)) return response;
+
+  const headers = new Headers(response.headers);
+  headers.delete("content-encoding");
+  headers.set("content-length", String(decompressed.byteLength));
+
+  return new Response(decompressed, {
+    status: response.status,
+    statusText: response.statusText,
+    headers
+  });
+}
+
+async function maybeNormalizeAmpManagementResponse(requestUrl, response) {
+  if (!shouldNormalizeAmpManagementResponse(requestUrl) || !response?.ok || !isJsonLikeResponse(response)) {
+    return null;
+  }
+
+  try {
+    const payload = await response.clone().json();
+    const normalizedPayload = normalizeAmpManagementPayload(requestUrl, payload);
+    if (!normalizedPayload) return null;
+    return passthroughResponseWithCors(new Response(JSON.stringify(normalizedPayload), {
+      status: response.status,
+      statusText: response.statusText,
+      headers: response.headers
+    }));
+  } catch {
+    return null;
+  }
+}
+
+export async function proxyAmpUpstreamRequest({ request, config, bodyOverride } = {}) {
+  const upstreamBase = String(config?.amp?.upstreamUrl || "").trim();
+  if (!upstreamBase) {
+    return jsonResponse({
+      type: "error",
+      error: {
+        type: "configuration_error",
+        message: "AMP upstream proxy is not configured. Set config.amp.upstreamUrl to enable AMP management and fallback proxying."
+      }
+    }, 503);
+  }
+
+  const requestUrl = new URL(request.url);
+  const upstreamUrl = new URL(`${requestUrl.pathname}${requestUrl.search}`, upstreamBase);
+  stripClientCredentialsFromAmpUpstreamUrl(upstreamUrl, request);
+  const headers = buildAmpProxyHeaders(request, config);
+  const init = {
+    method: request.method,
+    headers,
+    redirect: "manual"
+  };
+
+  if (request.method !== "GET" && request.method !== "HEAD") {
+    if (bodyOverride !== undefined) {
+      init.body = bodyOverride;
+    } else {
+      const rawBody = await request.arrayBuffer();
+      if (rawBody.byteLength > 0) init.body = rawBody;
+    }
+  }
+
+  try {
+    const response = await fetch(upstreamUrl, init);
+    const preparedResponse = await maybeDecompressAmpProxyResponse(response);
+    const normalizedResponse = await maybeNormalizeAmpManagementResponse(requestUrl, preparedResponse);
+    if (normalizedResponse) return normalizedResponse;
+    return passthroughResponseWithCors(preparedResponse);
+  } catch (error) {
+    return jsonResponse({
+      type: "error",
+      error: {
+        type: "api_error",
+        message: `AMP upstream proxy failed: ${error instanceof Error ? error.message : String(error)}`
+      }
+    }, 503);
+  }
+}

package/src/runtime/handler/auth.js

@@ -2,8 +2,23 @@ function parseAuthToken(request) {
   const authHeader = request.headers.get("Authorization");
   if (authHeader?.startsWith("Bearer ")) return authHeader.slice(7).trim();
   if (authHeader && !authHeader.startsWith("Bearer ")) return authHeader.trim();
-  const apiKey = request.headers.get("x-api-key");
-  return apiKey ? apiKey.trim() : "";
+
+  for (const headerName of ["x-api-key", "x-goog-api-key"]) {
+    const headerValue = request.headers.get(headerName);
+    if (headerValue) return headerValue.trim();
+  }
+
+  try {
+    const url = new URL(request.url);
+    for (const key of ["key", "api_key", "auth_token"]) {
+      const value = url.searchParams.get(key);
+      if (value) return value.trim();
+    }
+  } catch {
+    // Ignore malformed request URLs and continue with empty token.
+  }
+
+  return "";
 }
 
 function timingSafeStringEqual(left, right) {