diff --git a/CHANGELOG.md b/CHANGELOG.md index 034359bb..16a03423 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Record mode no longer silently drops fixtures when an SDK closes its socket immediately after `data: [DONE]` (e.g. the OpenAI Python SDK); the completed upstream response is now persisted. Genuine mid-stream aborts (client disconnect before `[DONE]`) still abort upstream and write no fixture (#288) + ## [1.35.0] - 2026-06-27 ### Added diff --git a/src/__tests__/recorder-early-client-close.test.ts b/src/__tests__/recorder-early-client-close.test.ts new file mode 100644 index 00000000..07be2169 --- /dev/null +++ b/src/__tests__/recorder-early-client-close.test.ts @@ -0,0 +1,131 @@ +import { describe, it, expect } from "vitest"; +import * as http from "node:http"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import type { FixtureFile } from "../types.js"; +import { proxyAndRecord } from "../recorder.js"; +import { Logger } from "../logger.js"; + +// =========================================================================== +// Regression: record mode must persist the fixture when the client closes its +// socket immediately after consuming `data: [DONE]`, before upstream has fired +// `res.end`. Some SDKs (notably the OpenAI Python SDK) close the response the +// instant they read `[DONE]`; if the recorder destroyed the upstream request +// and discarded buffered chunks on that client-close signal, record mode would +// silently produce NO fixture even though the full response body was received. +// +// This exercises the REAL recorder path (proxyAndRecord -> collapse -> +// persistFixture) against a live upstream and a real client that disconnects +// on `[DONE]` — not a stub of the recorder. +// =========================================================================== +describe("recorder — early client socket close after [DONE]", () => { + it("persists the full fixture when the client closes right after data: [DONE], before upstream end", async () => { + // Upstream streams a complete OpenAI chat-completions SSE turn ending with + // `data: [DONE]`, then holds the connection open briefly before calling + // `res.end()`. This reproduces the real race: the client sees `[DONE]` and + // closes its socket while upstream has NOT yet emitted `end`. + const UPSTREAM_END_DELAY_MS = 150; + const upstream = http.createServer((_upReq, upRes) => { + upRes.writeHead(200, { "Content-Type": "text/event-stream" }); + upRes.write( + `data: ${JSON.stringify({ choices: [{ delta: { role: "assistant", content: "Hello" } }] })}\n\n`, + ); + upRes.write(`data: ${JSON.stringify({ choices: [{ delta: { content: " world" } }] })}\n\n`); + upRes.write( + `data: ${JSON.stringify({ choices: [{ delta: {}, finish_reason: "stop" }] })}\n\n`, + ); + upRes.write("data: [DONE]\n\n"); + // Deliberately defer `end` so it fires AFTER the client has closed. + setTimeout(() => { + if (!upRes.writableEnded) upRes.end(); + }, UPSTREAM_END_DELAY_MS); + }); + await new Promise((resolve) => upstream.listen(0, "127.0.0.1", () => resolve())); + const upstreamPort = (upstream.address() as { port: number }).port; + + const fixturePath = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-early-close-")); + const logger = new Logger("silent"); + + const recorderServer = http.createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (c: Buffer) => chunks.push(c)); + req.on("end", async () => { + const rawBody = Buffer.concat(chunks).toString(); + await proxyAndRecord( + req, + res, + JSON.parse(rawBody), + "openai", + "/v1/chat/completions", + [], + { + record: { providers: { openai: `http://127.0.0.1:${upstreamPort}` }, fixturePath }, + logger, + }, + rawBody, + ); + }); + }); + await new Promise((resolve) => recorderServer.listen(0, "127.0.0.1", () => resolve())); + const recorderPort = (recorderServer.address() as { port: number }).port; + + try { + // Client that closes its socket the instant it observes `data: [DONE]`, + // mimicking the OpenAI Python SDK's behavior. + await new Promise((resolve) => { + const req = http.request( + { + hostname: "127.0.0.1", + port: recorderPort, + path: "/v1/chat/completions", + method: "POST", + headers: { "Content-Type": "application/json" }, + }, + (res) => { + let seen = ""; + res.on("data", (c: Buffer) => { + seen += c.toString(); + if (seen.includes("data: [DONE]")) { + // Close immediately, before upstream has fired `end`. + req.destroy(); + resolve(); + } + }); + res.on("error", () => resolve()); + res.on("end", () => resolve()); + }, + ); + req.on("error", () => resolve()); + req.end( + JSON.stringify({ + model: "gpt-4", + stream: true, + messages: [{ role: "user", content: "hi" }], + }), + ); + }); + + // Allow upstream `end` to fire and the recorder to persist the fixture. + await new Promise((resolve) => setTimeout(resolve, UPSTREAM_END_DELAY_MS + 150)); + + // The completed response must be recorded despite the early client close. + const all = fs.existsSync(fixturePath) ? fs.readdirSync(fixturePath) : []; + const jsons = all.filter((f) => f.endsWith(".json")); + const tmps = all.filter((f) => f.includes(".tmp.")); + expect(jsons).toHaveLength(1); + expect(tmps).toHaveLength(0); + + const fixture = JSON.parse( + fs.readFileSync(path.join(fixturePath, jsons[0]), "utf-8"), + ) as FixtureFile; + const saved = fixture.fixtures[0].response as { content?: string }; + // The full streamed content must be captured, not a truncated fragment. + expect(saved.content).toBe("Hello world"); + } finally { + await new Promise((resolve) => upstream.close(() => resolve())); + await new Promise((resolve) => recorderServer.close(() => resolve())); + fs.rmSync(fixturePath, { recursive: true, force: true }); + } + }); +}); diff --git a/src/recorder.ts b/src/recorder.ts index d98598e7..eb299688 100644 --- a/src/recorder.ts +++ b/src/recorder.ts @@ -447,6 +447,7 @@ export async function proxyAndRecord( // skip the final res.writeHead/res.end relay at the bottom of this fn. let streamedToClient = false; let clientDisconnected = false; + let sawDone = false; let frameTimestamps: number[] = []; let streamStartTime = 0; const maxProxyBufferBytes = clampMaxBufferBytes(record.maxProxyBufferBytes); @@ -469,6 +470,7 @@ export async function proxyAndRecord( rawBuffer = result.rawBuffer; streamedToClient = result.streamedToClient; clientDisconnected = result.clientDisconnected; + sawDone = result.sawDone; frameTimestamps = result.frameTimestamps; streamStartTime = result.streamStartTime; bufferTruncated = result.bufferTruncated; @@ -779,14 +781,24 @@ export async function proxyAndRecord( } } - // If the client disconnected mid-stream, the collected data is likely - // truncated. Saving a partial fixture is worse than saving none — skip - // fixture persistence entirely. + // Client may have closed its socket before upstream fired `end`. + // Distinguish two cases based on whether the SSE `[DONE]` marker was seen: + // - sawDone=true: client closed after consuming `data: [DONE]` (e.g. the + // OpenAI Python SDK closes the socket the moment it reads `[DONE]`). + // Upstream ran to completion; the buffered body is intact. Log and + // proceed to persist the full fixture. + // - sawDone=false: genuine mid-stream abort. The buffered body is partial; + // saving it would produce a corrupt fixture. Skip fixture persistence. if (clientDisconnected) { + if (!sawDone) { + defaults.logger.warn( + "Client disconnected mid-stream — skipping fixture save to avoid truncated data", + ); + return "relayed"; + } defaults.logger.warn( - "Client disconnected mid-stream — skipping fixture save to avoid truncated data", + "Client closed connection before upstream end — upstream response completed, recording full fixture", ); - return "relayed"; } // Build RecordedTimings from frame timestamps captured during streaming. @@ -1001,6 +1013,15 @@ function makeUpstreamRequest( * loud (502) rather than relay this truncated body as a success status. */ hardCeilingExceeded: boolean; + /** + * True when the SSE terminal marker `data: [DONE]` was observed in the + * upstream stream before the client disconnected. When `clientDisconnected` + * is true AND `sawDone` is true, the client closed after consuming `[DONE]` + * (e.g. the OpenAI Python SDK) — the stream was logically complete, so the + * fixture SHOULD be persisted. When `sawDone` is false the disconnect was a + * genuine mid-stream abort and the fixture MUST NOT be persisted. + */ + sawDone: boolean; }> { return new Promise((resolve, reject) => { const transport = target.protocol === "https:" ? https : http; @@ -1054,6 +1075,16 @@ function makeUpstreamRequest( let streamedToClient = false; let clientDisconnected = false; + // True once the SSE terminal marker `data: [DONE]` has been seen in the + // upstream stream. Used to distinguish two client-close scenarios: + // 1. Close AT/AFTER `[DONE]`: the stream is logically complete (e.g. + // the OpenAI Python SDK closes the socket the moment it reads + // `[DONE]`, before upstream fires `res.end`). The upstream should + // NOT be torn down — let it finish and persist the full fixture. + // 2. Close BEFORE `[DONE]`: genuine mid-stream abort. Restore the + // original teardown so upstream is destroyed and no partial fixture + // is persisted. + let sawDone = false; if (isProgressiveStream && clientRes && !clientRes.headersSent) { const relayHeaders: Record = { "Cache-Control": "no-cache, no-transform", @@ -1071,25 +1102,27 @@ function makeUpstreamRequest( // before the first data chunk arrives. if (typeof clientRes.flushHeaders === "function") clientRes.flushHeaders(); streamedToClient = true; - // Stop relaying if the client disconnects mid-stream. - // Check writableFinished to distinguish normal completion (where - // "close" also fires) from premature client disconnects. clientRes.on("close", () => { if (!clientRes.writableFinished) { clientDisconnected = true; - req.destroy(); - // Stop in-flight buffering immediately rather than draining the - // upstream socket under backpressure: detach the data listener so - // no further chunks accumulate, and free what's already buffered. - // The promise still settles via the 'end'/'error'/'close' the - // destroyed request emits; collapse/recording is skipped because - // the client disconnected. - res.removeListener("data", onUpstreamData); - chunks.length = 0; - bufferedBytes = 0; - frameTimestamps.length = 0; - frameBuffer = ""; - binaryFrameBuffer = Buffer.alloc(0); + if (!sawDone) { + // Genuine mid-stream abort — tear down upstream so it stops + // producing data, and clear all parse state so no partial + // fixture is persisted. Mirrors the pre-#288 behavior. + req.destroy(); + res.removeListener("data", onUpstreamData); + chunks.length = 0; + bufferedBytes = 0; + frameTimestamps.length = 0; + frameBuffer = ""; + binaryFrameBuffer = Buffer.alloc(0); + } + // If sawDone is true: client closed after consuming `[DONE]` + // (e.g. the OpenAI Python SDK). The upstream is logically + // complete; let it run to `res.end` so the full body is + // buffered and the fixture can be persisted. The + // `!clientDisconnected` guard inside `onUpstreamData` already + // prevents further writes to the now-closed client socket. } }); } @@ -1221,8 +1254,12 @@ function makeUpstreamRequest( tripTruncation("frame"); break; } - if (parts[fi].trim().length > 0) { + const frame = parts[fi].trim(); + if (frame.length > 0) { frameTimestamps.push(Date.now()); + // Track the SSE terminal marker so the client-close handler + // can distinguish a post-[DONE] close from a mid-stream abort. + if (frame === "data: [DONE]") sawDone = true; } } // Last part stays in buffer (may be incomplete). Skip when the @@ -1343,6 +1380,7 @@ function makeUpstreamRequest( truncationCap, totalBytes, hardCeilingExceeded, + sawDone, }); }); },