fix(proxy): bill partial streamed spend when the client disconnects mid-stream

[yassin-berriai]

Relevant issues

Follow-up to the security review finding on #32438 ("Client disconnects bypass sub-call billing")

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• 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)

Screenshots / Proof of Fix

All runs against a live proxy (python litellm/proxy/proxy_cli.py --config config.yaml --port 4245) backed by a real Postgres, with real gpt-5.5 calls costing real money. Key setup:

KEY=$(curl -s -X POST http://127.0.0.1:4245/key/generate \
  -H "Authorization: Bearer sk-litdisc-1234" -H "Content-Type: application/json" \
  -d '{"key_alias": "disconnect-verify"}' | jq -r .key)

Before (base, staging a7d01cb)

A streaming client that disconnects mid-stream pays nothing. The request streams real tokens, curl aborts partway, and no SpendLogs row is ever written:

curl -s -N --max-time 3 -X POST http://127.0.0.1:4245/v1/chat/completions \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model": "gpt-5.5", "messages": [{"role": "user", "content": "Write a 3000 word story about a lighthouse keeper. Do not stop early."}], "stream": true, "max_tokens": 8000}'
# curl: exit 28 (timed out mid-transfer)

 call_type | model | spend | prompt_tokens | completion_tokens | total_tokens
-----------+-------+-------+---------------+-------------------+--------------
(0 rows)

     key_alias     | spend
-------------------+-------
 disconnect-verify |     0

After (commit b6a8582)

Same curl, aborted after 261 streamed chunks (--max-time 10). The proxy bills exactly the partial delivered and attributes it to the key:

  call_type  |     model      |  spend   | prompt_tokens | completion_tokens | total_tokens
-------------+----------------+----------+---------------+-------------------+--------------
 acompletion | openai/gpt-5.5 | 0.008045 |            25 |               264 |          289

     key_alias     |  spend
-------------------+----------
 disconnect-verify | 0.008045

A second run aborted before the first content token (--max-time 3, gpt-5.5 still in TTFT) bills just the prompt side:

  call_type  |     model      |  spend   | prompt_tokens | completion_tokens | total_tokens
-------------+----------------+----------+---------------+-------------------+--------------
 acompletion | openai/gpt-5.5 | 0.000125 |            25 |                 0 |           25

A fully drained stream on the same proxy still bills exactly once (single row, no disconnect double-bill), pinned by the regression test as well

Type

🐛 Bug Fix

Changes

When a streaming client disconnects mid-stream, Starlette throws GeneratorExit/CancelledError into the proxy's streaming generator. Those are BaseExceptions, so neither the success nor the failure logging callback ever fires for the request; the code already acknowledges this (_release_max_parallel_requests_on_disconnect exists precisely to compensate for the leaked concurrency slot). The result is that every token streamed before the disconnect, plus any sub-call cost folded into the logging object (for example the RAG pipeline's vector search and rerank costs from #32438), never reaches SpendLogs, key spend, or budget enforcement. An authenticated caller can stream to 99% completion, disconnect before [DONE], and pay nothing

The fix finalizes streamed spend at disconnect time in _finalize_streaming_generator_cleanup, the shielded cleanup that already runs on every disconnect and records the 499 metadata. After the disconnect is recorded, the new _bill_partial_streamed_spend_on_disconnect assembles a partial response from the stream wrapper's collected chunks via the existing stream_chunk_builder and dispatches success logging for it through dispatch_success_handlers. Cost flows through the normal _response_cost_calculator, so additional_response_cost (sub-call cost) is included, and the standard proxy spend tracking attributes the row to the key with the already-stamped client_disconnected metadata. Double billing is impossible by construction: a natural end-of-stream schedules its dispatch before the generator can observe stream_completed, and dispatch_success_handlers dedups assembled-stream dispatches via has_dispatched_final_stream_success whichever side runs first. When post-call guardrails have armed deferred stream logging, that existing path fires instead and the new helper stands down

Live verification surfaced a second gap this fix depends on: the router wraps streamed responses in FallbackStreamWrapper, whose __anext__ bypasses the base CustomStreamWrapper iteration, so the wrapper's own chunks list stayed permanently empty and the disconnect path saw nothing to bill for router requests, which is every proxy request. The wrapper now aliases the inner stream's chunks list. Responses without a chunks attribute (raw passthrough generators) are skipped, and litellm.disable_streaming_logging is honored

This mirrors the intent of the existing _record_partial_usage_for_failure, which already recovers partial usage for streams that break with a provider error; client disconnect was the one abnormal termination with no billing path at all

Tests in tests/test_litellm/proxy/test_common_request_processing.py: a disconnect after two streamed chunks produces exactly one billing event carrying real usage plus a seeded additional_response_cost; a fully drained stream followed by a late disconnect cleanup still produces exactly one event; and a router-wrapped stream (the FallbackStreamWrapper case) bills on disconnect. The first and third fail on the unfixed code

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

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[greptile-apps]

Greptile Summary

This PR bills partial streamed spend when a client disconnects mid-stream by assembling the already-delivered chunks via stream_chunk_builder and dispatching a success logging event from the shielded cleanup block. It also coordinates the max_parallel_requests slot release so exactly one release fires regardless of which billing path wins.

• _bill_partial_streamed_spend_on_disconnect (common_request_processing.py): awaited directly in the shielded cleanup (not fire-and-forget), uses has_dispatched_final_stream_success to dedup against a natural end-of-stream dispatch, and returns a boolean so the caller knows whether the slot release is already owned by the success handler callback.
• Slot-release coordination (common_request_processing.py, proxy_server.py, utils.py): the eager create_task release in the except-block is removed; _arelease_max_parallel_requests_on_disconnect is now an async def awaited only when no success event fired, preventing the double-decrement race that was possible under the limiter's in-memory fallback.
• FallbackStreamWrapper chunks alias (router.py): three lines alias model_response.chunks onto the wrapper so the disconnect path sees accumulated chunks for every router-wrapped (i.e., every proxy) request, without which the billing path would always see an empty list and skip billing entirely.

Confidence Score: 5/5

Safe to merge — the fix is well-scoped, the critical paths have direct unit test coverage, and the live-proxy screenshots in the PR description confirm the end-to-end billing result.

The billing logic is awaited synchronously inside the existing shielded cleanup block so it cannot be GC'd before it completes, the dedup guard (has_dispatched_final_stream_success) prevents double-billing on a late disconnect after a fully-drained stream, and the slot-release now flows through exactly one owner in every reachable code path. The five new tests cover the disconnect billing case, the no-double-bill case, the router-wrapped case, and both directions of the slot-release coordination. No pre-existing mocks were weakened and no real network calls were added to the test file.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/common_request_processing.py	Adds _bill_partial_streamed_spend_on_disconnect and _deferred_stream_logging_is_armed; wires them into _finalize_streaming_generator_cleanup so partial spend is billed on disconnect and the max_parallel_requests slot is released exactly once regardless of which billing path fires.
litellm/proxy/proxy_server.py	Removes the eager _release_max_parallel_requests_on_disconnect call from the CancelledError/GeneratorExit except-block and instead passes user_api_key_dict/proxy_logging_obj to the shielded cleanup, letting it own the single slot release.
litellm/proxy/utils.py	Converts _release_max_parallel_requests_on_disconnect from fire-and-forget create_task to a proper async def (_arelease_max_parallel_requests_on_disconnect) that is awaited inside the shielded cleanup block, eliminating the old RuntimeError guard and the unrooted-task slot-leak risk.
litellm/router.py	Adds three lines in FallbackStreamWrapper.init to alias the inner stream's chunks list onto the wrapper, so the disconnect billing path can see accumulated chunks for every router-proxied request.
tests/test_litellm/proxy/test_common_request_processing.py	Adds five focused async tests covering disconnect billing (partial spend + additional_response_cost), no-double-bill on late disconnect after a completed stream, router FallbackStreamWrapper billing, no-double-slot-release when billing fires, and explicit release when nothing is billable.

_{Reviews (4): Last reviewed commit: "fix(proxy): use union syntax for disconn..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 81.39535% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/common_request_processing.py	78.37%	8 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yassin-berriai]

Note on CI: misc / Run tests also fails test_get_model_info_reports_realtime_mode (expects mode realtime, gets chat). That failure is inherited from the current staging tip; it reproduces locally on a clean checkout of the base, this PR touches nothing near model pricing, and the same test was already red on the PR that introduced it (#33728) before it merged. The test_acompletion_streaming_iterator failure from the first run was real and is fixed in 5d05119 (the chunks alias now guards for non-CustomStreamWrapper inner streams)

[yassin-berriai]

@greptileai please review the current head 759c6cb

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_stream_disconnect_billing (1c92dfd) with litellm_internal_staging (7015bd2)}

[yassin-berriai]

@greptileai please review the current head 047c9d2

[yassin-berriai]

@greptileai please review the current head 1c92dfd

[nuernber]

Thanks for fixing this -- just a quick note: does anyone know if AWS Bedrock charges for partial usage or the full usage in this scenario of interrupted/disconnected mid-streams? From my brief testing so far, I noticed that the AWS Bedrock invocation logs do appear to record a large amount of output tokens even though the interrupted stream may only output a few tokens. I have yet to inspect CUR logs though to confirm spend in AWS.

Is it possible that we add a flag in the config.yaml file where we can choose whether to let the stream bill on the full usage amount (i.e., drain the upstream on client disconnect mid-stream) instead of just the partial usage?

[WellMaxWang]

Related gap: streaming /v1/responses container ownership also breaks through FallbackResponsesStreamWrapper

While testing this fix, I hit an analogous issue on the Responses API path. Streaming /v1/responses through the proxy (i.e. through Router, so effectively every proxy deployment) logs:

LiteLLM Proxy:WARNING: common_request_processing.py:2073 - Container ownership recording skipped on streaming /v1/responses: no completed_response on stream iterator FallbackResponsesStreamWrapper. If this stream created any tool container (e.g. code_interpreter), follow-up /v1/containers//files calls will 403 for non-admin keys.

This looks like the same class of bug this PR fixed for chat completions — Router._aresponses_streaming_iterator wraps the real stream in FallbackResponsesStreamWrapper, which sets self.completed_response = None in init and never updates it, so _wrap_responses_stream_for_container_ownership → _extract_completed_responses_response always sees None and silently (well, now with a warning) skips recording.

This was previously reported as #30210 and supposedly fixed by #30213 (Propagate completed_response through FallbackResponsesStreamWrapper for streaming /v1/responses container ownership, released in v1.90.0). I'm seeing the exact same warning reproduce on the current build, so either:

It's a regression of #30213, or
#30213 only propagates completed_response on the natural end-of-stream path, and there's still a gap on early/disconnect termination (mirroring exactly the disconnect gap this PR just closed for billing) — in which case the fix here for chunks aliasing might need a parallel completed_response treatment for the Responses API side.

Given the security implication (non-admin keys 403'ing on their own container files created via streaming /v1/responses), could you clarify whether this is in scope for a follow-up here, or should I open a fresh issue referencing #30210?

Repro: streaming /v1/responses call through the proxy that implicitly creates a code_interpreter container, followed by a /v1/containers//files call with a non-admin key → 403, with the warning above in the proxy logs.

[nuernber]

I just tested this and it appears that it only works for /chat/completions endpoints, not /v1/messages, unless I'm testing incorrectly.