fix(mitm): route Claude Code standalone MITM traffic (#7574)

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

src/mitm/_internal/standaloneRouting.cjs

@@ -0,0 +1,78 @@
+"use strict";
+
+const AGENT_ROUTE_CONFIG = {
+  antigravity: {
+    aliasKey: "antigravity",
+    chatUrlPatterns: [":generateContent", ":streamGenerateContent"],
+    routerPath: "/v1/chat/completions",
+  },
+  "claude-code": {
+    aliasKey: "claude-code",
+    chatUrlPatterns: ["/v1/messages"],
+    routerPath: "/v1/messages",
+  },
+  kiro: {
+    aliasKey: "kiro",
+    chatUrlPatterns: ["/v1/messages"],
+    routerPath: "/v1/messages",
+  },
+};
+
+function getAgentRouteConfig(agentId) {
+  return AGENT_ROUTE_CONFIG[agentId] || AGENT_ROUTE_CONFIG.antigravity;
+}
+
+function resolveForwardTargetForAgent({
+  routerBaseUrl,
+  routerMessagesUrl,
+  body,
+  agentId,
+  fallbackResolver,
+}) {
+  const config = getAgentRouteConfig(agentId);
+  if (config.routerPath === "/v1/messages") {
+    return { format: "anthropic", url: routerMessagesUrl };
+  }
+  return fallbackResolver(routerBaseUrl, body);
+}
+
+function resolveMappedOverride(model, agentId, deps) {
+  if (!model) return null;
+
+  const config = getAgentRouteConfig(agentId);
+  const { fs, dbFile, getSqliteDb, aliasConfigShim } = deps;
+
+  try {
+    const db = getSqliteDb();
+    if (db) {
+      const row = db
+        .prepare("SELECT value FROM key_value WHERE namespace = 'mitmAlias' AND key = ?")
+        .get(config.aliasKey);
+      if (row) {
+        const mappings = aliasConfigShim.normalizeAliasMappings(JSON.parse(row.value));
+        return mappings[model] || null;
+      }
+    }
+  } catch {
+    // Fall through to JSON fallback.
+  }
+
+  try {
+    if (fs.existsSync(dbFile)) {
+      const db = JSON.parse(fs.readFileSync(dbFile, "utf-8"));
+      const mappings = aliasConfigShim.normalizeAliasMappings(db.mitmAlias?.[config.aliasKey]);
+      return mappings[model] || null;
+    }
+  } catch {
+    // Ignore malformed legacy state.
+  }
+
+  return null;
+}
+
+module.exports = {
+  AGENT_ROUTE_CONFIG,
+  getAgentRouteConfig,
+  resolveForwardTargetForAgent,
+  resolveMappedOverride,
+};

src/mitm/server.cjs

@@ -47,6 +47,7 @@ const ROUTER_BASE_URL = (
   .trim()
   .replace(/\/+$/, "");
 const ROUTER_URL = `${ROUTER_BASE_URL}/v1/chat/completions`;
+const ROUTER_MESSAGES_URL = `${ROUTER_BASE_URL}/v1/messages`;
 const API_KEY = process.env.ROUTER_API_KEY;
 const DATA_DIR = getDataDir();
 const DB_FILE = path.join(DATA_DIR, "db.json");
@@ -101,17 +102,18 @@ const SANITIZE_SOURCE_EXT = ["ts", "tsx", "js", "jsx", "mjs", "cjs"];
 function looksLikeAbsolutePath(tok) {
   if (tok.length < 4 || tok.length > 2048) return false;
   const isPosix = tok.charCodeAt(0) === 0x2f;
-  const isWindows =
-    tok.length > 2 && tok.charCodeAt(1) === 0x3a && /[A-Za-z]/.test(tok[0]);
+  const isWindows = tok.length > 2 && tok.charCodeAt(1) === 0x3a && /[A-Za-z]/.test(tok[0]);
   if (!isPosix && !isWindows) return false;
   const dot = tok.lastIndexOf(".");
   if (dot <= 0 || dot === tok.length - 1) return false;
-  const ext = tok.slice(dot + 1).split(":", 1)[0].toLowerCase();
+  const ext = tok
+    .slice(dot + 1)
+    .split(":", 1)[0]
+    .toLowerCase();
   return SANITIZE_SOURCE_EXT.includes(ext);
 }
 function sanitizeErrorMessage(message) {
-  let str =
-    typeof message === "string" ? message : String(message == null ? "" : message);
+  let str = typeof message === "string" ? message : String(message == null ? "" : message);
   if (str.length > SANITIZE_MAX_LEN) str = str.slice(0, SANITIZE_MAX_LEN);
   const nl = str.indexOf("\n");
   const firstLine = nl >= 0 ? str.slice(0, nl) : str;
@@ -137,6 +139,7 @@ const bypassShim = require("./_internal/bypass.cjs");
 const ingestShim = require("./_internal/ingest.cjs");
 const forwardShim = require("./_internal/forwardTarget.cjs");
 const aliasConfigShim = require("./_internal/aliasConfig.cjs");
+const standaloneRoutingShim = require("./_internal/standaloneRouting.cjs");
 
 // Inspector capture (D4 fallback). The standalone proxy intercepts AgentBridge
 // traffic inline (no MitmHandlerBase / agentBridgeHook), so it posts captured
@@ -199,9 +202,7 @@ function routeBypass(hostname) {
 
 const _bypassLoaded = loadUserBypassPatterns();
 if (_bypassLoaded > 0) {
-  console.log(
-    `[MITM] Loaded ${_bypassLoaded} user bypass pattern(s) from bypass.json`
-  );
+  console.log(`[MITM] Loaded ${_bypassLoaded} user bypass pattern(s) from bypass.json`);
 }
 
 let _sqliteDb = null;
@@ -239,9 +240,6 @@ const sslOptions = {
   cert: fs.readFileSync(path.join(certDir, "server.crt")),
 };
 
-// Chat endpoints that should be intercepted
-const CHAT_URL_PATTERNS = [":generateContent", ":streamGenerateContent"];
-
 // Log directory for request/response dumps
 const LOG_DIR = path.join(__dirname, "../../logs/mitm");
 if (ENABLE_FILE_LOG && !fs.existsSync(LOG_DIR)) fs.mkdirSync(LOG_DIR, { recursive: true });
@@ -344,39 +342,13 @@ function getSqliteDb() {
  * into the structured shape, so both old and new saves resolve consistently. Returns
  * `null` when there is no override at all for this model (passthrough).
  */
-function getMappedOverride(model) {
-  if (!model) return null;
-
-  // Primary: read from SQLite key_value table
-  try {
-    const db = getSqliteDb();
-    if (db) {
-      const row = db
-        .prepare(
-          "SELECT value FROM key_value WHERE namespace = 'mitmAlias' AND key = 'antigravity'"
-        )
-        .get();
-      if (row) {
-        const mappings = aliasConfigShim.normalizeAliasMappings(JSON.parse(row.value));
-        return mappings[model] || null;
-      }
-    }
-  } catch {
-    // Fall through to JSON fallback
-  }
-
-  // Fallback: read from db.json (legacy installs not yet migrated)
-  try {
-    if (fs.existsSync(DB_FILE)) {
-      const db = JSON.parse(fs.readFileSync(DB_FILE, "utf-8"));
-      const mappings = aliasConfigShim.normalizeAliasMappings(db.mitmAlias?.antigravity);
-      return mappings[model] || null;
-    }
-  } catch {
-    // Ignore
-  }
-
-  return null;
+function getMappedOverride(model, agentId = "antigravity") {
+  return standaloneRoutingShim.resolveMappedOverride(model, agentId, {
+    fs,
+    dbFile: DB_FILE,
+    getSqliteDb,
+    aliasConfigShim,
+  });
 }
 
 async function passthrough(req, res, bodyBuffer) {
@@ -465,7 +437,9 @@ async function intercept(req, res, bodyBuffer, override, sourceModel) {
   // other inbound clients and to record the originating IDE agent id.
   // Resolve agent id from the Host header against the target map; defensive
   // fallback to "unknown" when the host is somehow not in the map.
-  const reqHost = String(req.headers.host || "").split(":")[0].toLowerCase();
+  const reqHost = String(req.headers.host || "")
+    .split(":")[0]
+    .toLowerCase();
   const agentId = TARGET_HOST_AGENT.get(reqHost) || "unknown";
   const startedAt = Date.now();
   let upstreamStartedAt = startedAt;
@@ -492,7 +466,13 @@ async function intercept(req, res, bodyBuffer, override, sourceModel) {
     // the IDE gets its own format back; plain OpenAI bodies still go to
     // chat/completions. Without this, cloudcode hits chat/completions and 400s
     // on the missing `messages` field.
-    const forward = forwardShim.resolveForwardTarget(ROUTER_BASE_URL, body);
+    const forward = standaloneRoutingShim.resolveForwardTargetForAgent({
+      routerBaseUrl: ROUTER_BASE_URL,
+      routerMessagesUrl: ROUTER_MESSAGES_URL,
+      body,
+      agentId,
+      fallbackResolver: forwardShim.resolveForwardTarget,
+    });
     vlog(1, `[MITM] → forward ${forward.format} ${forward.url}`);
 
     upstreamStartedAt = Date.now();
@@ -578,10 +558,15 @@ const server = https.createServer(sslOptions, async (req, res) => {
   writeStats();
 
   const bodyBuffer = await collectBodyRaw(req);
-  const host = String(req.headers.host || "").split(":")[0].toLowerCase();
+  const host = String(req.headers.host || "")
+    .split(":")[0]
+    .toLowerCase();
   const model = bodyBuffer.length > 0 ? extractModel(bodyBuffer) : null;
 
-  vlog(1, `[MITM] ${req.method} ${host}${req.url} | body: ${bodyBuffer.length}B | model: ${model || "N/A"}`);
+  vlog(
+    1,
+    `[MITM] ${req.method} ${host}${req.url} | body: ${bodyBuffer.length}B | model: ${model || "N/A"}`
+  );
 
   if (bodyBuffer.length > 0) saveRequestLog(req.url, bodyBuffer);
 
@@ -595,14 +580,16 @@ const server = https.createServer(sslOptions, async (req, res) => {
     return passthrough(req, res, bodyBuffer);
   }
 
-  const isChatRequest = CHAT_URL_PATTERNS.some((p) => req.url.includes(p));
+  const agentId = TARGET_HOST_AGENT.get(host) || "antigravity";
+  const routeConfig = standaloneRoutingShim.getAgentRouteConfig(agentId);
+  const isChatRequest = routeConfig.chatUrlPatterns.some((p) => req.url.includes(p));
 
   if (!isChatRequest) {
     vlog(1, `[MITM] → PASSTHROUGH (URL ${req.url} does not match chat patterns)`);
     return passthrough(req, res, bodyBuffer);
   }
 
-  const mappedOverride = getMappedOverride(model);
+  const mappedOverride = getMappedOverride(model, agentId);
 
   if (!mappedOverride) {
     vlog(1, `[MITM] → PASSTHROUGH (model "${model}" has no MITM alias mapping)`);
@@ -615,7 +602,7 @@ const server = https.createServer(sslOptions, async (req, res) => {
 
   vlog(
     1,
-    `[MITM] INTERCEPTED ${model} → ${mappedOverride.model || model}` +
+    `[MITM] INTERCEPTED ${agentId} ${model} → ${mappedOverride.model || model}` +
       (mappedOverride.reasoningEffort ? ` (reasoningEffort=${mappedOverride.reasoningEffort})` : "")
   );
   return intercept(req, res, bodyBuffer, mappedOverride, model);
@@ -741,10 +728,7 @@ server.on("connect", (req, clientSocket, head) => {
     // https.createServer request handler can decrypt and route. We write the
     // 200 response ourselves and then `emit("connection")` so the TLS layer
     // picks the socket up.
-    vlog(
-      1,
-      `[MITM] CONNECT ${connectHost}:${connectPort} → TARGET (TLS terminate locally)`
-    );
+    vlog(1, `[MITM] CONNECT ${connectHost}:${connectPort} → TARGET (TLS terminate locally)`);
     clientSocket.write("HTTP/1.1 200 Connection Established\r\n\r\n");
     if (head && head.length > 0) clientSocket.unshift(head);
     server.emit("connection", clientSocket);
@@ -752,10 +736,7 @@ server.on("connect", (req, clientSocket, head) => {
   }
 
   // decision === "passthrough"
-  vlog(
-    1,
-    `[MITM] CONNECT ${connectHost}:${connectPort} → PASSTHROUGH (TCP tunnel)`
-  );
+  vlog(1, `[MITM] CONNECT ${connectHost}:${connectPort} → PASSTHROUGH (TCP tunnel)`);
   rawTcpForward(clientSocket, head, connectHost, connectPort, "passthrough");
 });
 

tests/unit/mitm-server-claude-code-routing.test.ts

@@ -0,0 +1,113 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import { createRequire } from "node:module";
+
+const require = createRequire(import.meta.url);
+const routing = require("../../src/mitm/_internal/standaloneRouting.cjs");
+const aliasConfigShim = require("../../src/mitm/_internal/aliasConfig.cjs");
+
+test("standalone MITM route config maps Claude Code and Kiro messages to /v1/messages", () => {
+  const claude = routing.getAgentRouteConfig("claude-code");
+  const kiro = routing.getAgentRouteConfig("kiro");
+  const antigravity = routing.getAgentRouteConfig("antigravity");
+
+  assert.equal(claude.aliasKey, "claude-code");
+  assert.deepEqual(claude.chatUrlPatterns, ["/v1/messages"]);
+  assert.equal(claude.routerPath, "/v1/messages");
+  assert.equal(kiro.aliasKey, "kiro");
+  assert.deepEqual(kiro.chatUrlPatterns, ["/v1/messages"]);
+  assert.equal(kiro.routerPath, "/v1/messages");
+  assert.deepEqual(antigravity.chatUrlPatterns, [":generateContent", ":streamGenerateContent"]);
+  assert.equal(routing.getAgentRouteConfig("unknown"), antigravity);
+});
+
+test("standalone MITM forwards agents from routerPath, not hardcoded agent ids", () => {
+  const fallbackCalls = [];
+  const fallbackResolver = (baseUrl, body) => {
+    fallbackCalls.push({ baseUrl, body });
+    return { format: "openai", url: `${baseUrl}/v1/chat/completions` };
+  };
+
+  const claudeForward = routing.resolveForwardTargetForAgent({
+    routerBaseUrl: "http://router",
+    routerMessagesUrl: "http://router/v1/messages",
+    body: { model: "claude-sonnet" },
+    agentId: "claude-code",
+    fallbackResolver,
+  });
+  const antigravityForward = routing.resolveForwardTargetForAgent({
+    routerBaseUrl: "http://router",
+    routerMessagesUrl: "http://router/v1/messages",
+    body: { model: "gemini-pro" },
+    agentId: "antigravity",
+    fallbackResolver,
+  });
+
+  assert.deepEqual(claudeForward, {
+    format: "anthropic",
+    url: "http://router/v1/messages",
+  });
+  assert.deepEqual(antigravityForward, {
+    format: "openai",
+    url: "http://router/v1/chat/completions",
+  });
+  assert.equal(fallbackCalls.length, 1);
+});
+
+test("standalone MITM resolves structured aliases from the agent-specific namespace", () => {
+  const rows = {
+    "claude-code": JSON.stringify({
+      "claude-source": { model: "anthropic/claude-sonnet-5", reasoningEffort: "high" },
+    }),
+    antigravity: JSON.stringify({
+      "claude-source": "antigravity-should-not-win",
+    }),
+  };
+  const db = {
+    prepare(sql) {
+      assert.equal(sql, "SELECT value FROM key_value WHERE namespace = 'mitmAlias' AND key = ?");
+      return {
+        get(key) {
+          return rows[key] ? { value: rows[key] } : undefined;
+        },
+      };
+    },
+  };
+
+  const override = routing.resolveMappedOverride("claude-source", "claude-code", {
+    fs: { existsSync: () => false },
+    dbFile: "/unused/db.json",
+    getSqliteDb: () => db,
+    aliasConfigShim,
+  });
+
+  assert.deepEqual(override, {
+    model: "anthropic/claude-sonnet-5",
+    reasoningEffort: "high",
+  });
+});
+
+test("standalone MITM resolves legacy JSON aliases from the agent-specific namespace", () => {
+  const legacyDb = {
+    mitmAlias: {
+      kiro: {
+        "kiro-source": "kiro/claude-sonnet-5",
+      },
+      antigravity: {
+        "kiro-source": "antigravity-should-not-win",
+      },
+    },
+  };
+
+  const override = routing.resolveMappedOverride("kiro-source", "kiro", {
+    fs: {
+      existsSync: () => true,
+      readFileSync: () => JSON.stringify(legacyDb),
+    },
+    dbFile: "/legacy/db.json",
+    getSqliteDb: () => null,
+    aliasConfigShim,
+  });
+
+  assert.deepEqual(override, { model: "kiro/claude-sonnet-5" });
+});