diff --git a/src/lib/db/inspectorCustomHosts.ts b/src/lib/db/inspectorCustomHosts.ts index 9874def829..a4b371eb60 100644 --- a/src/lib/db/inspectorCustomHosts.ts +++ b/src/lib/db/inspectorCustomHosts.ts @@ -75,3 +75,16 @@ export function touchLastSeen(host: string): void { const now = new Date().toISOString(); db.prepare("UPDATE inspector_custom_hosts SET last_seen_at = ? WHERE host = ?").run(now, host); } + +/** + * Returns true when `host` is present in inspector_custom_hosts with enabled=1. + * Used by agentBridgeHook to distinguish custom-host intercepts from agent-bridge + * intercepts so that Mode 2 (Custom Hosts) entries appear in the "Custom" profile. + */ +export function isCustomHost(host: string): boolean { + const db = getDbInstance(); + const row = db + .prepare("SELECT 1 AS found FROM inspector_custom_hosts WHERE host = ? AND enabled = 1") + .get(host) as { found: number } | undefined; + return row !== undefined; +} diff --git a/src/mitm/inspector/agentBridgeHook.ts b/src/mitm/inspector/agentBridgeHook.ts index e247f38628..00d3e37243 100644 --- a/src/mitm/inspector/agentBridgeHook.ts +++ b/src/mitm/inspector/agentBridgeHook.ts @@ -14,6 +14,7 @@ import { sanitizeHeaders } from "../sanitizeHeaders.ts"; import type { AgentId } from "../types.ts"; import { globalTrafficBuffer } from "./buffer.ts"; import type { InterceptedRequest } from "./types.ts"; +import { isCustomHost } from "@/lib/db/inspectorCustomHosts"; export interface RecordRequestStartOpts { req: IncomingMessage; @@ -42,13 +43,28 @@ export async function recordRequestStart( opts: RecordRequestStartOpts ): Promise { const requestBody = opts.body.length > 0 ? maskSecret(opts.body.toString("utf8")) : null; + + // Determine whether this request originates from a custom-host intercept + // (Mode 2 / Custom Hosts) or a standard agent-bridge intercept (Mode 1). + // + // Both modes reach this hook via the same MITM server path: custom hosts are + // added to inspector_custom_hosts by the Mode 2 UI and are spoofed to + // 127.0.0.1 by /etc/hosts entries, so they arrive here just like agent + // targets. The DB lookup below is the cheapest reliable way to distinguish + // them without touching server.cjs — it costs one SQLite read per request. + // + // If the host resolves as a custom-host entry, source="custom-host" and + // agent is left undefined so the "Custom" profile filter matches correctly. + const host = opts.req.headers.host ?? ""; + const customHost = isCustomHost(host); + const intercepted: InterceptedRequest = { id: randomUUID(), - source: "agent-bridge", - agent: opts.agentId, + source: customHost ? "custom-host" : "agent-bridge", + agent: customHost ? undefined : opts.agentId, timestamp: new Date().toISOString(), method: opts.req.method ?? "GET", - host: opts.req.headers.host ?? "", + host, path: opts.req.url ?? "/", requestHeaders: sanitizeHeaders(opts.req.headers), requestBody, diff --git a/tests/unit/inspector-agent-bridge-hook.test.ts b/tests/unit/inspector-agent-bridge-hook.test.ts new file mode 100644 index 0000000000..a1a0c75b69 --- /dev/null +++ b/tests/unit/inspector-agent-bridge-hook.test.ts @@ -0,0 +1,98 @@ +/** + * Unit tests: agentBridgeHook — source and agent field assignment + * + * Verifies that recordRequestStart() sets: + * - source="custom-host" + agent=undefined when the request host is in + * inspector_custom_hosts with enabled=1 (R5-8) + * - source="agent-bridge" + agent=agentId otherwise + */ + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import type { IncomingMessage } from "node:http"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-ab-hook-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const { resetDbInstance, getDbInstance } = await import("../../src/lib/db/core.ts"); +const { addCustomHost, toggleCustomHost } = await import( + "../../src/lib/db/inspectorCustomHosts.ts" +); +const { recordRequestStart } = await import( + "../../src/mitm/inspector/agentBridgeHook.ts" +); + +async function resetStorage() { + resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); + getDbInstance(); +} + +function makeFakeReq(host: string): IncomingMessage { + return { + method: "POST", + url: "/v1/chat/completions", + headers: { host, "content-type": "application/json" }, + } as unknown as IncomingMessage; +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(() => { + resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("recordRequestStart: custom-host entry → source=custom-host, agent=undefined", async () => { + addCustomHost("my-app.example.com", "app", "My App"); + + const entry = await recordRequestStart({ + req: makeFakeReq("my-app.example.com"), + body: Buffer.from("{}"), + agentId: "codex" as any, + mappedModel: "gpt-4o", + }); + + assert.equal(entry.source, "custom-host", "source should be custom-host"); + assert.equal(entry.agent, undefined, "agent should be undefined for custom-host"); + assert.equal(entry.host, "my-app.example.com"); +}); + +test("recordRequestStart: non-custom host → source=agent-bridge, agent=agentId", async () => { + // Do NOT add the host to inspector_custom_hosts + const entry = await recordRequestStart({ + req: makeFakeReq("api.openai.com"), + body: Buffer.from("{}"), + agentId: "codex" as any, + mappedModel: "gpt-4o", + }); + + assert.equal(entry.source, "agent-bridge", "source should be agent-bridge"); + assert.equal(entry.agent, "codex", "agent should be the provided agentId"); + assert.equal(entry.host, "api.openai.com"); +}); + +test("recordRequestStart: disabled custom-host → source=agent-bridge (not matched)", async () => { + addCustomHost("disabled-app.example.com"); + toggleCustomHost("disabled-app.example.com", false); + + const entry = await recordRequestStart({ + req: makeFakeReq("disabled-app.example.com"), + body: Buffer.from("{}"), + agentId: "codex" as any, + mappedModel: "gpt-4o", + }); + + assert.equal( + entry.source, + "agent-bridge", + "disabled custom-host should not be treated as custom-host source" + ); + assert.equal(entry.agent, "codex"); +});