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 e3f8786a39a..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,9 +36,6 @@ else: CLIENT_CONNECTION_CLASS = Any -REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" - - @dataclass(frozen=True, slots=True) class BackendClose: code: int @@ -1153,6 +1151,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/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 7ab5a14dfc2..2c1ce6068b2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -20,6 +20,7 @@ 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 @@ -346,6 +347,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 @@ -382,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/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dd7a39eb94c..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, @@ -11899,6 +11901,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None: + 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( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -11918,6 +11927,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") @@ -12021,6 +12031,9 @@ async def realtime_websocket_endpoint( 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. try: @@ -12050,12 +12063,10 @@ 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_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): + 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..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,9 +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_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, ) @@ -3399,6 +3399,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/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 2eadaee9e1a..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 @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm +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 @@ -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) @@ -684,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() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a5e6cf285ad..e09dddfec5b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -33,6 +33,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 @@ -9999,6 +10000,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: @@ -10006,8 +10008,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 @@ -10017,7 +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_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") @@ -10029,6 +10032,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 @@ -10037,7 +10042,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) ) @@ -10155,6 +10166,114 @@ 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, + 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 + 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( # 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() + ) + with hooks, expected_exit: + await _lit6973_drive_realtime_session( + 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) + 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( + 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, 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.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, phase_one_exit=phase_one_exit + ) + + 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_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),