[Fix] OpenRouter Streaming Usage Cost - #16162
Conversation
|
@dhruvyad is attempting to deploy a commit to the CLERKIEAI Team on Vercel. A member of the Team first needs to authorize it. |
d1f460f to
96d2c46
Compare
|
Is this PR still active? I'd love to have this feature |
I'm using a forked version cause this PR didn't get reviewed. Happy to rebase and address comments if there's interest in merging this. |
|
Pretty please review this. |
I'm also waiting for streaming support. In the meantime, the only workaround I've found to solve this problem is to manually define the list of models and their costs in |
Without this, the provider-reported cost (e.g. from OpenRouter) was available on usage.cost but never reached litellm's cost calculator, which reads from _hidden_params["additional_headers"]. Also cleans up setattr usage in tests since Usage already has a cost field.
4e7457e to
387a94e
Compare
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes OpenRouter streaming cost tracking by handling the provider-specific pattern where a usage chunk (containing token counts and cost) arrives after the Key changes:
Remaining concerns:
Confidence Score: 4/5Safe to merge after addressing the calculate_total_usage cost-loss edge case; the core fix is correct and well-tested. One P1 finding: calculate_total_usage can lose a previously-seen cost if any later usage chunk carries cost=None, because latest_usage_chunk is overwritten unconditionally. This is not triggered by the current OpenRouter flow (cost chunk is always last), but is a latent bug. The two P2 findings are minor defensive/style issues. The core fix is logically correct and backed by new unit and async integration tests. litellm/litellm_core_utils/streaming_handler.py — specifically calculate_total_usage and _propagate_usage_cost_to_hidden_params
|
| Filename | Overview |
|---|---|
| litellm/litellm_core_utils/streaming_handler.py | Core fix: instead of raising StopIteration on post-finish usage chunks, return them so cost data is accumulated in self.chunks; adds _propagate_usage_cost_to_hidden_params helper; extends calculate_total_usage to carry cost from the latest usage chunk. |
| litellm/litellm_core_utils/streaming_chunk_builder_utils.py | Threads cost through UsagePerChunk and calculate_usage; adds hasattr-first usage access for ModelResponseStream objects; setattr cost before recreating Usage via model_dump(). |
| litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py | Adds cost: Optional[float] to UsagePerChunk TypedDict to carry provider-reported cost through the pipeline. |
| litellm/types/utils.py | Workaround: explicitly re-assigns self.usage after super().init() to prevent Pydantic from silently dropping the Usage object (including its cost field) during ModelResponseStream initialisation. |
| tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py | Adds test_cost_field_in_usage_chunks covering ChunkProcessor.calculate_usage cost extraction; remaining changes are formatting/style only. |
| tests/test_litellm/litellm_core_utils/test_streaming_handler.py | Adds three new tests: calculate_total_usage cost unit test, async end-to-end cost propagation test, and a lighter hidden-params propagation test; remaining changes are formatting only. |
Sequence Diagram
sequenceDiagram
participant OR as OpenRouter Stream
participant CSW as CustomStreamWrapper
participant CC as chunk_creator
participant RPC as return_processed_chunk_logic
participant SCB as stream_chunk_builder
participant CCC as ChunkProcessor.calculate_usage
OR->>CSW: chunk1 (content delta)
CSW->>CC: chunk1
CC->>RPC: model_response (no usage)
RPC-->>CC: return model_response
CC-->>CSW: model_response
CSW->>CSW: chunks.append(chunk1)
CSW-->>OR: yield chunk1
OR->>CSW: chunk2 (finish_reason=stop)
CSW->>CC: chunk2
CC->>RPC: model_response (finish_reason)
RPC->>RPC: sent_last_chunk = True
RPC-->>CC: return model_response
CC-->>CSW: model_response
CSW->>CSW: chunks.append(chunk2)
CSW-->>OR: yield chunk2
OR->>CSW: chunk3 (usage with cost=0.00025)
CSW->>CC: chunk3
CC->>CC: model_response.usage = chunk3.usage
CC->>RPC: model_response (sent_last_chunk=True, has usage)
Note over RPC: NEW: return instead of StopIteration
RPC-->>CC: return model_response
CC-->>CSW: model_response with usage+cost
CSW->>CSW: chunks.append(copy with usage+cost)
CSW->>CSW: strip usage, is_empty=True, continue
OR->>CSW: StopIteration
CSW->>SCB: stream_chunk_builder(chunks)
SCB->>CCC: calculate_usage(chunks)
CCC->>CCC: extract cost=0.00025 from chunk3
CCC-->>SCB: Usage(tokens, cost=0.00025)
SCB-->>CSW: complete_response
CSW->>CSW: _propagate_usage_cost_to_hidden_params
Note over CSW: additional_headers[llm_provider-x-litellm-response-cost]=0.00025
CSW->>CSW: _last_returned_hidden_params[usage] = final_usage
CSW-->>OR: raise StopIteration
Reviews (2): Last reviewed commit: "address review feedback: remove provider..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if self.custom_llm_provider != "openrouter": | ||
| raise StopIteration | ||
| else: | ||
| # OpenRouter: continue processing - usage will come in later chunks | ||
| pass |
There was a problem hiding this comment.
llms/ directory
The hardcoded "openrouter" check in the core streaming_handler.py violates the project's architecture rule to keep provider-specific logic inside litellm/llms/. This check adds conditional branching that will grow over time as more providers share this late-usage pattern.
The more general fix already in this PR — having return_processed_chunk_logic return the usage chunk instead of raising StopIteration — works for ALL providers by design. This provider guard should be removed; if the chunk carries no content and generic_chunk_has_all_required_fields already accepted it, silently passing is the correct behavior for any provider, not just OpenRouter.
| if self.custom_llm_provider != "openrouter": | |
| raise StopIteration | |
| else: | |
| # OpenRouter: continue processing - usage will come in later chunks | |
| pass | |
| if not _chunk_has_content and ( | |
| not isinstance(chunk, dict) | |
| or "provider_specific_fields" not in chunk | |
| ): | |
| pass # no content / special fields – continue to build the response_obj |
Rule Used: What: Avoid writing provider-specific code outside... (source)
| # Don't raise StopIteration here - some providers (like OpenRouter) | ||
| # send usage/cost data in chunks after the finish_reason chunk | ||
| if ( | ||
| hasattr(model_response, "usage") | ||
| and model_response.usage is not None | ||
| ): | ||
| self.chunks.append(model_response) | ||
| raise StopIteration | ||
| return model_response | ||
| return |
There was a problem hiding this comment.
self.chunks
return_processed_chunk_logic now appends the usage chunk to self.chunks (line 1033) and then returns it (line 1034). Control passes back to __next__, which appends the same response again at line 1884 (self.chunks.append(response)).
Old code raised StopIteration here, so __next__'s append was never reached. With the new return-path that append is still executed, leaving every post-finish usage chunk recorded twice in self.chunks.
Currently calculate_total_usage and stream_chunk_builder both use a "last wins" strategy for token counts, so the duplicated chunk doesn't produce wrong numbers today. But self.chunks grows unexpectedly, and any future caller that sums or counts chunks would double-count usage.
The simplest fix is to not append inside return_processed_chunk_logic and let __next__'s existing append be the single source of truth:
| # Don't raise StopIteration here - some providers (like OpenRouter) | |
| # send usage/cost data in chunks after the finish_reason chunk | |
| if ( | |
| hasattr(model_response, "usage") | |
| and model_response.usage is not None | |
| ): | |
| self.chunks.append(model_response) | |
| raise StopIteration | |
| return model_response | |
| return | |
| # Don't raise StopIteration here - some providers (like OpenRouter) | |
| # send usage/cost data in chunks after the finish_reason chunk | |
| if ( | |
| hasattr(model_response, "usage") | |
| and model_response.usage is not None | |
| ): | |
| return model_response | |
| return |
| if ( | ||
| latest_usage_chunk | ||
| and hasattr(latest_usage_chunk, "cost") | ||
| and latest_usage_chunk.cost is not None | ||
| ): | ||
| return latest_usage_chunk |
There was a problem hiding this comment.
When the most recent usage chunk carries a cost, the function returns that chunk object directly, bypassing the Usage(...) construction below where prompt_tokens and completion_tokens from the local accumulators are used. If a provider ever emits cost in a final chunk that has zero/missing token counts, the returned object will show incorrect totals. Consider merging the cost value into a freshly-constructed Usage built from the accumulated variables to be consistent with the non-cost path.
|
|
||
| response = self.model_response_creator() | ||
| if complete_streaming_response is not None: | ||
| # Propagate provider-reported cost (e.g. OpenRouter) | ||
| # to _hidden_params so the cost calculator picks it up | ||
| _final_usage = getattr(complete_streaming_response, "usage", None) | ||
| if ( | ||
| _final_usage is not None | ||
| and hasattr(_final_usage, "cost") | ||
| and _final_usage.cost is not None | ||
| ): | ||
| if ( | ||
| "additional_headers" | ||
| not in complete_streaming_response._hidden_params | ||
| ): | ||
| complete_streaming_response._hidden_params[ | ||
| "additional_headers" | ||
| ] = {} | ||
| complete_streaming_response._hidden_params[ | ||
| "additional_headers" | ||
| ]["llm_provider-x-litellm-response-cost"] = float( | ||
| _final_usage.cost | ||
| ) |
There was a problem hiding this comment.
The ~20-line block that reads the cost from _final_usage and writes it into the response hidden params is copy-pasted verbatim into both the sync except StopIteration handler (~line 1936) and the async except (StopAsyncIteration, StopIteration) handler (~line 2183). Any future change to the header name or logic must be applied in two places. Consider extracting it to a small private helper method.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| model="openrouter/claude", | ||
| choices=[ | ||
| StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) | ||
| ], | ||
| usage=Usage( | ||
| completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 | ||
| ), | ||
| ) | ||
|
|
||
| # Build the complete response as stream_chunk_builder does | ||
| complete_response = litellm.stream_chunk_builder( | ||
| chunks=[chunk1, chunk2, chunk3], | ||
| messages=[{"role": "user", "content": "test"}], | ||
| ) | ||
|
|
||
| assert complete_response is not None | ||
| assert hasattr(complete_response.usage, "cost") | ||
| assert complete_response.usage.cost == 0.00025 | ||
|
|
||
| # Simulate the propagation logic from streaming_handler | ||
| _final_usage = getattr(complete_response, "usage", None) | ||
| if ( | ||
| _final_usage is not None | ||
| and hasattr(_final_usage, "cost") | ||
| and _final_usage.cost is not None | ||
| ): | ||
| if "additional_headers" not in complete_response._hidden_params: | ||
| complete_response._hidden_params["additional_headers"] = {} | ||
| complete_response._hidden_params["additional_headers"][ | ||
| "llm_provider-x-litellm-response-cost" | ||
| ] = float(_final_usage.cost) | ||
|
|
||
| assert "additional_headers" in complete_response._hidden_params | ||
| assert ( | ||
| complete_response._hidden_params["additional_headers"][ | ||
| "llm_provider-x-litellm-response-cost" | ||
| ] | ||
| == 0.00025 | ||
| ) | ||
|
|
||
| # Verify the cost calculator would pick this up | ||
| from litellm.cost_calculator import get_response_cost_from_hidden_params | ||
|
|
||
| provider_cost = get_response_cost_from_hidden_params( | ||
| complete_response._hidden_params | ||
| ) |
There was a problem hiding this comment.
test_openrouter_streaming_cost_propagates_to_hidden_params manually copies the propagation block from streaming_handler.py into the test body and then asserts that the code it just executed works. This means the test will pass even if the production path is broken — it only verifies the test's own inline copy of the logic.
To make this a real regression detector, the test should drive the complete response through CustomStreamWrapper (as test_openrouter_streaming_cost_after_finish_reason does) and then inspect the hidden params of the wrapper's final state, rather than simulating the propagation inline.
Rule Used: # Code Review Rule: Mock Test Integrity
What:... (source)
…ppend, DRY cost propagation - Remove hardcoded openrouter check from core handler; the general return-instead-of-raise in return_processed_chunk_logic already handles post-finish-reason usage chunks for any provider. - Remove self.chunks.append in return_processed_chunk_logic to avoid double-appending (caller already appends). - Build proper Usage in calculate_total_usage instead of returning raw chunk early — preserves cost while ensuring token fields are always reconstructed. - Extract cost propagation into _propagate_usage_cost_to_hidden_params static method, used by both sync and async paths. - Test calls the real method instead of copy-pasting production logic.
| if ( | ||
| latest_usage_chunk | ||
| and hasattr(latest_usage_chunk, "cost") | ||
| and latest_usage_chunk.cost is not None | ||
| ): | ||
| returned_usage_chunk.cost = latest_usage_chunk.cost | ||
|
|
There was a problem hiding this comment.
calculate_total_usage loses cost if any later chunk has cost=None
latest_usage_chunk is overwritten on every chunk that carries non-None usage, not just on the chunk that carries cost. If a provider sends a final summary chunk after the cost-bearing chunk (e.g. a token-count-only reconciliation chunk with cost=None), latest_usage_chunk would end up pointing at the later chunk and the cost check at line 2412 would return False, silently discarding the cost.
ChunkProcessor.calculate_usage_per_chunk uses the safer "last non-None wins" pattern:
if usage_chunk_dict["cost"] is not None:
cost = usage_chunk_dict["cost"] # only overwrites when there is an actual valueA safer alternative that mirrors calculate_usage_per_chunk:
# Track cost separately so a later cost=None chunk can't erase an earlier cost
latest_cost: Optional[float] = None
for chunk in chunks:
if "usage" in chunk and chunk["usage"] is not None:
usage = chunk["usage"]
latest_usage_chunk = usage
if "prompt_tokens" in usage:
prompt_tokens = usage.get("prompt_tokens", 0) or 0
if "completion_tokens" in usage:
completion_tokens = usage.get("completion_tokens", 0) or 0
if hasattr(usage, "cost") and usage.cost is not None:
latest_cost = usage.cost
if latest_cost is not None:
returned_usage_chunk.cost = latest_costIn the current OpenRouter flow the cost chunk is always the last one, so this doesn't manifest today — but it's a latent bug for any provider that sends a trailing reconciliation chunk.
|
@paulchaum @arbv @mparsam @krrish-berri-2 @ishaan-berri @ishaan-jaff Rebased and tested this PR given the demand. Let me know if I can help with anything else. |
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. |
|
not stale |
|
Reviewed and tested this locally (checked out this branch, Python 3.9) — it correctly fixes the core of #11626 (which is still open, auto-closed as stale). Mechanism looks right: Tests pass locally: One gap: this threads # Preserve OpenRouter-style provider usage metadata (issue #11626).
if getattr(latest_usage_chunk, "cost_details", None) is not None:
returned_usage_chunk.cost_details = latest_usage_chunk.cost_details
if getattr(latest_usage_chunk, "is_byok", None) is not None:
returned_usage_chunk.is_byok = latest_usage_chunk.is_byokI have this plus a test passing locally (167 total) and am happy to send it as a follow-up once this merges, or you're welcome to fold it in here. Either way, solid fix — would be great to see it land. It resolves #11626. |
|
+ @mateo Wang ***@***.***> for review
…On Fri, Jul 3, 2026 at 8:16 PM blakeaa827 ***@***.***> wrote:
*blakeaa827* left a comment (BerriAI/litellm#16162)
<#16162 (comment)>
Reviewed and tested this locally (checked out this branch, Python 3.9) —
it correctly fixes the core of #11626
<#11626> (which is still open,
auto-closed as stale).
*Mechanism looks right:* _propagate_usage_cost_to_hidden_params writes
usage.cost into
_hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"]
— the same key the non-streaming path (OpenrouterConfig.transform_response)
uses, and which get_response_cost_from_hidden_params short-circuits on
before the static cost map. So streamed OpenRouter calls now report the
real provider cost instead of falling back to the map (which is $0 for
models that aren't in it).
*Tests pass locally:*
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
tests/test_litellm/llms/openrouter/
→ 166 passed
*One gap:* this threads cost but not cost_details / is_byok, and #11626
<#11626>'s title explicitly
names is_byok. It's a small addition right after returned_usage_chunk.cost
= latest_usage_chunk.cost in calculate_total_usage:
# Preserve OpenRouter-style provider usage metadata (issue #11626).
if getattr(latest_usage_chunk, "cost_details", None) is not None:
returned_usage_chunk.cost_details = latest_usage_chunk.cost_details
if getattr(latest_usage_chunk, "is_byok", None) is not None:
returned_usage_chunk.is_byok = latest_usage_chunk.is_byok
I have this plus a test passing locally (167 total) and am happy to send
it as a follow-up once this merges, or you're welcome to fold it in here.
Either way, solid fix — would be great to see it land. It resolves #11626
<#11626>.
—
Reply to this email directly, view it on GitHub
<#16162?email_source=notifications&email_token=CARFWGF6IPK5DDECQN76LJD5DBZH5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYGA2DMMJXHEY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4880461791>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/CARFWGEVJBFKGCGGTA6KEGT5DBZH5AVCNFSNUABFKJSXA33TNF2G64TZHM3DOMJSGY4TKMBVHNEXG43VMU5TGNJXHA2DCMRWGAY2C5QC>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
--
*Krrish Dholakia | *CEO
Book a meeting with me
<https://calendly.com/d/cx9p-5yf-2nm/litellm-introductions>
LinkedIn <https://www.linkedin.com/in/krish-d/> | (770) 878 - 3106
***@***.***
P.S. See how LiteLLM helps LLM Platform teams move fast and stay in control
<https://www.litellm.ai/#features>
|
|
I'll make a copy branch for review. I see a bunch of folks testing via unit tests. @dhruvyad do you mind uploading screenshots/video or an ordered list of command+command outputs ran showing it works e2e? |
Port of BerriAI#16162 by @dhruvyad onto litellm_internal_staging. OpenRouter sends a usage chunk (including a provider-reported cost field) after the finish_reason chunk. Previously the stream handler raised StopIteration on the first post-finish chunk, so that usage/cost never reached the assembled response and cost tracking fell back to token-based estimates. Carry usage.cost through chunk accumulation, preserve stripped usage in _hidden_params, and propagate the provider cost into _hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] so the cost calculator uses it.
7 participants
Title
Fix usage cost calculation for streaming models via OpenRouter.
Relevant issues
Extension of #13653
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unitType
🐛 Bug Fix
✅ Test
Changes
Fix usage cost tracking for streaming, which is currently broken for OpenRouter models.