Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 78 additions & 0 deletions src/mitm/_internal/standaloneRouting.cjs
Original file line number Diff line number Diff line change
@@ -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,
};
99 changes: 40 additions & 59 deletions src/mitm/server.cjs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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 });
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand All @@ -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();
Expand Down Expand Up @@ -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);

Expand All @@ -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)`);
Expand All @@ -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);
Expand Down Expand Up @@ -741,21 +728,15 @@ 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);
return;
}

// 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");
});

Expand Down
113 changes: 113 additions & 0 deletions tests/unit/mitm-server-claude-code-routing.test.ts
Original file line number Diff line number Diff line change
@@ -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" });
});