Skip to content

fix: track token usage for streaming (SSE) LLM responses - #214

Open
anwesha01b wants to merge 1 commit into
Ruthwik000:mainfrom
anwesha01b:fix/issue-91-output-token-count-zero-streaming
Open

fix: track token usage for streaming (SSE) LLM responses#214
anwesha01b wants to merge 1 commit into
Ruthwik000:mainfrom
anwesha01b:fix/issue-91-output-token-count-zero-streaming

Conversation

@anwesha01b

Copy link
Copy Markdown

Problem

When making streaming LLM API requests (stream: true), TokenFirewall silently skipped all token usage tracking and budget management. The standardFetch function cloned the response and called clonedResponse.json() on it — which throws a SyntaxError on Server-Sent Events (SSE) payloads because data: … lines are not valid JSON. The error was caught by a broad catch block and swallowed silently, so inputTokens, outputTokens, totalCost, and budgetManager.track() were never called.

Closes #91

Root Cause

Before: always tried to JSON-parse the response body, which throws on SSE streams.

What Changed

src/interceptors/fetchInterceptor.ts

  • In standardFetch, added a Content-Type check before attempting clonedResponse.json().
  • If the response is text/event-stream or application/x-ndjson, the new processStreamResponse() function is fired asynchronously in the background so the stream is returned to the caller immediately without blocking.
  • processStreamResponse() reads the cloned ReadableStream chunk-by-chunk via a TextDecoder, parses each data: line, and accumulates token counts per provider convention:
    • OpenAI / Grok / Kimiusage.prompt_tokens / completion_tokens
    • Anthropicmessage_start.message.usage.input_tokens + message_delta.usage.output_tokens
    • GeminiusageMetadata.promptTokenCount / candidatesTokenCount
  • If no usage data is found (provider omitted the usage chunk), the function exits silently with no side-effects.
  • All existing non-streaming code paths are unchanged.

src/index.ts

  • Exports unpatchGlobalFetch from the public API surface (previously internal only) to support clean teardown in tests and integrations.

Tests

Three test scripts added under tests/:

File What it verifies
test-stream.js SSE stream with a usage chunk is tracked; totalSpent > 0
test-stream-no-usage.js SSE stream with no usage chunk does not crash; totalSpent == 0
test-non-streaming-regression.js Existing JSON response path still tracks correctly

All three pass. TypeScript build (tsc) passes with zero errors. git diff --check passes with zero trailing-whitespace violations.

Notes

  • Users must pass stream_options: { include_usage: true } in their OpenAI requests for the API to emit a usage chunk. TokenFirewall does not rewrite request bodies to inject this flag.
  • The background async nature of processStreamResponse means budget enforcement for a streaming call takes effect on the next request, not mid-stream — consistent with standard quota-tracking conventions.

Resolves Ruthwik000#91. When a provider returns a text/event-stream or
application/x-ndjson response, standardFetch previously called
clonedResponse.json() which failed silently on SSE payloads,
causing token counts and budget tracking to be entirely skipped.

Add processStreamResponse() that asynchronously reads the cloned
ReadableStream chunk-by-chunk, parses each SSE data line, and
extracts usage fields per provider convention:
- OpenAI / Grok / Kimi: usage.prompt_tokens / completion_tokens
- Anthropic: message_start.message.usage + message_delta.usage
- Gemini: usageMetadata.promptTokenCount / candidatesTokenCount

The function fires in the background so the stream is returned to
the caller immediately without blocking. If no usage data is found
(e.g. provider omitted the usage chunk) it exits silently.

Also export unpatchGlobalFetch from the public API so tests and
integrations can cleanly restore the original fetch after use.

Tests added:
- tests/test-stream.js: full usage tracking via SSE with usage chunk
- tests/test-stream-no-usage.js: graceful no-op when usage absent
- tests/test-non-streaming-regression.js: existing JSON path unchanged
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.

[Bug] Output Token Count Is Zero for Streaming LLM Responses

1 participant