diff --git a/packages/junior-evals/src/eval-ai-gateway-dispatcher.ts b/packages/junior-evals/src/eval-ai-gateway-dispatcher.ts index ac895cbea3..4463433a0e 100644 --- a/packages/junior-evals/src/eval-ai-gateway-dispatcher.ts +++ b/packages/junior-evals/src/eval-ai-gateway-dispatcher.ts @@ -11,13 +11,36 @@ const AI_GATEWAY_ORIGIN = new URL( // Leaves time for provider retry backoff and a second attempt inside the // separate 60-second eval reply budget. export const EVAL_AI_GATEWAY_BODY_TIMEOUT_MS = 25_000; +// Roach takes the value of each header whose name has `session` for a +// credential, and it redacts that value in the response that the agent gets. +// This header carries the Conversation ID, so a model call that names a child +// agent's Conversation ID would get `<>`. AI Gateway only groups +// requests by it, so evals drop it. +const SESSION_AFFINITY_HEADER = "x-session-affinity"; function requestOrigin(origin: Dispatcher.DispatchOptions["origin"]): string { if (!origin) return ""; return typeof origin === "string" ? new URL(origin).origin : origin.origin; } -/** Install an eval-scoped body timeout for AI Gateway provider requests. */ +// `fetch` passes its headers as a plain object. +function withoutSessionAffinity( + headers: Dispatcher.DispatchOptions["headers"], +): Dispatcher.DispatchOptions["headers"] { + if (!headers || Array.isArray(headers) || Symbol.iterator in headers) { + return headers; + } + return Object.fromEntries( + Object.entries(headers).filter( + ([name]) => name.toLowerCase() !== SESSION_AFFINITY_HEADER, + ), + ); +} + +/** + * Install eval-scoped body timeouts for AI Gateway provider requests, and drop + * their session affinity header before they reach Roach. + */ export function installEvalAiGatewayDispatcher( bodyTimeoutMs = EVAL_AI_GATEWAY_BODY_TIMEOUT_MS, targetOrigin = AI_GATEWAY_ORIGIN, @@ -30,6 +53,7 @@ export function installEvalAiGatewayDispatcher( requestOrigin(options.origin) === normalizedTargetOrigin ? { ...options, + headers: withoutSessionAffinity(options.headers), headersTimeout: bodyTimeoutMs, bodyTimeout: bodyTimeoutMs, } diff --git a/packages/junior-evals/tests/integration/eval-ai-gateway-dispatcher.test.ts b/packages/junior-evals/tests/integration/eval-ai-gateway-dispatcher.test.ts index 2c59c0fb47..71d45f640e 100644 --- a/packages/junior-evals/tests/integration/eval-ai-gateway-dispatcher.test.ts +++ b/packages/junior-evals/tests/integration/eval-ai-gateway-dispatcher.test.ts @@ -8,12 +8,8 @@ import { installEvalAiGatewayDispatcher } from "../../src/eval-ai-gateway-dispat const openServers = new Set(); -async function startStalledServer(sendHeaders = true): Promise { - const server = http.createServer((_request, response) => { - if (!sendHeaders) return; - response.writeHead(200, { "content-type": "text/event-stream" }); - response.flushHeaders(); - }); +async function startServer(listener: http.RequestListener): Promise { + const server = http.createServer(listener); openServers.add(server); await new Promise((resolve) => server.listen(0, "127.0.0.1", () => resolve()), @@ -25,6 +21,14 @@ async function startStalledServer(sendHeaders = true): Promise { return `http://127.0.0.1:${address.port}`; } +async function startStalledServer(sendHeaders = true): Promise { + return await startServer((_request, response) => { + if (!sendHeaders) return; + response.writeHead(200, { "content-type": "text/event-stream" }); + response.flushHeaders(); + }); +} + afterEach(async () => { for (const server of openServers) { server.closeAllConnections(); @@ -34,6 +38,27 @@ afterEach(async () => { }); describe("eval AI Gateway dispatcher", () => { + it("drops the session affinity header", async () => { + const targetOrigin = await startServer((request, response) => { + response.end(JSON.stringify(request.headers)); + }); + const restore = installEvalAiGatewayDispatcher(1_000, targetOrigin); + try { + const response = await fetch(targetOrigin, { + headers: { + "X-Session-Affinity": "agent:212aaafaf81eee9e28b821a24fd743e7", + "x-api-key": "kept", + }, + }); + const headers = await response.json(); + + expect(headers).not.toHaveProperty("x-session-affinity"); + expect(headers).toMatchObject({ "x-api-key": "kept" }); + } finally { + await restore(); + } + }); + it("terminates a request that never receives headers", async () => { const targetOrigin = await startStalledServer(false); const restore = installEvalAiGatewayDispatcher(100, targetOrigin);