From 37722eba68149c5f3e59ed0f3c12798a84aa1bc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:35:26 -0700 Subject: [PATCH] fix(realtime): close a rejected client before releasing its budget reservation A slow or unreachable counter store made a pre-relay rejection wait behind the reservation release before the client saw the error event and the close. Close first and release in finally, mirroring the relay's own failure path, so a client that already hung up still gets its reservation released. --- litellm/proxy/proxy_server.py | 20 ++++---- tests/test_litellm/proxy/test_proxy_server.py | 47 ++++++++++++++++++- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 65fe3ede822..a5c60d8e976 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11471,15 +11471,17 @@ async def _reject_realtime_session( 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) + try: + 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) + finally: + await _release_realtime_budget_reservation(user_api_key_dict) @app.websocket("/openai/v1/realtime") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8b151d9e1f6..4d70c9d436f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9533,7 +9533,11 @@ def _lit6973_fake_realtime_ws() -> MagicMock: async def _lit6973_drive_realtime_session( - reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None + reservation: dict, + *, + backend_logged_success: bool, + phase_one_exit: str | None = None, + websocket: MagicMock | None = None, ) -> MagicMock: """Drive realtime_websocket_endpoint through one of its reservation-settling exits. @@ -9574,7 +9578,7 @@ async def _lit6973_drive_realtime_session( 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() + ws: Final = websocket if websocket is not None else _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 @@ -9635,6 +9639,45 @@ async def test_realtime_session_denied_model_access_releases_the_budget_reservat ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") +@pytest.mark.asyncio +async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation(): + """The counter release can block on a slow or unreachable store, and a + rejected client must not sit behind it: the relay's own failure path closes + the client first and releases in its finally, so the pre-relay rejection + has to close first as well. The fake close checks the reservation is still + open when the client is closed.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + + async def close_while_reservation_is_still_open(**_: object) -> None: + assert reservation["finalized"] is False, "client was closed only after the reservation release" + + ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open) + + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws + ) + + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone(): + """A client that hung up before the rejection makes the close raise; the + reservation must still be released, or the key stays pinned.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected")) + + with pytest.raises(RuntimeError, match="client already disconnected"): + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws + ) + + assert reservation["finalized"] is True + + @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