Repository navigation
fix(ai-gateway): read Bedrock invoke-with-response-stream usage from base64 chunks - #4078
Merged
oschwartz10612 merged 4 commits intoNov 9, 2026
Merged
Conversation
Fix the upsert
Gemini reports the tokens of tool results (Google Search, URL context) in usageMetadata.toolUsePromptTokenCount, separate from promptTokenCount. The extractor ignored it, so usage and cost were under-counted.
…prompt-tokens fix(ai): count Gemini toolUsePromptTokenCount as input tokens
…base64 chunks
invoke-with-response-stream wraps each model event as {"bytes":"<base64>"} inside event-stream frames, so the usage fields never appear in clear text. extractUsage returned null and the pipeline fell back to estimating tokens over the base64, losing real counts and cache tokens. Decode the chunks before scanning.
pcbeingused333
requested review from
miloschwartz and
oschwartz10612
as code owners
November 9, 2026 10:18
oschwartz10612
approved these changes
Nov 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Community Contribution License Agreement
By creating this pull request, I grant the project maintainers an unlimited,
perpetual license to use, modify, and redistribute these contributions under any terms they
choose, including both the AGPLv3 and the Fossorial Commercial license terms. I
represent that I have the right to grant this license for all contributed content.
AI Disclosure
This change was prepared with AI assistance (Claude Code): code, test and description. It was reviewed before submission.
Description
invoke-with-response-streamhas nox-amzn-bedrock-*token headers. Each model event arrives inside a binary event-stream frame whose payload is{"bytes":"<base64>","p":"..."}, so the usage fields (input_tokens,cache_read_input_tokens,amazon-bedrock-invocationMetrics, ...) are never present as clear text.extractBedrockModelInvokesearched the raw text, soextractUsagereturnednulland the pipeline fell back toestimateUsageover the base64 text: estimated tokens/cost and no cache counts. The comment there assumed a regex over the binary was enough, which is true forconverse-stream(plain JSON) but not for this API.The fix adds a small private helper that pulls each
"bytes"value, decodes it withBuffer.from(..., "base64")and re-expresses the events asdata: <json>lines.extractBedrockModelInvokescans that text (falling back to the original text when no chunks are found), for both the Anthropic-style and theinputTokenCount/outputTokenCountshapes. Headers still take priority;extractBedrockConverseis untouched.Note: no response from real AWS Bedrock was captured, and the behaviour has not been verified against a live endpoint. The frame format was checked offline: frames encoded with the AWS event-stream codec were accepted by
@aws-sdk/client-bedrock-runtimeaschunkevents, and binary frames in that format, split at arbitrary byte boundaries and decoded with a streamingTextDecoderas the gateway does, yield the exact counts with this change. The unit tests use a simplified framing around the same{"bytes":...}payload.How to test?
npx tsx server/lib/aiUsageExtraction.test.ts: the newtestInvokeStreamAnthropicChunksandtestInvokeStreamInvocationMetricsfail onmain(extractor returnsnull) and pass with this change, together with the existing tests and a header-priority check.