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

@@ -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.

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
 

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.