scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): enforce project ITPM/OTPM quota on every Responses WebSocket frame

Browse source 
The connection-level pre-call hook only ran once per WebSocket
connection, so a project caller could send unlimited high-token
response.create frames after a single minimal reservation. Adds
enforce_project_io_token_quota_for_frame to the v3 rate limiter and
wires it into both the native and managed WebSocket handlers via a
duck-typed litellm.callbacks lookup, so the SDK layer stays free of
proxy imports. A rejected frame gets an error event; the connection
stays open for the client to retry.

Also fixes the RET504 and BLE001 strict-lint-budget violations the
litellm_internal_staging merge introduced in
parallel_request_limiter_v3.py, which were failing the lint check.
This commit is contained in:
Shivi Jain 2026-08-14 21:39:34 +05:30
parent a9227057a1
commit 312d12fe0c
5 changed files with 396 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -252,6 +252,30 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool:
return False
def _collect_ws_project_quota_callbacks() -> list:
"""Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM
enforcement, so the Responses WebSocket loop can charge every
``response.create`` frame, not just the connection's first one.
Uses duck-typing on ``litellm.callbacks`` (rather than importing the
proxy hook directly) to avoid a layering violation (SDK importing from
the proxy layer).
"""
try:
import litellm as _litellm
return [
cb for cb in _litellm.callbacks if callable(getattr(cb, "enforce_project_io_token_quota_for_frame", None))
]
except Exception as exc: # noqa: BLE001 - discovery must not block the connection
verbose_logger.warning(
"Responses WebSocket: failed to collect project quota callbacks — "
"per-frame ITPM/OTPM enforcement will be skipped. Error: %s",
exc,
)
return []
class BaseLLMHTTPHandler:
async def _make_common_async_call(
self,
@ -6168,6 +6192,8 @@ class BaseLLMHTTPHandler:
- Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls
- Forwards events over the websocket connection
"""
_ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks()
if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket():
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
@ -6184,6 +6210,7 @@ class BaseLLMHTTPHandler:
timeout=timeout,
custom_llm_provider=custom_llm_provider,
first_message=first_message,
quota_callbacks=_ws_quota_callbacks,
**kwargs,
)
await handler.run()
@ -6304,6 +6331,7 @@ class BaseLLMHTTPHandler:
first_message=first_message,
guardrail_callbacks=_ws_guardrail_callbacks,
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
quota_callbacks=_ws_quota_callbacks,
authorized_model=model,
)
await streaming.bidirectional_forward()

40
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -680,7 +680,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
config = data.get("config") if "config" in data else data.get("generationConfig")
translated_request = GoogleGenAIAdapter().translate_generate_content_to_completion(
return GoogleGenAIAdapter().translate_generate_content_to_completion(
model=data.get("model") if isinstance(data.get("model"), str) else "",
contents=contents,
config=config if isinstance(config, dict) else None,
@ -690,7 +690,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
toolConfig=data.get("toolConfig"),
tool_config=data.get("tool_config"),
)
return translated_request
@staticmethod
def _get_explicit_output_cap(data: object, call_type: str | None) -> int | None:
@ -2171,6 +2170,41 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
assert itpm_response is not None
return itpm_response, itpm_reserved, 0
async def enforce_project_io_token_quota_for_frame(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: str | None,
estimated_input_tokens: int,
estimated_output_tokens: int,
) -> None:
"""Reserve one WebSocket ``response.create`` frame's tokens against
the caller's project ITPM/OTPM quota.
The Responses WebSocket connection-level pre-call hook only runs once
per connection, but a connection accepts many ``response.create``
frames over its lifetime. Without this, a project caller could send
unlimited high-token generations after a single minimal reservation.
There is no per-frame post-call hook to reconcile against, so --
like the batch rate limiter -- this charges the estimate immediately
and never refunds it.
"""
descriptors: Final[list[RateLimitDescriptor]] = []
self._add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
if not descriptors:
return
response, _itpm_reserved, _otpm_reserved = await self.reserve_io_tokens(
descriptors=descriptors,
estimated_input_tokens=estimated_input_tokens,
estimated_output_tokens=estimated_output_tokens,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, descriptors, requested_model)
def create_organization_rate_limit_descriptor(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None
) -> list[RateLimitDescriptor]:
@ -3858,7 +3892,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
],
)
continue
except Exception as e:
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the plain increment fallback, never a 500
verbose_proxy_logger.warning(
"Window-guarded token adjustment failed for %s: %s",
operation["key"],

112
litellm/responses/streaming_iterator.py
View file

@ -20,7 +20,7 @@ from litellm.constants import (
LITELLM_MAX_STREAMING_DURATION_SECONDS,
STREAM_SSE_DONE_STRING,
)
from litellm.exceptions import MidStreamFallbackError
from litellm.exceptions import MidStreamFallbackError, RateLimitError
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -1326,6 +1326,79 @@ def _build_synthetic_response_events(
from litellm._logging import verbose_logger
# Conservative per-frame output-token floor used when a response.create
# frame omits max_output_tokens, so a project OTPM quota can't be bypassed
# by simply never declaring an output cap.
_FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR: Final = 1024
# Rough chars-per-token ratio for estimating a frame's input tokens without
# resolving a real per-model tokenizer, matching the conservative estimate
# the proxy's own rate limiter uses for the same purpose.
_FRAME_CHARS_PER_TOKEN_ESTIMATE: Final = 4
def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple[int, int | None]:
"""Extract a rough input-token count and any explicit max_output_tokens
from a ``response.create`` frame, handling both wire shapes:
flat: {"type": "response.create", "input": ..., "max_output_tokens": ...}
nested: {"type": "response.create", "response": {"input": ..., "max_output_tokens": ...}}
"""
nested: Final = msg_obj.get("response")
params: Final[Mapping[str, object]] = (
nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
)
text_parts: list[str] = [] # mutable-ok: local accumulator built in one pass, not shared
def _collect_text(value: object) -> None:
if isinstance(value, str):
text_parts.append(value)
elif _is_json_array(value):
for item in value:
if isinstance(item, str):
text_parts.append(item)
elif _is_json_object(item):
_collect_text(item.get("content"))
_collect_text(item.get("text"))
_collect_text(params.get("input"))
_collect_text(params.get("instructions"))
total_chars: Final = sum(len(part) for part in text_parts)
estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0
max_output_tokens = params.get("max_output_tokens")
return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None
async def _enforce_frame_project_quota(
quota_callbacks: Sequence[Any],
user_api_key_dict: UserAPIKeyAuth | None,
model: str | None,
raw_message: str,
) -> None:
"""Charge one response.create frame's estimated tokens against every
registered project ITPM/OTPM quota callback, in isolation from PII
masking / logging so a malformed frame still reaches those callbacks."""
if not quota_callbacks:
return
try:
msg_obj = json.loads(raw_message)
except (json.JSONDecodeError, TypeError):
return
if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create":
return
estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj)
estimated_output_tokens = (
explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR
)
for callback in quota_callbacks:
await callback.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model=model,
estimated_input_tokens=estimated_input_tokens,
estimated_output_tokens=estimated_output_tokens,
)
RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
"response.created",
"response.completed",
@ -1360,6 +1433,7 @@ class ResponsesWebSocketStreaming:
first_message: str | None = None,
guardrail_callbacks: list[Any] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: list[Any] | None = None,
authorized_model: str | None = None,
):
self.websocket = websocket
@ -1372,6 +1446,7 @@ class ResponsesWebSocketStreaming:
self.first_message = first_message
self.guardrail_callbacks: list[Any] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
self.quota_callbacks: list[Any] = quota_callbacks or []
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: str | None = authorized_model
@ -1781,10 +1856,31 @@ 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.
On rejection, sends an ``error`` event to the client and reports that
the frame must be dropped instead of forwarded, so the connection
stays open for the client to retry once the window resets.
"""
try:
await _enforce_frame_project_quota(
self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message
)
except RateLimitError as e:
try:
await self.websocket.send_text(
json.dumps({"type": "error", "error": {"type": "rate_limit_exceeded", "message": str(e)}})
)
except Exception: # noqa: BLE001, S110 - best-effort notification, client may already be gone
pass
return False
return True
async def client_to_backend(self) -> None:
"""Forward response.create events from client to backend."""
try:
if self.first_message is not None:
if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message):
masked_first: Final = await self._mask_response_create(self.first_message)
self._store_input(masked_first)
self._store_event(masked_first)
@ -1792,6 +1888,8 @@ class ResponsesWebSocketStreaming:
while True:
message = await self.websocket.receive_text()
if not await self._enforce_or_reject_frame(message):
continue
masked = await self._mask_response_create(message)
self._store_input(masked)
self._store_event(masked)
@ -1871,6 +1969,7 @@ class ManagedResponsesWebSocketHandler:
timeout: float | None = None,
custom_llm_provider: str | None = None,
first_message: str | None = None,
quota_callbacks: list[Any] | None = None,
**kwargs: object,
) -> None:
self.websocket = websocket
@ -1887,6 +1986,7 @@ class ManagedResponsesWebSocketHandler:
self.custom_llm_provider = custom_llm_provider
self._connection_provider = self._resolve_provider(model) or custom_llm_provider
self.first_message = first_message
self.quota_callbacks: list[Any] = quota_callbacks or []
# 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.
@ -2292,6 +2392,14 @@ class ManagedResponsesWebSocketHandler:
verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc)
return
try:
await _enforce_frame_project_quota(
self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message
)
except RateLimitError as e:
await self._send_error(str(e), error_type="rate_limit_exceeded")
return
call_kwargs: Final = self._build_base_call_kwargs(msg_obj)
call_kwargs["stream"] = True

