fix(proxy): bill partial usage when a streaming request is cancelled

[Bytechoreographer]

Relevant issues

Follow-up to #30522. Raised by @jingyu-lin on that PR: when a streaming request is aborted mid-flight, the tokens already generated upstream are never written to SpendLogs.

Pre-Submission checklist

• I have added meaningful tests (mutation-verified)
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Type

🐛 Bug Fix

Changes

On a mid-stream client disconnect the stream never reaches normal completion. CancelledError / GeneratorExit are BaseException, so neither the success nor the failure logging branch in the streaming generator runs, and the assembled-response success logging that builds usage via stream_chunk_builder and writes SpendLogs never fires. The result is that the tokens already produced upstream are billed by the provider (Azure / OpenAI / Google) but never recorded on the proxy. Spend undercounts by roughly the abort rate, and the gap is invisible in SpendLogs; it only shows up when reconciling against provider invoices.

This builds on the existing shielded streaming cleanup (_finalize_streaming_generator_cleanup). When a client disconnect is recorded, it now assembles the partial usage from the chunks received so far (stream_chunk_builder over response.chunks, the same builder normal completion uses) and dispatches success logging for it. dispatch_success_handlers already de-dupes via has_dispatched_final_stream_success, so this is a no-op when normal completion already logged. It only runs on the cancellation path; the exception path (e.g. provider timeout) sets stream_completed and already emits a failure log, so there is no double-logging.

The billing runs after response.aclose(), so the upstream provider connection is released before the partial-usage logging (and any post-call guardrails or external success callbacks) runs; a disconnecting client cannot keep the upstream stream held open while those callbacks finish. aclose() only closes completion_stream and leaves response.chunks intact, so the partial usage is still assembled from the chunks already received.

This is the billing-side half of the same root cause that #30522 addresses on the budget-reservation side. With this in place, the cost callback also fires on cancellation, so the reservation reconciles to the real partial cost rather than an estimate.

Scope / what is not covered here

The exception/timeout path still emits a failure log without partial usage (a separate, smaller gap). This PR targets the cancellation path, which is the dominant abort case.

Tests

tests/test_litellm/proxy/test_common_request_processing.py: dispatches the assembled partial usage on disconnect; is a no-op when normal completion already logged (has_dispatched_final_stream_success); is a no-op when no chunks were received; routes through the deferred guardrail path when post-call guardrails are active; swallows builder and dispatch failures so they cannot escape the shielded cleanup; closes the upstream stream before the billing dispatch. Removing the dispatch fails the first test, and reverting the close-before-bill order fails the ordering test.

Screenshots / Proof of Fix

Live proxy run to be added (recommend validating SpendLogs against a real aborted stream).

[Bytechoreographer]

@jingyu-lin this implements the SpendLogs side you root-caused on #30522. On a mid-stream disconnect it assembles partial usage from the received chunks via stream_chunk_builder in the shielded cleanup and dispatches success logging, de-duped by has_dispatched_final_stream_success and gated to the cancellation path so it does not double-log with the exception/failure path. Since you have the production-validated version, your review would be very welcome, especially on edge cases I cannot exercise in unit tests (real SpendLogs row, cost accuracy, multi-retry aborts).

@Sameerlite this is the follow-up I mentioned on #30522, kept as a separate focused PR. #30522 releases the budget reservation on cancel; this one records the partial spend so the provider tokens stop being unbilled. With both in, the cost callback fires on cancel and the reservation reconciles to real partial cost. Happy to adjust scope or fold differently if you prefer.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR adds partial-usage billing on mid-stream client disconnect by extending _finalize_streaming_generator_cleanup to call the new _bill_partial_stream_on_disconnect helper. On cancellation, stream_chunk_builder assembles usage from whatever chunks have been received and dispatches success logging — routing through the deferred guardrail path when post-call guardrails are active, and falling back to direct dispatch_success_handlers otherwise.

• Core change (common_request_processing.py): adds _bill_partial_stream_on_disconnect as a shielded, best-effort billing dispatch that runs after response.aclose() to release the upstream connection before any logging callbacks start, with full error swallowing so failures cannot escape the CancelScope(shield=True) block.
• Tests (test_common_request_processing.py): seven new tests and one new TestStreamingClientDisconnectLogging method cover the happy path, the de-dup guard, empty-chunks short-circuit, builder/dispatch error swallowing, the guardrail path, and the aclose-before-bill ordering. The rest of the test changes are pure Black reformatting with no logic changes.

Confidence Score: 5/5

Safe to merge — the new billing path runs inside an existing shielded CancelScope, errors are fully swallowed, and the upstream connection is released before any logging callbacks run.

