fix(streaming): price partial-stream spend rows at the real model and keep prompt and cache fields - #37734
Conversation
… keep prompt and cache fields A streaming chat completion that ends early (client disconnect, or the proxy cutting the stream at LITELLM_MAX_STREAMING_DURATION_SECONDS) wrote a spend log row with spend 0.0, prompt_tokens 0 on the proxy-cut path, and no cache fields in usage_object. The proxy restamps chunk.model in place to the client-facing alias, so the partial response rebuilt from those chunks priced the unmapped alias and came out at 0. The failure path also rebuilt usage without the request messages, so prompt tokens counted to 0, and a cut stream never sees the final usage event that normally zero-fills the cache fields. Restamp the rebuilt partial response with the wrapper's real model before cost calculation on both the disconnect and the failure paths, pass the request messages when rebuilding usage on the failure path, and zero-fill missing cache usage fields the way completed streams already do.
Greptile SummaryThe PR improves accounting for interrupted streams by rebuilding prompt and cache usage and selecting a priceable deployment model.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| litellm/litellm_core_utils/streaming_handler.py | Rebuilds partial failure usage with request messages, deployment-model attribution, and cache-token metadata. |
| litellm/proxy/common_request_processing.py | Updates disconnect billing to select between the wrapper deployment model and a model recovered from routed chunks. |
| tests/test_litellm/litellm_core_utils/test_streaming_handler.py | Adds coverage for partial-stream pricing, prompt token reconstruction, and cache-field preservation. |
| tests/test_litellm/proxy/test_common_request_processing.py | Adds disconnect-billing tests for aliases, Azure Model Router attribution, and cache usage. |
Reviews (5): Last reviewed commit: "Match the client-name check to the name ..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
cache_read_input_tokens and cache_creation_input_tokens are pydantic extras on Usage, not declared fields, so filling them in created keys that were not there before rather than replacing a None. Readers that test for presence then took the new zero as authoritative: the spend log writer skipped its own copy from prompt_tokens_details, turning a real cache read of 500 into 0, and the prometheus provider cache counters stopped incrementing. Carry the prompt_tokens_details counts up before defaulting to zero, so a partial row reports the same cache numbers a complete one does. Renamed the helper to say what it now does.
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit de7dcbb. Configure here.
The disconnect billing path was stamping the wrapper's model over whatever stream_chunk_builder assembled. For Azure Model Router that throws away the routed model: the proxy deliberately leaves those chunks unrestamped so the builder can pick the real model off a later chunk, and overwriting it prices the row at the router alias instead. Only apply the wrapper's model when the builder did not find a model beyond the first chunk's, which is every case except Model Router.
|
bugbot run |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 03a253a. Configure here.
A chunk carrying usage is stored as a pre-restamp copy, so an alias-restamped stream reaches disconnect billing with its first chunk still on the deployment model and every later chunk on the client's name. That is the same shape Azure Model Router produces, and the previous guard read it as a routed model and left the alias on the row, which is the unpriced name this PR set out to stop. Compare the assembled model against the name the proxy stamps chunks with, so the alias goes back to the deployment's model and the routed model stays.
|
bugbot run |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Client-name check uses rewritten model
- Restricted _assembled_model_is_the_name_the_client_asked_for to key off the preserved _litellm_client_requested_model (falling back to request_data['model'] only when the preserved key is absent) so a rewritten deployment id no longer masquerades as the client-asked name and clobbers later-chunk router recovery.
Or push these changes by commenting:
@cursor push 20165c0e12
Preview (20165c0e12)
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -304,11 +304,16 @@
That stamp is what leaves an unpriced alias on the partial response, so the deployment's
own model has to go back on before the row is costed.
+
+ Prefer the preserved client-requested name over ``request_data["model"]`` because pre-call
+ alias and routing rewrites can replace the latter with the deployment target; matching that
+ rewritten id would misclassify an Azure Model Router style later-chunk recovery as an alias
+ restamp and stamp the wrapper's model over it.
"""
- return assembled_model in (
- request_data.get("_litellm_client_requested_model"),
- request_data.get("model"),
- )
+ client_requested: Final = request_data.get("_litellm_client_requested_model")
+ if isinstance(client_requested, str):
+ return assembled_model == client_requested
+ return assembled_model == request_data.get("model")
async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool:
diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py
--- a/tests/test_litellm/proxy/test_common_request_processing.py
+++ b/tests/test_litellm/proxy/test_common_request_processing.py
@@ -5593,6 +5593,32 @@
assert standard_logging_object["response_cost"] > 0.0
@pytest.mark.asyncio
+ async def test_disconnect_billing_router_recovery_survives_pre_call_model_rewrite(self):
+ """
+ Pre-call routing can rewrite ``request_data["model"]`` to the deployment id while
+ preserving the client alias in ``_litellm_client_requested_model``. The
+ client-asked-for gate must key off the preserved name so that an Azure Model
+ Router style later-chunk recovery is not misread as an alias restamp and
+ stomped by the wrapper's model.
+ """
+ def restamp_like_azure_model_router(response):
+ response.chunks[0].model = "azure-model-router"
+ for chunk in response.chunks[1:]:
+ chunk.model = "gpt-4.1-nano-2025-04-14"
+
+ event = await self._bill_and_collect_success_event(
+ restamp_like_azure_model_router,
+ request_data={
+ "model": "gpt-4.1-nano-2025-04-14",
+ "_litellm_client_requested_model": "azure-model-router",
+ },
+ )
+
+ assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14"
+ standard_logging_object = event["kwargs"]["standard_logging_object"]
+ assert standard_logging_object["response_cost"] > 0.0
+
+ @pytest.mark.asyncio
async def test_disconnect_billing_backfills_missing_cache_fields(self):
event = await self._bill_and_collect_success_event()You can send follow-ups to the cloud agent here.
Pre-call processing rewrites request_data["model"] for aliasing and routing, so matching either key let a routed model count as the client's own name and put the wrapper model back on an Azure Model Router row.
|
bugbot run |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 0801493. Configure here.
tin-berri
left a comment
tin-berri
left a comment
There was a problem hiding this comment.
Solid billing-correctness fix. The core subtlety — distinguishing "stream_chunk_builder recovered the real routed model from later chunks" (Azure Model Router) from "chunks got restamped to the client-facing alias" — is handled correctly via _assembled_model_came_from_a_later_chunk + _assembled_model_is_the_name_the_client_asked_for, and the trickiest edge case (pre-call processing already rewrote request_data["model"] to the routed name) is explicitly tested via the _litellm_client_requested_model fallback — without it the naive comparison against request_data["model"] would silently misattribute the routed model back to the alias. Cache-field backfill correctly carries OpenAI-style prompt_tokens_details.cached_tokens up to the Anthropic-style top-level keys, defaults to real zero only when genuinely absent, and is proven not to clobber values already recovered from chunks. Test coverage is thorough — both the helper-level cases and full disconnect-billing integration tests asserting real response_cost > 0 and correct model attribution. CI green. Approved.
2 participants
TLDR
Problem this solves:
"spend": 0.0despite billable tokensprompt_tokens: 0How it solves it:
User Flow
Before: a developer whose stream gets cut mid-response is billed nothing for it, and the spend row loses token detail too
"model": "bedrock-claude-opus-5"and"stream": true, and text chunks start arriving"spend": 0.0, and its usage object is missing the cache token fields every completed request reports at 0408error chunk instead, and that row is worse:"prompt_tokens": 0,"spend": 0.0, usage objectnullAfter: the same cut streams get priced at the model's real rates and keep their token detail
"model": "bedrock-claude-opus-5"and"stream": true408row now carries the real prompt token count and non-zero spend for the tokens streamed before the cutRelevant issues
Linear ticket
Resolves LIT-5840
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Live runs against Bedrock
us.anthropic.claude-opus-5, real spend, no mocks. Shared setup for every case below:config.yamlEvery spend row below is read back the way a user would, through
GET /spend/logs?request_id=<id>with the master key, piped tojqfor the fields in question. Rates for pricing checks: input $5.5e-06, output $2.75e-05 per token.Before (9432f40)
Case A: client hangs up mid-stream on /v1/chat/completions
{ "request_id": "chatcmpl-17f0bb81-d88d-4bb3-9a19-935db3d12e9c", "status": "success", "prompt_tokens": 25, "completion_tokens": 600, "spend": 0.0166375, "model": "bedrock/us.anthropic.claude-opus-5", "usage_object": { "total_tokens": 625, "prompt_tokens": 25, "completion_tokens": 600, "prompt_tokens_details": {"text_tokens": 25, "cached_tokens": 0, "cache_write_tokens": 0, "cache_creation_tokens": 0}, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0 } }{ "request_id": "chatcmpl-736193e3-a7d9-4d92-8a18-6f5757cf3b13", "status": "success", "prompt_tokens": 20, "completion_tokens": 296, "spend": 0.0, "model": "bedrock/us.anthropic.claude-opus-5", "usage_object": { "total_tokens": 316, "prompt_tokens": 20, "completion_tokens": 296, "prompt_tokens_details": null, "completion_tokens_details": null }, "error_information": {"error_code": "499", "error_class": "ClientDisconnected", "error_message": "Client disconnected the request"} }Case B: proxy cuts the stream at its duration cap on /v1/chat/completions
LITELLM_MAX_STREAMING_DURATION_SECONDS=6and send a long stream:{ "request_id": "c50380a4-d092-4625-8bfc-46281a8d1fd3", "status": "failure", "prompt_tokens": 0, "completion_tokens": 219, "spend": 0.0, "model": "bedrock/us.anthropic.claude-opus-5", "usage_object": null, "error_information": {"error_code": "408", "error_class": "Timeout", "llm_provider": "bedrock", "error_message": "litellm.Timeout: Stream exceeded max streaming duration of 6.0s (elapsed 6.3s)"} }Case C: proxy cuts the stream at its duration cap on /v1/responses
{ "request_id": "9af3d402-257c-4c88-9782-f3c79129de21", "status": "failure", "prompt_tokens": 0, "completion_tokens": 81, "spend": 0.0022275, "model": "bedrock-claude-opus-5", "model_group": "", "usage_object": null, "error_information": {"error_code": "408", "error_class": "Timeout", "llm_provider": "bedrock"} }After (0801493)
Case A: client hangs up mid-stream on /v1/chat/completions
25 * 5.5e-06 + 600 * 2.75e-05 = 0.0166375:{ "request_id": "chatcmpl-72ff40bb-b926-4774-8bd3-07c06e37ab83", "status": "success", "prompt_tokens": 25, "completion_tokens": 600, "spend": 0.0166375, "model": "bedrock/us.anthropic.claude-opus-5", "usage_object": { "total_tokens": 625, "prompt_tokens": 25, "completion_tokens": 600, "prompt_tokens_details": {"text_tokens": 25, "cached_tokens": 0, "cache_write_tokens": 0, "cache_creation_tokens": 0}, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0 }, "error_information": null }20 * 5.5e-06 + 316 * 2.75e-05 = 0.0088, and carries the same cache fields at 0:{ "request_id": "chatcmpl-80db78a2-cc2c-4048-bc54-e4906040d659", "status": "success", "prompt_tokens": 20, "completion_tokens": 316, "spend": 0.0088, "model": "bedrock/us.anthropic.claude-opus-5", "usage_object": { "total_tokens": 336, "prompt_tokens": 20, "completion_tokens": 316, "prompt_tokens_details": {"cached_tokens": 0}, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0 }, "error_information": {"error_code": "499", "error_class": "ClientDisconnected", "error_message": "Client disconnected the request"} }Case B: proxy cuts the stream at its duration cap on /v1/chat/completions
20 * 5.5e-06 + 232 * 2.75e-05 = 0.00649:{ "request_id": "79786692-c5dd-4d33-a31d-1bc367faae2e", "status": "failure", "prompt_tokens": 20, "completion_tokens": 232, "spend": 0.00649, "model": "bedrock/us.anthropic.claude-opus-5", "model_group": "bedrock-claude-opus-5", "usage_object": null, "error_information": {"error_code": "408", "error_class": "Timeout", "llm_provider": "bedrock", "error_message": "litellm.Timeout: Stream exceeded max streaming duration of 6.0s (elapsed 6.1s)"} }Case C: proxy cuts the stream at its duration cap on /v1/responses
20 * 5.5e-06 + 116 * 2.75e-05 = 0.0033:{ "request_id": "1f44b81d-c4d4-4a0a-82be-24ab1b52dcc6", "status": "failure", "prompt_tokens": 20, "completion_tokens": 116, "spend": 0.0033, "model": "bedrock-claude-opus-5", "model_group": "", "usage_object": null, "error_information": {"error_code": "408", "error_class": "Timeout", "llm_provider": "bedrock", "error_message": "litellm.Timeout: Stream exceeded max streaming duration of 6.0s (elapsed 6.0s)"} }Notes from the runs
Type
🐛 Bug Fix
Caveats (if any)
cost_calculator.pyhas a branch that reads the presence ofcache_read_input_tokensas "this usage object is Anthropic-shaped" and folds the cache counts back into prompt tokens. It only runs when the caller passescustom_cost_per_tokenorcustom_cost_per_second, and the partial-stream path passes neither, so it cannot fire here. I left it alone rather than change shared cost code every caller goes throughstreaming_handler.pyhas no access to the client's model name, so it always prices at the deployment's model. An Azure Model Router stream cut by the duration cap is billed at the router deployment rather than the routed model, which is what a completed stream on that route logs todaymessagesinto the partial rebuild changes the prompt token count on every interrupted async stream, SDK callers included, not only proxy rows.token_counterraises onmessages=NoneandChunkProcessor.calculate_usageturns that intoprompt_tokens = 0, so the old number was exactly zero whenever no chunk reported prompt tokens. It is a tiktoken estimate now, which can differ from the provider's own count for non-OpenAI tokenizers or when tools and system content sit outsidemessagesFinal Attestation
Note
Medium Risk
Touches spend tracking for interrupted streams (disconnect and mid-flight failure). Wrong model selection could misprice rows, but the change is isolated to partial-billing paths and is covered by unit tests.
Overview
Interrupted streams (client disconnect or mid-flight failure) now log non-zero spend at the deployment model, not an unpriced public alias, and keep prompt/cache token fields that complete rows already report.
Pricing.
_record_partial_usage_for_failurerestamps the rebuilt response withself.modelbefore cost calc, and passes requestmessagesso prompt tokens are counted. Disconnect billing in_bill_partial_streamed_spend_on_disconnectprefers the wrapper’s real model unless Azure Model Router recovered a later-chunk routed name (distinguished from alias restamping via_litellm_client_requested_model).Cache fields. New
backfill_missing_cache_usage_fieldscopies OpenAI-styleprompt_tokens_detailscounts onto Anthropic-stylecache_read_input_tokens/cache_creation_input_tokens(or zeros them) so partial rows match complete ones without overwriting real cache counts.Reviewed by Cursor Bugbot for commit 0801493. Bugbot is set up for automated code reviews on this repo. Configure here.