From abc85ba6078b57466fd7d7a0d20217c578844add Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:47:37 +0000 Subject: [PATCH] fix(proxy): leave the realtime max_parallel_requests slot to the limiter failure callback when a refusal was logged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/realtime_streaming.py | 2 + litellm/proxy/proxy_server.py | 9 ++++- .../test_realtime_streaming.py | 21 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++-- 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e3f8786a39a..2177c999804 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -36,6 +36,7 @@ else: REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" @dataclass(frozen=True, slots=True) @@ -1153,6 +1154,7 @@ class RealTimeStreaming: self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True @staticmethod def _detect_beta_header(websocket: ScopedWebSocket) -> bool: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f37744f4685..74783398456 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11894,7 +11894,10 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None: - await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses + release_like_http_disconnect: Final = ( + proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared + ) + await release_like_http_disconnect(user_api_key_dict) async def _reject_realtime_session( @@ -12053,12 +12056,14 @@ async def realtime_websocket_endpoint( verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): await _release_realtime_budget_reservation(user_api_key_dict) - await _release_realtime_max_parallel_slot(user_api_key_dict) + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY): + await _release_realtime_max_parallel_slot(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 2a33d84ec78..e1eb61b59b9 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -12,6 +12,7 @@ from websockets.frames import Close import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, @@ -3399,6 +3400,26 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details +@pytest.mark.asyncio +async def test_refused_session_stamps_the_failure_ownership_marker(): + """LIT-6463: the enqueued failure callback releases the key's max_parallel_requests + slot from the logging worker, so a refusal stamps REALTIME_SESSION_FAILURE_LOGGED_KEY. + The proxy endpoint reads it to leave the slot to that callback instead of racing it. + A session that relayed frames logs a success and must not carry the failure stamp.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + refused: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + relayed: Final = _relay_session( + _client_ws_that_never_sends(), _backend_ws_closing_with(session_created, upstream_close) + ) + + await refused.run() + await relayed.run() + + assert refused.logging.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY) is True + assert REALTIME_SESSION_FAILURE_LOGGED_KEY not in relayed.logging.model_call_details + + @pytest.mark.asyncio async def test_transformed_transcription_completion_never_sends_response_create(): from typing import Final diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 404deb3ca50..986cac2a18e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9913,6 +9913,7 @@ async def _lit6973_drive_realtime_session( reservation: dict, *, backend_logged_success: bool, + backend_logged_failure: bool = False, phase_one_exit: str | None = None, websocket: MagicMock | None = None, ) -> MagicMock: @@ -9932,7 +9933,10 @@ async def _lit6973_drive_realtime_session( logging object carries a real model_call_details dict so the stamp is observable, and the reservation has empty entries so the real release touches no counter store.""" - from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") @@ -9944,6 +9948,8 @@ async def _lit6973_drive_realtime_session( async def fake_llm_call() -> None: if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + if backend_logged_failure: + logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True from litellm.proxy._types import ProxyException @@ -10082,6 +10088,7 @@ _LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests" async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( *, backend_logged_success: bool, + backend_logged_failure: bool = False, phase_one_exit: str | None = None, ) -> tuple[DualCache, RequestRateLimiterStash]: """Run the realtime endpoint with a real v3 limiter registered and the request's @@ -10107,13 +10114,20 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( stash_token: Final = _request_stash.set(stash) try: - hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter + hooks: Final = patch.dict( # test-quality-ok: registers the real limiter the route's release reads + ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter} + ) expected_exit: Final = ( - pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext() + pytest.raises(asyncio.CancelledError) + if phase_one_exit == "pre_call_cancelled" + else contextlib.nullcontext() ) with hooks, expected_exit: await _lit6973_drive_realtime_session( - reservation, backend_logged_success=backend_logged_success, phase_one_exit=phase_one_exit + reservation, + backend_logged_success=backend_logged_success, + backend_logged_failure=backend_logged_failure, + phase_one_exit=phase_one_exit, ) finally: _request_stash.reset(stash_token) @@ -10158,6 +10172,24 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_ assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} +@pytest.mark.asyncio +async def test_refused_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_failure_callback(): + """An upstream refusal before any frame enqueues the failure callback instead, and + the limiter's failure handler releases the slot from the logging worker just like + the success handler does. The route sees no success stamp, so it still settles the + budget reservation, but it must leave the slot to that callback or the two releases + race on the same acquisition.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, backend_logged_failure=True + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { + "slot-1": 1.0, + "slot-2": 2.0, + } + assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down),