From 1b31be1a9c3b06543addec28c6342e14a53bd3db Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:23:23 +0000 Subject: [PATCH] fix(proxy): leave the realtime max_parallel slot to the success callback when one is enqueued Releasing the slot unconditionally from the route raced the limiter's own success handler on the logging worker: both could read the same stashed acquisition before either cleared it, and under the integer in-memory fallback that double-decrements the counter. The route now releases only on exits without a success callback (pre-call rejection, pre-call cancellation, and Phase 2 exits without the success stamp), matching the HTTP disconnect path's ownership rule. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 122 +++++++++--------- tests/test_litellm/proxy/test_proxy_server.py | 95 ++++++++++---- 2 files changed, 131 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f046f0892b..f37744f4685 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11893,6 +11893,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 + + async def _reject_realtime_session( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -11912,6 +11916,7 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) @app.websocket("/openai/v1/realtime") @@ -11991,68 +11996,69 @@ 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: - 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 + ( + 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 + except BaseException: + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise - # Phase 2: route to upstream LLM. + # 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: - 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 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"), ) - 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) + 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: - 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 + 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) + await _release_realtime_max_parallel_slot(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 11c2d301365..404deb3ca50 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -10075,6 +10076,50 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c assert reservation["finalized"] is False +_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, + phase_one_exit: str | None = None, +) -> tuple[DualCache, RequestRateLimiterStash]: + """Run the realtime endpoint with a real v3 limiter registered and the request's + stash already holding slot-1 of a two-slot counter, the state pre-call leaves + behind. Returns the limiter's cache and the stash so the test can read what the + endpoint did to the slot.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + ) + from litellm.proxy.utils import InternalUsageCache + + dual_cache: Final = DualCache() + await dual_cache.async_set_cache( + key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True + ) + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) + stash: Final = RequestRateLimiterStash( + parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + ) + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + 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 + 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=backend_logged_success, phase_one_exit=phase_one_exit + ) + finally: + _request_stash.reset(stash_token) + return dual_cache, stash + + @pytest.mark.asyncio @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( @@ -10086,39 +10131,33 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa 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 ( - RequestRateLimiterStash, - _PROXY_MaxParallelRequestsHandler_v3, - _request_stash, + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, phase_one_exit=phase_one_exit ) - from litellm.proxy.utils import InternalUsageCache - counter_key: Final = "{api_key:hashed-token}:max_parallel_requests" - dual_cache: Final = DualCache() - await dual_cache.async_set_cache(key=counter_key, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True) - limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) - stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [counter_key]}) - reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - - 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 - 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 - ) - finally: - _request_stash.reset(stash_token) - - assert await dual_cache.async_get_cache(key=counter_key, local_only=True) == {"slot-2": 2.0} + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {"slot-2": 2.0} assert stash.parallel_slot is None +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_callback(): + """A session that enqueued its success callback hands the slot to the limiter's + own success handler, which runs on the logging worker. If the route also released + it, the two releases would race on the same stashed acquisition and, under the + limiter's integer in-memory fallback, double-decrement the counter so the key + admits more sessions than max_parallel_requests allows. With the success stamp + present the route leaves the slot and the stash alone.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=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),