fix(streaming): recover partial usage on sync mid-stream failure (#14457)

[nishant-uxs]

Relevant issues

Fixes #14457

Design note

Root cause

Sync CustomStreamWrapper.__next__ mid-stream failures called failure_handler without first recovering usage from chunks already delivered. The async path already did this via _log_stream_failure_and_raise → _record_partial_usage_for_failure. Separately, recovery called stream_chunk_builder(chunks=...) without messages / logging_obj, so when a provider only emits usage on the final chunk, mid-stream disconnect left prompt tokens unestimated.

Current data flow

1. Content chunks accumulate on self.chunks.
2. Successful end-of-stream: stream_chunk_builder(chunks, messages, logging_obj) → usage → success logging / spend.
3. Async mid-stream failure: stash combined_usage_object + response_cost, then failure handlers.
4. Sync mid-stream failure (before this PR): failure_handler zeros response_cost because combined_usage_object was never set.
5. Proxy disconnect billing (_bill_partial_streamed_spend_on_disconnect) is a separate path already covered by fix(proxy): bill partial streamed spend when the client disconnects mid-stream #33736 — out of scope here.

Where usage is lost

• Sync client disconnect / provider error before the final usage chunk.
• Estimation gap when delivered chunks have no usage and messages were not passed into the builder.

Why existing accounting fails

_failure_handler_helper_fn only preserves spend when combined_usage_object is present; otherwise it sets response_cost = 0. Sync never populated that stash.

Idempotency risks

Recovery overwrites (does not accumulate) combined_usage_object / response_cost. A later router-fallback success log on the same request id overwrites again. Covered by a regression that double-calls recovery.

Proposed minimal fix

1. Call _record_partial_usage_for_failure() in sync __next__ before failure_handler (mirror async).
2. Pass messages=self.messages and logging_obj=self.logging_obj into stream_chunk_builder inside recovery (same contract as end-of-stream assembly).
3. No public API changes. No unrelated streaming behavior changes.

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Type

🐛 Bug Fix
✅ Test

Changes

• Sync stream failure path now recovers partial usage/cost before failure_handler.
• Partial-usage recovery estimates tokens from request messages when the final provider usage chunk never arrives.
• Regression coverage for completed-stream usage, disconnect-before-usage, estimate-without-usage-chunk, no double-count, and failure_handler cost preservation.

Local verification

• ruff check litellm/litellm_core_utils/streaming_handler.py — pass
• pytest tests/test_litellm/litellm_core_utils/test_streaming_handler.py — 103 passed
• Focused regressions (partial usage / sync failure / idempotency / streaming_handler_with_usage) — 8 passed

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Test plan

• Normal completed stream still assembles usage (test_streaming_handler_with_usage)
• Sync mid-stream failure stashes partial usage before failure_handler
• Provider with usage only on final chunk → estimate from messages/content
• Re-running recovery does not double-count tokens/cost
• failure_handler preserves stashed response_cost
• No-chunks path remains a no-op

[greptile-apps]

Greptile Summary

This PR aligns synchronous stream-failure accounting with the asynchronous path.

• Recovers partial usage and cost before invoking synchronous failure handlers.
• Supplies request messages and logging context so missing final usage chunks can be estimated.
• Adds regression coverage for partial recovery, estimation, idempotency, and cost preservation.

Confidence Score: 4/5

The partial-usage recovery failure path must preserve raw-chunk usage before this PR is safe to merge.

An exception from partial stream assembly is suppressed without storing usage, after which failure logging resets the response cost to zero and loses billable partial spend.

Files Needing Attention: litellm/litellm_core_utils/streaming_handler.py

Important Files Changed

Filename	Overview
litellm/litellm_core_utils/streaming_handler.py	Adds synchronous partial-usage recovery, but an assembly exception still leaves failure logging without recovered usage and zeros the partial spend.
tests/test_litellm/litellm_core_utils/test_streaming_handler.py	Adds focused regression tests for missing provider usage, synchronous failures, cost preservation, and idempotent recovery.

_{Reviews (1): Last reviewed commit: "fix(streaming): recover partial usage on..." | Re-trigger Greptile}

[greptile-apps]

Recovery failure still zeros spend

When stream_chunk_builder raises while recovering an interrupted sync stream, the surrounding catch suppresses the error without stashing combined_usage_object; failure_handler then resets response_cost to zero, so the partial spend this path is intended to preserve is still lost. The completed-stream path already falls back to calculate_total_usage(chunks=self.chunks) for this condition.

Knowledge Base Used: Cost Tracking and Budget Enforcement

[nishant-uxs]

Good catch — addressed in the follow-up commit. _record_partial_usage_for_failure now mirrors the completed-stream path: if stream_chunk_builder raises, we fall back to calculate_total_usage(chunks=self.chunks) via model_response_creator before stashing combined_usage_object / response_cost. Added a regression that mocks the builder raise and asserts usage is still recovered.

[codecov]

Codecov Report

❌ Patch coverage is 90.90909% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/litellm_core_utils/streaming_handler.py	90.90%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing nishant-uxs:fix/streaming-partial-usage-on-disconnect-14457 (4dc7bfe) with litellm_internal_staging (c9292d3)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (6ba744b) during the generation of this report, so c9292d3 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[nuernber]

Does this fix it for /v1/messages as well? I posted at #33736 (comment) that a similar solution fixed it for /chat/completions for me, but not /v1/messages. I don't see any test cases for /v1/messages Anthropic endpoints in that PR nor this PR.

[nishant-uxs]

@nuernber Good question — short answer: this PR does not fix /v1/messages disconnect billing.

What this PR covers

#35349 is scoped to CustomStreamWrapper (sync mid-stream failure / recovery): stash partial usage before failure_handler, and estimate when the final provider usage chunk never arrives. That is the OpenAI-shaped streaming wrapper used by /chat/completions (and similar CSW consumers).

Why /v1/messages is different

Proxy route_type == anthropic_messages streams through async_sse_data_generator, not by iterating a billed CustomStreamWrapper:

• Native Anthropic: AnthropicMessagesStreamingResponse / pass-through SSE (PassThroughStreamingHandler) — no response.chunks.
• Bridge (non-Anthropic via /v1/messages): CSW may exist inside acompletion, but the object the proxy holds is an Anthropic SSE wrapper/generator, not CSW.

_bill_partial_streamed_spend_on_disconnect (#33736) is invoked from the shared streaming cleanup, but it only bills when getattr(response, chunks, None) is a non-empty list — a CSW field. For /v1/messages that check typically returns False immediately, which matches your observation that /chat/completions improved while /v1/messages did not.

Separate follow-up

A /v1/messages fix needs Anthropic-SSE / bridge-aware partial recovery (or exposing billable partial state on the object the proxy closes), not an extension of this CSW-only change. Happy to file that as a follow-up issue/PR if maintainers want it tracked separately — I deliberately kept this PR minimal so it stays reviewable for #14457.