From ebcd9bcb18a3cfae179df42ae0464384a7aa31e0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:29:11 +0000 Subject: [PATCH] fix(proxy): release max_parallel_requests slot when a realtime session ends without LLM callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 + tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..ed3d43800f7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11912,6 +11912,7 @@ 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") @@ -12050,6 +12051,7 @@ async def realtime_websocket_endpoint( if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): 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 ###################################################################### diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 03a24ec5e98..7b6a838aa3e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10068,6 +10068,46 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c assert reservation["finalized"] is False +@pytest.mark.asyncio +@pytest.mark.parametrize("phase_one_exit", [None, "pre_call"]) +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.""" + 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, + ) + 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 + with hooks: + 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 stash.parallel_slot is None + + @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),