scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-18 04:44:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(router): re-check deployment budget and RPM on each Responses WebSocket turn

Browse source
This commit is contained in:
mateo-berri 2026-10-15 22:51:48 -07:00
parent fbe9c3084c
commit 7a72a58353
9 changed files with 276 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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()

10
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -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(

4
litellm/responses/main.py
View file

@ -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,
)

27
litellm/responses/streaming_iterator.py
View file

@ -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)

31
litellm/router.py
View file

@ -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

32
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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
"""

22
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -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):

101
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -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

57
tests/unit/test_router/test_router_helper_utils.py
View file

@ -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():
"""