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..5d0bbffee69 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3466,6 +3466,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 +3719,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 +3768,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), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7f5983d5799..70b4326908f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3025,7 +3025,7 @@ async def websocket_passthrough_request( builtin_provider_route: bool = False, model: str | None = None, meter: WebSocketSessionMeter | None = None, -): +) -> None: """ WebSocket passthrough request handler. @@ -3044,7 +3044,8 @@ 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 apply to it, and a failure before the setup frame arrives is logged with it 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 +3255,37 @@ 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), + ) + ### 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 model is None else {"model": 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( @@ -3518,21 +3541,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 model is None else {"model": model}) if _client_socket_is_open(websocket): await websocket.close( @@ -3544,21 +3553,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 model is None else {"model": 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..5f89690ffd4 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, @@ -2312,7 +2316,7 @@ async def _websocket_http_exception_response(websocket: WebSocket, exc: Starlett 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)) + await websocket.close(code=_websocket_close_code(exc.status_code), reason=_websocket_close_reason(exc)) return None case WebSocketState.RESPONSE | WebSocketState.DISCONNECTED: return None @@ -2321,9 +2325,37 @@ async def _websocket_http_exception_response(websocket: WebSocket, exc: Starlett 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 {} + model: Final = exc.model if isinstance(exc.model, str) else "" + scope: Final = " ".join(part for part in (exc.descriptor_key, model) if part) + 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 {scope}" if scope 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..fbf44501a64 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 aura-2-thalia-en (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_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 824f04aee75..a85be48dca9 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 @@ -6,6 +6,7 @@ import json import logging import os import sys +import traceback import uuid import zlib from collections.abc import AsyncIterator, Callable, Mapping, Sequence @@ -32,7 +33,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 @@ -7650,6 +7657,164 @@ 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 _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 diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index 0f9d379afa1..9064c9c2bb7 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,81 @@ 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"} + + +@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:nova-3. Limit type: requests. Current limit: 1", + headers={"retry-after": "60", "rate_limit_type": "requests"}, + model="nova-3", + 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 nova-3 (requests), retry after 60s", + }, + ] def test_websocket_auth_rejection_reaches_the_client_as_a_denial_instead_of_a_server_error() -> None: