Codebase quality sweep: error handling, drift detection, test hygiene

CHANGELOG.md

@@ -1,5 +1,23 @@
 # @copilotkit/aimock
 
+## [Unreleased]
+
+### Fixed
+
+- **Progressive relay for NDJSON and Bedrock binary event streams** — Ollama NDJSON and Bedrock binary event streams were fully buffered before relay, triggering downstream idle timeouts; now relayed progressively as chunks arrive
+- **JSON.parse error detail in bare catch blocks** — capture and surface parse-error detail in all bare catch blocks across 25+ provider/WebSocket/stream-collapse handlers instead of swallowing context
+- **Unguarded stream write/end calls** — wrap stream write/end in try/catch (recorder.ts, agui-recorder.ts) to prevent unhandled exceptions on client disconnect
+- **Response termination for headers-already-sent paths** — add response termination in error paths where headers were already sent (server.ts, a2a-mock.ts, mcp-mock.ts), preventing connection hangs
+- **Vector-mock double body consumption** — fix route passthrough consuming the request body twice, causing empty-body forwarding
+- **Drift detection compared only first event per type** — `compareSSESequences` now compares ALL events per type, not just the first, catching previously invisible divergences
+- **Ollama drift tests used broken async describe.skipIf** — replaced with synchronous env-var gate so tests are correctly skipped or executed
+- **12 unrestored spy/mock leaks and misleading assertions** — fix spy/mock leaks across test files and correct assertions that passed for the wrong reasons
+
+### Changed
+
+- **Anti-buffering headers on all progressive stream relay paths** — standard headers (Cache-Control, Connection, X-Accel-Buffering) added to all progressive stream relay paths to prevent intermediate proxy buffering
+- **Stream-collapse returns firstDroppedSample** — stream-collapse functions now return the first dropped sample for forensic debugging of collapsed streams
+
 ## [1.21.0] - 2026-05-11
 
 ### Added

src/__tests__/api-conformance.test.ts

@@ -1037,7 +1037,7 @@ describe("OpenAI Embeddings API conformance", () => {
       );
       expect(res.status).toBe(400);
       const json = JSON.parse(res.body);
-      expect(json.error.message).toBe("Malformed JSON");
+      expect(json.error.message).toMatch(/^Malformed JSON body: /);
     });
 
     it("Content-Type is application/json", async () => {

src/__tests__/bedrock-stream.test.ts

@@ -1442,7 +1442,7 @@ describe("POST /model/{modelId}/converse (malformed JSON)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
   });
 });
 
@@ -1569,7 +1569,7 @@ describe("POST /model/{modelId}/converse-stream (malformed JSON)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
   });
 });
 

src/__tests__/bedrock.test.ts

@@ -243,7 +243,7 @@ describe("POST /model/{modelId}/invoke (error handling)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
   });
 });
 

src/__tests__/chaos.test.ts

@@ -590,6 +590,10 @@ describe("fixture-level chaos on non-OpenAI provider", () => {
 // ---------------------------------------------------------------------------
 
 describe("chaos with logLevel silent: invalid header is ignored gracefully", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
   it("proceeds normally and does not throw when x-aimock-chaos-drop is not a number", async () => {
     const fixtures: Fixture[] = [
       { match: { userMessage: "hello" }, response: { content: "Hi there" } },

src/__tests__/cohere.test.ts

@@ -575,7 +575,7 @@ describe("POST /v2/chat (validation)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON body: /);
   });
 
   it("returns 404 when no fixture matches", async () => {

src/__tests__/drift/ollama.drift.ts

@@ -1,10 +1,10 @@
 /**
  * Ollama drift tests.
  *
- * Compares aimock's Ollama endpoint output shapes against a real local
- * Ollama instance. Skips automatically if Ollama is not reachable.
+ * Compares aimock's Ollama endpoint output shapes against a real Ollama
+ * instance. Skips unless OLLAMA_HOST is set in the environment.
  *
- * Requires: local Ollama running at http://localhost:11434
+ * Requires: OLLAMA_HOST env var (e.g. http://localhost:11434)
  */
 
 import { describe, it, expect, beforeAll, afterAll } from "vitest";
@@ -13,21 +13,10 @@ import { extractShape, triangulate, formatDriftReport } from "./schema.js";
 import { httpPost, startDriftServer, stopDriftServer } from "./helpers.js";
 
 // ---------------------------------------------------------------------------
-// Connectivity check
+// Environment-based opt-in (consistent with other drift files)
 // ---------------------------------------------------------------------------
 
-let OLLAMA_REACHABLE = false;
-
-async function checkOllamaConnectivity(): Promise<boolean> {
-  try {
-    const res = await fetch("http://localhost:11434/api/tags", {
-      signal: AbortSignal.timeout(3000),
-    });
-    return res.ok;
-  } catch {
-    return false;
-  }
-}
+const OLLAMA_HOST = process.env.OLLAMA_HOST ?? "http://localhost:11434";
 
 // ---------------------------------------------------------------------------
 // Server lifecycle
@@ -36,7 +25,6 @@ async function checkOllamaConnectivity(): Promise<boolean> {
 let instance: ServerInstance;
 
 beforeAll(async () => {
-  OLLAMA_REACHABLE = await checkOllamaConnectivity();
   instance = await startDriftServer();
 });
 
@@ -119,7 +107,7 @@ function parseNDJSON(body: string): object[] {
     .map((line) => JSON.parse(line) as object);
 }
 
-describe.skipIf(!OLLAMA_REACHABLE)("Ollama drift", () => {
+describe.skipIf(!process.env.OLLAMA_HOST)("Ollama drift", () => {
   it("/api/chat response shape matches", async () => {
     const sdkShape = ollamaChatResponseShape();
 
@@ -130,7 +118,7 @@ describe.skipIf(!OLLAMA_REACHABLE)("Ollama drift", () => {
     };
 
     const [realRes, mockRes] = await Promise.all([
-      httpPost("http://localhost:11434/api/chat", body),
+      httpPost(`${OLLAMA_HOST}/api/chat`, body),
       httpPost(`${instance.url}/api/chat`, body),
     ]);
 
@@ -161,7 +149,7 @@ describe.skipIf(!OLLAMA_REACHABLE)("Ollama drift", () => {
     };
 
     const [realRes, mockRes] = await Promise.all([
-      httpPost("http://localhost:11434/api/chat", body),
+      httpPost(`${OLLAMA_HOST}/api/chat`, body),
       httpPost(`${instance.url}/api/chat`, body),
     ]);
 
@@ -199,7 +187,7 @@ describe.skipIf(!OLLAMA_REACHABLE)("Ollama drift", () => {
     };
 
     const [realRes, mockRes] = await Promise.all([
-      httpPost("http://localhost:11434/api/generate", body),
+      httpPost(`${OLLAMA_HOST}/api/generate`, body),
       httpPost(`${instance.url}/api/generate`, body),
     ]);
 

src/__tests__/drift/schema.ts

@@ -424,25 +424,23 @@ export function compareSSESequences(
     }
   }
 
-  // Compare shapes of matching event types
+  // Compare shapes of matching event types — collect ALL events per type and
+  // merge their shapes so that variant payloads (e.g. text_delta vs
+  // thinking_delta on the same Anthropic event type) are all represented.
   for (const type of realTypeSet) {
     if (!mockTypeSet.has(type)) continue;
-    const realEvent = real.find((e) => e.type === type);
-    const mockEvent = mock.find((e) => e.type === type);
-    const sdkEvent = sdk.find((e) => e.type === type);
-
-    if (realEvent && mockEvent) {
-      const eventDiffs = triangulate(
-        sdkEvent?.dataShape ?? null,
-        realEvent.dataShape,
-        mockEvent.dataShape,
-      );
-      for (const d of eventDiffs) {
-        diffs.push({
-          ...d,
-          path: `SSE:${type}.${d.path}`,
-        });
-      }
+
+    const realEvents = real.filter((e) => e.type === type);
+    const mockEvents = mock.filter((e) => e.type === type);
+    const sdkEvents = sdk.filter((e) => e.type === type);
+
+    const realMerged = mergeShapes(realEvents.map((e) => e.dataShape));
+    const mockMerged = mergeShapes(mockEvents.map((e) => e.dataShape));
+    const sdkMerged = sdkEvents.length > 0 ? mergeShapes(sdkEvents.map((e) => e.dataShape)) : null;
+
+    const eventDiffs = triangulateAt(`SSE:${type}`, sdkMerged, realMerged, mockMerged);
+    for (const d of eventDiffs) {
+      diffs.push(d);
     }
   }
 

src/__tests__/embeddings.test.ts

@@ -425,7 +425,7 @@ describe("POST /v1/embeddings (error handling)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON body: /);
     expect(body.error.code).toBe("invalid_json");
   });
 });

src/__tests__/fixture-loader.test.ts

@@ -49,6 +49,12 @@ describe("loadFixtureFile", () => {
   });
 
   afterEach(() => {
+    // Restore console spies that individual tests create via vi.spyOn(console, "warn").
+    // We cannot use vi.restoreAllMocks() here because the top-level vi.mock("node:fs")
+    // overrides would also be wiped, breaking subsequent tests.
+    if ("mockRestore" in console.warn) {
+      (console.warn as ReturnType<typeof vi.spyOn>).mockRestore();
+    }
     rmSync(tmpDir, { recursive: true, force: true });
   });
 
@@ -360,6 +366,12 @@ describe("loadFixturesFromDir", () => {
   });
 
   afterEach(() => {
+    // Restore console spies that individual tests create via vi.spyOn(console, "warn").
+    // We cannot use vi.restoreAllMocks() here because the top-level vi.mock("node:fs")
+    // overrides would also be wiped, breaking subsequent tests.
+    if ("mockRestore" in console.warn) {
+      (console.warn as ReturnType<typeof vi.spyOn>).mockRestore();
+    }
     rmSync(tmpDir, { recursive: true, force: true });
   });
 

src/__tests__/gemini.test.ts

@@ -610,7 +610,7 @@ describe("Gemini error handling", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON body: /);
   });
 
   it("returns 500 for unknown response type", async () => {

src/__tests__/mcp-mock.test.ts

@@ -874,6 +874,10 @@ describe("MCPMock", () => {
   // ---- Lifecycle edge cases ----
 
   describe("lifecycle", () => {
+    afterEach(() => {
+      vi.restoreAllMocks();
+    });
+
     it("start() when already started throws", async () => {
       mcp = new MCPMock();
       await mcp.start();

src/__tests__/messages.test.ts

@@ -703,7 +703,7 @@ describe("POST /v1/messages (error handling)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
   });
 
   it("returns 500 for unknown response type", async () => {

src/__tests__/metrics.test.ts

@@ -472,6 +472,7 @@ describe("MetricsRegistry: status label in counter output", () => {
 let instance: ServerInstance | null = null;
 
 afterEach(async () => {
+  vi.restoreAllMocks();
   if (instance) {
     await new Promise<void>((resolve) => instance!.server.close(() => resolve()));
     instance = null;

src/__tests__/ollama.test.ts

@@ -672,7 +672,7 @@ describe("POST /api/chat (error handling)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON body: /);
   });
 });
 
@@ -1102,7 +1102,7 @@ describe("POST /api/generate (malformed JSON)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON body: /);
   });
 });
 

src/__tests__/recorder.test.ts

@@ -2659,7 +2659,8 @@ describe("recorder content + toolCalls coexistence", () => {
         response: { content?: string; toolCalls?: Array<{ name: string; arguments: string }> };
       }>;
     };
-    // toolCalls should win
+    // both content and toolCalls should be preserved
+    expect(fixtureContent.fixtures[0].response.content).toBeDefined();
     expect(fixtureContent.fixtures[0].response.toolCalls).toBeDefined();
     expect(fixtureContent.fixtures[0].response.toolCalls).toHaveLength(1);
     expect(fixtureContent.fixtures[0].response.toolCalls![0].name).toBe("search");
@@ -3533,7 +3534,7 @@ describe("recorder streaming edge cases", () => {
     expect(savedResponse.content).toBe("Hello World");
   });
 
-  it("streaming with content + toolCalls: fixture saves toolCalls (not content)", async () => {
+  it("streaming with content + toolCalls: fixture saves both content and toolCalls", async () => {
     // Create a raw upstream that returns SSE with both text and tool call deltas
     const rawServer = http.createServer((_req, res) => {
       res.writeHead(200, { "Content-Type": "text/event-stream" });
@@ -3595,7 +3596,8 @@ describe("recorder streaming edge cases", () => {
       toolCalls?: Array<{ name: string; arguments: string }>;
       content?: string;
     };
-    // When toolCalls exist, they win over content
+    // Both content and toolCalls should be preserved
+    expect(savedResponse.content).toBeDefined();
     expect(savedResponse.toolCalls).toBeDefined();
     expect(savedResponse.toolCalls).toHaveLength(1);
     expect(savedResponse.toolCalls![0].name).toBe("get_weather");
@@ -3954,8 +3956,11 @@ async function setupUpstreamAndRecorder(
 
 describe("makeUpstreamRequest body timeout", () => {
   let fastRawServer: http.Server | undefined;
+  let setTimeoutSpy: ReturnType<typeof vi.spyOn> | undefined;
 
   afterEach(async () => {
+    setTimeoutSpy?.mockRestore();
+    setTimeoutSpy = undefined;
     if (fastRawServer) {
       await new Promise<void>((resolve) => fastRawServer!.close(() => resolve()));
       fastRawServer = undefined;
@@ -3975,7 +3980,7 @@ describe("makeUpstreamRequest body timeout", () => {
     await new Promise<void>((resolve) => fastRawServer!.listen(0, "127.0.0.1", resolve));
     const { port } = fastRawServer!.address() as { port: number };
 
-    const setTimeoutSpy = vi.spyOn(http.IncomingMessage.prototype, "setTimeout");
+    setTimeoutSpy = vi.spyOn(http.IncomingMessage.prototype, "setTimeout");
 
     tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "aimock-record-timeout-"));
     const record: RecordConfig = {
@@ -4009,7 +4014,6 @@ describe("makeUpstreamRequest body timeout", () => {
 
     // Verify res.setTimeout was called with the 30-second body accumulation timeout
     expect(setTimeoutSpy).toHaveBeenCalledWith(30_000, expect.any(Function));
-    setTimeoutSpy.mockRestore();
   });
 });
 

src/__tests__/responses.test.ts

@@ -790,7 +790,7 @@ describe("POST /v1/responses (error handling)", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
     expect(body.error.code).toBe("invalid_json");
   });
 

src/__tests__/server.test.ts

@@ -345,7 +345,7 @@ describe("POST /v1/chat/completions", () => {
 
     expect(res.status).toBe(400);
     const body = JSON.parse(res.body);
-    expect(body.error.message).toBe("Malformed JSON");
+    expect(body.error.message).toMatch(/^Malformed JSON: /);
     expect(body.error.code).toBe("invalid_json");
   });
 

src/__tests__/services.test.ts

@@ -244,7 +244,7 @@ describe("POST /search — edge cases", () => {
 
     expect(status).toBe(400);
     const data = json as { error: { message: string; type: string; code: string } };
-    expect(data.error.message).toBe("Malformed JSON");
+    expect(data.error.message).toMatch(/^Malformed JSON:/);
     expect(data.error.type).toBe("invalid_request_error");
     expect(data.error.code).toBe("invalid_json");
   });
@@ -318,7 +318,7 @@ describe("POST /v2/rerank — edge cases", () => {
 
     expect(status).toBe(400);
     const data = json as { error: { message: string; type: string; code: string } };
-    expect(data.error.message).toBe("Malformed JSON");
+    expect(data.error.message).toMatch(/^Malformed JSON:/);
     expect(data.error.type).toBe("invalid_request_error");
     expect(data.error.code).toBe("invalid_json");
   });
@@ -402,7 +402,7 @@ describe("POST /v1/moderations — edge cases", () => {
 
     expect(status).toBe(400);
     const data = json as { error: { message: string; type: string; code: string } };
-    expect(data.error.message).toBe("Malformed JSON");
+    expect(data.error.message).toMatch(/^Malformed JSON:/);
     expect(data.error.type).toBe("invalid_request_error");
     expect(data.error.code).toBe("invalid_json");
   });