From 8855c1385eb40e1a7fa766132b5d09bbd06cb4c3 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Mon, 19 Oct 2026 23:59:53 -0500 Subject: [PATCH] fix(realtime): enforce model budgets and preserve translation routing --- .../litellm_core_utils/realtime_streaming.py | 27 +++++++--- litellm/llms/azure/audio_transcriptions.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 9 +++- litellm/realtime_api/main.py | 2 + .../test_realtime_streaming.py | 44 ++++++++------- .../llms/azure/test_audio_transcriptions.py | 2 +- .../test_user_api_key_auth_request_flow.py | 53 +++++++++++++++++++ tests/unit/realtime_api/test_main.py | 41 ++++++++++++++ 8 files changed, 151 insertions(+), 29 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index ab8e19eb4ca..8b866666c6a 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -727,10 +727,25 @@ class RealTimeStreaming: backend, False if the provider transformation produced no output and the message was effectively dropped. """ - message = await self._apply_nested_transcription_model_policy(message) + from litellm.proxy._types import ProxyException + + try: + authorized_message: Final = await self._apply_nested_transcription_model_policy(message) + except ProxyException as denied: + denial_message: Final = redact_internal_details_from_client_message(denied.message) + try: + await self.websocket.send_text(realtime_error_event(denial_message, error_type="permission_error")) + finally: + await self.websocket.close( + code=1008, + reason=websocket_close_reason( + denial_message, fallback="Nested transcription model is not authorized" + ), + ) + return False if self.provider_config: transformed: Final = self.provider_config.transform_realtime_request( - message, self.model, self.session_configuration_request + authorized_message, self.model, self.session_configuration_request ) sent = False for msg in transformed: @@ -766,14 +781,14 @@ class RealTimeStreaming: self._content_sent_after_setup = True sent = True if sent: - self._capture_translation_input_audio(message) + self._capture_translation_input_audio(authorized_message) return sent - await self.backend_ws.send(message) + await self.backend_ws.send(authorized_message) if self._usage_meter is not None: - sent_event: Final = self._parse_backend_event(message) + sent_event: Final = self._parse_backend_event(authorized_message) if sent_event is not None: self._usage_meter.observe_client_event(sent_event) - self._capture_translation_input_audio(message) + self._capture_translation_input_audio(authorized_message) return True async def _apply_nested_transcription_model_policy(self, message: str) -> str: diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index c47e6832a4e..90c240a611c 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -60,7 +60,7 @@ class AzureAudioTranscription(AzureChatCompletion): provider_specific_entry: Final = ( as_str_mapping(model_info.get("provider_specific_entry")) if model_info is not None else None ) - requires_deployment_api: Final = model_info is None or ( + requires_deployment_api: Final = ( provider_specific_entry is not None and provider_specific_entry.get("transcription_deployment_api") == 1 ) resolved_api_version: Final = ( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d81083c5d8b..b51da241318 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -144,6 +144,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg populate_request_with_path_params, read_raw_json_body, read_request_body, + resolve_realtime_route_model, rewrite_request_model, safe_get_request_headers, safe_get_request_query_params, @@ -1180,7 +1181,13 @@ def _apply_budget_limits_to_end_user_params( async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: - return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model")) + route_path: Final[object] = websocket.scope.get("path") + model: Final = resolve_realtime_route_model( + model=websocket.query_params.get("model"), + intent=websocket.query_params.get("intent"), + is_translation=isinstance(route_path, str) and route_path.endswith("/realtime/translations"), + ) + return await user_api_key_auth_websocket_for_model(websocket, model=model) async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAuth: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f6e63e515cb..c5add19291d 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -289,6 +289,7 @@ async def acreate_realtime_translation_client_secret( model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( model=requested_model_name, + custom_llm_provider=litellm_params.custom_llm_provider, api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) @@ -492,6 +493,7 @@ async def arealtime_translation_calls( model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( model=requested_model_name, + custom_llm_provider=litellm_params.custom_llm_provider, api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 5e8d19fd8a2..4c91b0b1707 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -879,9 +879,7 @@ async def test_usage_meter_snapshot_stored_in_model_call_details(): backend_ws.send = AsyncMock() backend_ws.recv = AsyncMock( side_effect=[ - json.dumps( - {"type": "session.created", "session": {"turn_detection": {"type": None}}} - ).encode(), + json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(), json.dumps( {"type": "response.output_audio.delta", "delta": base64.b64encode(b"x" * 24000).decode()} ).encode(), @@ -1211,10 +1209,11 @@ async def test_transcription_session_update_enforces_authorized_nested_model(): @pytest.mark.asyncio async def test_translation_session_update_rejects_disallowed_nested_transcription_model() -> None: - backend_ws = MagicMock() + client_ws: Final = MagicMock(send_text=AsyncMock(), close=AsyncMock()) + backend_ws: Final = MagicMock() backend_ws.send = AsyncMock() - streaming = RealTimeStreaming( - MagicMock(), + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, MagicMock(), model="gpt-realtime-translate", @@ -1222,24 +1221,29 @@ async def test_translation_session_update_rejects_disallowed_nested_transcriptio translation_session=True, ) - with pytest.raises(Exception, match=r"gpt-live-transcribe.*not available"): - await streaming._send_to_backend( - json.dumps( - { - "type": "session.update", - "session": { - "type": "translation", - "audio": { - "input": { - "transcription": {"model": "gpt-live-transcribe"}, - } - }, + sent: Final = await streaming._send_to_backend( + json.dumps( + { + "type": "session.update", + "session": { + "type": "translation", + "audio": { + "input": { + "transcription": {"model": "gpt-live-transcribe"}, + } }, - } - ) + }, + } ) + ) + assert sent is False backend_ws.send.assert_not_awaited() assert streaming._is_transcription_session is False + error: Final = json.loads(client_ws.send_text.await_args.args[0]) + assert error["type"] == "error" + assert error["error"]["type"] == "permission_error" + assert "gpt-live-transcribe" in error["error"]["message"] + assert client_ws.close.await_args.kwargs["code"] == 1008 @pytest.mark.asyncio diff --git a/tests/unit/llms/azure/test_audio_transcriptions.py b/tests/unit/llms/azure/test_audio_transcriptions.py index d8f1a290d5e..f8a294c7953 100644 --- a/tests/unit/llms/azure/test_audio_transcriptions.py +++ b/tests/unit/llms/azure/test_audio_transcriptions.py @@ -172,7 +172,7 @@ async def test_azure_transcription_parses_response_formats( [ ("whisper-1", "/openai/v1/audio/transcriptions"), ("gpt-transcribe", "/openai/deployments/gpt-transcribe/audio/transcriptions"), - ("custom-transcribe-deployment", "/openai/deployments/custom-transcribe-deployment/audio/transcriptions"), + ("custom-transcribe-deployment", "/openai/v1/audio/transcriptions"), ], ) def test_azure_transcription_alias_uses_model_route( diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 4cdd7fefd28..6ea5693bbe6 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -11936,6 +11936,59 @@ async def test_websocket_auth_hands_the_reservation_to_the_socket_state(): assert websocket.scope["state"]["budget_reservation"] is reservation +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path", ("/realtime/translations", "/v1/realtime/translations", "/openai/v1/realtime/translations") +) +@pytest.mark.parametrize("query", (b"", b"model=gpt-realtime-translate"), ids=("default_model", "explicit_model")) +async def test_translation_websocket_auth_enforces_the_default_model_budget( + monkeypatch: pytest.MonkeyPatch, path: str, query: bytes +) -> None: + from fastapi import WebSocket + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + model: Final = "gpt-realtime-translate" + token: Final = UserAPIKeyAuth( + token="test-key", + models=[model], + model_max_budget={model: {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={model: 10.0}) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=token, model=model) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return token + + attrs: Final = { + **_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate), + "prisma_client": None, + "model_max_budget_limiter": limiter, + } + for name, value in attrs.items(): + monkeypatch.setattr(proxy_server, name, value) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + websocket: Final = WebSocket( + scope={ + "type": "websocket", + "path": path, + "headers": [(b"authorization", b"Bearer sk-test")], + "query_string": query, + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + with pytest.raises(HTTPException) as denied: + await user_api_key_auth_websocket(websocket) + + assert denied.value.status_code == 403 + assert model in str(denied.value.detail) + assert "budget" in str(denied.value.detail).lower() + + @pytest.mark.asyncio async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_one_redis_mget(): from fastapi import Request diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index dac63bf1d8e..26280ad4523 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -1222,6 +1222,47 @@ async def test_openai_realtime_http_routes_send_the_federated_bearer_for_a_keyle assert gateway.authorization == "Bearer federated-bearer" +@pytest.mark.asyncio +@pytest.mark.parametrize("flow", ("client_secret", "calls")) +async def test_translation_http_routes_preserve_explicit_azure_provider(flow: str) -> None: + expected_path: Final = "/openai/v1/realtime/translations/" + { + "client_secret": "client_secrets", + "calls": "calls", + }[flow] + + def send_response(request: httpx.Request) -> httpx.Response: + assert request.url.host == "azure-deployment.example" + assert request.url.path == expected_path + return httpx.Response(200, json={"value": "ek-test", "routed_path": request.url.path}) + + params: Final[Mapping[str, object]] = { + "model": "translation-deployment", + "custom_llm_provider": "azure", + "api_base": "https://azure-deployment.example/openai/v1/", + "api_key": "test-key", + "client": AsyncHTTPHandler(transport=httpx.MockTransport(send_response)), + "litellm_logging_obj": Logging( + model="translation-deployment", + messages=[], + stream=False, + call_type="acreate_realtime_translation_client_secret", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-translation-routing", + function_id="test-translation-routing", + ), + } + response: Final = await ( + realtime_main.acreate_realtime_translation_client_secret.__wrapped__(**params) + if flow == "client_secret" + else realtime_main.arealtime_translation_calls.__wrapped__( + openai_ephemeral_key="ek-test", sdp_body=b"v=0", **params + ) + ) + + assert response.status_code == 200 + assert response.json() == {"value": "ek-test", "routed_path": expected_path} + + @pytest.mark.asyncio @respx.mock async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_keyless_deployment(