fix(proxy): release budget reservation when a request is cancelled mid-flight

[Bytechoreographer]

Relevant issues

Intermittent 429 Budget has been exceeded! Key=... Current cost: 805.8, Max budget: 800.0 on keys whose real spend is far below budget (~$280 of an $800 budget), self-healing after a pause and recurring under load.

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

The pre-call budget reservation (reserve_budget_for_request) increments the cross-pod spend counter by a request's worst-case cost up front, then gives it back by reconciling to the actual cost on success (cost callback) or on error (async_post_call_failure_hook). Neither runs when a request is cancelled mid-flight: a client disconnect or timeout surfaces as asyncio.CancelledError (a BaseException, so except Exception misses it) or GeneratorExit in the streaming generator, so the reservation leaks. Under an agent retry storm the leaked reservations accumulate, pin the counter above the key's real spend, and 429 subsequent requests; the counter only recovers when its TTL lapses, which is why it is intermittent and self-healing.

This releases the reservation on the cancellation paths, mirroring the success/error paths. release_budget_reservation_on_cancel wraps release_budget_reservation in asyncio.shield so it completes despite the surrounding cancellation, and short-circuits on the reservation's finalized flag so it is a no-op once success/error handling already reconciled. It is wired into async_streaming_data_generator (which async_sse_data_generator delegates to, covering chat, Anthropic /v1/messages, and Google streaming) and into the non-streaming base_process_llm_request.

A blanket finally release does not work for streaming, where the handler returns the StreamingResponse before the stream is consumed and the success reconcile happens at stream end; a finally would release early and suppress the real reconcile. Catching only the cancellation signals targets the leak without touching the success or error paths. The reservation aggressiveness (worst-case max_tokens) is a separate policy question and is intentionally left unchanged.

Tests in tests/test_litellm/proxy/test_budget_reservation.py: cancel-release gives the counter back and is idempotent; a finalized reservation is left untouched; and driving async_streaming_data_generator to cancel mid-stream hands the reservation back (neutering the fix flips this test red).

Screenshots / Proof of Fix

Live proxy run to be added.

[greptile-apps]

Greptile Summary

This PR fixes intermittent spurious 429s caused by pre-call budget reservations leaking when a request is cancelled mid-flight (client disconnect or timeout), because asyncio.CancelledError and GeneratorExit are BaseException subclasses that bypass the existing except Exception success/failure handlers. The fix introduces release_budget_reservation_on_cancel, wired into both the streaming and non-streaming request paths, using asyncio.shield so the release completes even when the surrounding task is being torn down.

• Core fix: release_budget_reservation_on_cancel wraps release_budget_reservation in asyncio.shield and is called from except (asyncio.CancelledError, GeneratorExit) in async_streaming_data_generator and from except asyncio.CancelledError in base_process_llm_request; a finalized guard makes it idempotent.
• Ancillary fixes: llm_router is now threaded through get_model_from_request so router-managed model IDs (video/character) resolve correctly for budget checks; the pass-through route registry lookup is corrected to strip rather than prepend the server root path; and sanitize_vertex_anthropic_output_params becomes model-aware to avoid 400s on Haiku 4.5.
• Key management: Adds a is_ui_session_team_key exemption from the GHSA-q775-qw9r-2r4g budget ceiling and an _org_inherited_from_team exemption from org-assignment validation, both of which affect the security-sensitive key generation path.

Confidence Score: 3/5

The budget-reservation cancel fix is directionally correct, but the test suite verifies a synchronous proxy for cancellation rather than actual asyncio task cancellation, leaving the most important production scenario untested. The key management changes also touch the GHSA-q775-qw9r-2r4g security boundary.

The streaming and non-streaming cancel paths look sound, but the three new cancel-release tests all raise CancelledError from inside the iterator rather than from the outer task scheduler. Under real task cancellation asyncio.shield returns immediately while the inner task runs in the background, so finalized is not set before the function returns — a scenario the tests never exercise. Additionally, key_management_endpoints.py introduces two new exemptions near the GHSA-q775-qw9r-2r4g delegation ceiling: the is_ui_session_team_key exemption fires on any non-None team_id before downstream team validation confirms the team is real and the caller is a member.

tests/test_litellm/proxy/test_budget_reservation.py (cancel-release tests do not simulate actual outer-task cancellation) and litellm/proxy/management_endpoints/key_management_endpoints.py (session-token budget ceiling exemption guard).

Important Files Changed

