diff --git a/src/mitm/_internal/standaloneRouting.cjs b/src/mitm/_internal/standaloneRouting.cjs new file mode 100644 index 0000000000..cebd41e74c --- /dev/null +++ b/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, +}; diff --git a/src/mitm/server.cjs b/src/mitm/server.cjs index 9ae165f5d7..e7fe31c46d 100644 --- a/src/mitm/server.cjs +++ b/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"); }); diff --git a/tests/unit/mitm-server-claude-code-routing.test.ts b/tests/unit/mitm-server-claude-code-routing.test.ts new file mode 100644 index 0000000000..450a1daf0b --- /dev/null +++ b/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" }); +});