fix(aibridge): use latest streaming chat usage instead of cross-chunk sum - #27967
Conversation
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
nit: not sure if recorder should also be checked here. usage is already checked.
There was a problem hiding this comment.
Agreed, removed in dbe03ce. recordTokenUsage is unchanged by this PR and the check just re-fed the already-asserted lastUsage value into the recorder, so it added no coverage of the changed logic. Also dropped the mock recorder setup that existed only for it.
Posted by Mux on Mike's behalf.
…ertion recordTokenUsage is unchanged by this PR and the recorder check re-fed the already-asserted lastUsage value, adding no coverage of the last-wins tracking or the relayed payload.
2 participants
Problem
CODAGT-906: chats using OpenAI-compatible backends (e.g. poolside) through the AI Bridge persist token usage inflated 105x-640x, which falsely triggers automatic chat compaction on every turn.
The chat-completions streaming interceptor summed usage across every SSE chunk of one upstream stream and rewrote each relayed usage-bearing chunk with that running sum. Spec-compliant OpenAI emits usage once (final chunk with
stream_options.include_usage), so the sum equals the final value. vLLM-style backends emit cumulative usage snapshots on every chunk, so the relayed final usage becomes roughlyN_chunks x prompt_tokens(e.g. 417,012 persisted for a ~6,000-token context). chatd persists that value per assistant message and its compaction trigger reads it as context occupancy.Fix
Track the latest usage-bearing chunk's raw usage (last-wins) in the stream processor, updating only when a chunk actually carries usage so a trailing usage-less chunk cannot zero it.
marshalChunkrelays that value andrecordTokenUsagerecords the same value, unifying relayed and recorded usage. Last-wins is correct for both shapes: a single final usage chunk, and cumulative snapshots where each snapshot already includes all prior tokens.Per-iteration semantics are unchanged: each tool-loop iteration has its own processor, and the final iteration's usage is what the client sees.
Tests
TestStreamProcessorUsage(internal): cumulative snapshots with a trailing usage-less chunk, and the spec-compliant final-only shape; asserts relayed and recorded usage equal the last snapshot.streaming_cumulative_usage_injected_tool.txtarwith per-chunk cumulative usage plus an injected tool call; asserts client-visible final usage through the full interceptor.The blocking (non-streaming) path deliberately keeps its cross-iteration summation for external clients and is untouched. Remote dogfood UAT validated chat streaming, tool calls, plausible usage numbers, and zero spurious compactions.