Filename	Overview
litellm/proxy/spend_tracking/budget_reservation.py	Adds release_budget_reservation_on_cancel using asyncio.shield to release budget reservations on cancellation; also passes llm_router to get_model_from_request. Design is sound but the asyncio.shield creates a background task under real task cancellation, leaving finalized unset until the background task completes.
litellm/proxy/common_request_processing.py	Wires release_budget_reservation_on_cancel into both the non-streaming (base_process_llm_request) and streaming (async_streaming_data_generator) cancellation paths; streaming path also catches GeneratorExit alongside CancelledError.
tests/test_litellm/proxy/test_budget_reservation.py	Adds three tests for the cancel-release path; all three raise CancelledError from within the iterator rather than from outer-task cancellation, so asyncio.shield always completes synchronously in the tests — the production scenario where the shield fires asynchronously is not exercised.
litellm/proxy/management_endpoints/key_management_endpoints.py	Adds is_ui_session_team_key exemption from the GHSA-q775-qw9r-2r4g budget ceiling and _org_inherited_from_team exemption from org-assignment validation; exemption guard uses only _requested_team_id is not None before downstream team validation runs.
litellm/proxy/auth/auth_utils.py	Threads llm_router through get_model_from_request and _extract_models_from_managed_resource_id to resolve router-managed model IDs to their canonical names; new _resolve_model_id_with_router helper is well-guarded with try/except and fallback.
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py	Replaces _build_full_path_with_root with _route_for_registry_lookup which strips the root from the incoming route, fixing registry lookups when a non-root server root path is configured; also sets logging_obj.stream = True for streaming pass-through responses.
litellm/proxy/spend_tracking/spend_management_endpoints.py	Adds _resolve_request_response_payload to transparently serve prompt/response content from cold storage when the DB row holds only placeholders; correctly falls back to PG content when cold storage returns nothing.
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py	Makes sanitize_vertex_anthropic_output_params model-aware: output_config.effort is dropped for models that don't advertise support and forwarded for those that do, preventing upstream 400s.
litellm/proxy/spend_tracking/cold_storage_handler.py	Refactors ColdStorageHandler to support constructor-injected logger for testing; no behavior change in production.

Comments Outside Diff (1)

1. litellm/proxy/management_endpoints/key_management_endpoints.py, line 734-753 (link)

   Session-token budget-ceiling exemption relies entirely on downstream team validation

   is_ui_session_team_key evaluates to True whenever the calling token is a UI session token (team_id == UI_SESSION_TOKEN_TEAM_ID) and the request includes any non-None team_id. At this point the supplied team_id has not been validated (team existence and caller membership are checked later in _common_key_generation_helper). This means the GHSA-q775-qw9r-2r4g budget ceiling is bypassed for any key creation that includes a team_id field, before the system confirms the team is real or the caller belongs to it.

   If downstream team validation is strict and always rejects unknown/inaccessible teams, the exemption is safe. But if there is any path that allows a key to be stored with an unvalidated team_id (e.g., a team that has no max_budget), a session token holder could create a key with an unbounded budget by providing a bogus or unlimited-budget team_id. A guard like team_table is not None (which is only set after a successful DB lookup) would make the exemption self-evidently safe without relying on ordering.

_{Reviews (1): Last reviewed commit: "fix(proxy): release budget reservation w..." | Re-trigger Greptile}

[greptile-apps]

asyncio.shield fires asynchronously in the production scenario, but tests verify only the synchronous path

When the outer asyncio Task is genuinely cancelled (e.g. client disconnect), await asyncio.shield(release_budget_reservation(...)) immediately raises CancelledError in the outer coroutine; the inner release_budget_reservation Task is still running in the background. The except asyncio.CancelledError: pass silences the raise and the function returns while finalized is still False — the budget release and flag update happen later on the event loop.

The three tests in test_budget_reservation.py trigger CancelledError from inside the iterator, not from actual task cancellation. Without outer-task cancellation asyncio.shield completes synchronously (no outer CancelledError to surface), so the tests always see finalized = True immediately after the call. The production case (outer task cancelled) is not covered and would leave finalized = False until the background task finishes.

The background task still runs the release correctly on a live event loop, so no budget is permanently leaked; but the finalized guard being unset between the function returning and the background task completing means a hypothetical concurrent reconcile call (e.g. a late streaming success callback) could double-reconcile before the flag is written.

[veria-ai]

High: Budget bypass on cancelled streams

A caller can start an expensive streaming request, receive generated chunks, then disconnect before the stream completes; this branch releases the reservation with actual_cost=0.0, and the success cost callback never runs. Reconcile the generated partial stream cost on cancellation, or otherwise keep a bounded reservation/charge path for streams that have already produced output instead of unconditionally releasing it.

[veria-ai]

PR overview

This pull request changes proxy request handling so budget reservations are released when an in-flight request is cancelled. The touched logic is in the common request-processing path for proxied requests, including streaming flows.

There is still an open issue in the cancellation path for streaming requests: a caller can receive generated output and then disconnect, causing the reservation to be released without charging for the partial work. That leaves the PR vulnerable to budget enforcement bypass and potential cost/resource abuse for streaming calls. No issues have been fixed or addressed yet, so the current posture still depends on correcting that accounting behavior before merge.

Open issues (1)

• High: Budget bypass on cancelled streams — litellm/proxy/common_request_processing.py:2071

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

[Bytechoreographer]

Superseded by #30522 (rebased onto the latest litellm_internal_staging so the diff is a single isolated commit).