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/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:

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)
 
 
 ######################################################################

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

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),