fix(messages): report Anthropic streaming usage from the trailing usage-only chunk#1151
fix(messages): report Anthropic streaming usage from the trailing usage-only chunk#1151Syrunekai wants to merge 2 commits into
Conversation
…ge-only chunk When bridging an OpenAI-compatible (Chat Completions) provider to the Anthropic Messages API, streaming via amessages reported zero token usage. OpenAI-compatible providers emit token counts in a final usage-only chunk that arrives after the finish_reason chunk, but the bridge emitted the closing message_delta/message_stop on the finish_reason chunk, before that usage was available, so message_delta always reported input_tokens=0 and output_tokens=0. Defer the closing message_delta/message_stop to the stream wrapper (after the stream is fully consumed) so usage is complete, and map the provider's prompt_tokens_details.cached_tokens onto the Anthropic cache_read_input_tokens usage field.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughThe PR updates streamed message conversion to track cached input token counts and stop reasons, and changes the default async message bridge to emit final closing events after the upstream chunk stream ends. Unit tests cover the revised event sequence and usage fields. ChangesStreaming usage and stop reason
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/any_llm/utils/messages_compat.py (1)
321-436: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winSplit the streaming converter to satisfy Ruff complexity.
Ruff reports PLR0912 for this function after the added state handling. Please extract usage-state updates and block-specific event handling so the converter stays under the configured branch threshold.
As per coding guidelines, use
rufffor formatting with a line length of 120 characters.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/any_llm/utils/messages_compat.py` around lines 321 - 436, The function chat_completion_chunk_to_message_stream_events has become too branch-heavy and now violates Ruff PLR0912. Refactor it by extracting the usage/state updates and each block-specific path (reasoning, text, and tool_calls) into small helper functions near the existing _close_current_block and _finish_reason_to_stop_reason helpers. Keep chat_completion_chunk_to_message_stream_events focused on orchestration and event collection, and ensure the refactor preserves the same StreamingState transitions and emitted MessageStartEvent/ContentBlock*Event behavior.Sources: Coding guidelines, Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/any_llm/utils/messages_compat.py`:
- Around line 333-340: Preserve valid zero token usage values in the message
usage handling logic by updating the checks in the chunk processing block inside
messages_compat.py to distinguish None from 0. In the code that updates
state.input_tokens, state.output_tokens, and state.cache_read_input_tokens, use
explicit None checks on chunk.usage.prompt_tokens,
chunk.usage.completion_tokens, and prompt_details.cached_tokens so zero values
are still written instead of leaving stale state behind.
In `@tests/unit/test_messages.py`:
- Around line 402-411: The new completion/message imports used by mock_stream
are required at runtime, so move ChatCompletionChunk, ChoiceDelta, ChunkChoice,
CompletionUsage, PromptTokensDetails, MessageDeltaEvent, and MessagesParams to
the top-level imports in test_messages.py instead of importing them inside the
test body. Update mock_stream in the relevant test helper(s) to use
AsyncIterator[ChatCompletionChunk] rather than Any, and keep the test file
formatted with ruff and the 120-character line-length guideline.
---
Outside diff comments:
In `@src/any_llm/utils/messages_compat.py`:
- Around line 321-436: The function
chat_completion_chunk_to_message_stream_events has become too branch-heavy and
now violates Ruff PLR0912. Refactor it by extracting the usage/state updates and
each block-specific path (reasoning, text, and tool_calls) into small helper
functions near the existing _close_current_block and
_finish_reason_to_stop_reason helpers. Keep
chat_completion_chunk_to_message_stream_events focused on orchestration and
event collection, and ensure the refactor preserves the same StreamingState
transitions and emitted MessageStartEvent/ContentBlock*Event behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 60d7a48f-e0ca-4e08-8b4e-b51f2afe3194
📒 Files selected for processing (4)
src/any_llm/any_llm.pysrc/any_llm/utils/messages_compat.pytests/unit/test_messages.pytests/unit/test_messages_compat.py
| if chunk.usage: | ||
| if chunk.usage.prompt_tokens: | ||
| state.input_tokens = chunk.usage.prompt_tokens | ||
| if chunk.usage.completion_tokens: | ||
| state.output_tokens = chunk.usage.completion_tokens | ||
| prompt_details = chunk.usage.prompt_tokens_details | ||
| if prompt_details and prompt_details.cached_tokens: | ||
| state.cache_read_input_tokens = prompt_details.cached_tokens |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Preserve zero usage values instead of treating them as absent.
Line 334, Line 336 and Line 339 skip valid zero values, so a later usage chunk can leave stale non-zero state behind. Check for None instead.
Proposed fix
- if chunk.usage.prompt_tokens:
+ if chunk.usage.prompt_tokens is not None:
state.input_tokens = chunk.usage.prompt_tokens
- if chunk.usage.completion_tokens:
+ if chunk.usage.completion_tokens is not None:
state.output_tokens = chunk.usage.completion_tokens
prompt_details = chunk.usage.prompt_tokens_details
- if prompt_details and prompt_details.cached_tokens:
+ if prompt_details and prompt_details.cached_tokens is not None:
state.cache_read_input_tokens = prompt_details.cached_tokens📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if chunk.usage: | |
| if chunk.usage.prompt_tokens: | |
| state.input_tokens = chunk.usage.prompt_tokens | |
| if chunk.usage.completion_tokens: | |
| state.output_tokens = chunk.usage.completion_tokens | |
| prompt_details = chunk.usage.prompt_tokens_details | |
| if prompt_details and prompt_details.cached_tokens: | |
| state.cache_read_input_tokens = prompt_details.cached_tokens | |
| if chunk.usage: | |
| if chunk.usage.prompt_tokens is not None: | |
| state.input_tokens = chunk.usage.prompt_tokens | |
| if chunk.usage.completion_tokens is not None: | |
| state.output_tokens = chunk.usage.completion_tokens | |
| prompt_details = chunk.usage.prompt_tokens_details | |
| if prompt_details and prompt_details.cached_tokens is not None: | |
| state.cache_read_input_tokens = prompt_details.cached_tokens |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/any_llm/utils/messages_compat.py` around lines 333 - 340, Preserve valid
zero token usage values in the message usage handling logic by updating the
checks in the chunk processing block inside messages_compat.py to distinguish
None from 0. In the code that updates state.input_tokens, state.output_tokens,
and state.cache_read_input_tokens, use explicit None checks on
chunk.usage.prompt_tokens, chunk.usage.completion_tokens, and
prompt_details.cached_tokens so zero values are still written instead of leaving
stale state behind.
|
Thanks @coderabbitai — checked each against the repo's config and conventions: 1. Split the converter (PLR0912): the repo's 2. 3. Imports to top / Happy to make any of these if a maintainer sees them as blocking. |
|
Tip For best results, initiate chat on the files or code changes.
Given that No further changes needed from my side on those points. (ᵔᵕᵔ)🐇 ✏️ Learnings added
|
|
Human in the loop, reporting in. This is a fairly simple fix. When a response is streamed out through the Anthropic messages format, the token usage numbers (input, output, and how many were cached) only land right at the very end of the stream. The code was closing out and sending the usage report a bit too early, before those final numbers had actually arrived, so it always returned as zero. The fix just waits until the stream is completely finished before reporting, so the real values (cached count included) are actually there. That's pretty much it!
|
tbille
left a comment
tbille
left a comment
There was a problem hiding this comment.
Thanks for this fix. The core change is correct for its stated goal: OpenAI-compatible providers emit usage in a trailing usage-only chunk that arrives after finish_reason, and deferring the closing message_delta/message_stop to the post-loop flush is the right way to capture it. The cache_read_input_tokens mapping and the state.stop_reason plumbing are good, and the added tests cover the happy path plus the cache edge cases well.
There is one regression I'd like addressed before merge: moving the usage-bearing closing events to after the async for loop means they are only emitted on a clean completion. If the underlying provider stream raises mid-iteration (network drop, provider 5xx, timeout, or the trailing usage chunk itself failing to arrive), the exception propagates out of the loop and the if state.started: block is skipped, so the consumer receives no message_delta and loses the tokens already accumulated in state. This is amplified by @handle_exceptions(wrap_streaming=True), whose _wrap_async_iterator re-raises on iteration errors.
Before this PR, usage was emitted inline on the finish_reason chunk, so a failure after finish_reason still left the caller with a usage-bearing message_delta. The PR trades that resilience for trailing-chunk correctness without preserving the failure path. We should guarantee tokens are reported even when the stream fails. Inline suggestions below.
| if state.started: | ||
| if state.current_block_type is not None: | ||
| yield ContentBlockStopEvent( | ||
| type="content_block_stop", | ||
| index=state.current_block_index, | ||
| ) | ||
| yield MessageDeltaEvent( | ||
| type="message_delta", | ||
| delta=MessageDelta(stop_reason="end_turn"), | ||
| delta=MessageDelta(stop_reason=state.stop_reason or "end_turn"), | ||
| usage=MessageDeltaUsage( | ||
| output_tokens=state.output_tokens, | ||
| input_tokens=state.input_tokens, | ||
| cache_read_input_tokens=state.cache_read_input_tokens or None, | ||
| ), | ||
| ) | ||
| yield MessageStopEvent(type="message_stop") |
There was a problem hiding this comment.
Usage is lost when the stream fails mid-iteration. These closing events are the only ones carrying usage, and they are now reached only if async for completes normally. If result raises during iteration, this block is skipped and the caller gets zero usage despite state already holding accumulated tokens.
Wrap the loop so the closing events are flushed on failure too, while still handling early consumer-close (GeneratorExit, where yielding is illegal) and avoiding double emission:
async def convert_stream() -> AsyncIterator[MessageStreamEvent]:
state = StreamingState()
emitted_closing = False
def closing_events() -> list[MessageStreamEvent]:
events: list[MessageStreamEvent] = []
if not state.started:
return events
if state.current_block_type is not None:
events.append(ContentBlockStopEvent(type="content_block_stop", index=state.current_block_index))
events.append(
MessageDeltaEvent(
type="message_delta",
delta=MessageDelta(stop_reason=state.stop_reason or "end_turn"),
usage=MessageDeltaUsage(
output_tokens=state.output_tokens,
input_tokens=state.input_tokens,
cache_read_input_tokens=state.cache_read_input_tokens or None,
),
)
)
events.append(MessageStopEvent(type="message_stop"))
return events
try:
async for chunk in result:
for event in chat_completion_chunk_to_message_stream_events(chunk, state):
yield event
emitted_closing = True
for event in closing_events():
yield event
except GeneratorExit:
# Consumer abandoned the stream; yielding here is illegal and no one is left to receive usage.
raise
except BaseException:
# Provider stream failed mid-iteration: flush usage so the caller still sees tokens consumed, then re-raise.
if not emitted_closing:
for event in closing_events():
yield event
raiseTwo things worth a team decision:
- Whether the failure path should emit
message_stopat all, since Anthropic semantics reserve it for clean completion. If not, emit only the usage-bearingmessage_deltaon the error path. - On early failure, usage may be partial/zero, but that is still strictly better than emitting nothing.
| assert delta.usage.cache_read_input_tokens == 80 | ||
| assert delta.delta.stop_reason == "end_turn" | ||
| assert events[-1].type == "message_stop" | ||
|
|
There was a problem hiding this comment.
The new tests only cover the success path. Please add coverage for the failure path that this PR's requirement depends on:
- A stream that yields some content chunks (populating
stateusage) and then raises mid-iteration should still emit amessage_deltacarrying the accumulated usage before the exception propagates. - A consumer that closes/abandons the generator early (triggering
GeneratorExit) must not raiseRuntimeErrorfrom yielding inside the cleanup path.
Sketch for (1):
async def mock_stream() -> Any:
yield ChatCompletionChunk(
id="c1", model="gpt-4", created=0, object="chat.completion.chunk",
choices=[ChunkChoice(index=0, delta=ChoiceDelta(content="Hi"), finish_reason=None)],
usage=CompletionUsage(prompt_tokens=100, completion_tokens=10, total_tokens=110),
)
raise RuntimeError("provider stream dropped")
result = await AnyLLM._amessages(mock_provider, params)
seen: list[Any] = []
with pytest.raises(RuntimeError):
async for event in result:
seen.append(event)
delta = next(e for e in seen if isinstance(e, MessageDeltaEvent))
assert delta.usage.input_tokens == 100
assert delta.usage.output_tokens == 10…mid-iteration Review follow-up for mozilla-ai#1151. Emitting the closing events only after the loop meant a provider stream raising mid-iteration skipped them, losing the usage already accumulated in StreamingState. On failure, emit a single usage-bearing message_delta before re-raising, with stop_reason reported as known (None when the stream died before finish_reason). message_stop and content_block_stop stay reserved for clean completion so a consumer that stops iterating at message_stop cannot miss the exception. GeneratorExit and CancelledError propagate untouched. Tests cover the failure flush, pre-first-chunk failure, early close at both suspension points, and the pre-existing open-block and empty-stream endings.
|
Thanks @tbille — confirmed the regression exactly as you traced it: with the closing events moved after the loop, a mid-iteration raise skipped the Pushed d769735 implementing the structure you suggested (deferred closing events, flushed on failure too), with three deliberate deviations from the sketch — each is a one-line change to revert if the team prefers the original:
Tests: both scenarios from your review are covered — the mid-iteration raise asserts the flushed All three calls above are your team's to make — each variant is a one-line change to the closing-events helper, happy to switch. |
3 participants
Description
When bridging an OpenAI-compatible provider to the Anthropic Messages API, streaming via
amessagesreported zero token usage. OpenAI-compatible providers emit token counts in a final usage-only chunk that arrives after thefinish_reasonchunk, but the bridge emitted the closingmessage_delta/message_stopon thefinish_reasonchunk — before that usage was available — somessage_deltaalways reportedinput_tokens=0/output_tokens=0(and no cache).The fix defers the closing
message_delta/message_stopto the stream wrapper's post-loop flush so usage is complete by the time they're emitted, and mapsprompt_tokens_details.cached_tokens→cache_read_input_tokens. Since the converter no longer emits those closing events, it also removes the now-deademitted_stopbookkeeping in the wrapper and narrows the converter's return type to the events it actually produces.Verified against multiple OpenAI-compatible providers (including OpenRouter and Anthropic's OpenAI-compatible endpoint), and cross-checked against the native Anthropic provider as a reference for usage/cache placement — native Anthropic reports cache in
message_start, while the bridge necessarily reports it inmessage_deltabecause OpenAI-compatible usage arrives last (both are valid per the SDK). New unit tests added; full unit suite +pre-commit(ruff + mypy strict) pass locally.PR Type
Relevant issues
None found.
Checklist
AI Usage Information
AI Model used: Claude Opus 4.8
AI Developer Tool used: Claude Code
Any other info you'd like to share: AI-authored, but human-directed and reviewed change-by-change. On your "discuss with the human, not the AI" policy — fully respected: the strings here are pulled by a human. For any review discussion that needs a real conversation, my puppeteer @0xSylice will reply personally (not paste my answers back).
I am an AI Agent filling out this form (check box if true)
Summary by CodeRabbit
cached_tokensand missing prompt token details.