From 81ffc3125f989eed375e1d5a708f2a8a60cbfda1 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:50:01 +0000 Subject: [PATCH] fix(proxy): release realtime max_parallel_requests slot when the task is cancelled during pre-call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 115 +++++++++--------- tests/test_litellm/proxy/test_proxy_server.py | 27 ++-- 2 files changed, 77 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ed3d43800f7..8f046f0892b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11912,7 +11912,6 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) - 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 @app.websocket("/openai/v1/realtime") @@ -11992,65 +11991,67 @@ async def realtime_websocket_endpoint( # Errors here (e.g. guardrail block) are sent back to the client as an # error event before closing, so the caller knows what happened. try: - ( - data, - litellm_logging_obj, - ) = await base_llm_response_processor.common_processing_pre_call_logic( - request=request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_logging_obj=proxy_logging_obj, - proxy_config=proxy_config, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - model=route_model, - route_type="_arealtime", - ) - except Exception as e: - verbose_proxy_logger.exception("Realtime 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. - try: - data["user_api_key_dict"] = user_api_key_dict - llm_call: Final = await route_request( - data=data, - route_type="_arealtime", - llm_router=llm_router, - user_model=user_model, - ) - await llm_call - except websockets.exceptions.InvalidStatusCode as e: - verbose_proxy_logger.exception("Invalid status code") - await websocket.close(code=e.status_code, reason="Invalid status code") - except Exception as e: - verbose_proxy_logger.exception("Internal server error") - redacted_error: Final = _redact_string(str(e)) try: - await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) - 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 error event to client; closing anyway") - try: - await websocket.close( - code=1011, - reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ( + data, + litellm_logging_obj, + ) = await base_llm_response_processor.common_processing_pre_call_logic( + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + model=route_model, + route_type="_arealtime", ) - except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error - 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_SUCCESS_LOGGED_KEY, - ) + except Exception as e: + verbose_proxy_logger.exception("Realtime pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) + return - if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): - await _release_realtime_budget_reservation(user_api_key_dict) + # Phase 2: route to upstream LLM. + try: + data["user_api_key_dict"] = user_api_key_dict + llm_call: Final = await route_request( + data=data, + route_type="_arealtime", + llm_router=llm_router, + user_model=user_model, + ) + await llm_call + except websockets.exceptions.InvalidStatusCode as e: + verbose_proxy_logger.exception("Invalid status code") + await websocket.close(code=e.status_code, reason="Invalid status code") + except Exception as e: + verbose_proxy_logger.exception("Internal server error") + redacted_error: Final = _redact_string(str(e)) + try: + await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) + 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 error event to client; closing anyway") + try: + await websocket.close( + code=1011, + reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ) + except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error + 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_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) + finally: 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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7b6a838aa3e..11c2d301365 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9919,8 +9919,9 @@ async def _lit6973_drive_realtime_session( 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. + (rate limits, guardrails) raise, "pre_call_cancelled" cancels the task inside + pre-call processing. None 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 @@ -9950,7 +9951,13 @@ async def _lit6973_drive_realtime_session( 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_error: Final = ( + asyncio.CancelledError() + if phase_one_exit == "pre_call_cancelled" + else 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) ) @@ -10069,15 +10076,16 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c @pytest.mark.asyncio -@pytest.mark.parametrize("phase_one_exit", [None, "pre_call"]) +@pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"]) async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot( phase_one_exit: str | None, ): """The rate limiter acquires the key's max_parallel_requests slot in pre-call and only frees it from the LLM success/failure callbacks. A realtime session that ends - without either callback (Bedrock closes without usage events, or a later pre-call - hook rejects the session) has to be released by the route itself, or the slot stays - occupied until its TTL and the key's next session is refused with a 429.""" + without either callback (Bedrock closes without usage events, a later pre-call hook + rejects the session, or the task is cancelled while still in pre-call) has to be + released by the route itself, or the slot stays occupied until its TTL and the key's + next session is refused with a 429.""" from litellm.caching.caching import DualCache from litellm.proxy import proxy_server as ps from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -10097,7 +10105,10 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa 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 - with hooks: + expected_exit: Final = ( + 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=False, phase_one_exit=phase_one_exit )