The implementation correctly gates partial billing on a confirmed client disconnect, assembles usage via the same stream_chunk_builder path used for normal completion, preserves the guardrail routing for guarded streams, and de-dupes through the existing has_dispatched_final_stream_success flag. All new code paths are covered by purpose-built unit tests, including ordering, error swallowing, and the deferred-guardrail branch. No existing test assertions were weakened.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/common_request_processing.py	Adds _bill_partial_stream_on_disconnect (76 lines) and wires it into _finalize_streaming_generator_cleanup; ordering, de-dup, guardrail routing, and error-swallowing are all handled correctly.
tests/test_litellm/proxy/test_common_request_processing.py	Seven new unit tests cover every branch of the new billing helper; all other changes are Black-driven line-length reformatting with no assertion logic altered.

_{Reviews (5): Last reviewed commit: "fix(proxy): close upstream stream before..." | Re-trigger Greptile}

[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: 2 · PR risk: 0/10

[Bytechoreographer]

Addressed all three in 602241e:

veria (post-call guardrails bypassed): the partial response now routes through _on_deferred_stream_complete when it is set, which runs _run_deferred_stream_guardrails (the same audit-only post-call guardrail path normal completion uses) before logging. It only dispatches success directly when no deferred guardrail closure is set. Added a test asserting a guarded stream goes through the deferred path and does not dispatch success directly; removing the routing fails it.

greptile (redundant dedup guard): kept as an optimization to skip the stream_chunk_builder assembly, with a comment clarifying dispatch_success_handlers is the authoritative de-dup via has_dispatched_final_stream_success.

greptile (start_time=None): added an inline comment noting None falls back to self.start_time (original request start) and end_time=None to now, which is the intended behavior for a partial record.

@greptileai

[Sameerlite]

Thanks for the PR! A couple of things to get this over the finish line:

• Could you add proof of the change working (screenshots, test output, or a sample request/response)? Even a quick curl before/after really speeds up the review.

Triggering Greptile for a code review in the meantime:

@greptileai

[Bytechoreographer]

Good catch. Fixed in 131cf4f: the dispatch (both the deferred-guardrail path and the direct dispatch_success_handlers path) is now wrapped in a try/except that logs and continues, so a logging or callback failure can no longer escape _bill_partial_stream_on_disconnect, exit the shielded scope, and skip response.aclose(). Added a regression test asserting a raising dispatch_success_handlers does not propagate out of the helper.

@greptileai

[Bytechoreographer]

Thanks @Sameerlite. Here is a before/after on this branch. The repro drives the exact streaming-cleanup path (async_streaming_data_generator -> _finalize_streaming_generator_cleanup -> _bill_partial_stream_on_disconnect) against a real streaming completion, consuming a few chunks and then simulating a client disconnect mid-stream:

resp = await litellm.acompletion(model=..., stream=True, max_tokens=4000, messages=[...])
i = 0
async for _ in resp:        # client reads a few chunks...
    i += 1
    if i >= 6:
        break               # ...then disconnects mid-stream

# exactly what the proxy now runs in the shielded cleanup on disconnect:
await ProxyBaseLLMRequestProcessing._bill_partial_stream_on_disconnect(
    resp, {"litellm_logging_obj": resp.logging_obj}
)

A CustomLogger capturing async_log_success_event shows:

BEFORE this fix (disconnect; cleanup never assembles/logs):          logged rows = 0
AFTER  this fix (disconnect -> _bill_partial_stream_on_disconnect):  logged rows = 1
    -> completion_tokens=26, prompt_tokens=18

So the tokens the provider already generated (and billed us for) now produce a SpendLogs entry on a cancelled stream, whereas before they were silently dropped (total_tokens=0, no row). On a full proxy with the model priced in the cost map the row also carries a non-zero spend; in this repro response_cost was None only because the upstream model name was not in the local cost map, so I am showing the usage rather than the dollar figure.

Equivalent curl-level repro for reference: start a streaming /v1/chat/completions with stream:true, kill the client after ~2s, then GET /spend/logs?api_key=...&summarize=false and the most recent row for the key is the aborted request with non-zero total_tokens and metadata.client_disconnected=true.

[Sameerlite]

Thanks for the contribution! A couple of things to get this ready:

• Merge conflicts — The branch has conflicts with the base branch. Could you rebase or merge from the base?

Once that's in, we'll take another look — appreciate the work on this! 🙏

[Bytechoreographer]

Addressed in a0d1dd8.

veria (upstream stream held open during disconnect billing): the cleanup now calls response.aclose() before _bill_partial_stream_on_disconnect, so the provider connection is released first and the partial-usage success logging (including any post-call guardrails and external success callbacks) runs off the path that holds the stream. aclose() only closes completion_stream and leaves response.chunks intact, so the partial billing still assembles usage from the chunks already received. Added a regression test asserting the cleanup awaits aclose before the billing dispatch; reverting the order fails it.

This also subsumes the earlier 131cf4f reasoning: the billing dispatch is still wrapped so a logging/callback failure cannot propagate out of the shielded cleanup, and now even a failure there cannot skip aclose since aclose already ran.

@greptileai

[greptile-apps]

The fix in a0d1dd8 looks correct and well-reasoned. The ordering is sound: aclose() releases the upstream connection first, then _bill_partial_stream_on_disconnect does the billing work. Since aclose() only closes completion_stream and doesn't touch response.chunks, the partial usage assembly is unaffected.

The key invariants hold:

1. The provider connection is released before any slow external callbacks run
2. response.chunks remains intact for stream_chunk_builder
3. The existing has_dispatched_final_stream_success de-dup still guards against double-logging
4. A failure in the billing dispatch cannot skip aclose since it already ran

The regression test asserting aclose is awaited before the billing dispatch is the right way to lock in that ordering. LGTM.

[Bytechoreographer]

@Sameerlite the merge conflict is resolved (rebased onto litellm_internal_staging), and the upstream-stream-held-open concern veria raised is fixed in a0d1dd8. Cleanup now closes the upstream response before the partial-usage billing runs, so post-call guardrails and async success callbacks can no longer hold the provider connection open on a disconnect. aclose() only closes completion_stream and leaves response.chunks intact, so the partial usage is still assembled from the chunks already received. Added a regression test that fails if the close-before-bill order is reverted. This should be ready for another look whenever you have a moment

While you're here, I have a few other PRs that have been sitting in the review queue for a while; would you mind taking a look when you get a chance:

• feat(proxy): read cold-storage prompts back in the logs detail view #30364 feat(proxy): read cold-storage prompts back in the logs detail view
• fix(ui): wrap selected models in model picker instead of single respo… #30279 fix(ui): wrap selected models in the model picker instead of one per row
• fix(ui): usage charts clip Y-axis labels at large token/request counts #27464 fix(ui): usage charts clip Y-axis labels at large token/request counts
• fix(ui): team members tab — O(1) lookup + client-side pagination & search #26667 fix(ui): team members tab, O(1) lookup plus client-side pagination and search
• fix(ui): prevent infinite re-render loop in VirtualKeysTable #26601 fix(ui): prevent infinite re-render loop in VirtualKeysTable
• Fix/credential list stale on add model #26600 fix(ui): credential list goes stale after adding a model

Thanks for taking the time

[WellMaxWang]

Tested this against a request routed through a Router model group with fallbacks configured, and found a gap this PR doesn't cover yet: _bill_partial_stream_on_disconnect no-ops in that case because response.chunks is always empty.

Root cause: Router._acompletion_streaming_iterator wraps the real stream in an inline FallbackStreamWrapper(CustomStreamWrapper) (in router.py) whose __anext__ is overridden to delegate straight to self._async_generator.__anext__():

async def __anext__(self):
    return await self._async_generator.__anext__()

The inherited CustomStreamWrapper.__anext__ — the implementation that does self.chunks.append(...) per chunk — never runs on this instance. self.chunks is still initialized to [] via super().__init__(), so getattr(response, "chunks", None) returns a real-but-permanently-empty list instead of None, and _bill_partial_stream_on_disconnect hits the early return silently.

The actual chunk history accumulates on model_response (the underlying CustomStreamWrapper), which is only a closed-over local inside stream_with_fallbacks() — there's no reference to it on the FallbackStreamWrapper instance for outside code to reach.

Repro: configured a model group with a fallback, started a streaming request against it, disconnected after ~6 chunks. Added a log right before the early-return:

chunks = getattr(response, "chunks", None)
verbose_proxy_logger.info(
    "billing_debug: response_type=%s chunks_len=%s",
    type(response).__name__,
    len(chunks) if chunks else chunks,
)

Output:

billing_debug: response_type=FallbackStreamWrapper chunks_len=[]

So any disconnect on a request that went through a fallback-enabled model group still drops the partial usage from SpendLogs — same symptom this PR fixes for the non-fallback path.

Possible fix: have FallbackStreamWrapper expose a reference to whichever stream actually holds the live chunk history (e.g. self._underlying_stream, updated if a fallback switch happens mid-stream), and have _bill_partial_stream_on_disconnect prefer that over response.chunks when present. Alternatively, have FallbackStreamWrapper.__anext__ also append to self.chunks alongside delegating to self._async_generator.

Happy to test against a patch if useful — not sure if this should block this PR or be tracked as a fast-follow.

[Bytechoreographer]

@WellMaxWang confirmed, and thanks for root-causing it. I checked against litellm/router.py:2181-2199 and you're right: a request through a fallback-enabled model group comes back as FallbackStreamWrapper, whose __anext__ just forwards to self._async_generator.__anext__(), so the inherited CustomStreamWrapper.__anext__ that does self.chunks.append(...) never runs on that instance. super().__init__() initializes self.chunks to [], so response.chunks is a real-but-permanently-empty list rather than None, and the if not chunks: return in _bill_partial_stream_on_disconnect no-ops silently. Same symptom this PR fixes on the non-fallback path. The real chunk history lives on the underlying model_response (the true CustomStreamWrapper), which is a closure local inside stream_with_fallbacks() with no handle exposed on the wrapper

One detail worth flagging before anyone reaches for the quick fix: appending to self.chunks inside FallbackStreamWrapper.__anext__ (your option 2) would undercount usage. CustomStreamWrapper.__anext__ stores a copy with usage into self.chunks (streaming_handler.py around line 2141), then strips usage off the chunk it actually yields, so the wrapper only ever sees the usage-stripped chunk; stream_chunk_builder over those would miss the usage entirely. The usage-bearing chunk only exists on model_response.chunks

Your option 1 (expose the underlying stream and read its chunks) is correct for the common case, but a mid-stream fallback switch (MidStreamFallbackError) accumulates the post-switch chunks on a different stream, so a single _underlying_stream reference misses the tail unless it's repointed when the switch happens

So this is a real gap, but the usage-correct fix lives inside the router fallback machinery (router.py), not in the billing helper, and it isn't a one-liner. Since this PR already fixes the dominant non-fallback path correctly, I'd track the fallback case as a fast-follow rather than fold it in here, which is how you framed it too. Clean shape: keep self._underlying_stream = model_response on FallbackStreamWrapper, repoint it when stream_with_fallbacks switches streams, and have _bill_partial_stream_on_disconnect fall back to it when response.chunks is empty. Happy to credit you on that follow-up since you found and root-caused it

@Sameerlite flagging this for your call: I'd keep it out of this PR and take the fallback case as a fast-follow, since this PR already fixes the dominant non-fallback path and the usage-correct fix lives in the router fallback machinery. Happy to do it either way.

[Sameerlite]

Thanks for the contribution! Kicking off a Greptile code review on this one.

@greptileai

[WellMaxWang]

Deployed the latest commit and confirmed: /v1/chat/completions now bills correctly on disconnect through a fallback-enabled model group. Thanks for the quick turnaround.

/v1/responses still has a gap though — looks like the same underlying pattern, just on the responses-API side. Got this on a streaming /v1/responses call routed through a fallback-enabled model group:

common_request_processing.py:1828 - 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/<id>/files calls will 403 for non-admin keys.

Checked router.py's FallbackResponsesStreamWrapper and it's the same shape as FallbackStreamWrapper: __init__ explicitly sets self.completed_response = None, and __anext__ just delegates to self._async_generator.__anext__(). Nothing on this instance ever populates completed_response with the real terminal stream event — that only happens on the real BaseResponsesAPIStreamingIterator instance closed over inside the router's fallback generator, same as model_response was for the chat-completions wrapper.

Downstream effect: _wrap_responses_stream_for_container_ownership (in common_request_processing.py) does _extract_completed_responses_response(original_stream_response) after the stream ends, which reads getattr(stream_response, "completed_response", None) — on a FallbackResponsesStreamWrapper that's always None, so container-ownership recording is skipped for any tool that creates a container (e.g. code_interpreter) on a fallback-routed /v1/responses stream. Practical impact: non-admin keys get a 403 on follow-up /v1/containers/<id>/files calls for containers created during a fallback stream.

Given the usage-accounting caveat you raised for the chat-completions wrapper (mid-stream provider switch moves the live state to a different underlying object), I'd guess the same caveat applies here for completed_response — worth confirming whether FallbackResponsesStreamWrapper needs the same "repoint on switch" handling you described for _underlying_stream, or whether completed_response is simpler since it's only set once at stream end rather than accumulated per-chunk.

Letting you decide whether this rides along with the fast-follow for the chat-completions fallback case or gets tracked separately — flagging it now since it's the same root pattern and you're already in this code. Happy to help repro/verify either way.