From 5a35e6d41f76d2b09a258b3e2b051e3e7e745c79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:18:33 -0700 Subject: [PATCH] fix(realtime): release the budget reservation when a session is rejected before the relay starts The three pre-relay exits of realtime_websocket_endpoint (missing model, key/model access denied, pre-call rejection such as a rate limit or a guardrail) returned before the finally that releases the auth-time budget reservation, so a rejected session pinned the key at the reserved amount until the counter TTL expired and its next requests got budget_exceeded while /key/info showed spend 0. A single _reject_realtime_session helper now releases the reservation before sending the error event and closing, and release_or_invalidate_budget_reservation shields the release from a second cancellation and logs, rather than raises, a failing invalidate fallback so it can never mask the session's own outcome. --- litellm/proxy/proxy_server.py | 43 ++++++---- .../spend_tracking/budget_reservation.py | 4 +- tests/test_litellm/proxy/test_proxy_server.py | 79 +++++++++++++++++-- 3 files changed, 103 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d59dfa86c4..65fe3ede822 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11463,6 +11463,25 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _reject_realtime_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + code: int, + reason: str, + error_message: str | None = None, +) -> None: + await _release_realtime_budget_reservation(user_api_key_dict) + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11488,7 +11507,9 @@ async def realtime_websocket_endpoint( if intent == "transcription": route_model = "gpt-realtime-whisper" else: - await websocket.close(code=1008, reason="model query parameter is required") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1008, reason="model query parameter is required" + ) return assert route_model is not None try: @@ -11499,7 +11520,7 @@ async def realtime_websocket_endpoint( llm_router=llm_router, ) except ProxyException as e: - await websocket.close(code=1008, reason=e.message[:120]) + await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return await websocket.accept(**accept_kwargs) @@ -11558,21 +11579,9 @@ async def realtime_websocket_endpoint( ) except Exception as e: verbose_proxy_logger.exception("Realtime pre-call error") - try: - await websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_error", - "message": str(e), - }, - } - ) - ) - except Exception: - pass - await websocket.close(code=1011, reason="Pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) return # Phase 2: route to upstream LLM. diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2ee7320b82c..ed2bc87597c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -389,11 +389,13 @@ async def release_or_invalidate_budget_reservation( if budget_reservation is None or budget_reservation.get("finalized") is True: return try: - await release_budget_reservation(budget_reservation=budget_reservation) + await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation)) except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed + verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed") finally: budget_reservation["finalized"] = True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 697d4c182c7..8b151d9e1f6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,8 +9532,15 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: - """Drive realtime_websocket_endpoint to just before its budget-reservation finally. +async def _lit6973_drive_realtime_session( + reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None +) -> MagicMock: + """Drive realtime_websocket_endpoint through one of its reservation-settling exits. + + phase_one_exit picks a rejection before the relay: "model_access" makes the + key/model check raise ProxyException, "pre_call" makes pre-call processing + (rate limits, guardrails) raise. Neither reaches route_request, so no success + log can own the reservation and the endpoint has to release it on that exit. route_request resolves normally in both cases: the relay owns the session once route_request returns. A successful session enqueues its success cost @@ -9556,18 +9563,30 @@ async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_s if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) - can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test + from litellm.proxy._types import ProxyException + + model_access_error: Final = ( + ProxyException(message="key cannot access model", type="auth_error", param="model", code=401) + if phase_one_exit == "model_access" + else None + ) + pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None + pre_call: Final = AsyncMock( + side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) + ) + ws: Final = _lit6973_fake_realtime_ws() + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( - websocket=_lit6973_fake_realtime_ws(), + websocket=ws, model="vertex_ai/gemini-live-2.5-flash", intent=None, guardrails=None, user_api_key_dict=user_api_key_dict, ) + return ws @pytest.mark.asyncio @@ -9583,6 +9602,39 @@ async def test_refused_realtime_session_releases_the_budget_reservation(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation(): + """A rate-limit or guardrail rejection happens before route_request, so the + relay never runs and no success log can own the reservation. The endpoint + must release it on that exit too, or the key stays pinned at the reserved + amount and its next requests 429 with budget_exceeded while /key/info shows + spend 0 (reproduced live with rpm_limit=1). The client still gets the + pre-call error event and the 1011 close it got before.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call" + ) + + assert reservation["finalized"] is True + assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded" + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + + +@pytest.mark.asyncio +async def test_realtime_session_denied_model_access_releases_the_budget_reservation(): + """The key/model access check rejects before the socket is even accepted; + that exit skipped the release as well, pinning the reservation.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access" + ) + + assert reservation["finalized"] is True + ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") + + @pytest.mark.asyncio async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): """A billable realtime session settles its reservation through the enqueued @@ -9625,6 +9677,23 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails(): + """Both counter-store calls failing must not raise out of the realtime + endpoint's finally (it would mask the session's own outcome) and must still + stamp finalized so nothing retries the same reservation.""" + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + + with failing_release, failing_invalidate: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints.