From fd137c47f3b4f666279f1703bf802d1b9b6a8ab4 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 2 Nov 2026 10:53:26 -0800 Subject: [PATCH] fix(passthrough): apply per-model rate limits to the Deepgram and Vertex Live WebSocket routes The WebSocket relay ran the pre-call hooks with empty data, so the rate limiter never saw a model and skipped a key's model_rpm_limit and model_tpm_limit. The relay already took a model parameter, which only labeled failure logs, so it now puts that model in the hook data too The Deepgram listen and speak routes now pass the model their auth checked the key against Vertex Live serves the model its first client frame, the setup frame, names, whatever the optional model query param says. So the relay now reads that frame before it runs the hooks or dials Vertex, applies the hooks and failure logs to that model, then forwards the frame as before. The model is named the way the setup rewriter resolves it, the gateway alias the frame names or else the Vertex model id, so models/x, a full resource path and a LiteLLM id count against the same limit as the bare id. A session whose first frame names no model falls back to the query param. Vertex already keeps a socket open while it waits for setup, so waiting for the frame before dialing holds nothing new The shared WebSocket error handler no longer tries to close a socket whose client already left, which uvicorn rejects with an error A refused session used to close with 1008 and no reason, and the refusal never reached the failure hooks. The relay now reports a pre-call refusal to post_call_failure_hook the way HTTP pass-through does, and raises the replacement error a failure hook returns, if any The app's WebSocket error handler, which every relay route shares, now closes an accepted socket with a reason. That reason is the standard detail message, or for a rate limit the kind of limit and the retry time, run through the client redaction helper and cut to fit the close frame. A 429 there now closes with 1013 Try Again Later instead of 1008, matching how an upstream 429 already closes. This changes what the OpenAI, Vertex Live, Deepgram and configured WebSocket pass-through routes send on a refusal after accept. The realtime, live sessions and responses WebSocket routes handle their own errors and are unchanged Resolves LIT-10354 --- litellm/proxy/common_request_processing.py | 2 +- .../llm_passthrough_endpoints.py | 38 ++- .../pass_through_endpoints.py | 117 ++++--- litellm/proxy/proxy_server.py | 36 +- .../test_deepgram_ws_passthrough_routes.py | 185 +++++++++- .../test_llm_pass_through_endpoints.py | 7 +- .../test_pass_through_endpoints.py | 318 +++++++++++++++++- .../proxy_server/test_exception_handlers.py | 106 +++++- 8 files changed, 744 insertions(+), 65 deletions(-) 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: