Skip to content

Commit 1f7cdd4

Browse files
jpr5jpr5
committed
fix: normalize proxy relay status codes to prevent upstream leaks
Raw upstream HTTP status codes (429, 503, 401, etc.) were leaking through to aimock clients in 5 recorder proxy relay paths, exposing provider implementation details. Normalize to 200 for success, 502 for errors — aimock is the gateway. Fixture recording preserves the original upstream status for fidelity. Paths fixed: - recorder.ts non-SSE relay (writeHead with raw upstreamStatus) - recorder.ts SSE streaming relay (res.statusCode passthrough) - agui-recorder.ts success relay (exact 201/204 leak) - agui-recorder.ts error relay (raw 401/429/503 passthrough) - agui-recorder.ts timeout handler (implicit 200 instead of 502) Add 5 new status normalization tests with fixture preservation assertions, update 2 existing tests, fix test cleanup with awaited server.close in try/finally.
1 parent 773bd5d commit 1f7cdd4

4 files changed

Lines changed: 216 additions & 11 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,16 @@
11
# @copilotkit/aimock
22

3+
## [Unreleased]
4+
5+
### Fixed
6+
7+
- **Drift tests passed vacuously with zero assertions** — the `shouldFail` guard silently skipped all `expect` calls when no critical diffs were found, so broken extraction logic or warning-level drift went completely undetected. Replaced every guarded assertion across all 21 drift test files (89 instances) with unconditional `expect(diffs.filter(...)).toEqual([])`
8+
- **Proxy relay leaked raw upstream HTTP status codes** — 5 recorder relay paths in `recorder.ts` and `agui-recorder.ts` forwarded raw upstream codes (429, 503, 401, 201, etc.) to aimock clients, exposing provider implementation details. Normalized to 200 for success and 502 for errors; fixture recording preserves the original status for fidelity
9+
10+
### Added
11+
12+
- **Status code normalization tests** — 5 tests verifying proxy relay normalization (201→200, 429→502, 503→502, 401→502, SSE 429→502) with fixture preservation assertions; 2 existing tests updated to expect normalized 502
13+
314
## [1.19.5] - 2026-05-09
415

516
### Fixed

‎src/__tests__/recorder.test.ts‎

Lines changed: 184 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1199,9 +1199,10 @@ describe("recorder edge cases", () => {
11991199
messages: [{ role: "user", content: "trigger error" }],
12001200
});
12011201

1202-
expect(resp.status).toBe(500);
1202+
// Proxy relay normalizes upstream errors to 502 (Bad Gateway)
1203+
expect(resp.status).toBe(502);
12031204

1204-
// Fixture file created with error response
1205+
// Fixture file created with error response — preserves real upstream status
12051206
const files = fs.readdirSync(tmpDir);
12061207
const fixtureFiles = files.filter((f) => f.endsWith(".json"));
12071208
expect(fixtureFiles).toHaveLength(1);
@@ -1217,6 +1218,7 @@ describe("recorder edge cases", () => {
12171218
expect(savedResponse.status).toBe(500);
12181219

12191220
// Replay: second identical request matches the recorded error fixture
1221+
// (served by aimock's fixture serving logic, which uses the fixture's status field)
12201222
const resp2 = await post(`${recorder.url}/v1/chat/completions`, {
12211223
model: "gpt-4",
12221224
messages: [{ role: "user", content: "trigger error" }],
@@ -2258,6 +2260,183 @@ describe("recorder upstream connection failure", () => {
22582260
});
22592261
});
22602262

2263+
// ---------------------------------------------------------------------------
2264+
// Status code normalization — proxy relay normalizes upstream codes
2265+
// ---------------------------------------------------------------------------
2266+
2267+
describe("recorder status code normalization", () => {
2268+
it("normalizes upstream 201 to 200 for non-SSE responses", async () => {
2269+
// Create a raw upstream that returns 201 (e.g. Anthropic messages API)
2270+
const rawServer = http.createServer((_req, res) => {
2271+
res.writeHead(201, { "Content-Type": "application/json" });
2272+
res.end(
2273+
JSON.stringify({
2274+
content: [{ type: "text", text: "created response" }],
2275+
model: "claude-3",
2276+
role: "assistant",
2277+
}),
2278+
);
2279+
});
2280+
await new Promise<void>((resolve) => rawServer.listen(0, "127.0.0.1", resolve));
2281+
const rawAddr = rawServer.address() as { port: number };
2282+
const rawUrl = `http://127.0.0.1:${rawAddr.port}`;
2283+
2284+
try {
2285+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-"));
2286+
recorder = await createServer([], {
2287+
port: 0,
2288+
record: { providers: { anthropic: rawUrl }, fixturePath: tmpDir },
2289+
});
2290+
2291+
const resp = await post(`${recorder.url}/v1/messages`, {
2292+
model: "claude-3",
2293+
messages: [{ role: "user", content: "hello 201" }],
2294+
max_tokens: 100,
2295+
});
2296+
2297+
// Client sees 200, not 201
2298+
expect(resp.status).toBe(200);
2299+
} finally {
2300+
await new Promise<void>((r) => rawServer.close(() => r()));
2301+
}
2302+
});
2303+
2304+
it("normalizes upstream 429 to 502 for non-SSE responses", async () => {
2305+
const rawServer = http.createServer((_req, res) => {
2306+
res.writeHead(429, { "Content-Type": "application/json" });
2307+
res.end(
2308+
JSON.stringify({
2309+
error: { message: "Rate limited", type: "rate_limit_error" },
2310+
}),
2311+
);
2312+
});
2313+
await new Promise<void>((resolve) => rawServer.listen(0, "127.0.0.1", resolve));
2314+
const rawAddr = rawServer.address() as { port: number };
2315+
const rawUrl = `http://127.0.0.1:${rawAddr.port}`;
2316+
2317+
try {
2318+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-"));
2319+
recorder = await createServer([], {
2320+
port: 0,
2321+
record: { providers: { openai: rawUrl }, fixturePath: tmpDir },
2322+
});
2323+
2324+
const resp = await post(`${recorder.url}/v1/chat/completions`, {
2325+
model: "gpt-4",
2326+
messages: [{ role: "user", content: "rate limit me" }],
2327+
});
2328+
2329+
// Client sees 502, not 429
2330+
expect(resp.status).toBe(502);
2331+
2332+
// Fixture preserves the real upstream 429 status
2333+
const files = fs.readdirSync(tmpDir);
2334+
const fixtureFiles = files.filter((f) => f.endsWith(".json"));
2335+
expect(fixtureFiles).toHaveLength(1);
2336+
const fixtureContent = JSON.parse(
2337+
fs.readFileSync(path.join(tmpDir, fixtureFiles[0]), "utf-8"),
2338+
);
2339+
expect(fixtureContent.fixtures[0].response.status).toBe(429);
2340+
} finally {
2341+
await new Promise<void>((r) => rawServer.close(() => r()));
2342+
}
2343+
});
2344+
2345+
it("normalizes upstream 503 to 502 for non-SSE responses", async () => {
2346+
const rawServer = http.createServer((_req, res) => {
2347+
res.writeHead(503, { "Content-Type": "application/json" });
2348+
res.end(
2349+
JSON.stringify({
2350+
error: { message: "Service unavailable", type: "server_error" },
2351+
}),
2352+
);
2353+
});
2354+
await new Promise<void>((resolve) => rawServer.listen(0, "127.0.0.1", resolve));
2355+
const rawAddr = rawServer.address() as { port: number };
2356+
const rawUrl = `http://127.0.0.1:${rawAddr.port}`;
2357+
2358+
try {
2359+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-"));
2360+
recorder = await createServer([], {
2361+
port: 0,
2362+
record: { providers: { openai: rawUrl }, fixturePath: tmpDir },
2363+
});
2364+
2365+
const resp = await post(`${recorder.url}/v1/chat/completions`, {
2366+
model: "gpt-4",
2367+
messages: [{ role: "user", content: "service unavailable" }],
2368+
});
2369+
2370+
// Client sees 502, not 503
2371+
expect(resp.status).toBe(502);
2372+
} finally {
2373+
await new Promise<void>((r) => rawServer.close(() => r()));
2374+
}
2375+
});
2376+
2377+
it("normalizes upstream 401 to 502 for non-SSE responses", async () => {
2378+
const rawServer = http.createServer((_req, res) => {
2379+
res.writeHead(401, { "Content-Type": "application/json" });
2380+
res.end(
2381+
JSON.stringify({
2382+
error: { message: "Invalid API key", type: "authentication_error" },
2383+
}),
2384+
);
2385+
});
2386+
await new Promise<void>((resolve) => rawServer.listen(0, "127.0.0.1", resolve));
2387+
const rawAddr = rawServer.address() as { port: number };
2388+
const rawUrl = `http://127.0.0.1:${rawAddr.port}`;
2389+
2390+
try {
2391+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-"));
2392+
recorder = await createServer([], {
2393+
port: 0,
2394+
record: { providers: { openai: rawUrl }, fixturePath: tmpDir },
2395+
});
2396+
2397+
const resp = await post(`${recorder.url}/v1/chat/completions`, {
2398+
model: "gpt-4",
2399+
messages: [{ role: "user", content: "bad auth" }],
2400+
});
2401+
2402+
// Client sees 502, not 401
2403+
expect(resp.status).toBe(502);
2404+
} finally {
2405+
await new Promise<void>((r) => rawServer.close(() => r()));
2406+
}
2407+
});
2408+
2409+
it("normalizes SSE streaming upstream errors to 502", async () => {
2410+
// Upstream returns 429 with SSE content-type
2411+
const rawServer = http.createServer((_req, res) => {
2412+
res.writeHead(429, { "Content-Type": "text/event-stream" });
2413+
res.end("data: rate limited\n\n");
2414+
});
2415+
await new Promise<void>((resolve) => rawServer.listen(0, "127.0.0.1", resolve));
2416+
const rawAddr = rawServer.address() as { port: number };
2417+
const rawUrl = `http://127.0.0.1:${rawAddr.port}`;
2418+
2419+
try {
2420+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-"));
2421+
recorder = await createServer([], {
2422+
port: 0,
2423+
record: { providers: { openai: rawUrl }, fixturePath: tmpDir },
2424+
});
2425+
2426+
const resp = await post(`${recorder.url}/v1/chat/completions`, {
2427+
model: "gpt-4",
2428+
messages: [{ role: "user", content: "sse rate limit" }],
2429+
stream: true,
2430+
});
2431+
2432+
// Client sees 502, not 429 — even for SSE content-type
2433+
expect(resp.status).toBe(502);
2434+
} finally {
2435+
await new Promise<void>((r) => rawServer.close(() => r()));
2436+
}
2437+
});
2438+
});
2439+
22612440
// ---------------------------------------------------------------------------
22622441
// Filesystem write failure — response still relayed
22632442
// ---------------------------------------------------------------------------
@@ -3204,12 +3383,14 @@ describe("buildFixtureResponse format detection", () => {
32043383
messages: [{ role: "user", content: "rate limit test" }],
32053384
});
32063385

3207-
expect(resp.status).toBe(429);
3386+
// Proxy relay normalizes upstream errors to 502 (Bad Gateway)
3387+
expect(resp.status).toBe(502);
32083388

32093389
const files = fs.readdirSync(tmpDir);
32103390
const fixtureFiles = files.filter((f) => f.endsWith(".json"));
32113391
expect(fixtureFiles).toHaveLength(1);
32123392

3393+
// Fixture preserves real upstream status for fidelity
32133394
const fixtureContent = JSON.parse(
32143395
fs.readFileSync(path.join(tmpDir, fixtureFiles[0]), "utf-8"),
32153396
) as {

‎src/agui-recorder.ts‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -115,17 +115,22 @@ function teeUpstreamStream(
115115
(upstreamRes) => {
116116
const upstreamStatus = upstreamRes.statusCode ?? 200;
117117

118-
// Set appropriate headers on the client response
118+
// Normalize status codes: aimock acts as a gateway, so upstream
119+
// provider details (429, 503, etc.) should not leak.
120+
// Successes → 200, errors → 502 (Bad Gateway).
121+
const clientStatus = upstreamStatus >= 200 && upstreamStatus < 300 ? 200 : 502;
122+
123+
// Set appropriate headers on the client response.
119124
if (!clientRes.headersSent) {
120-
if (upstreamStatus >= 200 && upstreamStatus < 300) {
121-
clientRes.writeHead(upstreamStatus, {
125+
if (clientStatus === 200) {
126+
clientRes.writeHead(200, {
122127
"Content-Type": "text/event-stream",
123128
"Cache-Control": "no-cache",
124129
Connection: "keep-alive",
125130
});
126131
} else {
127132
const ct = upstreamRes.headers["content-type"] || "application/json";
128-
clientRes.writeHead(upstreamStatus, { "Content-Type": ct });
133+
clientRes.writeHead(502, { "Content-Type": ct });
129134
}
130135
}
131136

@@ -218,13 +223,12 @@ function teeUpstreamStream(
218223
logger.info("Proxied AG-UI request (proxy-only mode)");
219224
}
220225

221-
resolve(upstreamStatus);
226+
resolve(clientStatus);
222227
});
223228
},
224229
);
225230

226231
upstreamReq.on("timeout", () => {
227-
if (!clientRes.writableEnded) clientRes.end();
228232
upstreamReq.destroy(
229233
new Error(`Upstream AG-UI request timed out after ${UPSTREAM_TIMEOUT_MS / 1000}s`),
230234
);

‎src/recorder.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -455,7 +455,11 @@ export async function proxyAndRecord(
455455
if (ctString) {
456456
relayHeaders["Content-Type"] = ctString;
457457
}
458-
res.writeHead(upstreamStatus, relayHeaders);
458+
// Normalize status codes for the client: aimock acts as a gateway, so
459+
// upstream provider details (429 rate-limits, 503 outages, etc.) should
460+
// not leak. Successes → 200, errors → 502 (Bad Gateway).
461+
const clientStatus = upstreamStatus >= 200 && upstreamStatus < 300 ? 200 : 502;
462+
res.writeHead(clientStatus, relayHeaders);
459463
const isAudioRelay = ctString.toLowerCase().startsWith("audio/");
460464
res.end(isBinaryStream || isAudioRelay ? rawBuffer : upstreamBody);
461465
}
@@ -511,7 +515,12 @@ function makeUpstreamRequest(
511515
if (isSSE && clientRes && !clientRes.headersSent) {
512516
const relayHeaders: Record<string, string> = {};
513517
if (ctStr) relayHeaders["Content-Type"] = ctStr;
514-
clientRes.writeHead(res.statusCode ?? 200, relayHeaders);
518+
// Normalize status codes for the client: aimock acts as a gateway,
519+
// so upstream provider details should not leak.
520+
// Successes → 200, errors → 502 (Bad Gateway).
521+
const rawStatus = res.statusCode ?? 200;
522+
const clientStatus = rawStatus >= 200 && rawStatus < 300 ? 200 : 502;
523+
clientRes.writeHead(clientStatus, relayHeaders);
515524
// Flush headers immediately so the client starts parsing frames
516525
// before the first data chunk arrives.
517526
if (typeof clientRes.flushHeaders === "function") clientRes.flushHeaders();

0 commit comments

Comments
 (0)