recorder: don't drop fixtures when SDK closes socket after data: [DONE]

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

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<void>((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<void>((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<void>((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<void>((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<void>((resolve) => upstream.close(() => resolve()));
+      await new Promise<void>((resolve) => recorderServer.close(() => resolve()));
+      fs.rmSync(fixturePath, { recursive: true, force: true });
+    }
+  });
+});

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<string, string> = {
             "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,
           });
         });
       },