Skip to content

feat(record): capture and replay OpenAI/OpenRouter stream usage incl. cost - #369

Open
jpr5 wants to merge 15 commits into
mainfrom
claude/pr-review-n6v600
Open

feat(record): capture and replay OpenAI/OpenRouter stream usage incl. cost#369
jpr5 wants to merge 15 commits into
mainfrom
claude/pr-review-n6v600

Conversation

@jpr5

@jpr5 jpr5 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Closes #368.

Problem

Collapsing a streaming OpenAI-compatible chat completion dropped the final usage frame. collapseOpenAISSE bailed on any chunk with empty choices — which is exactly the shape of the frame carrying usage:

{ "object": "chat.completion.chunk", "choices": [], "usage": { "prompt_tokens": 1234, "cost": 0.0042 } }

So a recorded fixture kept content / reasoning / tool calls / timings but no token counts, and 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, so an app that bills from real provider cost couldn't e2e-test its wallet/ledger path from a tape — the same gap #269 closed for fal's x-fal-billable-units.

Changes

RecordcollapseOpenAISSE captures the last non-null usage object into a new CollapseResult.usage. The capture runs before the empty-choices guard (and after the Responses/transcription event handlers, so a transcription stream's usage isn't mistaken for chat usage). Last-usage-wins if a provider emits several. The non-streaming recorder captures the completion envelope's usage too.

Persist — the recorder writes it to the fixture's response.usage via a new sanitizeRecordedUsage: 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 values are dropped, mirroring the load-time validator in fixture-loader.ts — a fixture the recorder writes must always re-load cleanly.

Replay — recorded counts already win over estimation through the existing resolveUsage precedence, and OpenRouter shaping already emitted cost/breakdowns from response.usage. Added here: resolveOpenRouterShaping now passes unknown usage keys through verbatim instead of dropping them, and ResponseOverrides.usage / OpenRouterUsageExtras accept forward-compat extra keys.

Back-compat

A stream that reported no usage records no usage key, so every pre-existing recorded fixture stays byte-identical and estimation behavior is unchanged. Plain OpenAI (/v1/...) replays still emit token counts only — cost remains OpenRouter-shaped output.

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. Documented in docs/record-replay under a new "Recording Token Usage & Cost" section.

Tests

New src/__tests__/openai-stream-usage-record-replay.test.ts covering capture, persistence, and replay:

  • collapser captures the final empty-choices frame verbatim incl. cost; captures alongside tool calls; last-frame-wins; stays undefined with no usage; doesn't shadow transcription usage
  • recorder persists streamed and non-streaming usage, omits it when absent, and drops odd-shaped fields (asserted against validateFixtures)
  • replay emits recorded counts + cost + extras on the streaming usage chunk and the non-streaming envelope, a full record→replay round trip, and unchanged estimation for a usage-less fixture

Full suite passes with zero failures; format:check, lint, tsc --noEmit, and build all clean.

Code review (this pass)

A full multi-round code review ran on this PR. Fixes applied on the branch, each with red-green proof:

  • Sanitizer field classificationsanitizeRecordedUsage classified usage fields with key in (walks the prototype chain), so a usage key named like an Object.prototype member (toString, constructor, …) was misclassified. Switched to an own-key Set.has(), and replaced the dead inner-field arrays. Output remains a strict subset of what validateFixtures accepts (asserted end-to-end).
  • Empty usage frame — a trailing bare usage: {} frame could clobber a previously-captured populated usage (last-wins) and drop real token counts. Now only non-empty usage is captured (last-non-empty-wins).
  • Prototype-safe passthrough — the OpenRouter forward-compat passthrough now skips __proto__/constructor/prototype own-keys, preventing prototype mutation from a crafted fixture usage object.
  • Test hygiene — the round-trip recorder server is now registered with the drained afterEach cleanup (no socket leak on early throw); the prototype-pollution suite exercises all three guard legs (removed a vacuous bystander assertion, added the previously-untested prototype leg).
  • Docs & comments — record-replay docs and CHANGELOG corrected (OpenRouter listed in the collapse table; cost example marked OpenRouter-only; usage documented as orthogonal to response shape). Code comments clarified: the audio branch carries no usage, and record/replay usage capture is OpenAI-compatible-only (both streaming and non-streaming, which the "at parity" comment reflects).

Pre-existing bugs fixed (found during review)

The review surfaced several pre-existing correctness bugs in the stream-collapse path; the real ones are fixed here (each reproduced before the fix and verified after):

  • Bedrock tool-use accounting — a native/Converse tool_use content_block_start with an undefined index was silently dropped; now counted as a dropped chunk, matching the sibling arg-delta path.
  • Cohere tool-call correlationlastStartKey advanced before the tool_calls payload guard, so a payload-less start stole correlation and caused a false drop; now advanced only after the payload is confirmed.
  • OpenAI tool-call fragmentation — tool deltas missing both index and id minted fresh keys and split one call's arguments into separate entries; now fall back to the last-open tool-call key.
  • Transcription usage guards — transcription usage capture lacked the !Array.isArray / non-empty guards; an array or {} could be captured as a bogus record. Now guarded like the chat capture.

Investigated and not a bug: the over-cap relay's Content-Type Array.join branch is unreachable — Node collapses a duplicate content-type to a single string and joinDuplicateHeaders is never set, so that array branch is dead defensive code, left untouched.

Deferred follow-ups

Genuine feature extensions (net-new capability, not defects) — out of this PR's OpenAI/OpenRouter chat-usage subject:

  • Extend usage capture to non-OpenAI providers (Anthropic / Gemini / Cohere / Bedrock / Ollama), streaming and non-streaming — currently OpenAI-compatible only, symmetrically for both paths.
  • Capture Responses-API stream usage (parsed.response.usage on response.completed) — a distinct API surface.
  • collapseGeminiInteractionsSSE capturing interaction.usage on interaction.completed.
  • Minor: OpenRouterUsageExtras is not re-exported from the package entry point (pre-existing; not introduced here).

Generated by Claude Code

… cost

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
@pkg-pr-new

pkg-pr-new Bot commented Aug 11, 2026

Copy link
Copy Markdown

Open in StackBlitz

npm i https://pkg.pr.new/@copilotkit/aimock@369

commit: 2bde195

claude and others added 2 commits August 12, 2026 00:48
…apser

`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).
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.
jpr5 added a commit that referenced this pull request Aug 12, 2026
jpr5 added a commit that referenced this pull request Aug 12, 2026
- 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.
jpr5 added 8 commits August 12, 2026 10:30
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.
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.
- 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
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.
- 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.
…uard

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.
@jpr5
jpr5 force-pushed the claude/pr-review-n6v600 branch from 3284e99 to cc4d574 Compare August 12, 2026 17:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Record & replay OpenAI/OpenRouter stream usage (incl. OpenRouter usage.cost)

2 participants