diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a175ca1c3f6..8f1b2ce1cc2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1820,6 +1820,7 @@ class Logging(LiteLLMLoggingBaseClass): and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True and litellm_params.get(CallTypes.agenerate_content.value, False) is not True and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True + and litellm_params.get(CallTypes.arealtime.value, False) is not True ) def _is_assembled_stream_success(self, result=None) -> bool: diff --git a/litellm/litellm_core_utils/realtime_errors.py b/litellm/litellm_core_utils/realtime_errors.py index e1b957f4325..3c064728a66 100644 --- a/litellm/litellm_core_utils/realtime_errors.py +++ b/litellm/litellm_core_utils/realtime_errors.py @@ -29,3 +29,11 @@ def websocket_close_reason(message: str, fallback: str) -> str: if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES: return message return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + + +def client_close_code(upstream_code: int) -> int: + from websockets.frames import EXTERNAL_CLOSE_CODES, CloseCode + + if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000: + return upstream_code + return int(CloseCode.INTERNAL_ERROR) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8479e108d17..75046f2cf87 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,12 +1,15 @@ import asyncio import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +import traceback +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass +from enum import Enum, auto +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast from typing_extensions import ReadOnly import litellm -from litellm._logging import verbose_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_logger 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 ( @@ -19,9 +22,11 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging +from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from websockets.exceptions import ConnectionClosed from litellm.types.guardrails import GuardrailEventHooks @@ -30,8 +35,30 @@ else: CLIENT_CONNECTION_CLASS = Any -class _ClientWebSocketExceptions(Protocol): - ConnectionClosed: type[Exception] +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" + + +@dataclass(frozen=True, slots=True) +class BackendClose: + code: int + reason: str + + @property + def message(self) -> str: + if not self.reason: + return f"upstream websocket closed with code {self.code}" + return f"upstream websocket closed with code {self.code}: {self.reason}" + + +class ClientLoopExit(Enum): + CLIENT_DISCONNECTED = auto() + BACKEND_CLOSED = auto() + + +def backend_close_from(error: "ConnectionClosed") -> BackendClose: + if error.rcvd is None: + return BackendClose(code=1006, reason=str(error)) + return BackendClose(code=error.rcvd.code, reason=error.rcvd.reason) class _ASGIScope(TypedDict, total=False): @@ -69,10 +96,13 @@ class _ScopedWebSocket(Protocol): class _ClientWebSocket(_ScopedWebSocket, Protocol): - exceptions: _ClientWebSocketExceptions - async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class _LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... def _decode_json_object(payload: str) -> Mapping[str, object]: @@ -108,11 +138,14 @@ class RealTimeStreaming: backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, + logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj + self._logging_worker = logging_worker self.messages: list[OpenAIRealtimeEvents] = [] + self._backend_sent_frames: bool = False self.input_message: dict = {} self.input_messages: list[dict[str, str]] = [] self.session_tools: list[dict] = [] @@ -388,9 +421,10 @@ class RealTimeStreaming: # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True async def _send_to_backend(self, message: str) -> bool: """Send a message to the backend WebSocket. @@ -1035,60 +1069,84 @@ class RealTimeStreaming: return True return False - async def backend_to_client_send_messages(self): + async def _relay_backend_messages(self) -> NoReturn: + while True: + try: + raw_response = await self.backend_ws.recv(decode=False) + except TypeError: + raw_response = await self.backend_ws.recv() + self._backend_sent_frames = True + + if isinstance(raw_response, bytes): + try: + raw_response = raw_response.decode("utf-8") + except UnicodeDecodeError: + verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") + continue + + if self.provider_config: + try: + await self._handle_provider_config_message(raw_response) + except Exception as e: + verbose_logger.exception("Error processing backend message, skipping: %s", e) + continue + else: + event = self._parse_backend_event(raw_response) + if event is None: + await self.websocket.send_text(raw_response) + continue + + if self._should_drop_event_from_client(event): + continue + + if await self._handle_raw_backend_message(event, raw_response): + continue + + event = self._normalize_event_for_ga_client(event) + self.store_message(event) + + if not self._client_wants_beta: + await self.websocket.send_text(json.dumps(event)) + continue + + translated = self._translate_event_to_beta(event) + if translated is None: + continue + await self.websocket.send_text(json.dumps(translated)) + + async def backend_to_client_send_messages(self) -> BackendClose: import websockets try: - while True: - try: - raw_response = await self.backend_ws.recv(decode=False) - except TypeError: - raw_response = await self.backend_ws.recv() - - if isinstance(raw_response, bytes): - try: - raw_response = raw_response.decode("utf-8") - except UnicodeDecodeError: - verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") - continue - - if self.provider_config: - try: - await self._handle_provider_config_message(raw_response) - except Exception as e: - verbose_logger.exception("Error processing backend message, skipping: %s", e) - continue - else: - event = self._parse_backend_event(raw_response) - if event is None: - await self.websocket.send_text(raw_response) - continue - - if self._should_drop_event_from_client(event): - continue - - if await self._handle_raw_backend_message(event, raw_response): - continue - - event = self._normalize_event_for_ga_client(event) - self.store_message(event) - - if not self._client_wants_beta: - await self.websocket.send_text(json.dumps(event)) - continue - - translated = self._translate_event_to_beta(event) - if translated is None: - continue - await self.websocket.send_text(json.dumps(translated)) - + await self._relay_backend_messages() except websockets.exceptions.ConnectionClosed as e: verbose_logger.exception("Connection closed in backend to client send messages - %s", e) - except Exception as e: - verbose_logger.exception("Error in backend to client send messages: %s", e) - finally: + close: Final = backend_close_from(e) + self._flush_unbilled_transcription_usage() + if self._backend_refused_session(close): + await self.log_backend_refusal(e) + else: + await self.log_messages() + return close + except asyncio.CancelledError: self._flush_unbilled_transcription_usage() await self.log_messages() + raise + except Exception as e: + verbose_logger.exception("Error in backend to client send messages: %s", e) + self._flush_unbilled_transcription_usage() + await self.log_messages() + return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") + + def _backend_refused_session(self, close: BackendClose) -> bool: + return close.code != 1000 and not self._backend_sent_frames + + async def log_backend_refusal(self, error: Exception) -> None: + if not self.logging_obj: + return + self._logging_worker.ensure_initialized_and_enqueue( + self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) + ) @staticmethod def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: @@ -1243,11 +1301,22 @@ class RealTimeStreaming: item["content"] = new_content return item - async def client_ack_messages(self): + async def _receive_client_message(self) -> str | None: + try: + return await self.websocket.receive_text() + except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone + verbose_logger.debug("Client disconnected: %s", e) + return None + + async def client_ack_messages(self) -> ClientLoopExit: + import websockets + client_event: _ClientEventFrame try: while True: - message = await self.websocket.receive_text() + message = await self._receive_client_message() + if message is None: + return ClientLoopExit.CLIENT_DISCONNECTED ## GUARDRAIL: intercept conversation.item.create for text-based injection. guardrail_turn_detection_injected = False @@ -1481,23 +1550,38 @@ class RealTimeStreaming: if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True + except websockets.exceptions.ConnectionClosed as e: + verbose_logger.debug("Backend closed while forwarding a client message: %s", e) + return ClientLoopExit.BACKEND_CLOSED except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) + return ClientLoopExit.CLIENT_DISCONNECTED - async def bidirectional_forward(self): + async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) + client_task: Final = asyncio.create_task(self.client_ack_messages()) try: - await self.client_ack_messages() - except self.websocket.exceptions.ConnectionClosed: - verbose_logger.debug("Connection closed") - forward_task.cancel() + await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) + if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED: + return + await self._close_client(await forward_task) finally: - if not forward_task.done(): - forward_task.cancel() - try: - await forward_task - except asyncio.CancelledError: - pass + forward_task.cancel() + client_task.cancel() + await asyncio.gather(forward_task, client_task, return_exceptions=True) + + async def _close_client(self, close: BackendClose) -> None: + redacted_message: Final = redact_internal_details_from_client_message(close.message) + redacted_reason: Final = redact_internal_details_from_client_message(close.reason) + try: + if close.code != 1000: + await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) + await self.websocket.close( + code=client_close_code(close.code), + reason=websocket_close_reason(redacted_reason, fallback=redacted_message), + ) + except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way + verbose_logger.debug("Could not relay the upstream close to the client: %s", e) def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 026d8d2228c..0f0abd0eeae 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11458,6 +11458,37 @@ def _realtime_query_params_template(model: str | None, intent: str | None) -> tu return tuple(params) +async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth) -> None: + from litellm.proxy.spend_tracking.budget_reservation import ( + release_or_invalidate_budget_reservation, + ) + + await release_or_invalidate_budget_reservation( + budget_reservation=user_api_key_dict.budget_reservation, + ) + + +async def _reject_realtime_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + code: int, + reason: str, + error_message: str | None = None, +) -> None: + try: + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + 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 pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + finally: + await _release_realtime_budget_reservation(user_api_key_dict) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11483,7 +11514,9 @@ async def realtime_websocket_endpoint( if intent == "transcription": route_model = "gpt-realtime-whisper" else: - await websocket.close(code=1008, reason="model query parameter is required") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1008, reason="model query parameter is required" + ) return assert route_model is not None try: @@ -11494,7 +11527,7 @@ async def realtime_websocket_endpoint( llm_router=llm_router, ) except ProxyException as e: - await websocket.close(code=1008, reason=e.message[:120]) + await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return await websocket.accept(**accept_kwargs) @@ -11553,21 +11586,9 @@ async def realtime_websocket_endpoint( ) except Exception as e: verbose_proxy_logger.exception("Realtime pre-call error") - try: - await websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_error", - "message": str(e), - }, - } - ) - ) - except Exception: - pass - await websocket.close(code=1011, reason="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. @@ -11597,6 +11618,13 @@ 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) ###################################################################### diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 91d2ece7a51..ed2bc87597c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -373,6 +373,33 @@ async def invalidate_budget_reservation_counters( await _invalidate_spend_counter(counter_key=counter_key) +async def release_or_invalidate_budget_reservation( + budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict +) -> None: + """Reconcile a still-open reservation on a terminal path that settles no cost. + + A failed or upstream-refused request never runs the success cost callback, so + its pre-call reservation stays open and keeps the spend counter pinned above + real spend until the counter's TTL expires, 429ing later requests on the same + key. Release it to zero; if the release itself fails (e.g. the counter store is + unreachable) drop the reserved counters directly and mark the reservation + finalized so nothing reprocesses it. Idempotent: the finalized guard makes a + second call a no-op once success or failure handling already reconciled. + """ + if budget_reservation is None or budget_reservation.get("finalized") is True: + return + try: + await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation)) + except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead + verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") + try: + await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed + verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed") + finally: + budget_reservation["finalized"] = True + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f1a23eba6c4..b824a5928c6 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -27,7 +27,7 @@ from litellm.types.realtime import ( RealtimeTranscriptionSessionRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params @@ -360,7 +360,7 @@ async def _arealtime( user: Final = kwargs.get("user", None) litellm_params: Final = GenericLiteLLMParams(**kwargs) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.arealtime.value: True} model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( model=model, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f1de7390b5b..af75691eb10 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -995,6 +995,35 @@ async def test_anthropic_messages_marks_litellm_params_async(): litellm.callbacks = original_callbacks +@pytest.mark.asyncio +async def test_arealtime_marks_litellm_params_async(monkeypatch): + """LIT-6973: ``_arealtime`` must plant ``_arealtime`` in ``litellm_params`` so + ``_is_sync_litellm_request`` classifies the session async and a failed session + reaches a CustomLogger's failure hook once, through the async path only, even + though the sync ``failure_handler`` still runs ahead of the async one.""" + captured = {} + async_logged = asyncio.Event() + + class CaptureLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + captured["litellm_params"] = kwargs.get("litellm_params", {}) + async_logged.set() + + logger = CaptureLogger() + logger.log_failure_event = MagicMock() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + with pytest.raises(ValueError, match="Unsupported model"): + await litellm._arealtime(model="anthropic/claude-x", websocket=MagicMock()) + await asyncio.wait_for(async_logged.wait(), timeout=10) + logger.log_failure_event.assert_not_called() + assert captured["litellm_params"].get("_arealtime") is True + assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False + + @pytest.mark.asyncio async def test_agenerate_content_marks_litellm_params_async(): """LIT-4475: the async ``agenerate_content`` entrypoint must plant @@ -1180,6 +1209,7 @@ def test_is_sync_litellm_request(): assert LitellmLogging._is_sync_litellm_request({}) is True assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False + assert LitellmLogging._is_sync_litellm_request({"_arealtime": True}) is False assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py index 494d16b0b9b..1d2cf905f4e 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py @@ -1,8 +1,10 @@ import json +import pytest from litellm.litellm_core_utils.realtime_errors import ( WEBSOCKET_CLOSE_REASON_MAX_BYTES, + client_close_code, realtime_error_event, websocket_close_reason, ) @@ -42,3 +44,11 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes(): assert len(reason.encode("utf-8")) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES assert reason == "あ" * (WEBSOCKET_CLOSE_REASON_MAX_BYTES // 3) assert "�" not in reason + + +@pytest.mark.parametrize( + ("upstream_code", "expected"), + [(1000, 1000), (1008, 1008), (1011, 1011), (4001, 4001), (1005, 1011), (1006, 1011), (1015, 1011), (2999, 1011)], +) +def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected): + assert client_close_code(upstream_code) == expected 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 52e88db753a..9c0f6f59463 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -1,14 +1,20 @@ +import asyncio import json +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from websockets.exceptions import ConnectionClosed +from websockets.frames import Close import litellm 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, ) @@ -2941,13 +2947,11 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): realtime turn leaves a suspended task pinning its response in memory -> an unbounded leak. Regression for that fix.""" logging_obj = MagicMock() - streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj) + mock_worker = MagicMock() + streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj, logging_worker=mock_worker) streaming.messages = [{"type": "session.created"}] - with ( - patch("litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER") as mock_worker, - patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task, - ): + with patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task: await streaming.log_messages() mock_worker.ensure_initialized_and_enqueue.assert_called_once() @@ -3028,12 +3032,12 @@ async def test_session_close_flushes_unbilled_transcription_usage(): messages before log_messages runs, and never forwarded to the client.""" from typing import Final - from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage + from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeResponseTypedDict client_ws: Final = MagicMock() client_ws.send_text = AsyncMock() backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws.recv = AsyncMock(side_effect=[b'{"serverContent": {}}', ConnectionClosed(None, None)]) logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() @@ -3045,7 +3049,24 @@ async def test_session_close_flushes_unbilled_transcription_usage(): "total_tokens": 171, "input_token_details": {"text_tokens": 0, "audio_tokens": 153}, } + transcript_frame: Final[RealtimeResponseTypedDict] = { + "response": { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "transcript": "ahoy", + "item_id": "item_1", + "content_index": 0, + }, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } provider_config: Final = MagicMock() + provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame) provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage) streaming: Final = RealTimeStreaming( @@ -3077,7 +3098,9 @@ async def test_session_close_flushes_unbilled_transcription_usage(): ) assert len(flushed) == 1 assert flushed[0] in logged_snapshots[0] - assert not client_ws.send_text.called + forwarded: Final = tuple(json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + assert [event.get("transcript") for event in forwarded] == ["ahoy"] + assert all("usage" not in event for event in forwarded) @pytest.mark.asyncio @@ -3111,3 +3134,281 @@ async def test_session_close_flush_noop_without_unbilled_usage(): isinstance(message, dict) and message.get("type") == "conversation.item.input_audio_transcription.completed" for message in streaming.messages ) + + + +_UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found" + + +class _InlineLoggingWorker: + def __init__(self) -> None: + self.enqueued: tuple[Coroutine[object, object, None], ...] = () + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + self.enqueued = (*self.enqueued, async_coroutine) + + async def drain(self) -> None: + for coroutine in self.enqueued: + await coroutine + + +class _RecordingLogging: + def __init__(self) -> None: + self.model_call_details: dict[str, object] = {} + self.logged_sessions: tuple[tuple[dict, ...], ...] = () + self.logged_failures: tuple[Exception, ...] = () + + def pre_call(self, input: str | dict, api_key: str) -> None: + return None + + async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None: + self.logged_sessions = (*self.logged_sessions, tuple(result)) + + async def dispatch_failure_handlers( + self, exception: Exception, traceback_exception: str, prefer_async_handlers: bool = False + ) -> None: + self.logged_failures = (*self.logged_failures, exception) + + +@dataclass(frozen=True, slots=True) +class _RelaySession: + streaming: RealTimeStreaming + logging: _RecordingLogging + worker: _InlineLoggingWorker + + async def run(self) -> None: + await asyncio.wait_for(self.streaming.bidirectional_forward(), timeout=2) + await self.worker.drain() + + +async def _wait_forever() -> str: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def _client_ws_that_never_sends() -> MagicMock: + client_ws: Final = MagicMock() + client_ws.headers = {} + client_ws.receive_text = AsyncMock(side_effect=_wait_forever) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + return client_ws + + +def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=list(frames)) + return backend_ws + + +def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: + logging: Final = _RecordingLogging() + worker: Final = _InlineLoggingWorker() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker + ) + return _RelaySession(streaming=streaming, logging=logging, worker=worker) + + +def _error_events_sent_to(client_ws: MagicMock) -> list[dict]: + events: Final = (json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + return [event for event in events if event.get("type") == "error"] + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert error_event["error"]["type"] == "server_error" + assert "1008" in error_event["error"]["message"] + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.parametrize( + "leaked_detail", + ( + pytest.param("sk-live-abcdef0123456789abcdef0123", id="credential"), + pytest.param("vertex-int.svc.cluster.local", id="internal-hostname"), + pytest.param("/etc/litellm/service-account.json", id="filesystem-path"), + ), +) +@pytest.mark.asyncio +async def test_upstream_close_details_are_scrubbed_before_reaching_the_client(leaked_detail: str): + """LIT-6973: the relayed close goes through the proxy's client-facing redaction, so an upstream + error echoing a credential, an internal host, or a server path never reaches the client verbatim.""" + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, f"upstream rejected: {leaked_detail}"), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert leaked_detail not in error_event["error"]["message"] + relayed_reason: Final = client_ws.close.await_args.kwargs["reason"] + assert leaked_detail not in relayed_reason + assert "REDACTED" in relayed_reason + + +@pytest.mark.asyncio +async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(None, None))) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert "1006" in error_event["error"]["message"] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1011 + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_normal_upstream_close_without_error_event(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(Close(1000, ""), None))) + + await session.run() + + assert _error_events_sent_to(client_ws) == [] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1000 + + +@pytest.mark.asyncio +async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure(): + """LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before + the relay starts. It is not an upstream frame, so a refusal after it is still a refusal.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session.streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "sess_synthetic"}})) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + (logged_session,) = session.logging.logged_sessions + assert [event["type"] for event in logged_session] == ["session.created"] + assert session.logging.logged_failures == () + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.asyncio +async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + backend_closed: Final = asyncio.Event() + client_messages: Final = iter((json.dumps({"type": "response.create"}),)) + + async def receive_text() -> str: + message = next(client_messages, None) + return message if message is not None else await _wait_forever() + + async def send_to_backend(_message: str) -> None: + backend_closed.set() + raise upstream_close + + async def recv_from_backend() -> bytes: + await backend_closed.wait() + raise upstream_close + + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = receive_text + backend_ws: Final = MagicMock() + backend_ws.send = send_to_backend + backend_ws.recv = recv_from_backend + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + assert session.logging.logged_failures == (upstream_close,) + + +@pytest.mark.asyncio +async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client went away")) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the_backend_closing(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_success_logging_stamps_the_reservation_ownership_marker(): + """LIT-6973: only the success path enqueues the cost callback that settles the + session's budget reservation, so it stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on + the shared logging object. The proxy endpoint reads that stamp to decide whether to + release the reservation itself, so a logged-as-success session must carry it.""" + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1000, ""), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + assert session.logging.logged_sessions != () + assert session.logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + +@pytest.mark.asyncio +async def test_refused_session_does_not_stamp_the_reservation_ownership_marker(): + """A refused session logs a failure, not a success, so it must not stamp + REALTIME_SESSION_SUCCESS_LOGGED_KEY. If it did, the proxy endpoint would skip its + own reservation release and the refused session's reservation would stay pinned.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 34216f7e1b9..9393ec0f8e6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9521,6 +9521,222 @@ def test_realtime_websocket_route_aliases_registered(): ) +def _lit6973_fake_realtime_ws() -> MagicMock: + ws = MagicMock() + ws.headers = {} + ws.scope = {"headers": [], "type": "websocket"} + ws.url = "ws://testserver/v1/realtime" + ws.accept = AsyncMock() + ws.send_text = AsyncMock() + ws.close = AsyncMock() + return ws + + +async def _lit6973_drive_realtime_session( + reservation: dict, + *, + backend_logged_success: bool, + phase_one_exit: str | None = None, + websocket: MagicMock | None = None, +) -> MagicMock: + """Drive realtime_websocket_endpoint through one of its reservation-settling exits. + + 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. + + route_request resolves normally in both cases: the relay owns the session + once route_request returns. A successful session enqueues its success cost + callback and stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on the shared logging + object; a refused one does neither. The endpoint keys its reservation cleanup + off that stamp, so backend_logged_success reproduces both branches. The fake + 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.proxy import proxy_server as ps + + user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") + user_api_key_dict.budget_reservation = reservation + + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} + + async def fake_llm_call() -> None: + if backend_logged_success: + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + from litellm.proxy._types import ProxyException + + model_access_error: Final = ( + ProxyException(message="key cannot access model", type="auth_error", param="model", code=401) + 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: Final = AsyncMock( + side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) + ) + ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws() + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test + pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state + route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object + with can_call, pre, route: + await ps.realtime_websocket_endpoint( + websocket=ws, + model="vertex_ai/gemini-live-2.5-flash", + intent=None, + guardrails=None, + user_api_key_dict=user_api_key_dict, + ) + return ws + + +@pytest.mark.asyncio +async def test_refused_realtime_session_releases_the_budget_reservation(): + """LIT-6973: a refused realtime session enqueues no success cost callback, so + the pre-call reservation would stay open and pin the key/team/user spend + counters, locking the key after a couple of refusals. The endpoint sees no + success stamp and reconciles it: the reservation ends up finalized.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_realtime_session(reservation, backend_logged_success=False) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation(): + """A rate-limit or guardrail rejection happens before route_request, so the + relay never runs and no success log can own the reservation. The endpoint + must release it on that exit too, or the key stays pinned at the reserved + amount and its next requests 429 with budget_exceeded while /key/info shows + spend 0 (reproduced live with rpm_limit=1). The client still gets the + pre-call error event and the 1011 close it got before.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call" + ) + + assert reservation["finalized"] is True + assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded" + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + + +@pytest.mark.asyncio +async def test_realtime_session_denied_model_access_releases_the_budget_reservation(): + """The key/model access check rejects before the socket is even accepted; + that exit skipped the release as well, pinning the reservation.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access" + ) + + assert reservation["finalized"] is True + ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") + + +@pytest.mark.asyncio +async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation(): + """The counter release can block on a slow or unreachable store, and a + rejected client must not sit behind it: the relay's own failure path closes + the client first and releases in its finally, so the pre-relay rejection + has to close first as well. The fake close checks the reservation is still + open when the client is closed.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + + async def close_while_reservation_is_still_open(**_: object) -> None: + assert reservation["finalized"] is False, "client was closed only after the reservation release" + + ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open) + + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws + ) + + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone(): + """A client that hung up before the rejection makes the close raise; the + reservation must still be released, or the key stays pinned.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected")) + + with pytest.raises(RuntimeError, match="client already disconnected"): + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws + ) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): + """A billable realtime session settles its reservation through the enqueued + success cost callback, not the endpoint. The endpoint must not finalize it in + its finally, or it would reconcile the reservation to zero before the cost + callback applies real spend, so billable sessions stop counting against budget. + With the success stamp present, the endpoint leaves the reservation untouched.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_realtime_session(reservation, backend_logged_success=True) + + assert reservation["finalized"] is False + + +@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), + the reserved counters must be invalidated directly so the estimate does not + stay pinned, and the reservation is finalized so nothing reprocesses it.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = { + "reserved_cost": 0.55, + "input_cost": 0.0, + "finalized": False, + "entries": [{"counter_key": "spend:key:hashed-token"}], + } + invalidated: Final[list[str]] = [] + + async def _record(counter_key: str) -> None: + invalidated.append(counter_key) + + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated + sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable + with failing_release, sink: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert invalidated == ["spend:key:hashed-token"] + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails(): + """Both counter-store calls failing must not raise out of the realtime + endpoint's finally (it would mask the session's own outcome) and must still + stamp finalized so nothing retries the same reservation.""" + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + + with failing_release, failing_invalidate: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints.