diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9282ee7f703..d3d8c242ff4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -687,7 +687,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch def serialize_http_exception_detail( detail: object, -) -> tuple[str, dict | None]: +) -> tuple[str, Mapping[str, object] | None]: """ Convert an HTTPException.detail value into (message, structured_fields) for ProxyException / SSE error frames. diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index be9f80694b8..1cbfd31e7d6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -153,6 +153,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( create_websocket_passthrough_route, refuse_websocket, truncated_close_reason, + vertex_live_setup_frame_model, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType @@ -178,7 +179,7 @@ from litellm.types.passthrough_endpoints.tinyfish import ( is_allowed_tinyfish_endpoint, ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials -from litellm.types.router import LiteLLMParamsTypedDict +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict from litellm.types.utils import CallTypes, LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -3466,6 +3467,7 @@ class _WebsocketRelay(Protocol): client_frame_model_gate: WebsocketFrameModelGate | None = None, upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, meter: WebSocketSessionMeter | None = None, + model: str | None = None, ) -> None: ... @@ -3718,6 +3720,7 @@ async def deepgram_listen_websocket_route( forward_headers=False, endpoint=websocket.url.path, accept_websocket=False, + model=deepgram_listen_requested_model(websocket.url.query), ) @@ -3766,6 +3769,7 @@ async def deepgram_speak_websocket_route( endpoint=websocket.url.path, accept_websocket=False, meter=DeepgramSpeakMeter(target, character_rate), + model=deepgram_speak_requested_model(websocket.url.query), ) @@ -4171,21 +4175,26 @@ def _build_vertex_live_setup_model_rewriter( return rewrite -def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None) -> str: +def _setup_model_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: """ The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ if llm_router is None: - return setup_model + return None candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1]) - upstream: Final = next( + return next( ( - deployment["litellm_params"].get("model") + deployment for deployment in (llm_router.registry().deployments() or ()) if deployment.get("model_name") in candidates ), None, ) + + +def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None) -> str: + deployment: Final = _setup_model_deployment(setup_model, llm_router) + upstream: Final = None if deployment is None else deployment["litellm_params"].get("model") if upstream is None: return setup_model try: @@ -4195,13 +4204,29 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None return upstream.removeprefix(f"{provider}/") +def _build_vertex_live_session_model(llm_router: Router | None) -> Callable[[str | bytes], str | None]: + """ + Name a Live session's model the way the setup rewriter resolves it, the gateway alias its setup frame names or + else the publisher model id Vertex serves + """ + + def session_model(frame: str | bytes) -> str | None: + setup_model: Final = vertex_live_setup_frame_model(frame) + if setup_model is None: + return None + deployment: Final = _setup_model_deployment(setup_model, llm_router) + return setup_model.rsplit("/", 1)[-1] if deployment is None else deployment["model_name"] + + return session_model + + async def vertex_ai_live_websocket_passthrough( websocket: WebSocket, model: str | None = None, vertex_project: str | None = None, vertex_location: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, -): +) -> None: """ Vertex AI Live API WebSocket Pass-through Function @@ -4301,6 +4326,7 @@ async def vertex_ai_live_websocket_passthrough( accept_websocket=False, builtin_provider_route=True, model=model, + opening_frame_model=_build_vertex_live_session_model(_get_llm_router()), setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project=resolved_project, vertex_location=resolved_location, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7f5983d5799..756f5835d70 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -34,7 +34,7 @@ from fastapi.routing import APIRoute from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Match, Route -from starlette.types import Scope +from starlette.types import Message, Scope from starlette.websockets import WebSocketState from typing_extensions import ReadOnly from websockets.asyncio.client import ClientConnection, connect @@ -2581,6 +2581,23 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def vertex_live_setup_frame_model(frame: str | bytes) -> str | None: + """ + The model a Vertex AI Live ``setup`` frame names, as the client wrote it + """ + message: Final = _json_object_frame_or_none(frame) + setup: Final = None if message is None else as_str_mapping(message.get("setup")) + setup_model: Final = None if setup is None else setup.get("model") + return setup_model if isinstance(setup_model, str) else None + + +def _json_object_frame_or_none(frame: str | bytes) -> Mapping[str, object] | None: + try: + return _json_object_frame(frame) + except (ValueError, RecursionError): + return None + + def _decoded_frame(frame: str | bytes) -> object: try: return _DECODED_JSON.validate_python(orjson.loads(frame)) @@ -2998,6 +3015,16 @@ _EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0 +def _opening_session_model( + message: Message | None, opening_frame_model: Callable[[str | bytes], str | None] | None +) -> str | None: + if message is None or opening_frame_model is None: + return None + text: Final = message.get("text") + frame: Final = text if isinstance(text, str) else message.get("bytes") + return opening_frame_model(frame) if isinstance(frame, str | bytes) else None + + def _metered_client_text(meter: WebSocketSessionMeter | None, client_text: str) -> int: return 0 if meter is None else meter.measure(client_text) @@ -3024,8 +3051,9 @@ async def websocket_passthrough_request( walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS, builtin_provider_route: bool = False, model: str | None = None, + opening_frame_model: Callable[[str | bytes], str | None] | None = None, meter: WebSocketSessionMeter | None = None, -): +) -> None: """ WebSocket passthrough request handler. @@ -3044,7 +3072,10 @@ async def websocket_passthrough_request( walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client left or was refused, so the provider's usage for it still reaches spend builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider - model: The model the caller named before connecting, logged on a failure before the setup frame arrives + model: The model the route checked the key against before connecting. Pre-call hooks such as per-model rate + limits and failure logs use it unless the opening frame names another + opening_frame_model: Reads the model the session's first client frame selects, for an upstream that serves + that model whatever the URL names. The relay waits for that frame before running hooks or dialing meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage """ from litellm.litellm_core_utils.litellm_logging import Logging @@ -3254,15 +3285,42 @@ async def websocket_passthrough_request( log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row) ledger.close_turn(lane, end_time) + async def report_failure(exception: Exception, request_body: Mapping[str, object] | None) -> HTTPException | None: + return await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exception, + request_data=_build_passthrough_failure_request_payload( + parsed_body=None if request_body is None else dict(request_body), + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), + traceback_str=traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG), + ) + + opening_message: Final = None if opening_frame_model is None else await websocket.receive() + session_model: Final = _opening_session_model(opening_message, opening_frame_model) or model + ### CALL HOOKS ### - modify incoming data / reject request before calling the model websocket_hook_data: Final = guardrail_request_data_with_streaming( - MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True - ) - await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, - data=websocket_hook_data, - call_type="pass_through_endpoint", + MappingProxyType( + {"litellm_call_id": litellm_call_id, **({} if session_model is None else {"model": session_model})} + ), + is_streaming=True, ) + try: + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=websocket_hook_data, + call_type="pass_through_endpoint", + ) + except Exception as refusal: + replacement: Final = await report_failure(refusal, websocket_hook_data) + if replacement is not None: + raise replacement from refusal + raise try: verbose_proxy_logger.debug( @@ -3276,12 +3334,17 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Upstream connection established successfully", endpoint ) + async def client_messages() -> AsyncIterator[Message]: + if opening_message is not None: + yield opening_message + while True: + yield await websocket.receive() + async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" nonlocal metered_total try: - while True: - message = await websocket.receive() + async for message in client_messages(): message_type = message.get("type") if message_type == "websocket.disconnect": break @@ -3518,21 +3581,7 @@ async def websocket_passthrough_request( except InvalidStatus as exc: verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint) - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=exc, - request_data=_build_passthrough_failure_request_payload( - parsed_body=None if model is None else {"model": model}, - kwargs=kwargs, - logging_obj=logging_obj, - custom_llm_provider=None, - url=httpx.URL(target), - builtin_provider_route=builtin_provider_route, - ), - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) + await report_failure(exc, None if session_model is None else {"model": session_model}) if _client_socket_is_open(websocket): await websocket.close( @@ -3544,21 +3593,7 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint ) - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=_build_passthrough_failure_request_payload( - parsed_body=None if model is None else {"model": model}, - kwargs=kwargs, - logging_obj=logging_obj, - custom_llm_provider=None, - url=httpx.URL(target), - builtin_provider_route=builtin_provider_route, - ), - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) + await report_failure(e, None if session_model is None else {"model": session_model}) if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb08272e687..a940b0fa51d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -35,6 +35,7 @@ from collections.abc import ( ) from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone +from http.client import responses as http_status_phrases from itertools import chain from types import MappingProxyType, TracebackType, UnionType from typing import ( @@ -168,6 +169,7 @@ from litellm.proxy.common_utils.callback_utils import ( process_callback, strip_callback_config, ) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # legacy module exports _realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export realtime_request_body, @@ -319,6 +321,7 @@ from litellm._internal_context import service_target, with_service_target from litellm._logging import ( STARTUP_OUTPUT_SETTINGS, print_startup_message, + redact_internal_details_from_client_message, redact_string, verbose_proxy_logger, verbose_router_logger, @@ -485,6 +488,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l open_sse_before_first_byte, proxy_exception_from_route_error, resolve_litellm_call_id, + serialize_http_exception_detail, should_return_raw_model_name, sse_error_payload, ttft_keepalive_interval, @@ -2311,19 +2315,45 @@ async def _websocket_http_exception_response(websocket: WebSocket, exc: Starlett match state: case WebSocketState.CONNECTING: return JSONResponse({"detail": exc.detail}, status_code=exc.status_code, headers=exc.headers) - case WebSocketState.CONNECTED: - await websocket.close(code=_websocket_close_code(exc.status_code)) + case WebSocketState.CONNECTED if websocket.client_state is not WebSocketState.DISCONNECTED: + await websocket.close(code=_websocket_close_code(exc.status_code), reason=_websocket_close_reason(exc)) return None - case WebSocketState.RESPONSE | WebSocketState.DISCONNECTED: + case WebSocketState.CONNECTED | WebSocketState.RESPONSE | WebSocketState.DISCONNECTED: return None case _: assert_never(state) def _websocket_close_code(status_code: int) -> int: + if status_code == status.HTTP_429_TOO_MANY_REQUESTS: + return status.WS_1013_TRY_AGAIN_LATER return status.WS_1011_INTERNAL_ERROR if status_code >= 500 else status.WS_1008_POLICY_VIOLATION +def _websocket_close_reason(exc: StarletteHTTPException) -> str: + message: Final = ( + _rate_limit_close_message(exc) + if isinstance(exc, ProxyRateLimitError) + else serialize_http_exception_detail(exc.detail)[0] + ) + return websocket_close_reason( + redact_internal_details_from_client_message(message), + fallback=http_status_phrases.get(exc.status_code, "WebSocket error"), + ) + + +def _rate_limit_close_message(exc: ProxyRateLimitError) -> str: + headers: Final = exc.headers or {} + limit_type: Final = exc.rate_limit_type or headers.get("rate_limit_type") + retry_after: Final = headers.get("retry-after") + return ( + "Rate limit exceeded" + + (f" for {exc.descriptor_key}" if exc.descriptor_key else "") + + (f" ({limit_type})" if limit_type else "") + + (f", retry after {retry_after}s" if retry_after else "") + ) + + def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Exception | None = None) -> None: parent_otel_span: Final[_Span | None] = getattr(request.state, "parent_otel_span", None) if parent_otel_span is None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 5e31906a4b9..dcabe8a530b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -2,33 +2,41 @@ target URL.""" import asyncio -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.routing import WebSocketRoute -from starlette.websockets import WebSocketDisconnect +from starlette.types import Message +from starlette.websockets import WebSocket, WebSocketDisconnect +from websockets.exceptions import ConnectionClosedOK +from websockets.frames import Close import litellm from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import WEBSOCKET_CLOSE_REASON_MAX_BYTES from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map from litellm.llms.deepgram.common_utils import DeepgramSpeakMeter from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import cache_key_object +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _websocket_relay, deepgram_listen_websocket_route, + deepgram_speak_user_api_key_auth, deepgram_speak_websocket_route, router, ) -from litellm.proxy.utils import hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionMeter GET_CREDENTIALS: Final = ( @@ -87,6 +95,7 @@ class _FakeRelay: endpoint: str, accept_websocket: bool, meter: WebSocketSessionMeter | None = None, + model: str | None = None, ) -> None: self.calls.append( _RelayCall( @@ -536,6 +545,174 @@ async def test_deepgram_speak_relays_with_provider_auth_and_meters_the_session( assert websocket.closed is None +def _client_that_hangs_up(path: str, query: str) -> WebSocket: + handshake: Final = iter(({"type": "websocket.connect"},)) + + async def receive() -> Message: + return next(handshake, {"type": "websocket.disconnect", "code": 1000}) + + async def send(message: Message) -> None: + return None + + return WebSocket({"type": "websocket", "path": path, "query_string": query.encode(), "headers": []}, receive, send) + + +class _UpstreamThatClosesAtOnce: + async def recv(self, decode: bool | None = None) -> str: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + + async def send(self, message: str | bytes) -> None: + return None + + async def close(self, code: int = 1000, reason: str = "") -> None: + return None + + +class _UpstreamConnect: + async def __aenter__(self) -> _UpstreamThatClosesAtOnce: + return _UpstreamThatClosesAtOnce() + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _UpstreamDials: + def __init__(self) -> None: + self.targets: tuple[str, ...] = () + + def dial(self, target: str, additional_headers: Mapping[str, str]) -> _UpstreamConnect: + self.targets = (*self.targets, target) + return _UpstreamConnect() + + +def _upstream_dials_under_the_real_rate_limiter(monkeypatch: pytest.MonkeyPatch) -> _UpstreamDials: + from litellm.proxy import proxy_server + + fixed_now: Final = datetime.fromisoformat("2026-10-31T12:00:00+00:00") + limiter: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache(in_memory_cache=InMemoryCache(clock=fixed_now.timestamp))), + time_provider=lambda: fixed_now, + model_group_resolver=lambda _model: None, + ) + monkeypatch.setattr(litellm, "callbacks", [limiter]) + monkeypatch.setattr( + proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=proxy_server.user_api_key_cache) + ) + dials: Final = _UpstreamDials() + monkeypatch.setattr("litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", dials.dial) + return dials + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "path", "limited_model", "first_query", "second_query", "first_target"), + [ + pytest.param( + deepgram_listen_websocket_route, + "/deepgram/v1/listen", + "nova-3", + "model=nova-3", + "model=nova-3&model=nova-2", + "wss://api.deepgram.com/v1/listen?model=nova-3", + id="listen repeated model", + ), + pytest.param( + deepgram_listen_websocket_route, + "/deepgram/v1/listen", + "nova-3", + "", + "", + "wss://api.deepgram.com/v1/listen?model=nova-3", + id="listen default model", + ), + pytest.param( + deepgram_speak_websocket_route, + "/deepgram/v1/speak", + "aura-2-thalia-en", + "model=aura-2-thalia-en", + "model=aura-2-thalia-en&model=aura-asteria-en", + "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", + id="speak repeated model", + ), + pytest.param( + deepgram_speak_websocket_route, + "/deepgram/v1/speak", + "aura-asteria-en", + "", + "", + "wss://api.deepgram.com/v1/speak?model=aura-asteria-en", + id="speak default model", + ), + ], +) +async def test_deepgram_routes_refuse_a_second_session_over_the_keys_model_rpm_limit( + route: Callable[..., Awaitable[None]], + path: str, + limited_model: str, + first_query: str, + second_query: str, + first_target: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + dialed: Final = _upstream_dials_under_the_real_rate_limiter(monkeypatch) + caller: Final = UserAPIKeyAuth( + api_key=hash_token("sk-deepgram-model-rpm"), metadata={"model_rpm_limit": {limited_model: 1}} + ) + + async def session(query: str) -> None: + await route( + websocket=_client_that_hangs_up(path, query), + user_api_key_dict=caller, + relay=_websocket_relay(), + ) + + await asyncio.create_task(session(first_query)) + with pytest.raises(HTTPException) as refused: + await asyncio.create_task(session(second_query)) + + assert (refused.value.status_code, refused.value.detail, dialed.targets) == ( + 429, + f"Rate limit exceeded for model_per_key: {caller.api_key}:{limited_model}. Limit type: requests. " + "Current limit: 1, Remaining: 0. Limit resets at: 2026-10-31 12:01:00 UTC", + (first_target,), + ) + + +@pytest.mark.timeout(10) +def test_deepgram_speak_tells_a_client_over_its_model_rpm_limit_why_it_was_refused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.proxy_server import otlp_http_exception_handler + + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + _upstream_dials_under_the_real_rate_limiter(monkeypatch) + caller: Final = UserAPIKeyAuth( + api_key=hash_token("sk-deepgram-model-rpm"), metadata={"model_rpm_limit": {"aura-2-thalia-en": 1}} + ) + app: Final = FastAPI() + app.include_router(router) + app.exception_handler(StarletteHTTPException)(otlp_http_exception_handler) + app.dependency_overrides[deepgram_speak_user_api_key_auth] = lambda: caller + client: Final = TestClient(app) + + with client.websocket_connect("/deepgram/v1/speak?model=aura-2-thalia-en") as first: + first_close: Final = first.receive() + with client.websocket_connect("/deepgram/v1/speak?model=aura-2-thalia-en") as second: + second_close: Final = second.receive() + + assert (first_close, second_close) == ( + {"type": "websocket.close", "code": 1000, "reason": ""}, + { + "type": "websocket.close", + "code": 1013, + "reason": "Rate limit exceeded for model_per_key (requests), retry after 60s", + }, + ) + + @pytest.mark.asyncio async def test_deepgram_speak_bills_at_the_connect_time_rate_after_a_cost_map_reload_drops_the_row( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4a76aa77f14..853a1cd34b3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6324,7 +6324,8 @@ class TestVertexAILiveWebsocketPassthrough: "vertex_location": "global", "vertex_credentials": '{"type": "service_account"}', }, - } + }, + {"model_name": "live/gemini", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, ] ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) @@ -6358,6 +6359,10 @@ class TestVertexAILiveWebsocketPassthrough: assert rewriter("gemini-live") == ( "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" ) + session_model: Final = passthrough_kwargs["opening_frame_model"] + assert tuple( + session_model(json.dumps({"setup": {"model": model}})) for model in ("models/gemini-live", "live/gemini") + ) == ("gemini-live", "live/gemini") websocket.close.assert_not_awaited() @pytest.mark.parametrize( diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 824f04aee75..1b533d9e906 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2,10 +2,12 @@ import asyncio import base64 import contextvars import gzip +import itertools import json import logging import os import sys +import traceback import uuid import zlib from collections.abc import AsyncIterator, Callable, Mapping, Sequence @@ -32,7 +34,13 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, ANTHROPIC_MAX_USER_ID_LENGTH, DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.constants import ( + ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, + ANTHROPIC_MAX_USER_ID_LENGTH, + DEFAULT_REQUEST_TIMEOUT_SECONDS, + SERVER_STREAMING_CLASSIFICATION_KEY, + SERVER_STREAMING_CLASSIFICATION_MARKER, +) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -44,6 +52,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _build_vertex_live_session_model from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, @@ -7650,6 +7659,283 @@ async def test_websocket_passthrough_relays_upstream_policy_close_to_client(): websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +class _UpstreamDials: + def __init__(self) -> None: + self.targets: tuple[str, ...] = () + + def dial(self, target: str, additional_headers: Mapping[str, str]) -> FakeUpstreamConnect: + self.targets = (*self.targets, target) + return FakeUpstreamConnect(FakeUpstreamWebSocket()) + + +def _record_upstream_dials(monkeypatch: pytest.MonkeyPatch) -> _UpstreamDials: + dials: Final = _UpstreamDials() + monkeypatch.setattr("litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", dials.dial) + return dials + + +async def _relay_one_client_session(caller: UserAPIKeyAuth, model: str | None) -> None: + await websocket_passthrough_request( + websocket=_client_websocket(AsyncMock(return_value={"type": "websocket.disconnect"})), + target="wss://upstream.example.test/v1/speak", + custom_headers={}, + user_api_key_dict=caller, + forward_headers=False, + endpoint="/ws-model-rpm", + accept_websocket=False, + model=model, + ) + + +class _PreCallDataRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.data: tuple[Mapping[str, object], ...] = () + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> None: + self.data = (*self.data, MappingProxyType(dict(data))) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "hook_data"), + [ + pytest.param( + None, + {SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER}, + id="route authorized no model", + ), + pytest.param( + "nova-3", + {"model": "nova-3", SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER}, + id="route authorized nova-3", + ), + ], +) +async def test_websocket_passthrough_hands_pre_call_hooks_exactly_the_model_it_was_given( + model: str | None, hook_data: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import user_api_key_cache + + recorder: Final = _PreCallDataRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=user_api_key_cache)) + _record_upstream_dials(monkeypatch) + + await _relay_one_client_session(UserAPIKeyAuth(), model) + + assert tuple( + {key: value for key, value in data.items() if key != "litellm_call_id"} for data in recorder.data + ) == (hook_data,) + + +class _VertexLiveDials: + def __init__(self) -> None: + self.upstreams: tuple[RecordingUpstreamWebSocket, ...] = () + + def dial(self, target: str, additional_headers: Mapping[str, str]) -> FakeUpstreamConnect: + upstream: Final = RecordingUpstreamWebSocket() + self.upstreams = (*self.upstreams, upstream) + return FakeUpstreamConnect(upstream) + + def frames_sent(self) -> tuple[str | bytes, ...]: + calls: Final = itertools.chain.from_iterable(upstream.send.await_args_list for upstream in self.upstreams) + return tuple(call.args[0] for call in calls) + + +def _vertex_live_setup(model: str) -> str: + return json.dumps({"setup": {"model": model, "generationConfig": {"responseModalities": ["AUDIO"]}}}) + + +def _client_frame_message(frame: str | bytes) -> Mapping[str, object]: + return ( + {"type": "websocket.receive", "text": frame} + if isinstance(frame, str) + else {"type": "websocket.receive", "bytes": frame} + ) + + +_TOO_DEEP_TO_PARSE: Final = b"[" * 100_000 + b"]" * 100_000 +_TOO_LONG_A_NUMBER_TO_PARSE: Final = b'{"setup": NaN, "n": ' + b"1" * 5_000 + b"}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("url_model", "first_frame", "hooked_model"), + [ + pytest.param(None, _vertex_live_setup("gemini-live-2.5-flash"), "gemini-live-2.5-flash", id="no url model"), + pytest.param( + "gemini-2.0-flash-live-001", + _vertex_live_setup("models/gemini-live-2.5-flash"), + "gemini-live-2.5-flash", + id="setup model overrides a different url model", + ), + pytest.param( + None, + _vertex_live_setup("projects/p/locations/us-central1/publishers/google/models/gemini-live-2.5-flash"), + "gemini-live-2.5-flash", + id="full resource path", + ), + pytest.param( + None, + _vertex_live_setup("vertex_ai/gemini-live-2.5-flash"), + "gemini-live-2.5-flash", + id="litellm model id", + ), + pytest.param(None, _vertex_live_setup("models/gemini-live"), "gemini-live", id="gateway alias"), + pytest.param(None, _vertex_live_setup("team/gemini-live"), "gemini-live", id="gateway alias behind a prefix"), + pytest.param(None, _vertex_live_setup("live/gemini"), "live/gemini", id="gateway alias with a slash"), + pytest.param( + None, + _vertex_live_setup("gemini-live-2.5-flash").encode(), + "gemini-live-2.5-flash", + id="setup sent as bytes", + ), + pytest.param( + "gemini-live-2.5-flash", + json.dumps({"realtimeInput": {"text": "hi"}}), + "gemini-live-2.5-flash", + id="first frame names no model", + ), + pytest.param( + "gemini-live-2.5-flash", _TOO_DEEP_TO_PARSE, "gemini-live-2.5-flash", id="first frame too deep to parse" + ), + pytest.param( + "gemini-live-2.5-flash", + _TOO_LONG_A_NUMBER_TO_PARSE, + "gemini-live-2.5-flash", + id="first frame with a number too long to parse", + ), + pytest.param("gemini-live-2.5-flash", None, "gemini-live-2.5-flash", id="client leaves before setup"), + ], +) +async def test_vertex_live_relay_runs_pre_call_hooks_on_the_model_the_setup_frame_selects( + url_model: str | None, first_frame: str | bytes | None, hooked_model: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import user_api_key_cache + + recorder: Final = _PreCallDataRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=user_api_key_cache)) + dials: Final = _VertexLiveDials() + monkeypatch.setattr("litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", dials.dial) + client_frames: Final = () if first_frame is None else (first_frame,) + llm_router: Final = litellm.Router( + model_list=[ + {"model_name": "gemini-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + {"model_name": "live/gemini", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ] + ) + + await websocket_passthrough_request( + websocket=_client_websocket( + AsyncMock(side_effect=[*map(_client_frame_message, client_frames), {"type": "websocket.disconnect"}]) + ), + target="wss://aiplatform.example.test/ws/live", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + model=url_model, + opening_frame_model=_build_vertex_live_session_model(llm_router), + ) + + assert (tuple(data.get("model") for data in recorder.data), dials.frames_sent()) == ( + (hooked_model,), + client_frames, + ) + + +class _RefusingPreCallHook(CustomLogger): + def __init__(self, refusal: HTTPException) -> None: + super().__init__() + self.refusal = refusal + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> None: + data["guardrail_seen"] = "pii-guard" # rebind-ok: the guardrail-style write this test traces + raise self.refusal + + +@pytest.mark.asyncio +async def test_websocket_passthrough_reports_a_pre_call_refusal_to_the_failure_hooks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import user_api_key_cache + + refusal: Final = HTTPException(status_code=429, detail="Rate limit exceeded") + recorder: Final = _FailureHookRecorder() + monkeypatch.setattr(litellm, "callbacks", [_RefusingPreCallHook(refusal), recorder]) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=user_api_key_cache)) + dialed: Final = _record_upstream_dials(monkeypatch) + + with pytest.raises(HTTPException) as raised: + await _relay_one_client_session(UserAPIKeyAuth(), "nova-3") + + assert raised.value is refusal + assert ( + [failure.original_exception for failure in recorder.failures], + [(failure.request_data["model"], failure.request_data["guardrail_seen"]) for failure in recorder.failures], + [str(failure.traceback_str).splitlines()[-1] for failure in recorder.failures], + dialed.targets, + ) == ( + [refusal], + [("nova-3", "pii-guard")], + ["".join(traceback.format_exception_only(refusal)).strip()], + (), + ) + + +class _RewritingFailureHook(CustomLogger): + def __init__(self, replacement: HTTPException) -> None: + super().__init__() + self.replacement = replacement + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> HTTPException: + return self.replacement + + +@pytest.mark.asyncio +async def test_websocket_passthrough_raises_the_failure_hooks_replacement_for_a_refusal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import user_api_key_cache + + refusal: Final = HTTPException(status_code=429, detail="Rate limit exceeded for model_per_key: internal detail") + replacement: Final = HTTPException(status_code=429, detail="Too many sessions, try again later") + monkeypatch.setattr(litellm, "callbacks", [_RefusingPreCallHook(refusal), _RewritingFailureHook(replacement)]) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=user_api_key_cache)) + _record_upstream_dials(monkeypatch) + + with pytest.raises(HTTPException) as raised: + await _relay_one_client_session(UserAPIKeyAuth(), "nova-3") + + assert raised.value is replacement + assert raised.value.__cause__ is refusal + + @pytest.mark.asyncio async def test_websocket_passthrough_keeps_normal_upstream_close_normal(): from websockets.exceptions import ConnectionClosedOK @@ -10856,16 +11142,29 @@ async def test_websocket_passthrough_connect_failure_logs_the_provider_of_the_up @pytest.mark.asyncio @pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"]) -@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")]) +@pytest.mark.parametrize( + "model, setup_model, expected_model", + [ + ("gemini-3.8-live", None, "gemini-3.8-live"), + (None, None, ""), + ("gemini-3.8-live", "models/gemini-live-2.5-flash", "gemini-live-2.5-flash"), + ], +) async def test_websocket_passthrough_connect_failure_logs_the_model_the_caller_named( - failure_kind: str, model: str | None, expected_model: str -): - failure: Final = await _run_websocket_connect_failure(failure_kind, _VERTEX_LIVE_TARGET, True, model=model) + failure_kind: str, model: str | None, setup_model: str | None, expected_model: str +) -> None: + failure: Final = await _run_websocket_connect_failure( + failure_kind, _VERTEX_LIVE_TARGET, True, model=model, setup_model=setup_model + ) assert failure.request_data["model"] == expected_model async def _run_websocket_connect_failure( - failure_kind: str, target: str, builtin_provider_route: bool, model: str | None = None + failure_kind: str, + target: str, + builtin_provider_route: bool, + model: str | None = None, + setup_model: str | None = None, ) -> _RecordedFailure: recorder: Final = _FailureHookRecorder() with ( @@ -10880,7 +11179,11 @@ async def _run_websocket_connect_failure( ): litellm_patch.setattr(litellm, "callbacks", [recorder]) await websocket_passthrough_request( - websocket=_client_websocket(_pending_receive), + websocket=_client_websocket( + _pending_receive + if setup_model is None + else AsyncMock(return_value={"type": "websocket.receive", "text": _vertex_live_setup(setup_model)}) + ), target=target, custom_headers={}, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), @@ -10888,6 +11191,7 @@ async def _run_websocket_connect_failure( accept_websocket=False, builtin_provider_route=builtin_provider_route, model=model, + opening_frame_model=None if setup_model is None else _build_vertex_live_session_model(None), ) assert len(recorder.failures) == 1 return recorder.failures[0] diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index 0f9d379afa1..4832bd7a924 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -25,6 +25,7 @@ from starlette.types import Message from starlette.websockets import WebSocket, WebSocketDisconnect from litellm.proxy._types import ProxyException +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.proxy_server import ( _close_dangling_otel_server_span, openai_exception_handler, @@ -565,7 +566,7 @@ async def test_http_exception_on_a_connecting_websocket_denies_the_upgrade_with_ @pytest.mark.asyncio -@pytest.mark.parametrize(("status_code", "close_code"), [(403, 1008), (429, 1008), (500, 1011), (503, 1011)]) +@pytest.mark.parametrize(("status_code", "close_code"), [(403, 1008), (429, 1013), (500, 1011), (503, 1011)]) async def test_http_exception_on_an_accepted_websocket_closes_it_with_a_matching_code( status_code: int, close_code: int ) -> None: @@ -576,7 +577,108 @@ async def test_http_exception_on_an_accepted_websocket_closes_it_with_a_matching response: Final = await otlp_http_exception_handler(websocket, HTTPException(status_code, "late failure")) assert response is None - assert sent[-1] == {"type": "websocket.close", "code": close_code, "reason": ""} + assert sent[-1] == {"type": "websocket.close", "code": close_code, "reason": "late failure"} + + +class _SentMessages: + def __init__(self) -> None: + self.messages: tuple[Message, ...] = () + + async def send(self, message: Message) -> None: + self.messages = (*self.messages, message) + + +@pytest.mark.asyncio +async def test_http_exception_after_the_client_left_an_accepted_websocket_sends_nothing_more() -> None: + client_messages: Final = iter(({"type": "websocket.connect"}, {"type": "websocket.disconnect", "code": 1000})) + sent: Final = _SentMessages() + + async def receive() -> Message: + return next(client_messages) + + websocket: Final = WebSocket( + {"type": "websocket", "path": "/vertex_ai/live", "headers": [], "query_string": b""}, receive, sent.send + ) + await websocket.accept() + await websocket.receive() + + response: Final = await otlp_http_exception_handler(websocket, HTTPException(429, "Rate limit exceeded")) + + assert (response, sent.messages) == (None, ({"type": "websocket.accept", "subprotocol": None, "headers": []},)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "detail", "close_code", "reason"), + [ + pytest.param( + 429, "é" * 100, 1013, "é" * 61, id="long detail cut to the close frame without splitting a character" + ), + pytest.param( + 400, + {"error": "Violated content safety policy", "guardrail_response": {"flagged": True}}, + 1008, + "Violated content safety policy", + id="structured detail closes with its error message", + ), + pytest.param( + 400, {"error": {"message": "blocked by policy"}}, 1008, "blocked by policy", id="nested error message" + ), + pytest.param( + 503, + "guardrail unreachable at http://guardrail.internal.svc:8080/v1/check", + 1011, + "guardrail unreachable at http://REDACTED.svc:8080/v1/check", + id="internal host in a server error is redacted", + ), + pytest.param( + 400, + "x" * 112 + " sk-Abc123Def456Ghi789Jkl012", + 1008, + "x" * 112 + " REDACTED", + id="key across the close frame limit is redacted before the cut", + ), + pytest.param(429, "", 1013, "Too Many Requests", id="empty detail falls back to the status phrase"), + pytest.param(499, "", 1008, "WebSocket error", id="non-standard status falls back to a generic reason"), + ], +) +async def test_http_exception_on_an_accepted_websocket_closes_it_with_its_detail_as_the_reason( + status_code: int, detail: object, close_code: int, reason: str +) -> None: + sent: Final[list[Message]] = [] + websocket: Final = _websocket(sent) + await websocket.accept() + + await otlp_http_exception_handler(websocket, HTTPException(status_code, detail)) + + assert sent == [ + {"type": "websocket.accept", "subprotocol": None, "headers": []}, + {"type": "websocket.close", "code": close_code, "reason": reason}, + ] + + +@pytest.mark.asyncio +async def test_rate_limit_on_an_accepted_websocket_closes_with_the_limit_and_retry_time() -> None: + sent: Final[list[Message]] = [] + websocket: Final = _websocket(sent) + await websocket.accept() + refusal: Final = ProxyRateLimitError( + detail="Rate limit exceeded for model_per_key: 5f2c:live/gemini. Limit type: requests. Current limit: 1", + headers={"retry-after": "60", "rate_limit_type": "requests"}, + model="gemini-3.8-live", + descriptor_key="model_per_key", + ) + + await otlp_http_exception_handler(websocket, refusal) + + assert sent == [ + {"type": "websocket.accept", "subprotocol": None, "headers": []}, + { + "type": "websocket.close", + "code": 1013, + "reason": "Rate limit exceeded for model_per_key (requests), retry after 60s", + }, + ] def test_websocket_auth_rejection_reaches_the_client_as_a_denial_instead_of_a_server_error() -> None: