From ebcd9bcb18a3cfae179df42ae0464384a7aa31e0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:29:11 +0000 Subject: [PATCH 1/7] 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), From 81ffc3125f989eed375e1d5a708f2a8a60cbfda1 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:50:01 +0000 Subject: [PATCH 2/7] 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 ) From 1b31be1a9c3b06543addec28c6342e14a53bd3db Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:23:23 +0000 Subject: [PATCH 3/7] 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), From abc85ba6078b57466fd7d7a0d20217c578844add Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:47:37 +0000 Subject: [PATCH 4/7] 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), From d2342f06ceb52f429d5a405b68adb090d76b012c Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:13:29 +0000 Subject: [PATCH 5/7] fix(bedrock): stamp the realtime success ownership marker when Nova Sonic spend is logged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 6 ++++- .../realtime/test_bedrock_realtime_handler.py | 27 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 7ab5a14dfc2..e83cb743aa0 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -24,7 +24,10 @@ from litellm.constants import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes +from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + DefaultLoggedRealTimeEventTypes, +) from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput @@ -346,6 +349,7 @@ class BedrockRealtime(BaseAWSLLM): prefer_async_handlers=True, ) ) + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True if outcome.provider_failure is None: return diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 2eadaee9e1a..7e418817e74 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm +from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -77,6 +78,7 @@ class UnavailableBedrockStream: class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id + self.model_call_details = {} class DisconnectingClientWS: @@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation: await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + @pytest.mark.asyncio + async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + idle_logging = FakeLogging() + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=idle_logging, + **self.AWS_PARAMS, + ) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details + @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) From b4d0f4ad2658fa2e3740d30acbd6c0ad687a3b60 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:33:48 +0000 Subject: [PATCH 6/7] refactor(realtime): move session ownership marker keys into constants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ litellm/litellm_core_utils/realtime_streaming.py | 5 +---- litellm/llms/bedrock/realtime/handler.py | 6 ++---- litellm/proxy/proxy_server.py | 7 ++----- .../litellm_core_utils/test_realtime_streaming.py | 3 +-- .../llms/bedrock/realtime/test_bedrock_realtime_handler.py | 2 +- tests/test_litellm/proxy/test_proxy_server.py | 5 +---- 7 files changed, 10 insertions(+), 20 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e8aafda1797..c106be688e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 2177c999804..fa567bdf4c9 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -10,6 +10,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -35,10 +36,6 @@ else: CLIENT_CONNECTION_CLASS = Any -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) class BackendClose: code: int diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index e83cb743aa0..43138b2c526 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -20,14 +20,12 @@ from litellm.constants import ( BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - DefaultLoggedRealTimeEventTypes, -) +from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f499ee04506..1c931863a2f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -274,6 +274,8 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + REALTIME_SESSION_FAILURE_LOGGED_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, @@ -12061,11 +12063,6 @@ async def realtime_websocket_endpoint( 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_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) if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY): 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 e1eb61b59b9..7e6d4d24905 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,10 +10,9 @@ from websockets.exceptions import ConnectionClosed from websockets.frames import Close import litellm +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY 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, ) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 7e418817e74..6c7659d84cd 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -8,7 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm -from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY +from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 64828ae3cb4..e09dddfec5b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10020,10 +10020,7 @@ 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_FAILURE_LOGGED_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) + from litellm.constants 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") From fbc1011d277b4630b4b20b061b82f5e0c0032130 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:42:42 +0000 Subject: [PATCH 7/7] fix(bedrock): end the realtime session when the client disconnects instead of waiting for Nova Sonic Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 2 +- .../realtime/test_bedrock_realtime_handler.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 43138b2c526..2c1ce6068b2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -384,7 +384,7 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) - await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION) + await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED) client_disconnected: Final = ( client_task.done() and not client_task.cancelled() and client_task.exception() is None ) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 6c7659d84cd..ac3a43b742f 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -711,6 +711,25 @@ class TestBedrockRealtimeProviderFailurePropagation: assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_client_disconnect_ends_the_session_while_bedrock_output_stays_open(self, stub_aws_sdk_client): + receiver = DrainedThenOpenBedrockReceiver([]) + stream = ScriptedBedrockStream([], receiver_type=lambda _payloads: receiver) + stub_aws_sdk_client["streams"] = [stream] + + await asyncio.wait_for( + BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=FakeLogging(), + **self.AWS_PARAMS, + ), + timeout=1, + ) + + assert receiver.drained.is_set(), "the handler must have been waiting on the open provider stream" + assert stream.input_stream.closed + @pytest.mark.asyncio async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models): handler = BedrockRealtime()