56
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -3246,6 +3246,62 @@ async def test_project_model_itpm_and_tpm_limits_coexist_v3():
assert "model_per_project_otpm" in descriptor_keys
@pytest.mark.asyncio
async def test_enforce_project_io_token_quota_for_frame_blocks_over_limit_otpm():
"""VERIA regression: the Responses WebSocket connection-level pre-call
hook only runs once, but a connection accepts many response.create
frames. enforce_project_io_token_quota_for_frame is the per-frame check
that closes that gap; it must reserve against the caller's project OTPM
limit and reject once a frame's estimated output tokens exceed it."""
_api_key = hash_token("sk-ws-frame-otpm")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key,
project_id="proj-mantle-ws",
project_metadata={"model_otpm_limit": {"gpt-4o": 50}},
)
await handler.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model="gpt-4o",
estimated_input_tokens=1,
estimated_output_tokens=30,
)
with pytest.raises(HTTPException) as exc:
await handler.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model="gpt-4o",
estimated_input_tokens=1,
estimated_output_tokens=30,
)
assert exc.value.status_code == 429
assert "model_per_project_otpm" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_enforce_project_io_token_quota_for_frame_noop_without_project_limits():
"""A key with no project ITPM/OTPM configured must never be blocked by
the per-frame check (no descriptors to reserve against)."""
_api_key = hash_token("sk-ws-frame-no-limits")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
await handler.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model="gpt-4o",
estimated_input_tokens=10_000_000,
estimated_output_tokens=10_000_000,
)
@pytest.mark.asyncio
async def test_pre_call_hook_keeps_internal_stash_out_of_request_body():
"""Regression for #27001 / #35197: the limiter's per-request bookkeeping

165
tests/test_litellm/responses/test_responses_websocket_all_providers.py
View file

@ -1030,6 +1030,171 @@ class TestWebSocketErrorHandling:
assert "Invalid JSON" in error_event
class TestWebSocketProjectQuotaEnforcement:
"""VERIA regression: the connection-level pre-call hook only runs once,
but a WebSocket connection accepts many response.create frames. Every
frame must be checked against any registered project ITPM/OTPM quota
callback, not just the first one."""
@pytest.mark.asyncio
async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm
from litellm.exceptions import RateLimitError
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
)
aresponses_called = False
async def fake_aresponses(*args, **kwargs):
nonlocal aresponses_called
aresponses_called = True
monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
quota_callback = MagicMock()
quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(
side_effect=RateLimitError(message="project OTPM exceeded", llm_provider="", model="")
)
mock_websocket = MagicMock()
mock_websocket.send_text = AsyncMock()
mock_logging_obj = Logging(
model="test-model",
messages=[],
stream=True,
call_type="aresponses",
start_time=0,
litellm_call_id="test-id",
function_id="test-func",
)
handler = ManagedResponsesWebSocketHandler(
websocket=mock_websocket,
model="test-model",
logging_obj=mock_logging_obj,
quota_callbacks=[quota_callback],
)
await handler._process_response_create(json.dumps({"type": "response.create", "input": "hi"}))
quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
assert aresponses_called is False
mock_websocket.send_text.assert_called_once()
error_event = mock_websocket.send_text.call_args[0][0]
assert "rate_limit_exceeded" in error_event
@pytest.mark.asyncio
async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
)
aresponses_called = False
async def fake_aresponses(*args, **kwargs):
nonlocal aresponses_called
aresponses_called = True
async def _empty():
return
yield
return _empty()
monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
quota_callback = MagicMock()
quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(return_value=None)
mock_websocket = MagicMock()
mock_websocket.send_text = AsyncMock()
mock_logging_obj = Logging(
model="test-model",
messages=[],
stream=True,
call_type="aresponses",
start_time=0,
litellm_call_id="test-id",
function_id="test-func",
)
handler = ManagedResponsesWebSocketHandler(
websocket=mock_websocket,
model="test-model",
logging_obj=mock_logging_obj,
quota_callbacks=[quota_callback],
)
await handler._process_response_create(json.dumps({"type": "response.create", "input": "hi"}))
quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
assert aresponses_called is True
@pytest.mark.asyncio
async def test_native_handler_blocks_frame_rejected_by_quota_callback(self):
from unittest.mock import AsyncMock, MagicMock
from litellm.exceptions import RateLimitError
from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming
quota_callback = MagicMock()
quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(
side_effect=RateLimitError(message="project OTPM exceeded", llm_provider="", model="")
)
mock_backend_ws = MagicMock()
mock_backend_ws.send = AsyncMock()
mock_websocket = MagicMock()
mock_websocket.send_text = AsyncMock()
handler = ResponsesWebSocketStreaming(
websocket=mock_websocket,
backend_ws=mock_backend_ws,
logging_obj=MagicMock(),
authorized_model="gpt-4o",
quota_callbacks=[quota_callback],
)
allowed = await handler._enforce_or_reject_frame(
json.dumps({"type": "response.create", "input": "hi"})
)
assert allowed is False
mock_backend_ws.send.assert_not_called()
mock_websocket.send_text.assert_called_once()
assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0]
@pytest.mark.asyncio
async def test_native_handler_forwards_frame_allowed_by_quota_callback(self):
from unittest.mock import AsyncMock, MagicMock
from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming
quota_callback = MagicMock()
quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(return_value=None)
handler = ResponsesWebSocketStreaming(
websocket=MagicMock(),
backend_ws=MagicMock(),
logging_obj=MagicMock(),
authorized_model="gpt-4o",
quota_callbacks=[quota_callback],
)
allowed = await handler._enforce_or_reject_frame(
json.dumps({"type": "response.create", "input": "hi"})
)
assert allowed is True
quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
class TestNativeWebSocketGuardrails:
@pytest.mark.asyncio
async def test_response_create_injects_authorized_model(self):