From 7a72a58353aaa7962ccb62a5175fce48e2bb3a64 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:51:48 -0700 Subject: [PATCH] fix(router): re-check deployment budget and RPM on each Responses WebSocket turn --- litellm/llms/custom_httpx/llm_http_handler.py | 5 +- .../proxy/response_api_endpoints/endpoints.py | 10 +- litellm/responses/main.py | 4 +- litellm/responses/streaming_iterator.py | 27 ++++- litellm/router.py | 31 +++++- .../hooks/test_proxy_track_cost_callback.py | 32 ++++++ .../response_api_endpoints/test_endpoints.py | 22 +++- .../test_responses_websocket_all_providers.py | 101 ++++++++++++++++++ .../test_router/test_router_helper_utils.py | 57 ++++++++++ 9 files changed, 276 insertions(+), 13 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 610321fe1f6..ba5947abf44 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2,7 +2,7 @@ import asyncio import inspect import json import ssl -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache, partial from types import MappingProxyType @@ -6806,6 +6806,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs: object, ) -> Exception | None: """ @@ -6838,6 +6839,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + turn_admission=turn_admission, **kwargs, ) await handler.run() @@ -6962,6 +6964,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + turn_admission=turn_admission, ) return await streaming.bidirectional_forward() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b7f57d93a9b..cdaf54e10ce 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1731,12 +1732,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec def _responses_ws_failure_frame(failure: Exception) -> str: - raw_status: Final = getattr(failure, "status_code", None) - status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 + proxy_exception: Final = handle_exception_on_proxy(failure) + status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500 error_type: Final = ( "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error" ) - return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}}) + return json.dumps( + {"type": "error", "status": status, "error": {"type": error_type, "message": proxy_exception.message}} + ) async def _enforce_responses_ws_first_frame_model_auth( @@ -1896,7 +1899,6 @@ async def _responses_websocket_session( request_data=routed_data, ) except Exception as e: - verbose_proxy_logger.exception("Responses WebSocket error") with contextlib.suppress(Exception): await websocket.send_text(_responses_ws_failure_frame(e)) await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f566c8f4e02..269b7367f34 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import partial @@ -2475,6 +2475,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs, ) -> Exception | None: """ @@ -2567,5 +2568,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + turn_admission=turn_admission, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cf68332c07..58f55ac73fd 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1892,6 +1892,13 @@ async def _enforce_frame_project_quota( ) +async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> None: + frame: Final = _json_object_or_none(raw_message) + if turn_admission is None or frame is None or frame.get("type") != "response.create": + return + await turn_admission() + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2062,6 +2069,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -2081,6 +2089,7 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission def _start_turn(self, client_frame: str | None) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time @@ -2522,8 +2531,10 @@ class ResponsesWebSocketStreaming: return json.dumps(evt_obj) if modified else response_str - async def _enforce_or_reject_frame(self, message: str) -> bool: - """Run the per-frame project quota check. + async def _enforce_or_reject_frame( + self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> bool: + """Run the per-frame project quota check, then the deployment's own admission for a turn after the first. On rejection, sends an ``error`` event to the client and reports that the frame must be dropped instead of forwarded, so the connection @@ -2533,6 +2544,7 @@ class ResponsesWebSocketStreaming: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) + await _admit_frame_turn(turn_admission, message) except RateLimitError as e: try: await self.websocket.send_text( @@ -2561,7 +2573,7 @@ class ResponsesWebSocketStreaming: while True: message = await self.websocket.receive_text() - if not await self._enforce_or_reject_frame(message): + if not await self._enforce_or_reject_frame(message, self.turn_admission): continue masked = await self._mask_response_create(message) self._record_client_frame(masked) @@ -2642,6 +2654,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2660,6 +2673,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3027,7 +3041,9 @@ class ManagedResponsesWebSocketHandler: # Core request handler # ------------------------------------------------------------------ - async def _process_response_create(self, raw_message: str) -> None: + async def _process_response_create( + self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, and forward every streaming event to the client. @@ -3067,6 +3083,7 @@ class ManagedResponsesWebSocketHandler: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message ) + await _admit_frame_turn(turn_admission, raw_message) except RateLimitError as e: await self._send_error(str(e), error_type="rate_limit_exceeded") return @@ -3123,7 +3140,7 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc) break - await self._process_response_create(message) + await self._process_response_create(message, self.turn_admission) except Exception as exc: verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) diff --git a/litellm/router.py b/litellm/router.py index 130cdfab309..94debd5b03d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5632,6 +5632,7 @@ class Router: ) custom_llm_provider: Final = provider_for_generic_call(data) + generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") response_kwargs: Final = { **data, @@ -5639,6 +5640,11 @@ class Router: **kwargs, "model": model_name, **_with_router_resolved_session_model(kwargs.get("session"), model_name), + **( + {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)} + if generic_function_name == "_aresponses_websocket" + else {} + ), } # Only set custom_llm_provider if it's not None if custom_llm_provider is not None: @@ -5648,7 +5654,7 @@ class Router: self, deployment, response_kwargs, - surface_for_call(getattr(original_generic_function, "__name__", "")), + surface_for_call(generic_function_name), ) async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): response = await original_generic_function(**compacted_input) @@ -9247,6 +9253,29 @@ class Router: ) raise e + async def async_admit_deployment_turn( + self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> None: + budget_limiter: Final = self.get_router_deployment_budget_limiter() + if budget_limiter is not None: + try: + await budget_limiter.async_filter_deployments( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts + model=model_group, + healthy_deployments=[dict(deployment)], + messages=None, + request_kwargs=dict(request_kwargs), + ) + except ValueError as e: + raise litellm.RateLimitError( + message=str(e), + llm_provider="", + model=model_group, + category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from e + await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict + deployment=dict(deployment), parent_otel_span=None + ) + @contextlib.asynccontextmanager async def _deployment_slot( self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 2f0e8a268c2..ba2af41bf8e 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1105,6 +1105,38 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +class _FailedTrackingAlerts: + def __init__(self) -> None: + self.messages: tuple[str, ...] = () + + async def failed_tracking_alert(self, error_message: str, failing_model: str) -> None: + self.messages = (*self.messages, error_message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "alerted"), [("_aresponses_websocket", True), ("_arealtime", False)]) +async def test_track_cost_callback_alerts_on_a_websocket_turn_with_no_logging_object(call_type: str, alerted: bool): + kwargs: Final = { + "call_type": call_type, + "model": "gpt-5-mini", + "litellm_call_id": "turn-call-id", + "litellm_params": {}, + "stream": False, + } + alerts: Final = _FailedTrackingAlerts() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", alerts): + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.sleep(0) + + assert any(f"call_type: {call_type}" in message for message in alerts.messages) is alerted + + @pytest.mark.asyncio async def test_track_cost_callback_defers_in_progress_background_interaction(): # test-quality-ok: writing no spend row and raising no alert is the whole observable contract of the deferral path """ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index dd7a0ec7a2e..f19a702dc13 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1596,7 +1596,7 @@ class TestManagedResponsesWSFirstMessage: processed: list = [] - async def fake_process(msg: str) -> None: + async def fake_process(msg: str, turn_admission: object = None) -> None: processed.append(msg) handler = ManagedResponsesWebSocketHandler( @@ -1641,6 +1641,26 @@ class TestResponsesWSStreamingFirstMessage: backend_ws.send.assert_awaited_once_with(first) +@pytest.mark.parametrize( + ("failure", "expected_status", "expected_type"), + [ + (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"), + (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"), + (RuntimeError("boom"), 500, "server_error"), + ], + ids=["crossed-budget", "bad-request", "unexpected"], +) +def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure( + failure: Exception, expected_status: int, expected_type: str +): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame + + frame: Final = json.loads(_responses_ws_failure_frame(failure)) + + assert frame["status"] == expected_status + assert frame["error"]["type"] == expected_type + + class TestWSSessionCostTracking: @pytest.mark.asyncio async def test_router_budget_limiter_skips_arealtime_call_type(self): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 2b78aa98b04..92c815aa4ab 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1344,6 +1344,107 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _turn_frames() -> tuple[str, str, str]: + return tuple(json.dumps({"type": "response.create", "input": word}) for word in ("one", "two", "three")) + + +def _admission_refusing_the_next_turn(): + from unittest.mock import AsyncMock + + from litellm.exceptions import RateLimitError + + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + + +class TestWebSocketTurnAdmission: + """A turn after the connection's first is admitted by the deployment again before it reaches the provider.""" + + @pytest.mark.asyncio + async def test_native_relay_drops_a_later_turn_the_deployment_refuses(self): + from unittest.mock import AsyncMock + + first, second, third = _turn_frames() + cancel: Final = json.dumps({"type": "response.cancel"}) + turn_admission: Final = _admission_refusing_the_next_turn() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, cancel, third, Exception("stop")]) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + turn_admission=turn_admission, + ) + + await handler.client_to_backend() + + assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [ + json.loads(first), + json.loads(cancel), + json.loads(third), + ] + assert turn_admission.await_count == 2 + refusal: Final = json.loads(websocket.send_text.await_args.args[0]) + assert refusal["error"]["type"] == "rate_limit_exceeded" + assert "crossed budget" in refusal["error"]["message"] + + @pytest.mark.asyncio + async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch): + from unittest.mock import AsyncMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args, **kwargs): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + turn_admission: Final = _admission_refusing_the_next_turn() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, third, Exception("stop")]) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + turn_admission=turn_admission, + ) + + await handler.run() + + assert len(answered) == 2 + assert "two" not in json.dumps(answered) + assert turn_admission.await_count == 2 + refusals: Final = [ + json.loads(call.args[0]) + for call in websocket.send_text.await_args_list + if json.loads(call.args[0]).get("type") == "error" + ] + assert [refusal["error"]["type"] for refusal in refusals] == ["rate_limit_exceeded"] + + def _deployment_defaults(): from types import MappingProxyType diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 39fa0970aa1..0e046ec48ff 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -526,6 +526,63 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 +async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]: + routed: Final[dict[str, object]] = {} + + async def call(**kwargs: object) -> None: + routed.update(kwargs) + + call.__name__ = function_name + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call) + return routed + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget(): + router = Router( + model_list=[ + { + "model_name": "ws-budget", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "max_budget": 0.001, + "budget_duration": "1d", + }, + "model_info": {"id": "ws-budget-dep"}, + } + ] + ) + assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses") + admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + await admit_turn() + await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002) + + with pytest.raises(litellm.RateLimitError, match="crossed budget"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): + router = Router( + model_list=[ + { + "model_name": "ws-rpm", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1}, + "model_info": {"id": "ws-rpm-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await admit_turn() + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """