From 2d39546fa938591dc2acc44f5387c92f365bc8a3 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 11 Aug 2026 22:11:37 +0000 Subject: [PATCH 01/15] feat(record): capture and replay OpenAI/OpenRouter stream usage incl. cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapsing a streaming OpenAI-compatible chat completion dropped the final usage frame — the `chat.completion.chunk` with an empty `choices` array and a populated `usage` — because the collapser skipped every chunk without choices. Recorded fixtures kept content / reasoning / tool calls / timings but no token counts, so replay could only serve the ceil(len/4) estimate or a hand-authored `response.usage`. OpenRouter's provider-reported `cost` was never captured at all, making consumer billing paths untestable from a tape (the gap #269 closed for fal's `x-fal-billable-units`). Record: `collapseOpenAISSE` captures the last non-null `usage` object into `CollapseResult.usage`; the non-streaming recorder captures the completion envelope's `usage`. Persist: the recorder writes it to `response.usage`, passing through the standard token fields plus `cost`, `cost_details`, `prompt_tokens_details`, `completion_tokens_details`, `is_byok`, and unmodelled provider extras such as `native_tokens_*`. Non-numeric / unknown-shaped fields are dropped so a recorded fixture always passes load-time validation. A stream that reported no usage records no `usage` key, keeping pre-existing fixtures byte-identical. Replay: recorded counts already win over estimation via `resolveUsage`; OpenRouter shaping now also passes forward-compat `usage` keys through verbatim instead of dropping them, and `ResponseOverrides.usage` accepts extra keys. Closes #368 --- CHANGELOG.md | 8 + docs/record-replay/index.html | 37 ++ .../openai-stream-usage-record-replay.test.ts | 481 ++++++++++++++++++ src/openrouter-chat.ts | 36 ++ src/recorder.ts | 98 +++- src/stream-collapse.ts | 32 ++ src/types.ts | 16 + 7 files changed, 706 insertions(+), 2 deletions(-) create mode 100644 src/__tests__/openai-stream-usage-record-replay.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 8016a060..8df70b8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- **Recorded OpenAI/OpenRouter token usage — including OpenRouter `usage.cost` (#368).** Collapsing a streaming OpenAI-compatible chat completion previously dropped the final usage frame (the `chat.completion.chunk` with an empty `choices` array and a populated `usage`), because the collapser skipped every chunk without choices. A recorded fixture therefore kept content / reasoning / tool calls / timings but no token counts at all, and replay could only ever serve the `ceil(length / 4)` estimate or a hand-authored `response.usage` override. OpenRouter's provider-reported `cost` was never captured, so an app that bills from real provider cost could not e2e-test its wallet/ledger path from a tape (the same gap #269 closed for fal's `x-fal-billable-units`). + - **Record:** `collapseOpenAISSE` now captures the last non-null `usage` object on the stream (`CollapseResult.usage`), and the non-streaming recorder captures the completion envelope's `usage`. + - **Persist:** the recorder writes it to the fixture's `response.usage`, passing through the standard token fields plus `cost`, `cost_details`, `prompt_tokens_details`, `completion_tokens_details`, `is_byok`, and unmodelled provider extras such as OpenRouter's `native_tokens_*`. Non-numeric / unknown-shaped fields are dropped so a recorded fixture always passes load-time validation. **Back-compatible:** a stream that reported no usage records no `usage` key and the fixture stays byte-identical to before. + - **Replay:** recorded counts win over estimation (existing `resolveUsage` precedence), and OpenRouter-shaped responses emit the recorded `cost` / breakdowns on both the final streaming usage chunk and the non-streaming envelope. `ResponseOverrides.usage` accepts forward-compat extra keys, and OpenRouter shaping now passes any such key through verbatim rather than dropping it. + - **Note:** capturing cost requires the recorded request to actually elicit a usage frame — `stream_options: { include_usage: true }` on OpenAI-compatible streams (OpenRouter sends it regardless), or a non-streaming response. Plain OpenAI (`/v1/...`) replays continue to emit token counts only; `cost` is OpenRouter-shaped output. + ### Changed - `POST /__aimock/reset` is now the canonical full reset and returns a plain `{ "reset": true }` with no deprecation header or body fields. `POST /__aimock/reset/journal` is unaffected. diff --git a/docs/record-replay/index.html b/docs/record-replay/index.html index 92ba9f05..d9e6c62d 100644 --- a/docs/record-replay/index.html +++ b/docs/record-replay/index.html @@ -449,6 +449,43 @@

Recording Block Order

faithfully block order is reconstructable on each provider's wire.

+

Recording Token Usage & Cost

+

+ OpenAI-compatible recordings (including OpenRouter) keep + the provider's reported usage on the fixture. For a non-streaming call that + is the envelope's usage object; for a stream it is the final + chat.completion.chunk — the one with an empty choices array — + that OpenAI emits when the request sets + stream_options: { include_usage: true } (OpenRouter emits it either way). +

+ 
{
+  "match": { "userMessage": "summarize this" },
+  "response": {
+    "content": "…",
+    "usage": {
+      "prompt_tokens": 1234,
+      "completion_tokens": 567,
+      "total_tokens": 1801,
+      "cost": 0.0042
+    }
+  }
+}
+

+ On replay those counts are served verbatim instead of aimock's + ceil(length / 4) estimate, and OpenRouter-shaped responses re-emit + usage.cost — so a test can assert a wallet or ledger deduction against the + amount the provider actually charged. Extra provider fields (cost_details, + prompt_tokens_details, completion_tokens_details, + native_tokens_*, …) round-trip too. +

+

+ Record with usage enabled to get cost. If the recorded request did not + ask for usage — and the provider therefore never sent a usage frame — the fixture is + written without a usage key and replay falls back to estimated token counts, + exactly as before. You can always hand-author response.usage on a fixture + instead. +

+

Header Forwarding

When proxying to upstream providers, aimock forwards the original request's headers except diff --git a/src/__tests__/openai-stream-usage-record-replay.test.ts b/src/__tests__/openai-stream-usage-record-replay.test.ts new file mode 100644 index 00000000..6177f240 --- /dev/null +++ b/src/__tests__/openai-stream-usage-record-replay.test.ts @@ -0,0 +1,481 @@ +/** + * Record & replay of provider-reported OpenAI/OpenRouter stream usage (#368). + * + * Before this, collapsing an SSE chat completion dropped the final usage frame + * (empty `choices` → `continue`), so a recorded fixture could only ever replay + * ESTIMATED token counts (`ceil(len/4)`) and OpenRouter's `usage.cost` was never + * captured at all — making consumer billing paths untestable from a tape. + * + * These tests pin the full round trip: capture (collapser) → persist (recorder) + * → replay (OpenRouter-shaped usage chunk / completion envelope). + */ + +import { describe, it, expect, afterEach } from "vitest"; +import http from "node:http"; +import { mkdtempSync, readdirSync, readFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { collapseOpenAISSE } from "../stream-collapse.js"; +import { createServer, type ServerInstance } from "../server.js"; +import { validateFixtures } from "../fixture-loader.js"; +import type { Fixture, FixtureFile, SSEChunk, TextResponse } from "../types.js"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function httpPost(url: string, body: object): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + const req = http.request( + url, + { method: "POST", headers: { "Content-Type": "application/json" } }, + (res) => { + const chunks: Buffer[] = []; + res.on("data", (c) => chunks.push(c)); + res.on("end", () => + resolve({ status: res.statusCode!, body: Buffer.concat(chunks).toString() }), + ); + }, + ); + req.on("error", reject); + req.write(JSON.stringify(body)); + req.end(); + }); +} + +/** Split an SSE body into data-frame JSON objects (skips comments + [DONE]). */ +function parseSSE(body: string): SSEChunk[] { + return body + .split("\n\n") + .map((block) => block.split("\n").find((l) => l.startsWith("data: "))) + .filter((l): l is string => !!l && !l.includes("[DONE]")) + .map((l) => JSON.parse(l.slice(6)) as SSEChunk); +} + +const OR = "/api/v1/chat/completions"; + +/** Stand-in upstream that replies with a fixed body + content-type. */ +function createUpstream( + servers: http.Server[], + body: string, + contentType: string, +): Promise { + return new Promise((resolve) => { + const srv = http.createServer((_req, res) => { + res.writeHead(200, { "Content-Type": contentType }); + res.end(body); + }); + srv.listen(0, "127.0.0.1", () => { + const addr = srv.address() as { port: number }; + servers.push(srv); + resolve(`http://127.0.0.1:${addr.port}`); + }); + }); +} + +/** The single fixture the recorder wrote into `dir`. */ +function readRecordedFixture(dir: string): Fixture { + const files = readdirSync(dir).filter((f) => f.endsWith(".json")); + expect(files).toHaveLength(1); + const file = JSON.parse(readFileSync(join(dir, files[0]), "utf-8")) as FixtureFile; + expect(file.fixtures).toHaveLength(1); + return file.fixtures[0]; +} + +/** + * A real OpenRouter streaming body: content deltas, a finish chunk, then the + * final EMPTY-`choices` usage frame carrying token counts AND `cost`. + */ +const OPENROUTER_SSE = [ + `data: ${JSON.stringify({ + id: "gen-1", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { role: "assistant", content: "Hello" } }], + })}`, + "", + `data: ${JSON.stringify({ + id: "gen-1", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { content: " world" } }], + })}`, + "", + `data: ${JSON.stringify({ + id: "gen-1", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}`, + "", + `data: ${JSON.stringify({ + id: "gen-1", + object: "chat.completion.chunk", + choices: [], + usage: { + prompt_tokens: 1234, + completion_tokens: 567, + total_tokens: 1801, + cost: 0.0042, + cost_details: { + upstream_inference_cost: 0.004, + upstream_inference_prompt_cost: 0.001, + upstream_inference_completions_cost: 0.003, + }, + prompt_tokens_details: { cached_tokens: 12 }, + completion_tokens_details: { reasoning_tokens: 34 }, + is_byok: false, + native_tokens_prompt: 1200, + native_tokens_completion: 550, + }, + })}`, + "", + "data: [DONE]", + "", +].join("\n"); + +let instance: ServerInstance | null = null; +let servers: http.Server[] = []; +let tmpDir: string | null = null; + +afterEach(async () => { + if (instance) { + await new Promise((resolve) => instance!.server.close(() => resolve())); + instance = null; + } + for (const s of servers) { + await new Promise((resolve) => s.close(() => resolve())); + } + servers = []; + if (tmpDir) { + rmSync(tmpDir, { recursive: true, force: true }); + tmpDir = null; + } +}); + +// --------------------------------------------------------------------------- +// 1. Capture — collapseOpenAISSE +// --------------------------------------------------------------------------- + +describe("collapseOpenAISSE usage capture (#368)", () => { + it("captures the final empty-choices usage frame verbatim, including cost", () => { + const result = collapseOpenAISSE(OPENROUTER_SSE); + expect(result.content).toBe("Hello world"); + expect(result.usage).toEqual({ + prompt_tokens: 1234, + completion_tokens: 567, + total_tokens: 1801, + cost: 0.0042, + cost_details: { + upstream_inference_cost: 0.004, + upstream_inference_prompt_cost: 0.001, + upstream_inference_completions_cost: 0.003, + }, + prompt_tokens_details: { cached_tokens: 12 }, + completion_tokens_details: { reasoning_tokens: 34 }, + is_byok: false, + native_tokens_prompt: 1200, + native_tokens_completion: 550, + }); + }); + + it("captures usage alongside tool calls", () => { + const body = [ + `data: ${JSON.stringify({ + choices: [ + { + delta: { + tool_calls: [{ index: 0, id: "call_1", function: { name: "f", arguments: "{}" } }], + }, + }, + ], + })}`, + "", + `data: ${JSON.stringify({ choices: [], usage: { prompt_tokens: 7, completion_tokens: 3 } })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + + const result = collapseOpenAISSE(body); + expect(result.toolCalls).toHaveLength(1); + expect(result.usage).toEqual({ prompt_tokens: 7, completion_tokens: 3 }); + }); + + it("keeps the LAST usage frame when a provider emits more than one", () => { + const body = [ + `data: ${JSON.stringify({ choices: [{ delta: { content: "x" } }], usage: { cost: 1 } })}`, + "", + `data: ${JSON.stringify({ choices: [], usage: { cost: 2 } })}`, + "", + ].join("\n"); + expect(collapseOpenAISSE(body).usage).toEqual({ cost: 2 }); + }); + + it("leaves usage undefined when the stream carries none (byte-identical to pre-#368)", () => { + const body = [ + `data: ${JSON.stringify({ choices: [{ delta: { content: "hi" } }] })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + const result = collapseOpenAISSE(body); + expect(result).toEqual({ content: "hi" }); + }); + + it("does not mistake a transcription stream's usage for chat usage", () => { + const result = collapseOpenAISSE( + `data: ${JSON.stringify({ + type: "transcript.text.done", + text: "hello", + usage: { type: "duration", seconds: 3 }, + })}\n\n`, + ); + expect(result.transcription?.usage).toEqual({ type: "duration", seconds: 3 }); + expect(result.usage).toBeUndefined(); + }); +}); + +// --------------------------------------------------------------------------- +// 2. Persist — the recorder writes response.usage +// --------------------------------------------------------------------------- + +describe("recorder persists provider usage (#368)", () => { + it("writes the streamed usage (incl. cost) onto the recorded fixture", async () => { + const upstreamUrl = await createUpstream(servers, OPENROUTER_SSE, "text/event-stream"); + tmpDir = mkdtempSync(join(tmpdir(), "aimock-usage-")); + instance = await createServer([], { + port: 0, + record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, + }); + + const res = await httpPost(`${instance.url}/v1/chat/completions`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + stream_options: { include_usage: true }, + }); + expect(res.status).toBe(200); + + const fixture = readRecordedFixture(tmpDir); + const response = fixture.response as TextResponse; + expect(response.content).toBe("Hello world"); + expect(response.usage).toEqual({ + prompt_tokens: 1234, + completion_tokens: 567, + total_tokens: 1801, + cost: 0.0042, + cost_details: { + upstream_inference_cost: 0.004, + upstream_inference_prompt_cost: 0.001, + upstream_inference_completions_cost: 0.003, + }, + prompt_tokens_details: { cached_tokens: 12 }, + completion_tokens_details: { reasoning_tokens: 34 }, + is_byok: false, + native_tokens_prompt: 1200, + native_tokens_completion: 550, + }); + // A fixture the recorder writes must always pass load-time validation. + expect(validateFixtures([fixture]).filter((r) => r.severity === "error")).toEqual([]); + }); + + it("writes usage from a NON-streaming completion envelope too", async () => { + const upstreamUrl = await createUpstream( + servers, + JSON.stringify({ + id: "gen-2", + object: "chat.completion", + choices: [{ index: 0, message: { role: "assistant", content: "hey" } }], + usage: { prompt_tokens: 11, completion_tokens: 2, total_tokens: 13, cost: 0.5 }, + }), + "application/json", + ); + tmpDir = mkdtempSync(join(tmpdir(), "aimock-usage-")); + instance = await createServer([], { + port: 0, + record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, + }); + + const res = await httpPost(`${instance.url}/v1/chat/completions`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + }); + expect(res.status).toBe(200); + + const response = readRecordedFixture(tmpDir).response as TextResponse; + expect(response.content).toBe("hey"); + expect(response.usage).toEqual({ + prompt_tokens: 11, + completion_tokens: 2, + total_tokens: 13, + cost: 0.5, + }); + }); + + it("omits usage entirely when upstream reported none", async () => { + const noUsageSSE = [ + `data: ${JSON.stringify({ choices: [{ index: 0, delta: { content: "plain" } }] })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + const upstreamUrl = await createUpstream(servers, noUsageSSE, "text/event-stream"); + tmpDir = mkdtempSync(join(tmpdir(), "aimock-usage-")); + instance = await createServer([], { + port: 0, + record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, + }); + + await httpPost(`${instance.url}/v1/chat/completions`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + }); + + const response = readRecordedFixture(tmpDir).response as TextResponse; + expect(response).toEqual({ content: "plain" }); + }); + + it("drops non-numeric usage fields so the recorded fixture still validates", async () => { + const oddSSE = [ + `data: ${JSON.stringify({ choices: [{ index: 0, delta: { content: "x" } }] })}`, + "", + `data: ${JSON.stringify({ + choices: [], + usage: { + prompt_tokens: 3, + // Shapes the fixture contract does not accept — must not be persisted. + model: "some/model", + buckets: [1, 2], + nested: { note: "text", tokens: 5 }, + cost: null, + }, + })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + const upstreamUrl = await createUpstream(servers, oddSSE, "text/event-stream"); + tmpDir = mkdtempSync(join(tmpdir(), "aimock-usage-")); + instance = await createServer([], { + port: 0, + record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, + }); + + await httpPost(`${instance.url}/v1/chat/completions`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + stream_options: { include_usage: true }, + }); + + const fixture = readRecordedFixture(tmpDir); + expect((fixture.response as TextResponse).usage).toEqual({ prompt_tokens: 3 }); + expect(validateFixtures([fixture]).filter((r) => r.severity === "error")).toEqual([]); + }); +}); + +// --------------------------------------------------------------------------- +// 3. Replay — recorded usage reaches the wire +// --------------------------------------------------------------------------- + +describe("replay of recorded usage (#368)", () => { + const recorded: Fixture = { + match: { userMessage: "hi" }, + response: { + content: "Hello world", + usage: { + prompt_tokens: 1234, + completion_tokens: 567, + total_tokens: 1801, + cost: 0.0042, + cost_details: { + upstream_inference_cost: 0.004, + upstream_inference_prompt_cost: 0.001, + upstream_inference_completions_cost: 0.003, + }, + prompt_tokens_details: { cached_tokens: 12 }, + completion_tokens_details: { reasoning_tokens: 34 }, + native_tokens_prompt: 1200, + }, + }, + }; + + it("emits the recorded token counts and cost on the final stream usage chunk", async () => { + instance = await createServer([recorded]); + const res = await httpPost(`${instance.url}${OR}`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + stream_options: { include_usage: true }, + }); + + const usageChunk = parseSSE(res.body).find((c) => c.usage !== undefined); + expect(usageChunk).toBeDefined(); + // Real recorded counts, NOT the ceil(len/4) estimate. + expect(usageChunk!.usage!.prompt_tokens).toBe(1234); + expect(usageChunk!.usage!.completion_tokens).toBe(567); + expect(usageChunk!.usage!.total_tokens).toBe(1801); + expect(usageChunk!.usage!.cost).toBe(0.0042); + expect(usageChunk!.usage!.cost_details).toEqual({ + upstream_inference_cost: 0.004, + upstream_inference_prompt_cost: 0.001, + upstream_inference_completions_cost: 0.003, + }); + expect(usageChunk!.usage!.prompt_tokens_details).toEqual({ cached_tokens: 12 }); + expect(usageChunk!.usage!.completion_tokens_details).toEqual({ reasoning_tokens: 34 }); + // Unmodelled provider fields pass through verbatim. + expect(usageChunk!.usage!.native_tokens_prompt).toBe(1200); + }); + + it("emits the recorded usage on the non-streaming completion envelope", async () => { + instance = await createServer([recorded]); + const res = await httpPost(`${instance.url}${OR}`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + }); + const json = JSON.parse(res.body) as { usage: Record }; + expect(json.usage.prompt_tokens).toBe(1234); + expect(json.usage.completion_tokens).toBe(567); + expect(json.usage.total_tokens).toBe(1801); + expect(json.usage.cost).toBe(0.0042); + expect(json.usage.native_tokens_prompt).toBe(1200); + }); + + it("round-trips record → replay end to end (cost survives the tape)", async () => { + const upstreamUrl = await createUpstream(servers, OPENROUTER_SSE, "text/event-stream"); + tmpDir = mkdtempSync(join(tmpdir(), "aimock-usage-")); + const recorder = await createServer([], { + port: 0, + record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, + }); + await httpPost(`${recorder.url}/v1/chat/completions`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + stream_options: { include_usage: true }, + }); + await new Promise((resolve) => recorder.server.close(() => resolve())); + + instance = await createServer([readRecordedFixture(tmpDir)]); + const res = await httpPost(`${instance.url}${OR}`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + stream: true, + stream_options: { include_usage: true }, + }); + const usageChunk = parseSSE(res.body).find((c) => c.usage !== undefined); + expect(usageChunk!.usage!.cost).toBe(0.0042); + expect(usageChunk!.usage!.prompt_tokens).toBe(1234); + expect(usageChunk!.usage!.completion_tokens).toBe(567); + }); + + it("still estimates tokens for a fixture without recorded usage", async () => { + instance = await createServer([{ match: { userMessage: "hi" }, response: { content: "yo" } }]); + const res = await httpPost(`${instance.url}${OR}`, { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "hi" }], + }); + const json = JSON.parse(res.body) as { usage: Record }; + expect(json.usage.completion_tokens).toBe(1); + expect(json.usage.cost).toBeUndefined(); + }); +}); diff --git a/src/openrouter-chat.ts b/src/openrouter-chat.ts index 57433fd7..5f91d449 100644 --- a/src/openrouter-chat.ts +++ b/src/openrouter-chat.ts @@ -92,6 +92,29 @@ export interface OpenRouterShaping { usageExtras: OpenRouterUsageExtras; } +/** + * `usage` keys that are already accounted for elsewhere and must therefore NOT + * be re-emitted through the forward-compat passthrough below: the canonical + * token counts (and their per-provider aliases), which `resolveUsage` folds into + * `prompt_tokens` / `completion_tokens` / `total_tokens` on the base usage + * object, plus every field this function shapes explicitly. + */ +const SHAPED_OR_CANONICAL_USAGE_KEYS = new Set([ + "prompt_tokens", + "completion_tokens", + "total_tokens", + "input_tokens", + "output_tokens", + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "cost", + "cost_details", + "prompt_tokens_details", + "completion_tokens_details", + "is_byok", +]); + /** * Resolve the shaping context from a fixture's response overrides and the * winning model slug. `provider` defaults to the winning slug's author; @@ -134,6 +157,19 @@ export function resolveOpenRouterShaping( if (u?.is_byok !== undefined) { usageExtras.is_byok = u.is_byok; } + // Forward-compat passthrough (#368): a fixture's `usage` may carry provider + // fields aimock does not model — typically written by the recorder straight + // from a real upstream usage frame (OpenRouter `native_tokens_prompt`, + // `native_tokens_completion`, …). Emit them verbatim rather than silently + // dropping them, so a recorded usage object round-trips on replay. The + // canonical token counts and every field shaped above are excluded: those are + // already resolved onto the base usage object (see helpers.ts `resolveUsage`) + // or handled with their own required-member defaults. + for (const [key, value] of Object.entries(u ?? {})) { + if (value === undefined) continue; + if (SHAPED_OR_CANONICAL_USAGE_KEYS.has(key)) continue; + usageExtras[key] = value; + } return { provider: overrides?.provider ?? deriveOpenRouterProvider(winningModel), ...(overrides?.nativeFinishReason !== undefined && { diff --git a/src/recorder.ts b/src/recorder.ts index 05ed6c52..151566bd 100644 --- a/src/recorder.ts +++ b/src/recorder.ts @@ -12,6 +12,7 @@ import type { RecordConfig, RecordedTimings, RecordProviderKey, + ResponseOverrides, ToolCall, } from "./types.js"; import { getLastMessageByRole, getTextContent, currentTurnHasToolResult } from "./router.js"; @@ -691,6 +692,15 @@ export async function proxyAndRecord( `Harmony tokens present but unparseable — content preserved verbatim${collapsed.harmonyNote ? ` (${collapsed.harmonyNote})` : ""}`, ); } + // Provider-reported usage captured from the stream's final usage frame + // (#368). Persisted on the text / toolCall fixture shapes so replay serves + // the REAL token counts instead of the `ceil(len/4)` estimate, and so + // OpenRouter's `usage.cost` round-trips to the replayed usage chunk / + // completion envelope. Omitted when the stream carried no usage, keeping + // every pre-#368 recorded fixture byte-identical. The transcription and + // audio shapes carry their own usage slots and are left untouched. + const collapsedUsage = sanitizeRecordedUsage(collapsed.usage); + const usageSpread = collapsedUsage ? { usage: collapsedUsage } : {}; // Audio from streamed inlineData (e.g. Gemini SSE with audio parts). // A single Gemini turn can interleave audio with a functionCall and/or // text/thought parts; preserve those companion modalities so the tool call @@ -753,6 +763,7 @@ export async function proxyAndRecord( ...reasoningSignatureSpread, ...redactedThinkingSpread, ...webSearchesSpread, + ...usageSpread, }; } else { const reasoningSpread = collapsed.reasoning ? { reasoning: collapsed.reasoning } : {}; @@ -802,6 +813,7 @@ export async function proxyAndRecord( ...reasoningSignatureSpread, ...redactedThinkingSpread, ...webSearchesSpread, + ...usageSpread, }; } else { fixtureResponse = { @@ -810,6 +822,7 @@ export async function proxyAndRecord( ...reasoningSignatureSpread, ...redactedThinkingSpread, ...webSearchesSpread, + ...usageSpread, }; } } else { @@ -819,6 +832,7 @@ export async function proxyAndRecord( ...reasoningSignatureSpread, ...redactedThinkingSpread, ...webSearchesSpread, + ...usageSpread, }; } } @@ -1512,6 +1526,69 @@ function toToolCallArguments(raw: unknown): string { return JSON.stringify(raw); } +/** + * `usage` sub-objects that carry numeric breakdowns rather than a scalar, and + * the inner fields aimock documents for each. Mirrors the load-time validator in + * fixture-loader.ts so anything this recorder writes loads cleanly. + */ +const USAGE_OBJECT_FIELDS: Record = { + cost_details: [ + "upstream_inference_cost", + "upstream_inference_prompt_cost", + "upstream_inference_completions_cost", + ], + prompt_tokens_details: ["cached_tokens", "cache_write_tokens", "audio_tokens"], + completion_tokens_details: ["reasoning_tokens"], +}; + +/** + * Sanitize a provider-reported `usage` object into the fixture + * `response.usage` override shape (#368). + * + * The recorder persists real upstream usage so replay can serve the provider's + * ACTUAL token counts (instead of the `ceil(len/4)` estimate) and OpenRouter's + * `cost` — which previously never survived recording at all. Upstream usage is + * provider-specific and forward-extending, so this passes fields through rather + * than whitelisting them, but only in the shapes the fixture contract accepts: + * + * - numeric scalars (`prompt_tokens`, `cost`, `native_tokens_prompt`, …) — kept + * - `is_byok` boolean — kept + * - the documented breakdown objects — kept, with non-numeric inner fields + * dropped (the load-time validator rejects those, and a fixture the recorder + * writes must always load) + * - anything else (strings, arrays, nulls, unknown objects) — DROPPED, because + * `fixture-loader.ts` errors on a non-numeric extra `usage` key and a recorded + * fixture that fails validation is worse than one missing an exotic field + * + * Returns `undefined` when nothing survives, so the caller omits `usage` + * entirely and keeps pre-#368 fixtures byte-identical. + */ +function sanitizeRecordedUsage(raw: unknown): ResponseOverrides["usage"] | undefined { + if (typeof raw !== "object" || raw === null || Array.isArray(raw)) return undefined; + const out: Record = {}; + for (const [key, val] of Object.entries(raw as Record)) { + if (val === undefined || val === null) continue; + if (key in USAGE_OBJECT_FIELDS) { + if (typeof val !== "object" || val === null || Array.isArray(val)) continue; + const inner: Record = {}; + for (const [innerKey, innerVal] of Object.entries(val as Record)) { + // Documented inner fields must be numeric (validator-enforced); undocumented + // ones are unvalidated, so pass them through only when numeric too — keeping + // the recorded object uniformly safe to re-load. + if (typeof innerVal === "number" && Number.isFinite(innerVal)) inner[innerKey] = innerVal; + } + if (Object.keys(inner).length > 0) out[key] = inner; + continue; + } + if (key === "is_byok") { + if (typeof val === "boolean") out[key] = val; + continue; + } + if (typeof val === "number" && Number.isFinite(val)) out[key] = val; + } + return Object.keys(out).length > 0 ? (out as ResponseOverrides["usage"]) : undefined; +} + /** * Detect the response format from the parsed upstream JSON and convert * it into an aimock FixtureResponse. @@ -1760,6 +1837,13 @@ function buildFixtureResponse( ? message.reasoning : undefined; + // Provider-reported usage from the non-streaming envelope (#368), at + // parity with the collapsed-stream path: replay then serves the real token + // counts rather than the `ceil(len/4)` estimate, and OpenRouter's + // `usage.cost` survives recording. Omitted when upstream sent no usage. + const openaiUsage = sanitizeRecordedUsage(obj.usage); + const usageSpread = openaiUsage ? { usage: openaiUsage } : {}; + if (hasToolCalls) { const toolCalls: ToolCall[] = (message.tool_calls as Array>).map( (tc) => { @@ -1776,19 +1860,29 @@ function buildFixtureResponse( content: message.content as string, toolCalls, ...(openaiReasoning ? { reasoning: openaiReasoning } : {}), + ...usageSpread, }; } - return { toolCalls, ...(openaiReasoning ? { reasoning: openaiReasoning } : {}) }; + return { + toolCalls, + ...(openaiReasoning ? { reasoning: openaiReasoning } : {}), + ...usageSpread, + }; } // Text content only if (hasContent) { return { content: message.content as string, ...(openaiReasoning ? { reasoning: openaiReasoning } : {}), + ...usageSpread, }; } // Recognized OpenAI shape but empty content (e.g. content filtering, zero max_tokens) - return { content: "", ...(openaiReasoning ? { reasoning: openaiReasoning } : {}) }; + return { + content: "", + ...(openaiReasoning ? { reasoning: openaiReasoning } : {}), + ...usageSpread, + }; } } diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index 88918d1e..38fdc973 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -122,6 +122,22 @@ export interface CollapseResult { */ redactedThinking?: string[]; webSearches?: string[]; + /** + * The provider-reported `usage` object from the stream's final usage frame, + * captured VERBATIM (#368). OpenAI-compatible streams (including OpenRouter) + * close with a `chat.completion.chunk` whose `choices` is empty and whose + * `usage` carries the real token counts — plus, on OpenRouter, + * `cost` / `cost_details` / `native_tokens_*`. Collapsing dropped that frame, + * so a recorded fixture could only ever replay estimated (`ceil(len/4)`) token + * counts and never a provider-reported cost. Last-usage-wins: a stream that + * emits several usage frames keeps the final one. Absent when the upstream + * stream carried no usage (e.g. `stream_options.include_usage` unset). + * + * Kept as a loose record because the field set is provider-specific and + * forward-extending; the recorder sanitizes it into the fixture + * `response.usage` override shape before persisting. + */ + usage?: Record; toolCalls?: ToolCall[]; droppedChunks?: number; firstDroppedSample?: string; @@ -372,6 +388,8 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { let transcriptUsage: Record | undefined; let reasoning = ""; const webSearchQueries: string[] = []; + // Provider-reported usage from the stream's final usage frame (#368). + let usage: Record | undefined; let droppedChunks = 0; let firstDroppedSample: string | undefined; let harmonyUnparsed = false; @@ -465,6 +483,17 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { continue; } + // Final usage frame (#368). OpenAI-compatible providers close an + // include_usage stream with a `chat.completion.chunk` carrying EMPTY + // `choices` and a populated `usage` — the only frame with the real token + // counts, and on OpenRouter the only one with `cost`. Capture it BEFORE the + // empty-`choices` guard below, which would otherwise skip the frame + // entirely. Some providers also attach `usage` to the finish chunk (non-empty + // choices), so this runs for every chat chunk, last-usage-wins. + if (parsed.usage && typeof parsed.usage === "object" && !Array.isArray(parsed.usage)) { + usage = parsed.usage as Record; + } + const choices = parsed.choices as Array> | undefined; if (!choices || choices.length === 0) continue; @@ -605,6 +634,8 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { ...(reasoning ? { reasoning } : {}), // webSearches parity with the text-only return branch. ...(webSearchQueries.length > 0 ? { webSearches: webSearchQueries } : {}), + // Provider-reported usage parity with the text-only return branch (#368). + ...(usage ? { usage } : {}), ...(droppedChunks > 0 ? { droppedChunks } : {}), ...(firstDroppedSample ? { firstDroppedSample } : {}), ...(harmonyUnparsed ? { harmonyUnparsed: true } : {}), @@ -626,6 +657,7 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { content, ...(reasoning ? { reasoning } : {}), ...(webSearchQueries.length > 0 ? { webSearches: webSearchQueries } : {}), + ...(usage ? { usage } : {}), ...(droppedChunks > 0 ? { droppedChunks } : {}), ...(firstDroppedSample ? { firstDroppedSample } : {}), ...(harmonyUnparsed ? { harmonyUnparsed: true } : {}), diff --git a/src/types.ts b/src/types.ts index eb049746..c8dcd985 100644 --- a/src/types.ts +++ b/src/types.ts @@ -232,6 +232,16 @@ export interface ResponseOverrides { * overridden (avoids asserting a fake BYOK state by default). */ is_byok?: boolean; + /** + * Forward-compat escape hatch for provider usage fields aimock does not + * model (OpenRouter `native_tokens_prompt` / `native_tokens_completion`, + * future breakdowns, …). The recorder persists such fields when an upstream + * usage frame carries them, and OpenRouter-shaped replays pass any extra key + * through verbatim onto the emitted `usage` object. The load-time validator + * still requires extra scalars to be numbers (see fixture-loader.ts), so a + * typo'd field is caught rather than silently emitted as a string. + */ + [key: string]: unknown; }; systemFingerprint?: string; finishReason?: string; @@ -790,6 +800,12 @@ export interface OpenRouterUsageExtras { }; completion_tokens_details?: { reasoning_tokens: number; [key: string]: unknown }; is_byok?: boolean; + /** + * Provider usage fields aimock does not model, passed through verbatim from a + * fixture's `response.usage` (typically written by the recorder from a real + * upstream usage frame — e.g. OpenRouter `native_tokens_prompt`). + */ + [key: string]: unknown; } export interface SSEChunk { From 6c3785124634a869f86c67586c99935e17a8ef45 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 12 Aug 2026 00:48:12 +0000 Subject: [PATCH 02/15] fix(record): route the openrouter provider key to the OpenAI SSE collapser MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `collapseStreamingResponse`'s provider switch had no case for `openrouter`, even though it is a first-class RecordProviderKey that server.ts sets on every `/api/v1/chat/completions` request. Recording a streaming OpenRouter completion fell to the `default` arm and logged [stream-collapse] unknown SSE provider "openrouter", falling back to OpenAI SSE format on every recorded stream. The collapse was already correct — the fallback IS the OpenAI collapser and OpenRouter speaks the OpenAI SSE wire format — so this is a diagnostics fix with no behavior change. The warning claimed aimock did not recognize a provider it ships first-class support for, which is actively misleading while debugging a recording (including the #368 cost-capture flow). --- CHANGELOG.md | 4 ++++ src/__tests__/stream-collapse.test.ts | 27 ++++++++++++++++++++++++++- src/stream-collapse.ts | 8 ++++++++ 3 files changed, 38 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8df70b8f..1ff686cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,10 @@ - **Replay:** recorded counts win over estimation (existing `resolveUsage` precedence), and OpenRouter-shaped responses emit the recorded `cost` / breakdowns on both the final streaming usage chunk and the non-streaming envelope. `ResponseOverrides.usage` accepts forward-compat extra keys, and OpenRouter shaping now passes any such key through verbatim rather than dropping it. - **Note:** capturing cost requires the recorded request to actually elicit a usage frame — `stream_options: { include_usage: true }` on OpenAI-compatible streams (OpenRouter sends it regardless), or a non-streaming response. Plain OpenAI (`/v1/...`) replays continue to emit token counts only; `cost` is OpenRouter-shaped output. +### Fixed + +- **`openrouter` is no longer treated as an unknown SSE provider when recording.** `collapseStreamingResponse`'s provider switch had cases for `openai` / `azure` / `anthropic` / `gemini` / `cohere` / `bedrock` but none for `openrouter`, even though it is a first-class `RecordProviderKey` that the server sets on every `/api/v1/chat/completions` request. Recording a streaming OpenRouter completion therefore hit the `default` arm and logged `[stream-collapse] unknown SSE provider "openrouter", falling back to OpenAI SSE format` on **every** recorded stream. The collapse itself was already correct (the fallback is the OpenAI collapser, and OpenRouter speaks the OpenAI SSE wire format), so this is a log-noise/diagnostics fix with no behavior change — but the warning claimed aimock did not recognize a provider it ships first-class support for, which is actively misleading while debugging a recording. + ### Changed - `POST /__aimock/reset` is now the canonical full reset and returns a plain `{ "reset": true }` with no deprecation header or body fields. `POST /__aimock/reset/journal` is unaffected. diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index 0cc50b8e..2250ab92 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi } from "vitest"; import { collapseOpenAISSE, collapseAnthropicSSE, @@ -751,6 +751,31 @@ describe("collapseStreamingResponse", () => { expect(result!.content).toBe("buf-hi"); }); + it("routes the openrouter provider key to the OpenAI collapser without warning", () => { + // OpenRouter is a first-class RecordProviderKey (set by every + // /api/v1/chat/completions record) and speaks the OpenAI SSE wire format. + // It must be a recognized case, not the unknown-provider fallback, which + // logged a misleading warning on every recorded OpenRouter stream. + const openrouterSse = [ + 'data: {"choices":[{"delta":{"content":"hi"}}]}', + "", + 'data: {"choices":[],"usage":{"prompt_tokens":4,"completion_tokens":1,"cost":0.001}}', + "", + "data: [DONE]", + "", + ].join("\n"); + const logger = { warn: vi.fn(), error: vi.fn(), info: vi.fn(), debug: vi.fn() }; + const result = collapseStreamingResponse( + "text/event-stream", + "openrouter", + openrouterSse, + logger as unknown as Parameters[3], + ); + expect(result?.content).toBe("hi"); + expect(result?.usage).toEqual({ prompt_tokens: 4, completion_tokens: 1, cost: 0.001 }); + expect(logger.warn).not.toHaveBeenCalled(); + }); + it("unknown SSE provider key falls back to OpenAI SSE format", () => { const openaiSse = 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\ndata: [DONE]\n\n'; // "unknown-provider" is not in RecordProviderKey; "as never" lets us test the runtime default branch diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index 38fdc973..b3bfb2ce 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -1895,8 +1895,16 @@ export function collapseStreamingResponse( if (ct.includes("text/event-stream")) { const str = typeof body === "string" ? body : body.toString("utf8"); switch (providerKey) { + // OpenRouter belongs here with OpenAI/Azure: it IS the OpenAI SSE wire + // format (it is the OpenAI-compatible gateway), and `openrouter` is a + // first-class RecordProviderKey set by every `/api/v1/chat/completions` + // record. Without this case it fell to the `default` arm, which collapsed + // it correctly but logged `unknown SSE provider "openrouter"` on EVERY + // recorded OpenRouter stream — telling users aimock does not recognize a + // provider it ships first-class support for. case "openai": case "azure": + case "openrouter": return collapseOpenAISSE(str); case "anthropic": return collapseAnthropicSSE(str); From 51503417e4f1ca1917c390dfec067e046ab96cf2 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 09:51:08 -0700 Subject: [PATCH 03/15] fix: own-key usage-field classification in sanitizeRecordedUsage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sanitizeRecordedUsage classified usage keys with `key in USAGE_OBJECT_FIELDS`, which walks the prototype chain. A usage key named like an Object.prototype member (toString, valueOf, constructor, ...) was misclassified: a numeric one was dropped, and an object-valued one was emitted under a bogus key that fixture-loader's validateFixtures then rejects — breaking the recorder->loader parity that guarantees a recorded fixture always validates. USAGE_OBJECT_FIELDS also declared per-field inner-name arrays that the code never used (the inner loop keeps every finite-numeric inner field), implying an allowlist that was not enforced. Replace it with a ReadonlySet of the three object-field names and classify via own-key `.has()`, mirroring the validator's own `new Set(Object.keys(...))` + escape-hatch behavior. This removes the dead arrays and fixes the prototype-chain misclassification in one change, keeping sanitizer output a subset of what the validator accepts. Adds src/__tests__/recorder-usage-sanitize.test.ts exercising the real exported function and asserting recorder->validator parity. --- src/__tests__/recorder-usage-sanitize.test.ts | 96 +++++++++++++++++++ src/recorder.ts | 29 +++--- 2 files changed, 111 insertions(+), 14 deletions(-) create mode 100644 src/__tests__/recorder-usage-sanitize.test.ts diff --git a/src/__tests__/recorder-usage-sanitize.test.ts b/src/__tests__/recorder-usage-sanitize.test.ts new file mode 100644 index 00000000..a7966a08 --- /dev/null +++ b/src/__tests__/recorder-usage-sanitize.test.ts @@ -0,0 +1,96 @@ +import { describe, it, expect } from "vitest"; +import { sanitizeRecordedUsage } from "../recorder.js"; +import { validateFixtures } from "../fixture-loader.js"; +import type { Fixture } from "../types.js"; + +// --------------------------------------------------------------------------- +// sanitizeRecordedUsage (#369): field classification must be OWN-KEY, and the +// object-field data structure must not imply a dead inner-field allowlist. +// +// The sanitizer's output must always be a SUBSET of what validateFixtures +// accepts, so "a recorded fixture always validates" holds. We assert that +// parity directly by feeding sanitized usage through validateFixtures. +// --------------------------------------------------------------------------- + +/** Build a minimal fixture whose response.usage is the sanitized override. */ +function fixtureWithUsage(usage: unknown): Fixture[] { + return [ + { + match: {}, + response: { content: "ok", usage: usage as never }, + }, + ] as unknown as Fixture[]; +} + +function usageErrors(usage: unknown): string[] { + return validateFixtures(fixtureWithUsage(usage)) + .filter((r) => r.severity === "error") + .map((r) => r.message); +} + +describe("sanitizeRecordedUsage — own-key classification (F1)", () => { + it("keeps a numeric prototype-named scalar and real fields alongside it", () => { + // `toString` collides with Object.prototype: a `key in USAGE_OBJECT_FIELDS` + // check misclassifies it as an object-field and drops the numeric value. + const out = sanitizeRecordedUsage({ toString: 5, prompt_tokens: 10 }); + expect(out).toEqual({ toString: 5, prompt_tokens: 10 }); + // And the validator accepts what the sanitizer produced (parity). + expect(usageErrors(out)).toEqual([]); + }); + + it("drops an object under a prototype-named key instead of emitting it (validator would reject it)", () => { + // `valueOf` collides with Object.prototype: `key in USAGE_OBJECT_FIELDS` is + // true, so an object value is emitted under `valueOf` — which the validator + // rejects ("usage.valueOf must be a number"), breaking recorder→loader parity. + const out = sanitizeRecordedUsage({ + valueOf: { reasoning_tokens: 3 }, + completion_tokens: 7, + }); + // `toEqual` compares OWN enumerable keys, so this proves no `valueOf` key + // was emitted (a `not.toHaveProperty` check would wrongly match the + // inherited Object.prototype.valueOf). + expect(out).toEqual({ completion_tokens: 7 }); + expect(usageErrors(out)).toEqual([]); + }); + + it("classifies the real object fields as objects, not scalars", () => { + const out = sanitizeRecordedUsage({ + prompt_tokens: 4, + prompt_tokens_details: { cached_tokens: 2 }, + }); + expect(out).toEqual({ + prompt_tokens: 4, + prompt_tokens_details: { cached_tokens: 2 }, + }); + expect(usageErrors(out)).toEqual([]); + }); +}); + +describe("sanitizeRecordedUsage — object-field passthrough (F3 characterization)", () => { + it("keeps numeric inner fields (documented and undocumented) — the design's forward-compat passthrough", () => { + // The object-field structure carries only the three field NAMES; the inner + // arrays were dead. All finite-numeric inner fields survive, matching the + // validator's index-signature escape hatch for forward-compat inner keys. + const out = sanitizeRecordedUsage({ + cost_details: { + upstream_inference_cost: 1.5, // documented + some_future_inner_cost: 0.25, // undocumented but numeric + }, + }); + expect(out).toEqual({ + cost_details: { + upstream_inference_cost: 1.5, + some_future_inner_cost: 0.25, + }, + }); + expect(usageErrors(out)).toEqual([]); + }); + + it("drops non-numeric inner fields", () => { + const out = sanitizeRecordedUsage({ + completion_tokens_details: { reasoning_tokens: 9, label: "nope" }, + }); + expect(out).toEqual({ completion_tokens_details: { reasoning_tokens: 9 } }); + expect(usageErrors(out)).toEqual([]); + }); +}); diff --git a/src/recorder.ts b/src/recorder.ts index 151566bd..43c5f6d6 100644 --- a/src/recorder.ts +++ b/src/recorder.ts @@ -1527,19 +1527,20 @@ function toToolCallArguments(raw: unknown): string { } /** - * `usage` sub-objects that carry numeric breakdowns rather than a scalar, and - * the inner fields aimock documents for each. Mirrors the load-time validator in - * fixture-loader.ts so anything this recorder writes loads cleanly. + * `usage` sub-objects that carry numeric breakdowns rather than a scalar. The + * inner fields are NOT allowlisted here on purpose: the recorder keeps every + * finite-numeric inner field (documented or forward-compat), matching the + * load-time validator in fixture-loader.ts, whose index-signature escape hatch + * accepts undocumented numeric inner keys too. A `Set` (own-key membership via + * `.has()`, no prototype-chain walk) mirrors that validator's + * `new Set(Object.keys(...))` classification so anything this recorder writes + * loads cleanly. */ -const USAGE_OBJECT_FIELDS: Record = { - cost_details: [ - "upstream_inference_cost", - "upstream_inference_prompt_cost", - "upstream_inference_completions_cost", - ], - prompt_tokens_details: ["cached_tokens", "cache_write_tokens", "audio_tokens"], - completion_tokens_details: ["reasoning_tokens"], -}; +const USAGE_OBJECT_FIELDS: ReadonlySet = new Set([ + "cost_details", + "prompt_tokens_details", + "completion_tokens_details", +]); /** * Sanitize a provider-reported `usage` object into the fixture @@ -1563,12 +1564,12 @@ const USAGE_OBJECT_FIELDS: Record = { * Returns `undefined` when nothing survives, so the caller omits `usage` * entirely and keeps pre-#368 fixtures byte-identical. */ -function sanitizeRecordedUsage(raw: unknown): ResponseOverrides["usage"] | undefined { +export function sanitizeRecordedUsage(raw: unknown): ResponseOverrides["usage"] | undefined { if (typeof raw !== "object" || raw === null || Array.isArray(raw)) return undefined; const out: Record = {}; for (const [key, val] of Object.entries(raw as Record)) { if (val === undefined || val === null) continue; - if (key in USAGE_OBJECT_FIELDS) { + if (USAGE_OBJECT_FIELDS.has(key)) { if (typeof val !== "object" || val === null || Array.isArray(val)) continue; const inner: Record = {}; for (const [innerKey, innerVal] of Object.entries(val as Record)) { From cb0b998e151fa8d166681e3b1e57bbd826378bc9 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 09:49:04 -0700 Subject: [PATCH 04/15] fix(record): ignore empty usage frames so a trailing usage:{} cannot clobber token counts (#369) --- src/__tests__/stream-collapse.test.ts | 21 +++++++++++++++++++++ src/stream-collapse.ts | 10 ++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index 2250ab92..0843fb3e 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -226,6 +226,27 @@ describe("collapseOpenAISSE", () => { expect(result.toolCalls![0].name).toBe("lookup"); expect(result.toolCalls![0].arguments).toBe('{"q":"test"}'); }); + + it("does not let a trailing empty usage frame clobber a populated one (#369 F4)", () => { + const body = [ + `data: ${JSON.stringify({ + choices: [], + usage: { prompt_tokens: 11, completion_tokens: 5, total_tokens: 16 }, + })}`, + "", + `data: ${JSON.stringify({ choices: [], usage: {} })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + + const result = collapseOpenAISSE(body); + expect(result.usage).toEqual({ + prompt_tokens: 11, + completion_tokens: 5, + total_tokens: 16, + }); + }); }); // --------------------------------------------------------------------------- diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index b3bfb2ce..db5a1a4c 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -489,8 +489,14 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { // counts, and on OpenRouter the only one with `cost`. Capture it BEFORE the // empty-`choices` guard below, which would otherwise skip the frame // entirely. Some providers also attach `usage` to the finish chunk (non-empty - // choices), so this runs for every chat chunk, last-usage-wins. - if (parsed.usage && typeof parsed.usage === "object" && !Array.isArray(parsed.usage)) { + // choices), so this runs for every chat chunk, last-NON-EMPTY-usage-wins: + // a trailing bare `usage: {}` frame must not clobber a good capture (#369). + if ( + parsed.usage && + typeof parsed.usage === "object" && + !Array.isArray(parsed.usage) && + Object.keys(parsed.usage).length > 0 + ) { usage = parsed.usage as Record; } From bdfcfbd022783b82be8a85740419e22ee190a168 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 09:50:18 -0700 Subject: [PATCH 05/15] fix(openrouter): make usage passthrough prototype-safe The forward-compat usage passthrough assigned every un-shaped key onto usageExtras verbatim. A JSON-parsed usage frame (fixtures mirror upstream provider responses) can carry a real own __proto__/constructor/prototype key; a plain usageExtras[key] = value for one of those hits the prototype setter and corrupts the emitted object's prototype chain. Skip those keys in the passthrough loop. Legitimate un-shaped keys still pass through. --- .../openrouter-passthrough-proto.test.ts | 65 +++++++++++++++++++ src/openrouter-chat.ts | 11 ++++ 2 files changed, 76 insertions(+) create mode 100644 src/__tests__/openrouter-passthrough-proto.test.ts diff --git a/src/__tests__/openrouter-passthrough-proto.test.ts b/src/__tests__/openrouter-passthrough-proto.test.ts new file mode 100644 index 00000000..797d562c --- /dev/null +++ b/src/__tests__/openrouter-passthrough-proto.test.ts @@ -0,0 +1,65 @@ +import { describe, it, expect } from "vitest"; +import { resolveOpenRouterShaping } from "../openrouter-chat.js"; +import type { ResponseOverrides } from "../types.js"; + +// --------------------------------------------------------------------------- +// Prototype-pollution guard for the forward-compat usage passthrough (#369 CR). +// +// The passthrough loop in resolveOpenRouterShaping copies every un-shaped usage +// key onto the emitted `usageExtras`. Recorded fixtures are attacker-adjacent — +// they come from upstream provider responses — and JSON.parse materializes a +// `"__proto__"` (or `"constructor"` / `"prototype"`) key as a REAL own, +// enumerable property. A plain `usageExtras[key] = value` assignment for such a +// key hits the prototype setter, corrupting the emitted object's prototype +// chain rather than emitting a data field. The passthrough must be +// prototype-safe: skip those keys while letting legitimate un-shaped keys +// (e.g. `native_tokens_prompt`) through verbatim. +// --------------------------------------------------------------------------- + +/** + * Build a usage override whose extras include a `__proto__` (and `constructor`) + * OWN key. A source-level object literal cannot express this — `{ __proto__: … }` + * sets the prototype instead of an own key — so parse it from JSON, exactly as + * the recorder would load a captured upstream usage frame. + */ +function usageWithProtoKeys(): ResponseOverrides { + const usage = JSON.parse( + '{"__proto__":{"polluted":true},"constructor":{"polluted":true},"native_tokens_prompt":1200}', + ) as ResponseOverrides["usage"]; + return { usage }; +} + +describe("openrouter usage passthrough — prototype safety", () => { + it("does not let a __proto__ usage key corrupt the emitted object's prototype", () => { + const shaping = resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); + const extras = shaping.usageExtras; + + // The emitted object's prototype chain must be intact — no bogus inherited prop. + expect(Object.getPrototypeOf(extras)).toBe(Object.prototype); + expect("polluted" in extras).toBe(false); + expect((extras as Record).polluted).toBeUndefined(); + }); + + it("does not emit __proto__ / constructor / prototype as own keys", () => { + const shaping = resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); + const extras = shaping.usageExtras; + + expect(Object.prototype.hasOwnProperty.call(extras, "__proto__")).toBe(false); + expect(Object.prototype.hasOwnProperty.call(extras, "constructor")).toBe(false); + expect(Object.prototype.hasOwnProperty.call(extras, "prototype")).toBe(false); + }); + + it("does not pollute a bystander plain object after shaping", () => { + resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); + const bystander: Record = {}; + expect("polluted" in bystander).toBe(false); + expect(bystander.polluted).toBeUndefined(); + }); + + it("still passes legitimate un-shaped usage keys through verbatim", () => { + const shaping = resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); + const extras = shaping.usageExtras as Record; + expect(extras.native_tokens_prompt).toBe(1200); + expect(Object.prototype.hasOwnProperty.call(extras, "native_tokens_prompt")).toBe(true); + }); +}); diff --git a/src/openrouter-chat.ts b/src/openrouter-chat.ts index 5f91d449..cbd99d24 100644 --- a/src/openrouter-chat.ts +++ b/src/openrouter-chat.ts @@ -115,6 +115,16 @@ const SHAPED_OR_CANONICAL_USAGE_KEYS = new Set([ "is_byok", ]); +/** + * Keys that must never flow through the passthrough assignment below: a + * JSON-parsed usage frame (fixtures are attacker-adjacent — they mirror upstream + * provider responses) can carry a real own `__proto__`/`constructor`/`prototype` + * key, and a plain `usageExtras[key] = value` for one of those hits the + * prototype setter, corrupting the emitted object's prototype chain instead of + * emitting a data field. Real providers never send these, so skipping is safe. + */ +const UNSAFE_PROTO_KEYS = new Set(["__proto__", "constructor", "prototype"]); + /** * Resolve the shaping context from a fixture's response overrides and the * winning model slug. `provider` defaults to the winning slug's author; @@ -168,6 +178,7 @@ export function resolveOpenRouterShaping( for (const [key, value] of Object.entries(u ?? {})) { if (value === undefined) continue; if (SHAPED_OR_CANONICAL_USAGE_KEYS.has(key)) continue; + if (UNSAFE_PROTO_KEYS.has(key)) continue; usageExtras[key] = value; } return { From f3e40e74a537a83f22a14adc7a05821cb1152ca2 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 09:49:26 -0700 Subject: [PATCH 06/15] test(record): register round-trip recorder with drained cleanup The round-trip record->replay test created a local recorder server and closed it inline, but never registered it with the afterEach-drained `servers` array. An assertion throw before the explicit close would leak the listening socket. Push `recorder.server` onto `servers` so it is torn down even on early throw, matching the idiom every other test in the file uses. --- src/__tests__/openai-stream-usage-record-replay.test.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/__tests__/openai-stream-usage-record-replay.test.ts b/src/__tests__/openai-stream-usage-record-replay.test.ts index 6177f240..aaecad9b 100644 --- a/src/__tests__/openai-stream-usage-record-replay.test.ts +++ b/src/__tests__/openai-stream-usage-record-replay.test.ts @@ -447,6 +447,9 @@ describe("replay of recorded usage (#368)", () => { port: 0, record: { providers: { openai: upstreamUrl }, fixturePath: tmpDir }, }); + // Register with the afterEach-drained cleanup so an early assertion throw + // before the explicit close below can't leak the listening socket. + servers.push(recorder.server); await httpPost(`${recorder.url}/v1/chat/completions`, { model: "openai/gpt-4o", messages: [{ role: "user", content: "hi" }], From 45b7f7ca165f01a45ce7a8a0066c79b4cae40a0e Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 09:51:25 -0700 Subject: [PATCH 07/15] docs(record): correct OpenRouter SSE/usage doc accuracy - record-replay page: OpenAI SSE row now lists OpenRouter alongside OpenAI/Azure (collapseOpenAISSE groups all three) - ResponseOverrides: replace stale "all 7" wording with the 7 common fields plus the 2 OpenRouter-only fields (provider, nativeFinishReason) - CHANGELOG: distinguish persist-time drop from the separate load-time validator so usage-field handling no longer reads as contradictory --- CHANGELOG.md | 2 +- docs/record-replay/index.html | 2 +- src/types.ts | 9 ++++++--- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ff686cf..9385f7f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - **Recorded OpenAI/OpenRouter token usage — including OpenRouter `usage.cost` (#368).** Collapsing a streaming OpenAI-compatible chat completion previously dropped the final usage frame (the `chat.completion.chunk` with an empty `choices` array and a populated `usage`), because the collapser skipped every chunk without choices. A recorded fixture therefore kept content / reasoning / tool calls / timings but no token counts at all, and replay could only ever serve the `ceil(length / 4)` estimate or a hand-authored `response.usage` override. OpenRouter's provider-reported `cost` was never captured, so an app that bills from real provider cost could not e2e-test its wallet/ledger path from a tape (the same gap #269 closed for fal's `x-fal-billable-units`). - **Record:** `collapseOpenAISSE` now captures the last non-null `usage` object on the stream (`CollapseResult.usage`), and the non-streaming recorder captures the completion envelope's `usage`. - - **Persist:** the recorder writes it to the fixture's `response.usage`, passing through the standard token fields plus `cost`, `cost_details`, `prompt_tokens_details`, `completion_tokens_details`, `is_byok`, and unmodelled provider extras such as OpenRouter's `native_tokens_*`. Non-numeric / unknown-shaped fields are dropped so a recorded fixture always passes load-time validation. **Back-compatible:** a stream that reported no usage records no `usage` key and the fixture stays byte-identical to before. + - **Persist:** the recorder writes it to the fixture's `response.usage`, passing through the standard token fields plus `cost`, `cost_details`, `prompt_tokens_details`, `completion_tokens_details`, `is_byok`, and unmodelled provider extras such as OpenRouter's `native_tokens_*`. At persist time the recorder drops non-numeric / unknown-shaped fields, so a recorded fixture always passes the load-time validator — a separate stage that independently rejects a non-numeric extra `usage` scalar in a hand-authored fixture (see `fixture-loader.ts`). **Back-compatible:** a stream that reported no usage records no `usage` key and the fixture stays byte-identical to before. - **Replay:** recorded counts win over estimation (existing `resolveUsage` precedence), and OpenRouter-shaped responses emit the recorded `cost` / breakdowns on both the final streaming usage chunk and the non-streaming envelope. `ResponseOverrides.usage` accepts forward-compat extra keys, and OpenRouter shaping now passes any such key through verbatim rather than dropping it. - **Note:** capturing cost requires the recorded request to actually elicit a usage frame — `stream_options: { include_usage: true }` on OpenAI-compatible streams (OpenRouter sends it regardless), or a non-streaming response. Plain OpenAI (`/v1/...`) replays continue to emit token counts only; `cost` is OpenRouter-shaped output. diff --git a/docs/record-replay/index.html b/docs/record-replay/index.html index d9e6c62d..b07eb978 100644 --- a/docs/record-replay/index.html +++ b/docs/record-replay/index.html @@ -395,7 +395,7 @@

Stream Collapsing

OpenAI SSE - OpenAI, Azure + OpenAI, Azure, OpenRouter text/event-stream diff --git a/src/types.ts b/src/types.ts index c8dcd985..fc84ac5a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -177,9 +177,12 @@ export interface FixtureMatch { * When total_tokens (or provider equivalent) is omitted, it is auto-computed * from the component fields. * - * Provider support: OpenAI Chat (all 7), Responses API (5: no role, - * systemFingerprint), Claude (5: no created, systemFingerprint), - * Gemini (2: only finishReason, usage). + * Provider support: seven fields are common (id, created, model, usage, + * systemFingerprint, finishReason, role); provider and nativeFinishReason are + * OpenRouter-only and ignored elsewhere. OpenAI Chat honors the seven common + * fields; OpenRouter additionally honors provider/nativeFinishReason; Responses + * API (5: no role, systemFingerprint); Claude (5: no created, + * systemFingerprint); Gemini (2: only finishReason, usage). */ export interface ResponseOverrides { id?: string; From dca6a94b65be97b6f1491e112c9daef10164c650 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:05:13 -0700 Subject: [PATCH 08/15] test(openrouter): remove vacuous bystander prototype-pollution test The 'does not pollute a bystander' case asserted a property that cannot fail: resolveOpenRouterShaping copies keys via plain assignment (usageExtras[key] = value), which for a '__proto__' key reparents only usageExtras itself and never mutates Object.prototype, so an unrelated plain object can never be polluted regardless of the UNSAFE_PROTO_KEYS guard. The test passed even against a guard-removed implementation, providing zero regression protection. The guard's real, observable effect is fully covered by the two genuine sibling tests, both of which go red without the guard: the emitted object's prototype chain stays intact (no inherited 'polluted') and __proto__/constructor/prototype never become own keys. --- src/__tests__/openrouter-passthrough-proto.test.ts | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/__tests__/openrouter-passthrough-proto.test.ts b/src/__tests__/openrouter-passthrough-proto.test.ts index 797d562c..06fb3efa 100644 --- a/src/__tests__/openrouter-passthrough-proto.test.ts +++ b/src/__tests__/openrouter-passthrough-proto.test.ts @@ -49,13 +49,6 @@ describe("openrouter usage passthrough — prototype safety", () => { expect(Object.prototype.hasOwnProperty.call(extras, "prototype")).toBe(false); }); - it("does not pollute a bystander plain object after shaping", () => { - resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); - const bystander: Record = {}; - expect("polluted" in bystander).toBe(false); - expect(bystander.polluted).toBeUndefined(); - }); - it("still passes legitimate un-shaped usage keys through verbatim", () => { const shaping = resolveOpenRouterShaping(usageWithProtoKeys(), "openai/gpt-4o"); const extras = shaping.usageExtras as Record; From a64adcfd7930e784ec8525b74f4d813788fba08d Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:06:55 -0700 Subject: [PATCH 09/15] docs: correct usage/cost record-replay docs and comments (#369) - record-replay page: stream-collapse summary now mentions the captured final usage frame and ordered blocks; usage/cost example annotated so cost is only claimed to round-trip on OpenRouter-endpoint replay. - CHANGELOG: reword the persist bullet so the modelled non-numeric fields (cost_details/*_details objects, is_byok boolean) are described as kept by shape and the drop rule applies only to off-shape/non-finite fields, matching sanitizeRecordedUsage + the fixture-loader validator. - recorder.ts comments: state that AudioResponse has no usage slot (audio branch cannot carry usage) and that stream/non-stream usage parity is OpenAI-compatible-only. --- CHANGELOG.md | 2 +- docs/record-replay/index.html | 25 ++++++++++++++++++------- src/recorder.ts | 17 ++++++++++++----- 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9385f7f0..42cb229f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - **Recorded OpenAI/OpenRouter token usage — including OpenRouter `usage.cost` (#368).** Collapsing a streaming OpenAI-compatible chat completion previously dropped the final usage frame (the `chat.completion.chunk` with an empty `choices` array and a populated `usage`), because the collapser skipped every chunk without choices. A recorded fixture therefore kept content / reasoning / tool calls / timings but no token counts at all, and replay could only ever serve the `ceil(length / 4)` estimate or a hand-authored `response.usage` override. OpenRouter's provider-reported `cost` was never captured, so an app that bills from real provider cost could not e2e-test its wallet/ledger path from a tape (the same gap #269 closed for fal's `x-fal-billable-units`). - **Record:** `collapseOpenAISSE` now captures the last non-null `usage` object on the stream (`CollapseResult.usage`), and the non-streaming recorder captures the completion envelope's `usage`. - - **Persist:** the recorder writes it to the fixture's `response.usage`, passing through the standard token fields plus `cost`, `cost_details`, `prompt_tokens_details`, `completion_tokens_details`, `is_byok`, and unmodelled provider extras such as OpenRouter's `native_tokens_*`. At persist time the recorder drops non-numeric / unknown-shaped fields, so a recorded fixture always passes the load-time validator — a separate stage that independently rejects a non-numeric extra `usage` scalar in a hand-authored fixture (see `fixture-loader.ts`). **Back-compatible:** a stream that reported no usage records no `usage` key and the fixture stays byte-identical to before. + - **Persist:** the recorder writes it to the fixture's `response.usage` via `sanitizeRecordedUsage`, which keeps each field by its documented shape: the numeric token scalars (`cost` among them) and any unmodelled scalar extra such as OpenRouter's `native_tokens_*`, but only when finite; the `cost_details` / `prompt_tokens_details` / `completion_tokens_details` objects, keeping their finite-numeric inner scalars; and the `is_byok` boolean. What it drops is anything that fits none of those shapes — an unmodelled extra that is not a finite number, or a modelled field of the wrong type — so a recorded fixture always passes the load-time validator (a separate stage in `fixture-loader.ts` that type-checks each documented field per its type and independently rejects, for example, a non-numeric `usage` scalar in a hand-authored fixture). **Back-compatible:** a stream that reported no usage records no `usage` key and the fixture stays byte-identical to before. - **Replay:** recorded counts win over estimation (existing `resolveUsage` precedence), and OpenRouter-shaped responses emit the recorded `cost` / breakdowns on both the final streaming usage chunk and the non-streaming envelope. `ResponseOverrides.usage` accepts forward-compat extra keys, and OpenRouter shaping now passes any such key through verbatim rather than dropping it. - **Note:** capturing cost requires the recorded request to actually elicit a usage frame — `stream_options: { include_usage: true }` on OpenAI-compatible streams (OpenRouter sends it regardless), or a non-streaming response. Plain OpenAI (`/v1/...`) replays continue to emit token counts only; `cost` is OpenRouter-shaped output. diff --git a/docs/record-replay/index.html b/docs/record-replay/index.html index b07eb978..5179fd1a 100644 --- a/docs/record-replay/index.html +++ b/docs/record-replay/index.html @@ -426,8 +426,14 @@

Stream Collapsing

- The collapse extracts text content and tool calls from streaming chunks and produces a - simple { content } or { toolCalls } fixture response. + The collapse extracts text content and tool calls from streaming chunks — plus, for + OpenAI-compatible streams, the provider-reported token usage from the final + usage frame — and produces a fixture response such as { content }, + { toolCalls }, or { content, toolCalls, usage }. Genuinely + tool-first or interleaved streams additionally gain an ordered + blocks array. See + Recording Block Order and + Recording Token Usage & Cost below.

Recording Block Order

@@ -472,11 +478,16 @@

Recording Token Usage & Cost

}

On replay those counts are served verbatim instead of aimock's - ceil(length / 4) estimate, and OpenRouter-shaped responses re-emit - usage.cost — so a test can assert a wallet or ledger deduction against the - amount the provider actually charged. Extra provider fields (cost_details, - prompt_tokens_details, completion_tokens_details, - native_tokens_*, …) round-trip too. + ceil(length / 4) estimate. The usage.cost shown above is + OpenRouter-shaped output: it is re-emitted only when the fixture is replayed through + aimock's OpenRouter endpoint + (/api/v1/chat/completions), so a test can assert a wallet or ledger deduction + against the amount the provider actually charged. On a plain OpenAI (/v1/…) + replay the cost key is still recorded and validated on the fixture but is + not served back (token counts round-trip while cost stays inert), so + copy this example into an OpenRouter fixture if you need the cost to appear. Extra + provider fields (cost_details, prompt_tokens_details, + completion_tokens_details, native_tokens_*, …) round-trip too.

Record with usage enabled to get cost. If the recorded request did not diff --git a/src/recorder.ts b/src/recorder.ts index 43c5f6d6..661f01e7 100644 --- a/src/recorder.ts +++ b/src/recorder.ts @@ -697,8 +697,13 @@ export async function proxyAndRecord( // the REAL token counts instead of the `ceil(len/4)` estimate, and so // OpenRouter's `usage.cost` round-trips to the replayed usage chunk / // completion envelope. Omitted when the stream carried no usage, keeping - // every pre-#368 recorded fixture byte-identical. The transcription and - // audio shapes carry their own usage slots and are left untouched. + // every pre-#368 recorded fixture byte-identical. Transcription fixtures + // carry their own `usage` slot, so the transcription branch below is left + // untouched. The audio shape (`AudioResponse` in types.ts) has NO usage + // field, so the audio branch cannot carry `usage` at all — inert today (a + // collapse never yields both `audioB64` and a stream usage frame), but a + // future collapser that set both would drop the usage on the audio path + // with no compiler signal. const collapsedUsage = sanitizeRecordedUsage(collapsed.usage); const usageSpread = collapsedUsage ? { usage: collapsedUsage } : {}; // Audio from streamed inlineData (e.g. Gemini SSE with audio parts). @@ -1839,9 +1844,11 @@ function buildFixtureResponse( : undefined; // Provider-reported usage from the non-streaming envelope (#368), at - // parity with the collapsed-stream path: replay then serves the real token - // counts rather than the `ceil(len/4)` estimate, and OpenRouter's - // `usage.cost` survives recording. Omitted when upstream sent no usage. + // parity with the collapsed-stream path (both capture usage for + // OpenAI-compatible providers only — only `collapseOpenAISSE` sets + // `CollapseResult.usage`): replay then serves the real token counts rather + // than the `ceil(len/4)` estimate, and OpenRouter's `usage.cost` survives + // recording. Omitted when upstream sent no usage. const openaiUsage = sanitizeRecordedUsage(obj.usage); const usageSpread = openaiUsage ? { usage: openaiUsage } : {}; From a1c359d2fea492a499982bd1f1ca7e818324cced Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:15:22 -0700 Subject: [PATCH 10/15] test(openrouter): exercise the prototype leg of the proto-pollution guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The proto-safety fixture carried __proto__, constructor, and native_tokens_prompt but never a prototype key, so the assertion that prototype is not passed through onto usageExtras was vacuously green — an impl regression dropping "prototype" from UNSAFE_PROTO_KEYS would still pass. Add a prototype own key to the JSON-parsed usage fixture so the guard's prototype leg is actually driven. --- src/__tests__/openrouter-passthrough-proto.test.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/__tests__/openrouter-passthrough-proto.test.ts b/src/__tests__/openrouter-passthrough-proto.test.ts index 06fb3efa..c24187ff 100644 --- a/src/__tests__/openrouter-passthrough-proto.test.ts +++ b/src/__tests__/openrouter-passthrough-proto.test.ts @@ -17,14 +17,15 @@ import type { ResponseOverrides } from "../types.js"; // --------------------------------------------------------------------------- /** - * Build a usage override whose extras include a `__proto__` (and `constructor`) - * OWN key. A source-level object literal cannot express this — `{ __proto__: … }` - * sets the prototype instead of an own key — so parse it from JSON, exactly as - * the recorder would load a captured upstream usage frame. + * Build a usage override whose extras include a `__proto__`, `constructor`, and + * `prototype` OWN key. A source-level object literal cannot express the first + * of these — `{ __proto__: … }` sets the prototype instead of an own key — so + * parse it from JSON, exactly as the recorder would load a captured upstream + * usage frame. All three keys exercise a distinct leg of `UNSAFE_PROTO_KEYS`. */ function usageWithProtoKeys(): ResponseOverrides { const usage = JSON.parse( - '{"__proto__":{"polluted":true},"constructor":{"polluted":true},"native_tokens_prompt":1200}', + '{"__proto__":{"polluted":true},"constructor":{"polluted":true},"prototype":{"polluted":true},"native_tokens_prompt":1200}', ) as ResponseOverrides["usage"]; return { usage }; } From cc4d574546fdbbc3f8007f52186b94b0a0423476 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:14:40 -0700 Subject: [PATCH 11/15] docs: clarify recorded usage is orthogonal to response shape (#369) --- docs/record-replay/index.html | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/record-replay/index.html b/docs/record-replay/index.html index 5179fd1a..f8d2e318 100644 --- a/docs/record-replay/index.html +++ b/docs/record-replay/index.html @@ -426,11 +426,13 @@

Stream Collapsing

- The collapse extracts text content and tool calls from streaming chunks — plus, for - OpenAI-compatible streams, the provider-reported token usage from the final - usage frame — and produces a fixture response such as { content }, - { toolCalls }, or { content, toolCalls, usage }. Genuinely - tool-first or interleaved streams additionally gain an ordered + The collapse extracts text content and tool calls from streaming chunks and produces a + fixture response such as { content }, { toolCalls }, or + { content, toolCalls }. For OpenAI-compatible streams it also captures the + provider-reported token usage from the final usage frame and attaches it to + whichever of those shapes results — usage is recorded independently of + whether the response carries content, tool calls, or both. Genuinely tool-first or + interleaved streams additionally gain an ordered blocks array. See Recording Block Order and Recording Token Usage & Cost below. From 7965cefef87f45533b96dc44c4222c025860ab08 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:55:10 -0700 Subject: [PATCH 12/15] fix(record): account unkeyable Bedrock tool_use starts as dropped chunks --- src/__tests__/stream-collapse.test.ts | 37 ++++++++++++++++ src/stream-collapse.ts | 62 +++++++++++++++++++-------- 2 files changed, 82 insertions(+), 17 deletions(-) diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index 0843fb3e..7c015aa1 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -4373,3 +4373,40 @@ describe("stream block-order instrumentation (#274)", () => { }); }); }); + +// --------------------------------------------------------------------------- +// Stream-collapse integrity: uncorrelated tool_use-start accounting (Bedrock) +// --------------------------------------------------------------------------- + +describe("collapseBedrockEventStream uncorrelated tool_use start accounting", () => { + // A Bedrock native tool_use content_block_start with an undefined block index + // cannot be keyed, so it must be accounted as a dropped chunk (matching the + // sibling arg-delta path), never silently lose the tool call's identity. + it("native: tool_use content_block_start with no index is accounted, not silently lost", () => { + const frame = encodeEventStreamMessage("chunk", { + type: "content_block_start", + // No `index` — the tool call's identity cannot be keyed. + content_block: { type: "tool_use", id: "toolu_x", name: "get_weather", input: {} }, + }); + const result = collapseBedrockEventStream(Buffer.from(frame)); + // Either captured or counted as dropped — but NOT silently vanished. + const captured = (result.toolCalls?.length ?? 0) > 0; + const accounted = (result.droppedChunks ?? 0) > 0; + expect(captured || accounted).toBe(true); + // Sibling accounting behavior: an unkeyable tool_use start is a dropped chunk. + expect(result.droppedChunks).toBe(1); + }); + + // The Converse contentBlockStart path shares the same `index !== undefined` + // guard and must account an unkeyable toolUse start too. + it("converse: toolUse contentBlockStart with no index is accounted, not silently lost", () => { + const frame = encodeEventStreamMessage("contentBlockStart", { + // No contentBlockIndex anywhere — the tool call cannot be keyed. + contentBlockStart: { + start: { toolUse: { toolUseId: "tu_1", name: "get_weather" } }, + }, + }); + const result = collapseBedrockEventStream(Buffer.from(frame)); + expect(result.droppedChunks).toBe(1); + }); +}); diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index db5a1a4c..70de81f6 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -1541,14 +1541,29 @@ export function collapseBedrockEventStream(rawBody: Buffer): CollapseResult { if (redactedData !== undefined) { redactedThinking.push(redactedData); } - if (block?.type === "tool_use" && index !== undefined) { - const created = { - id: (block.id as string) ?? "", - name: (block.name as string) ?? "", - arguments: "", - }; - toolCallMap.set(index, created); - orderAtoms.push({ kind: "toolCall", ref: created }); + if (block?.type === "tool_use") { + if (index !== undefined) { + const created = { + id: (block.id as string) ?? "", + name: (block.name as string) ?? "", + arguments: "", + }; + toolCallMap.set(index, created); + orderAtoms.push({ kind: "toolCall", ref: created }); + } else { + // A tool_use start with no block index cannot be keyed, so no later + // input_json_delta can ever correlate to it and the tool call's + // identity is silently lost. Account for it as a dropped chunk — + // matching the sibling uncorrelated arg-delta path below — rather + // than vanishing without a trace. + droppedChunks++; + if (droppedChunks === 1) { + firstDroppedSample = `tool_use content_block_start with no index — tool call identity lost: ${surrogateSafeSlice( + frameStr, + 200, + )}`; + } + } } continue; } @@ -1561,15 +1576,28 @@ export function collapseBedrockEventStream(rawBody: Buffer): CollapseResult { | number | undefined; const start = blockStart.start as Record | undefined; - if (start?.toolUse && index !== undefined) { - const toolUse = start.toolUse as Record; - const created = { - id: (toolUse.toolUseId as string) ?? "", - name: (toolUse.name as string) ?? "", - arguments: "", - }; - toolCallMap.set(index, created); - orderAtoms.push({ kind: "toolCall", ref: created }); + if (start?.toolUse) { + if (index !== undefined) { + const toolUse = start.toolUse as Record; + const created = { + id: (toolUse.toolUseId as string) ?? "", + name: (toolUse.name as string) ?? "", + arguments: "", + }; + toolCallMap.set(index, created); + orderAtoms.push({ kind: "toolCall", ref: created }); + } else { + // Same accounting as the native path: an unkeyable toolUse start + // would lose its identity, so count it as a dropped chunk rather + // than dropping it silently. + droppedChunks++; + if (droppedChunks === 1) { + firstDroppedSample = `contentBlockStart toolUse with no index — tool call identity lost: ${surrogateSafeSlice( + frameStr, + 200, + )}`; + } + } } } From 8dadd3b1dca3cde47a95364c2b6b097f20952d65 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:55:47 -0700 Subject: [PATCH 13/15] fix(record): advance Cohere lastStartKey only after a tool_calls payload --- src/__tests__/stream-collapse.test.ts | 40 +++++++++++++++++++++++++++ src/stream-collapse.ts | 10 +++++-- 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index 7c015aa1..1ae868b9 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -4410,3 +4410,43 @@ describe("collapseBedrockEventStream uncorrelated tool_use start accounting", () expect(result.droppedChunks).toBe(1); }); }); + +// --------------------------------------------------------------------------- +// Stream-collapse integrity: Cohere tool-call-start correlation ordering +// --------------------------------------------------------------------------- + +describe("collapseCohereSSE tool-call-start correlation ordering", () => { + // A payload-less tool-call-start must NOT advance lastStartKey; otherwise a + // following index-less tool-call-delta is stolen away from the prior valid + // start and miscounted as a dropped chunk. + it("payload-less tool-call-start does not steal correlation from a prior start", () => { + const body = [ + `event: tool-call-start`, + `data: ${JSON.stringify({ + type: "tool-call-start", + index: 0, + delta: { + message: { + tool_calls: { id: "call_1", type: "function", function: { name: "fn", arguments: "" } }, + }, + }, + })}`, + "", + // A start event carrying NO tool_calls payload and NO index. + `event: tool-call-start`, + `data: ${JSON.stringify({ type: "tool-call-start" })}`, + "", + // An index-less delta that should correlate to the prior valid start (0). + `event: tool-call-delta`, + `data: ${JSON.stringify({ + type: "tool-call-delta", + delta: { message: { tool_calls: { function: { arguments: '{"x":1}' } } } }, + })}`, + "", + ].join("\n"); + const result = collapseCohereSSE(body); + expect(result.toolCalls).toHaveLength(1); + expect(result.toolCalls![0].arguments).toBe('{"x":1}'); + expect(result.droppedChunks).toBeUndefined(); + }); +}); diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index 70de81f6..1979dd81 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -1234,9 +1234,6 @@ export function collapseCohereSSE(rawBody: string): CollapseResult { } else { index = nextSyntheticIndex++; } - // Track the most-recent start key (real OR synthetic) so a following - // index-less delta correlates to whichever call just opened. - lastStartKey = index; const delta = parsed.delta as Record | undefined; const message = delta?.message as Record | undefined; const toolCalls = message?.tool_calls as Record | undefined; @@ -1248,6 +1245,13 @@ export function collapseCohereSSE(rawBody: string): CollapseResult { arguments: "", }; toolCallMap.set(index, created); + // Track the most-recent start key (real OR synthetic) so a following + // index-less delta correlates to whichever call just opened. Advance it + // ONLY after confirming this start actually carried a tool_calls payload + // and created an entry; a payload-less start must not steal correlation + // from a prior valid start (which would miscount the next index-less + // delta as a dropped chunk). + lastStartKey = index; // Record the tool atom at the position its tool-call-start arrived; it // references `created` so later tool-call-delta args fill it in place. orderAtoms.push({ kind: "toolCall", ref: created }); From 20477649c10b1eb988f3cefc42f9b122279eedd1 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:56:32 -0700 Subject: [PATCH 14/15] fix(record): correlate index-and-id-less OpenAI tool deltas to last-open call --- src/__tests__/stream-collapse.test.ts | 57 +++++++++++++++++++++++++++ src/stream-collapse.ts | 16 ++++++++ 2 files changed, 73 insertions(+) diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index 1ae868b9..eb3106fc 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -4450,3 +4450,60 @@ describe("collapseCohereSSE tool-call-start correlation ordering", () => { expect(result.droppedChunks).toBeUndefined(); }); }); + +// --------------------------------------------------------------------------- +// Stream-collapse integrity: OpenAI index-and-id-less tool-call correlation +// --------------------------------------------------------------------------- + +describe("collapseOpenAISSE index-and-id-less tool-call correlation", () => { + // Arg deltas that omit BOTH index and id must fall back to the last-open + // tool-call key so one call's arguments streamed across multiple such deltas + // concatenate into one tool call instead of fragmenting. + it("arg deltas lacking both index and id concatenate into one tool call", () => { + const body = [ + `data: ${JSON.stringify({ + choices: [{ delta: { tool_calls: [{ function: { name: "foo", arguments: '{"a":' } }] } }], + })}`, + "", + `data: ${JSON.stringify({ + choices: [{ delta: { tool_calls: [{ function: { arguments: "1}" } }] } }], + })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + const result = collapseOpenAISSE(body); + expect(result.toolCalls).toHaveLength(1); + expect(result.toolCalls![0].name).toBe("foo"); + expect(result.toolCalls![0].arguments).toBe('{"a":1}'); + }); + + // Distinct calls still separate when each carries its own id, and a trailing + // bare delta correlates to whichever call opened most recently. + it("bare delta correlates to the last-open call, not a merge of all calls", () => { + const body = [ + `data: ${JSON.stringify({ + choices: [ + { delta: { tool_calls: [{ id: "a", function: { name: "fa", arguments: '{"x":1}' } }] } }, + ], + })}`, + "", + `data: ${JSON.stringify({ + choices: [ + { delta: { tool_calls: [{ id: "b", function: { name: "fb", arguments: '{"y":' } }] } }, + ], + })}`, + "", + `data: ${JSON.stringify({ + choices: [{ delta: { tool_calls: [{ function: { arguments: "2}" } }] } }], + })}`, + "", + "data: [DONE]", + "", + ].join("\n"); + const result = collapseOpenAISSE(body); + expect(result.toolCalls).toHaveLength(2); + expect(result.toolCalls![0]).toMatchObject({ name: "fa", arguments: '{"x":1}' }); + expect(result.toolCalls![1]).toMatchObject({ name: "fb", arguments: '{"y":2}' }); + }); +}); diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index 1979dd81..85336c8c 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -404,6 +404,14 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { // it (they are small per-stream counters), so synthetic keys never collide. let nextSyntheticIndex = 1_000_000; const idKeyMap = new Map(); + // The key of the most-recently-opened tool call (real, id-correlated, or + // synthetic). A delta that omits BOTH `index` and `id` carries no identity of + // its own; without a fallback each such fragment would mint a fresh synthetic + // key, splitting one call's arguments streamed across multiple bare deltas + // into separate tool-call entries with no drop/truncation signal. Correlate + // it to the last-open call instead (real OpenAI always streams `index`, so + // this only affects the degenerate index-and-id-less case). + let lastToolCallKey: number | undefined; // Cross-channel order atoms (#274), in stream arrival order. A toolCall atom // references the same accumulator object stored in toolCallMap, so later arg // deltas mutate the block in place. @@ -540,9 +548,17 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { index = nextSyntheticIndex++; idKeyMap.set(rawId, index); } + } else if (lastToolCallKey !== undefined) { + // No `index` AND no `id`: correlate to the last-open tool call so a + // call's arguments streamed across multiple bare deltas concatenate + // into one entry instead of fragmenting under fresh synthetic keys. + index = lastToolCallKey; } else { index = nextSyntheticIndex++; } + // Remember this key so a following index-and-id-less delta correlates + // to whichever call most recently opened. + lastToolCallKey = index; if (!toolCallMap.has(index)) { const created = { From 2bde1954bad3459e77351114f9c32c4789bba825 Mon Sep 17 00:00:00 2001 From: Jordan Ritter Date: Wed, 12 Aug 2026 10:56:58 -0700 Subject: [PATCH 15/15] fix(record): guard transcription usage against array and empty-object values --- src/__tests__/stream-collapse.test.ts | 32 +++++++++++++++++++++++++++ src/stream-collapse.ts | 10 ++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/__tests__/stream-collapse.test.ts b/src/__tests__/stream-collapse.test.ts index eb3106fc..d00a01fa 100644 --- a/src/__tests__/stream-collapse.test.ts +++ b/src/__tests__/stream-collapse.test.ts @@ -4507,3 +4507,35 @@ describe("collapseOpenAISSE index-and-id-less tool-call correlation", () => { expect(result.toolCalls![1]).toMatchObject({ name: "fb", arguments: '{"y":2}' }); }); }); + +// --------------------------------------------------------------------------- +// Stream-collapse integrity: transcription usage capture guards +// --------------------------------------------------------------------------- + +describe("collapseOpenAISSE transcription usage capture guards", () => { + // An array-valued usage is a type lie once cast to a Record — never capture it. + it("array-valued usage is not captured", () => { + const result = collapseOpenAISSE( + 'data: {"type":"transcript.text.done","text":"hi","usage":[]}\n\n', + ); + expect(result.transcription).toBeDefined(); + expect(result.transcription!.usage).toBeUndefined(); + }); + + // An empty usage object carries no information — never capture it. + it("empty usage object is not captured", () => { + const result = collapseOpenAISSE( + 'data: {"type":"transcript.text.done","text":"hi","usage":{}}\n\n', + ); + expect(result.transcription).toBeDefined(); + expect(result.transcription!.usage).toBeUndefined(); + }); + + // A real, non-empty usage object is still captured. + it("non-empty usage object is still captured", () => { + const result = collapseOpenAISSE( + 'data: {"type":"transcript.text.done","text":"hi","usage":{"total_tokens":5}}\n\n', + ); + expect(result.transcription!.usage).toEqual({ total_tokens: 5 }); + }); +}); diff --git a/src/stream-collapse.ts b/src/stream-collapse.ts index 85336c8c..142ceedd 100644 --- a/src/stream-collapse.ts +++ b/src/stream-collapse.ts @@ -455,7 +455,15 @@ export function collapseOpenAISSE(rawBody: string): CollapseResult { ) .map((language) => ({ code: language.code as string })); } - if (parsed.usage && typeof parsed.usage === "object") { + // Only capture a real usage OBJECT. `typeof === "object"` alone admits an + // array (a type lie once cast to Record) and an empty `{}` (meaningless); + // guard against both, exactly like any chat usage capture would. + if ( + parsed.usage && + typeof parsed.usage === "object" && + !Array.isArray(parsed.usage) && + Object.keys(parsed.usage).length > 0 + ) { transcriptUsage = parsed.usage as Record; } continue;