scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 312

fix(proxy): count native Responses WebSocket turn tokens toward key and tag TPM limits (#49257)

Browse source 
* fix(proxy): count Responses WebSocket turn tokens toward key and tag TPM limits

A native Responses WebSocket turn logs its raw events instead of a usage-carrying
response object, so the v3 rate limiter read no usage and never charged the key
or tag token window. Read the turn's usage from its events the way the cost
calculator does, and charge the TPM-limited tags the socket was opened under.

* fix(proxy): refuse socket turns past TPM and count managed, failed, and realtime usage

A later response.create turn on an open Responses WebSocket now gets a read-only
check of every TPM window it counts toward (key, team, user, org, end user, team
member, model-scoped limits and the socket's tags). A turn with no tokens left gets
the 429 rate_limit_exceeded frame and the socket stays open. The check never
reserves or increments a counter and leaves RPM, budget, PTU and project windows
to their own paths. It runs on native and managed sockets.

Managed turns carry the socket's call id in their metadata, so they charge the
socket's TPM-limited tags. A failed native turn charges the partial usage its
failure row already carries. A socket refused at open by a rate limit gets the
same refusal frame and closes 1008, logged once without a traceback.
/v1/realtime sessions now count toward TPM through the realtime usage combiner,
and the v1 limiter reads socket usage through the same helper

* fix(proxy): count failed socket turns, keep socket usage off team PTU, check agent windows per turn

A native socket turn that ends in response.failed now counts the usage that event reports toward spend and TPM, as the HTTP stream path does. Native socket and realtime usage no longer charge the team PTU counter, matching the merge base. The per-turn TPM check now also reads the agent and agent session windows the socket opened with.

* fix(proxy): start missing legacy user, team, and end-user TPM buckets at zero

* fix(proxy): stop charging a realtime session's whole usage into its close minute

Realtime and live sessions log once at close, so counting them put a long session's tokens into one TPM window and refused calls that stayed under the limit. They count no TPM again, as before this branch. Also covers the legacy v1 count-once fix with an HTTP regression case.

* test(proxy): check each per-turn TPM window on its own

* fix(proxy): send a fail-closed per-turn socket refusal with its 503 status

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules 2026-10-31 02:53:27 -07:00 • committed by GitHub
parent da71a03f6e
commit cee72ffdd6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1164 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -1691,6 +1691,7 @@ BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: Final = "responses_websocket_call_id"
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"

2
litellm/cost_calculator.py
View file

@ -3319,7 +3319,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
class _ResponsesWsEventResponse(LiteLLMBaseModel):

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

@ -131,6 +131,8 @@ from litellm.responses.streaming_iterator import (
ResponsesAPIStreamingIterator,
ResponsesWebSocketStreaming,
SyncResponsesAPIStreamingIterator,
TokenLimitCallback,
with_token_limit_checks,
)
from litellm.types.containers.main import (
ContainerFileListResponse,
@ -598,10 +600,10 @@ def _decoded_body_headers(response: httpx.Response) -> httpx.Headers:
)
def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
"""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.
def _ws_callbacks_exposing(method_name: str) -> tuple[object, ...]:
"""Duck-type discover proxy hooks exposing a per-frame check, so the
Responses WebSocket loop can run it on 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
@ -612,13 +614,28 @@ def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
callbacks: Final = cast( # cast-ok: callback registry is inspected before protocol use
Sequence[object], _litellm.callbacks
)
return tuple(callback for callback in callbacks if callable(getattr(callback, method_name, None)))
def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
return tuple(
cast(ProjectQuotaCallback, callback) # cast-ok: required callback method is callable
for callback in callbacks
if callable(getattr(callback, "enforce_project_io_token_quota_for_frame", None))
for callback in _ws_callbacks_exposing("enforce_project_io_token_quota_for_frame")
)
def _collect_ws_token_limit_callbacks() -> tuple[TokenLimitCallback, ...]:
return tuple(
cast(TokenLimitCallback, callback) # cast-ok: required callback method is callable
for callback in _ws_callbacks_exposing("enforce_token_limits_for_frame")
)
def _ws_model_group(litellm_metadata: Mapping[str, object] | None, model: str) -> str:
model_group: Final = (litellm_metadata or {}).get("model_group")
return model_group if isinstance(model_group, str) and model_group else model
def _unwrap_provider_error(error: Exception) -> Exception:
if not isinstance(error, httpx.TransportError):
return error
@ -7285,6 +7302,12 @@ class BaseLLMHTTPHandler:
- Forwards events over the websocket connection
"""
_ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks()
_ws_turn_admission: Final = with_token_limit_checks(
turn_admission,
_collect_ws_token_limit_callbacks(),
user_api_key_dict,
_ws_model_group(litellm_metadata, model),
)
_ws_guardrail_callbacks: list = []
_ws_output_guardrail_callbacks: list = []
@ -7332,7 +7355,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
first_message=first_message,
quota_callbacks=_ws_quota_callbacks,
turn_admission=turn_admission,
turn_admission=_ws_turn_admission,
guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks
frame_guard=frame_guard,
**kwargs,
@ -7429,7 +7452,7 @@ class BaseLLMHTTPHandler:
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
turn_admission=turn_admission,
turn_admission=_ws_turn_admission,
frame_guard=frame_guard,
)
return await streaming.bidirectional_forward()

29
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -21,7 +21,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.budget_throttle import throttled_limit
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit, socket_session_usage
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.types.utils import CallTypes, Usage
@ -41,10 +41,12 @@ def _precise_minute(now: datetime) -> str:
return now.strftime("%Y-%m-%d-%H-%M")
def _response_total_tokens(response_obj: object) -> int:
if not isinstance(response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse)):
return 0
response_usage: Final = getattr(response_obj, "usage", None)
def _response_total_tokens(kwargs: Mapping[str, object], response_obj: object) -> int:
response_usage: Final = (
getattr(response_obj, "usage", None)
if isinstance(response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse))
else socket_session_usage(kwargs, response_obj)
)
return response_usage.total_tokens if isinstance(response_usage, Usage) else 0
@ -556,7 +558,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
precise_minute: Final = _precise_minute(self._clock())
total_tokens: int = _response_total_tokens(response_obj)
total_tokens: Final = _response_total_tokens(
kwargs, # pyright: ignore[reportUnknownArgumentType] # legacy callback kwargs are untyped
response_obj,
)
# ------------
# Update usage - API Key
@ -630,8 +635,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# Update usage - User
# ------------
if user_api_key_user_id is not None:
total_tokens = _response_total_tokens(response_obj)
request_count_api_key = f"{user_api_key_user_id}::{precise_minute}::request_count"
current = await self.internal_usage_cache.async_get_cache(
@ -639,7 +642,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_tpm": 0,
"current_rpm": 1,
}
@ -656,8 +659,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# Update usage - Team
# ------------
if user_api_key_team_id is not None:
total_tokens = _response_total_tokens(response_obj)
request_count_api_key = f"{user_api_key_team_id}::{precise_minute}::request_count"
current = await self.internal_usage_cache.async_get_cache(
@ -665,7 +666,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_tpm": 0,
"current_rpm": 1,
}
@ -682,8 +683,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# Update usage - End User
# ------------
if user_api_key_end_user_id is not None:
total_tokens = _response_total_tokens(response_obj)
request_count_api_key = f"{user_api_key_end_user_id}::{precise_minute}::request_count"
current = await self.internal_usage_cache.async_get_cache(
@ -691,7 +690,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_tpm": 0,
"current_rpm": 1,
}

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

@ -46,8 +46,10 @@ from litellm.constants import (
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION,
RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY,
)
from litellm.cost_calculator import get_usage_object
from litellm.exceptions import ServiceUnavailableError
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
@ -85,7 +87,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
batch_response_view,
canonical_provider_batch_id,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit, socket_session_usage
from litellm.proxy.spend_tracking.input_tokens import (
inline_pdf_file_count,
models_counted_for,
@ -745,6 +747,7 @@ class RequestRateLimiterStash:
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
reservation_released: bool = False
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
socket_agent_windows: tuple[RateLimitDescriptor, ...] = ()
@dataclass(frozen=True, slots=True)
@ -864,6 +867,62 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
return call_id if isinstance(call_id, str) else None
def _socket_connection_call_id(kwargs: Mapping[str, object]) -> object:
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
litellm_metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("litellm_metadata"))
return None if litellm_metadata is None else litellm_metadata.get(RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY)
def _is_socket_turn(kwargs: Mapping[str, object]) -> bool:
if kwargs.get("call_type") == CallTypes.aresponses_websocket.value:
return True
connection_stash: Final = get_request_stash()
connection_call_id: Final = _socket_connection_call_id(kwargs)
return (
connection_stash is not None
and connection_call_id is not None
and connection_call_id == connection_stash.owner_litellm_call_id
)
_AGENT_DESCRIPTOR_KEYS: Final = frozenset({"agent", "agent_session"})
_SOCKET_TURN_UNCOUNTED_DESCRIPTOR_KEYS: Final = frozenset(
{PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, PTU_TEAM_DESCRIPTOR_KEY}
)
def _socket_turn_tpm_windows(descriptors: Sequence[RateLimitDescriptor]) -> tuple[RateLimitDescriptor, ...]:
return tuple(
RateLimitDescriptor(
key=descriptor["key"],
value=descriptor["value"],
rate_limit=RateLimitDescriptorRateLimitObject(
tokens_per_unit=tokens_per_unit, window_size=(descriptor.get("rate_limit") or {}).get("window_size")
),
)
for descriptor in descriptors
if descriptor["key"] not in _SOCKET_TURN_UNCOUNTED_DESCRIPTOR_KEYS
and (tokens_per_unit := (descriptor.get("rate_limit") or {}).get("tokens_per_unit")) is not None
)
def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None:
combined_usage: Final = kwargs.get("combined_usage_object")
if isinstance(combined_usage, Usage):
return combined_usage
return socket_session_usage(kwargs, response_obj)
def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None, kwargs: Mapping[str, object]) -> frozenset[str]:
if call_stash is not None:
return call_stash.tpm_limited_tags
connection_stash: Final = get_request_stash()
if connection_stash is None or not _is_socket_turn(kwargs):
return frozenset()
return connection_stash.tpm_limited_tags
def _parse_output_cap_value(raw_value: object) -> int | None:
if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float, str)):
return None
@ -2967,6 +3026,40 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, descriptors, requested_model)
@with_service_target("rate_limits")
async def enforce_token_limits_for_frame(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
) -> None:
connection_stash: Final = get_request_stash()
tags: Final = sorted(connection_stash.tpm_limited_tags) if connection_stash is not None else []
agent_windows: Final = connection_stash.socket_agent_windows if connection_stash is not None else ()
hidden_stash: Final = _request_stash.set(None)
try:
descriptors: Final = await self._build_request_rate_limit_descriptors(
user_api_key_dict,
MappingProxyType({"model": requested_model, "metadata": {"tags": tags}}),
CallTypes.aresponses_websocket.value,
)
finally:
_request_stash.reset(hidden_stash)
built_scopes: Final = frozenset((d["key"], d["value"]) for d in descriptors)
tpm_windows: Final = _socket_turn_tpm_windows(
(*descriptors, *(d for d in agent_windows if (d["key"], d["value"]) not in built_scopes))
)
try:
response: Final = await self.should_rate_limit(
tpm_windows,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True,
skip_parallel_check=True,
)
except RateLimitUnverifiableError as unverifiable:
raise ServiceUnavailableError(
message=RATE_LIMIT_UNVERIFIABLE_MESSAGE, llm_provider="litellm_proxy", model=requested_model or ""
) from unverifiable
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, tpm_windows, requested_model)
def _rate_limited_model(self, requested_model: str | None) -> RateLimitedModel | None:
if not requested_model:
return None
@ -3777,7 +3870,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _handle_rate_limit_error(
self,
response: RateLimitResponse,
descriptors: list[RateLimitDescriptor],
descriptors: Sequence[RateLimitDescriptor],
requested_model: str | None = None,
) -> None:
"""Handle rate limit exceeded by raising :class:`ProxyRateLimitError` (a 429)."""
@ -4329,6 +4422,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
for d in descriptors
if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None
)
stash.socket_agent_windows = tuple(d for d in descriptors if d["key"] in _AGENT_DESCRIPTOR_KEYS)
# Only check rate limits if we have descriptors with actual limits
if descriptors:
@ -5286,9 +5380,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Get total tokens from response. Responses LiteLLM does not model
# (e.g. pass-through, whose usage is reported by the upstream rather
# than parsed out of the body) carry their usage in
# ``combined_usage_object`` instead, and would otherwise never charge
# the TPM window.
_combined_usage: Final = kwargs.get("combined_usage_object")
# ``combined_usage_object`` instead (a Responses WebSocket turn in its
# raw events) and would otherwise never charge the TPM window.
_usage: Final[Usage | dict[str, object] | None] = (
getattr(response_obj, "usage", None)
if isinstance(
@ -5302,9 +5395,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
else get_usage_object(response_obj)
if isinstance(response_obj, ExtractionResponse)
else _combined_usage
if isinstance(_combined_usage, Usage)
else None
else _unmodeled_response_usage(kwargs, response_obj)
)
total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
if total_tokens == 0:
@ -5337,7 +5428,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
targets: Final = self._collect_tpm_scope_targets(
standard_logging_metadata=standard_logging_metadata,
kwargs=kwargs,
tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(),
tpm_limited_tags=_tpm_limited_tags(stash, kwargs),
model_group=reconcile_model.group if reconcile_model is not None else None,
)
targets.extend(
@ -5365,16 +5456,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
reserved_tokens=reserved_tokens,
)
)
ptu_skips_socket_usage: Final = kwargs.get("call_type") == CallTypes.aresponses_websocket.value
pipeline_operations.extend(
self._build_team_ptu_tpm_ops(
standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata
usage_source=response_obj if _usage is None else _usage,
usage_source=response_obj if _usage is None or ptu_skips_socket_usage else _usage,
reconcile_model=reconcile_model,
reserved_scopes=reserved_scopes,
reserved_ceiling=stash.ptu_ceiling if stash is not None else None,
reserved_tokens=stash.ptu_reserved_tokens if stash is not None else 0,
reserved_windows=stash.ptu_reserved_window_identities if stash is not None else frozenset(),
total_tokens=total_tokens,
total_tokens=0 if ptu_skips_socket_usage else total_tokens,
)
)
@ -5469,27 +5561,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else None,
)
pipeline_operations: Final = self._build_success_event_pipeline_operations(
kwargs=kwargs,
response_obj=response_obj,
rate_limit_type=rate_limit_type,
await self._apply_usage_operations(
self._build_success_event_pipeline_operations(
kwargs=kwargs,
response_obj=response_obj,
rate_limit_type=rate_limit_type,
),
litellm_parent_otel_span,
)
unguarded_operations: Final = [
operation for operation in pipeline_operations if _as_window_guarded(operation) is None
]
window_guarded_operations: Final = tuple(
guarded for operation in pipeline_operations if (guarded := _as_window_guarded(operation)) is not None
)
if unguarded_operations:
await self.async_increment_tokens_with_ttl_preservation(
pipeline_operations=unguarded_operations,
parent_otel_span=litellm_parent_otel_span,
)
if window_guarded_operations:
await self.async_increment_reservation_aware_tokens(
pipeline_operations=window_guarded_operations,
parent_otel_span=litellm_parent_otel_span,
)
io_token_operations: Final = self._build_io_token_reservation_ops(
kwargs=kwargs,
response_obj=response_obj,
@ -5509,6 +5588,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit success event: %s", e)
async def _apply_usage_operations(
self, pipeline_operations: Sequence[RedisPipelineIncrementOperation], parent_otel_span: Span | None
) -> None:
unguarded_operations: Final = [
operation for operation in pipeline_operations if _as_window_guarded(operation) is None
]
window_guarded_operations: Final = tuple(
guarded for operation in pipeline_operations if (guarded := _as_window_guarded(operation)) is not None
)
if unguarded_operations:
await self.async_increment_tokens_with_ttl_preservation(
pipeline_operations=unguarded_operations,
parent_otel_span=parent_otel_span,
)
if window_guarded_operations:
await self.async_increment_reservation_aware_tokens(
pipeline_operations=window_guarded_operations,
parent_otel_span=parent_otel_span,
)
@with_service_target("rate_limits")
async def async_logging_hook(
self,
@ -5622,11 +5721,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
try:
litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs)
failed_call: Final = _callback_kwargs_mapping(kwargs)
litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(failed_call)
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failed_call))
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True)
# Skip the reservation refund if async_post_call_failure_hook
@ -5638,7 +5738,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if stash is None or stash.reservation_released
else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(kwargs)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(failed_call)
if stash is not None and reserved_tokens > 0:
verbose_proxy_logger.debug(
@ -5722,6 +5822,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0):
stash.reservation_released = True
if stash is None and tpm_actual > 0 and _is_socket_turn(failed_call):
await self._apply_usage_operations(
self._build_success_event_pipeline_operations(
kwargs=dict(failed_call), response_obj=None, rate_limit_type=self.get_rate_limit_type()
),
litellm_parent_otel_span,
)
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit failure event: %s", e)

20
litellm/proxy/hooks/rate_limiter_utils.py
View file

@ -2,14 +2,32 @@
Shared utility functions for rate limiter hooks.
"""
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
from litellm.litellm_core_utils.core_helpers import iter_items
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import PriorityReservationDict
from litellm.types.utils import CallTypes, PriorityReservationDict, Usage
_SOCKET_SESSION_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def socket_session_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None:
if not isinstance(response_obj, list):
return None
match kwargs.get("call_type"):
case CallTypes.aresponses_websocket.value:
return ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
_SOCKET_SESSION_EVENTS.validate_python(response_obj)
)
case _:
return None
def resolve_llm_provider_for_rate_limit(

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

@ -36,7 +36,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
from litellm.caching.dual_cache import DualCache
from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES
from litellm.exceptions import EndpointNotSupportedByModelError
from litellm.exceptions import EndpointNotSupportedByModelError, RateLimitError
from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
@ -103,6 +103,7 @@ from litellm.responses.streaming_iterator import (
ResponsesWsFrameGuard,
ResponsesWsFrameRejection,
spawn_logging_task,
turn_refusal_frame,
)
from litellm.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
@ -2627,6 +2628,12 @@ async def _responses_websocket_session(
route_type="_aresponses_websocket",
skip_guardrails=True,
)
except RateLimitError as e:
verbose_proxy_logger.info("Responses WebSocket refused at open by a rate limit: %s", e)
with contextlib.suppress(WebSocketDisconnect, RuntimeError, OSError):
await websocket.send_text(turn_refusal_frame(e))
await websocket.close(code=1008, reason="Rate limit exceeded")
return
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket pre-call error")
try:

45
litellm/responses/streaming_iterator.py
View file

@ -35,9 +35,10 @@ import litellm
from litellm.constants import (
EMPTY_MAPPING,
LITELLM_MAX_STREAMING_DURATION_SECONDS,
RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY,
STREAM_SSE_DONE_STRING,
)
from litellm.exceptions import MidStreamFallbackError, RateLimitError
from litellm.exceptions import MidStreamFallbackError, RateLimitError, ServiceUnavailableError
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import (
process_response_headers,
@ -100,6 +101,12 @@ class ProjectQuotaCallback(Protocol):
) -> None: ...
class TokenLimitCallback(Protocol):
async def enforce_token_limits_for_frame(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
) -> None: ...
@lru_cache(maxsize=1)
def get_openai_response_types():
from litellm.types.llms import openai as openai_types
@ -2158,6 +2165,25 @@ class _RefusedTurn:
ResponsesWsTurnAdmission: TypeAlias = Callable[[], Awaitable[Mapping[str, object]]]
def with_token_limit_checks(
turn_admission: ResponsesWsTurnAdmission | None,
token_limit_callbacks: Sequence[TokenLimitCallback],
user_api_key_dict: UserAPIKeyAuth | None,
requested_model: str | None,
) -> ResponsesWsTurnAdmission | None:
if not token_limit_callbacks or user_api_key_dict is None:
return turn_admission
async def admit() -> Mapping[str, object]:
for callback in token_limit_callbacks:
await callback.enforce_token_limits_for_frame(
user_api_key_dict=user_api_key_dict, requested_model=requested_model
)
return MappingProxyType({}) if turn_admission is None else await turn_admission()
return admit
async def _admit_frame_turn(
turn_admission: ResponsesWsTurnAdmission | None, raw_message: str
) -> _AdmittedTurn | _RefusedTurn:
@ -2173,22 +2199,24 @@ async def _admit_frame_turn(
return _RefusedTurn(e)
def _turn_refusal_frame(refusal: RateLimitError) -> str:
def turn_refusal_frame(refusal: RateLimitError) -> str:
return json.dumps(
{"type": "error", "status": 429, "error": {"type": "rate_limit_exceeded", "message": str(refusal)}}
)
def _turn_admission_error_frame(failure: Exception) -> str:
return json.dumps({"type": "error", "status": 500, "error": {"type": "server_error", "message": str(failure)}})
def _turn_admission_error_frame(failure: Exception, status: int) -> str:
return json.dumps({"type": "error", "status": status, "error": {"type": "server_error", "message": str(failure)}})
def _turn_rejection_frame(rejection: Exception) -> str:
match rejection:
case RateLimitError():
return _turn_refusal_frame(rejection)
return turn_refusal_frame(rejection)
case ServiceUnavailableError():
return _turn_admission_error_frame(rejection, rejection.status_code)
case _:
return _turn_admission_error_frame(rejection)
return _turn_admission_error_frame(rejection, 500)
RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
@ -3015,7 +3043,7 @@ class ResponsesWebSocketStreaming:
async def _send_turn_refusal(self, refusal: RateLimitError) -> None:
try:
await self.websocket.send_text(_turn_refusal_frame(refusal))
await self.websocket.send_text(turn_refusal_frame(refusal))
except Exception: # noqa: BLE001, S110 # client may already be gone
pass
@ -3569,7 +3597,7 @@ class ManagedResponsesWebSocketHandler:
self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message
)
except RateLimitError as e:
await self._send_frame(_turn_refusal_frame(e))
await self._send_frame(turn_refusal_frame(e))
return
guarded: Final = await self._guard_frame(raw_message)
if guarded is None:
@ -3613,6 +3641,7 @@ class ManagedResponsesWebSocketHandler:
call_kwargs["litellm_metadata"] = {
**(turn_litellm_metadata if _is_json_object(turn_litellm_metadata) else {}),
**verdict.reservation,
RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: self.logging_obj.litellm_call_id,
}
try:

171
tests/integration/routing/test_responses_websocket_turn_admission.py
View file

@ -71,6 +71,13 @@ class FaultingTurnCheck(CustomLogger):
instance = FaultingTurnCheck()
"""
HEAVY_TURN_TOKENS: Final = 600
HEAVY_TURN_USAGE: Final[dict[str, JsonValue]] = {
"input_tokens": 400,
"output_tokens": 200,
"total_tokens": HEAVY_TURN_TOKENS,
}
SOCKET_TPM_LIMIT: Final = 1000
FrameReader = Callable[[], int]
@ -124,7 +131,16 @@ def _admitting_proxy(
(directory / f"{module}.py").write_text(source)
config: Final = _config(directory, router_settings, tuple(f"{module}.instance" for module in callback_sources))
with gateway_from_environment() as base:
with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=config) as proxy:
with owned_proxy(
base,
directory,
{
"SSL_CERT_FILE": str(cert[0]),
"REDIS_HOST": os.environ["REDIS_HOST"],
"REDIS_PORT": os.environ["REDIS_PORT"],
},
config=config,
) as proxy:
yield proxy
@ -659,6 +675,159 @@ def test_usage_based_routing_refuses_a_turn_only_once_earlier_turns_used_up_the_
assert deployment.forwarded_turns() == 2
def _heavy_turn(response_id: str) -> tuple[dict[str, JsonValue], ...]:
*started, done = events(response_id, "heavy")
return (*started, {**done, "response": {**object_value(done["response"]), "usage": HEAVY_TURN_USAGE}})
def _counted_tokens(counter: str) -> int:
with _redis() as redis:
return int(redis.get(counter) or 0)
def _chat(gateway: Gateway, key: str, model: str, headers: Mapping[str, str] | None = None) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"after the socket {uuid.uuid4().hex}"}]},
key=key,
headers=headers,
)
def _two_heavy_turns_then_chat(
gateway: Gateway,
cert: tuple[Path, Path],
counter: Callable[[str], str],
headers: Mapping[str, str],
**limits: JsonValue,
) -> httpx.Response:
with responses_peer(cert, _heavy_turn) as peer, gateway.scenario() as scenario:
deployment: Final = _native(scenario, peer)
chat: Final = scenario.model()
key: Final = scenario.key(models=[deployment.name, chat], **limits)
_fresh_minute()
async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
async with _socket(gateway, key, deployment.name, headers) as connection:
return (await _say(connection, deployment.name), await _say(connection, deployment.name))
turns: Final = asyncio.run(session())
assert [completed(received)["status"] for received in turns] == ["completed"] * 2
eventually(
partial(_counted_tokens, counter(key)),
lambda tokens: tokens >= 2 * HEAVY_TURN_TOKENS,
seconds=15,
)
return _chat(gateway, key, chat, headers)
def test_websocket_turn_tokens_count_toward_the_key_tpm_so_its_next_http_call_is_refused(
limited: Gateway, cert: tuple[Path, Path]
) -> None:
refused: Final = _two_heavy_turns_then_chat(
limited, cert, lambda key: f"{{api_key:{_key_hash(key)}}}:tokens", {}, tpm_limit=SOCKET_TPM_LIMIT
)
assert refused.status_code == 429, refused.text
assert "token" in refused.text.lower(), refused.text
def test_websocket_turn_tokens_count_toward_the_tag_tpm_so_its_next_http_call_is_refused(
limited: Gateway, cert: tuple[Path, Path]
) -> None:
tag: Final = f"ws-tag-tpm-{uuid.uuid4().hex}"
limited.post("/tag/new", {"name": tag, "tpm_limit": SOCKET_TPM_LIMIT})
try:
refused: Final = _two_heavy_turns_then_chat(
limited, cert, lambda _: f"{{tag:{tag}}}:tokens", {"x-litellm-tags": tag}
)
finally:
limited.post("/tag/delete", {"name": tag})
assert refused.status_code == 429, refused.text
assert "tag" in refused.text.lower(), refused.text
def _key_tokens(key: str) -> int:
return _counted_tokens(f"{{api_key:{_key_hash(key)}}}:tokens")
def test_a_turn_is_refused_on_the_open_socket_once_earlier_turns_used_up_the_key_tpm(
limited: Gateway, cert: tuple[Path, Path]
) -> None:
with responses_peer(cert, _heavy_turn) as peer, limited.scenario() as scenario:
deployment: Final = _native(scenario, peer)
key: Final = scenario.key(models=[deployment.name], tpm_limit=SOCKET_TPM_LIMIT)
async def session() -> tuple[tuple[tuple[dict[str, JsonValue], ...], ...], int | None]:
async with _socket(limited, key, deployment.name) as connection:
answered: Final = (await _say(connection, deployment.name), await _say(connection, deployment.name))
await _until(partial(_key_tokens, key), lambda tokens: tokens >= 2 * HEAVY_TURN_TOKENS)
refused: Final = (await _say(connection, deployment.name), await _say(connection, deployment.name))
return (*answered, *refused), connection.close_code
(first, second, third, fourth), close_code = asyncio.run(session())
assert [completed(received)["status"] for received in (first, second)] == ["completed"] * 2
_assert_refused(third)
_assert_refused(fourth)
assert close_code is None
assert deployment.forwarded_turns() == 2
def test_managed_socket_turns_count_toward_the_sockets_tag_tpm_like_the_key_tpm(limited: Gateway) -> None:
tag: Final = f"ws-managed-tag-tpm-{uuid.uuid4().hex}"
limited.post("/tag/new", {"name": tag, "tpm_limit": 1_000_000})
try:
with limited.scenario() as scenario:
deployment: Final = _managed(scenario)
key: Final = scenario.key(models=[deployment.name], tpm_limit=1_000_000)
async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
async with _socket(limited, key, deployment.name, {"x-litellm-tags": tag}) as connection:
return (await _say(connection, deployment.name), await _say(connection, deployment.name))
turns: Final = asyncio.run(session())
turn_tokens: Final = sum(
int(str(object_value(completed(received)["usage"])["total_tokens"])) for received in turns
)
assert turn_tokens > 0, turns
eventually(
lambda: (_key_tokens(key), _counted_tokens(f"{{tag:{tag}}}:tokens")),
lambda counters: counters[0] >= turn_tokens and counters[0] == counters[1],
seconds=20,
)
finally:
limited.post("/tag/delete", {"name": tag})
def test_a_socket_opened_after_turns_used_up_the_key_tpm_gets_the_turn_refusal_frame_and_closes_1008(
limited: Gateway, cert: tuple[Path, Path]
) -> None:
with responses_peer(cert, _heavy_turn) as peer, limited.scenario() as scenario:
deployment: Final = _native(scenario, peer)
key: Final = scenario.key(models=[deployment.name], tpm_limit=SOCKET_TPM_LIMIT)
async def spend_then_reopen() -> tuple[tuple[dict[str, JsonValue], ...], int | None]:
async with _socket(limited, key, deployment.name) as spending:
completed(await _say(spending, deployment.name))
completed(await _say(spending, deployment.name))
await _until(partial(_key_tokens, key), lambda tokens: tokens >= 2 * HEAVY_TURN_TOKENS)
async with _socket(limited, key, deployment.name) as refused:
received: Final = await _say(refused, deployment.name)
await asyncio.wait_for(refused.wait_closed(), timeout=20)
return received, refused.close_code
received, close_code = asyncio.run(spend_then_reopen())
_assert_refused(received)
assert close_code == 1008, close_code
assert deployment.forwarded_turns() == 2
def _assert_http_chat_counted_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
identity: Final = f"ws-admission-{uuid.uuid4().hex}"

32
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -47,6 +47,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
_collect_ws_project_quota_callbacks,
_collect_ws_token_limit_callbacks,
_google_genai_streaming_hidden_params,
_has_pre_call_deployment_hook,
_rust_responses_websocket_enabled,
@ -4043,26 +4044,29 @@ def test_a_callback_that_overrides_the_deployment_hook_is_detected(monkeypatch):
assert _has_pre_call_deployment_hook(logging_obj) is True
def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monkeypatch):
@pytest.mark.parametrize(
"collect, frame_check",
[
(_collect_ws_project_quota_callbacks, "enforce_project_io_token_quota_for_frame"),
(_collect_ws_token_limit_callbacks, "enforce_token_limits_for_frame"),
],
ids=["project-quota", "token-limits"],
)
def test_only_callbacks_that_can_check_a_frame_are_collected_for_ws_turns(monkeypatch, collect, frame_check):
from litellm.integrations.custom_logger import CustomLogger
class _PlainLogger(CustomLogger):
pass
async def check_frame(self, *args, **kwargs):
return None
class _QuotaLogger(CustomLogger):
async def enforce_project_io_token_quota_for_frame(self, *args, **kwargs):
return None
class _NotCallableAttribute:
enforce_project_io_token_quota_for_frame = "not a method"
plain, quota, decoy = _PlainLogger(), _QuotaLogger(), _NotCallableAttribute()
plain: Final = CustomLogger()
checker: Final = type("_FrameChecker", (CustomLogger,), {frame_check: check_frame})()
decoy: Final = type("_NotCallableAttribute", (), {frame_check: "not a method"})()
monkeypatch.setattr(litellm, "callbacks", [plain, decoy])
assert _collect_ws_project_quota_callbacks() == ()
assert collect() == ()
monkeypatch.setattr(litellm, "callbacks", [plain, quota, decoy])
assert _collect_ws_project_quota_callbacks() == (quota,)
monkeypatch.setattr(litellm, "callbacks", [plain, checker, decoy])
assert collect() == (checker,)
@pytest.mark.asyncio

72
tests/unit/proxy/hooks/test_parallel_request_limiter.py
View file

@ -110,8 +110,9 @@ async def test_pre_call_hook_counts_a_cli_session_under_the_per_user_alias_not_t
),
],
)
@pytest.mark.parametrize("opened_this_minute", [True, False], ids=["opened-this-minute", "opened-last-minute"])
@pytest.mark.asyncio
async def test_async_log_success_event_counts_non_chat_response_tokens(response_obj):
async def test_async_log_success_event_counts_non_chat_response_tokens(response_obj, opened_this_minute):
"""
Embedding and text completion responses must increment the per key, user,
team, and end user TPM counters, not just chat completion ModelResponse
@ -129,7 +130,7 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_
precise_minute = FROZEN_INSTANT.strftime("%Y-%m-%d-%H-%M")
scope_ids = [_api_key, user_id, team_id, end_user_id]
for scope_id in scope_ids:
for scope_id in scope_ids if opened_this_minute else ():
await parallel_request_handler.internal_usage_cache.async_set_cache(
key=f"{scope_id}::{precise_minute}::request_count",
value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1},
@ -166,6 +167,73 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_
)
def _socket_turn_events(input_tokens: int, output_tokens: int) -> list[dict[str, object]]:
usage: Final = {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
}
return [
{"type": "response.created", "response": {"id": "resp_ws_turn", "status": "in_progress", "output": []}},
{"type": "response.completed", "response": {"id": "resp_ws_turn", "status": "completed", "usage": usage}},
]
@pytest.mark.parametrize(
"call_type, events, tokens",
[
(CallTypes.aresponses_websocket.value, _socket_turn_events(input_tokens=590, output_tokens=10), 600),
(CallTypes.aresponses.value, _socket_turn_events(input_tokens=590, output_tokens=10), 0),
],
ids=["responses-websocket-turn", "events-of-a-non-socket-call"],
)
@pytest.mark.parametrize("opened_this_minute", [True, False], ids=["opened-this-minute", "opened-last-minute"])
@pytest.mark.asyncio
async def test_async_log_success_event_counts_socket_session_tokens(
call_type: str, events: list[dict[str, object]], tokens: int, opened_this_minute: bool
):
api_key: Final = hash_token("sk-socket-session")
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
)
precise_minute: Final = FROZEN_INSTANT.strftime("%Y-%m-%d-%H-%M")
scope_ids: Final = (api_key, "socket-user", "socket-team", "socket-end-user")
for scope_id in scope_ids if opened_this_minute else ():
await handler.internal_usage_cache.async_set_cache(
key=f"{scope_id}::{precise_minute}::request_count",
value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1},
litellm_parent_otel_span=None,
)
await handler.async_log_success_event(
kwargs={
"call_type": call_type,
"litellm_params": {
"metadata": {
"user_api_key": api_key,
"user_api_key_user_id": "socket-user",
"user_api_key_team_id": "socket-team",
"user_api_key_model_max_budget": {},
}
},
"user": "socket-end-user",
},
response_obj=events,
start_time=FROZEN_INSTANT,
end_time=FROZEN_INSTANT,
)
counted: Final = [
(
await handler.internal_usage_cache.async_get_cache(
key=f"{scope_id}::{precise_minute}::request_count", litellm_parent_otel_span=None
)
)["current_tpm"]
for scope_id in scope_ids
]
assert counted == [tokens, tokens, tokens, tokens]
@pytest.mark.asyncio
async def test_a_pre_call_across_a_minute_rollover_lands_in_the_bucket_of_its_first_clock_read():
internal_usage_cache: Final = InternalUsageCache(DualCache())

437
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -23,9 +23,11 @@ import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY
from litellm.exceptions import ServiceUnavailableError
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.litellm_core_utils import ptu_pricing
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import UserAPIKeyAuth
@ -50,9 +52,11 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.spend_tracking.input_tokens import models_counted_for
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
from litellm.responses.streaming_iterator import _record_usage_for_failure
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.extraction import ExtractionResponse
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -3746,6 +3750,401 @@ async def test_post_call_failure_hook_settles_reservation_at_recovered_partial_u
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
def _responses_websocket_turn_events(total_tokens: int) -> list[dict[str, object]]:
response: Final = {"id": "resp_ws_turn", "object": "response", "created_at": 1700000000, "model": "gpt-5.4"}
return [
{"type": "response.created", "response": {**response, "status": "in_progress", "output": []}},
{
"type": "response.completed",
"response": {
**response,
"status": "completed",
"output": [],
"usage": {"input_tokens": total_tokens - 10, "output_tokens": 10, "total_tokens": total_tokens},
},
},
]
async def _tokens_counted(handler, local_cache: DualCache, key: str, value: str) -> int:
counter: Final = handler.create_rate_limit_keys(key=key, value=value, rate_limit_type="tokens")
return int(await local_cache.async_get_cache(key=counter) or 0)
@pytest.mark.asyncio
async def test_responses_websocket_turns_charge_their_usage_to_the_key_and_the_sockets_tpm_limited_tags_v3():
_api_key = hash_token("sk-responses-ws-turns")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_request_stash.set(
RequestRateLimiterStash(owner_litellm_call_id="socket-open", tpm_limited_tags=frozenset({"ws-tag"}))
)
for turn_call_id in ("turn-1", "turn-2"):
await handler.async_log_success_event(
kwargs={
"litellm_call_id": turn_call_id,
"call_type": CallTypes.aresponses_websocket.value,
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=_responses_websocket_turn_events(total_tokens=600),
start_time=None,
end_time=None,
)
assert await _tokens_counted(handler, local_cache, "api_key", _api_key) == 1200
assert await _tokens_counted(handler, local_cache, "tag", "ws-tag") == 1200
@pytest.mark.asyncio
async def test_a_nested_call_charges_its_own_usage_but_not_the_tags_of_the_request_it_runs_under_v3():
_api_key = hash_token("sk-nested-call")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_request_stash.set(
RequestRateLimiterStash(owner_litellm_call_id="outer-call", tpm_limited_tags=frozenset({"outer-tag"}))
)
await handler.async_log_success_event(
kwargs={
"litellm_call_id": "nested-call",
"call_type": CallTypes.acompletion.value,
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)),
start_time=None,
end_time=None,
)
assert await _tokens_counted(handler, local_cache, "api_key", _api_key) == 15
assert await _tokens_counted(handler, local_cache, "tag", "outer-tag") == 0
def _socket_turn_kwargs(
call_id: str, metadata: dict[str, object], call_type: str = CallTypes.aresponses_websocket.value
) -> dict[str, object]:
return {
"litellm_call_id": call_id,
"call_type": call_type,
"standard_logging_object": {"metadata": metadata},
}
def _logged_http_responses_call(model: str, total_tokens: int) -> ResponsesAPIResponse:
response: Final = ResponsesAPIResponse(
id="resp_http_call",
created_at=0,
model=model,
object="response",
output=[],
status="completed",
usage=ResponseAPIUsage(input_tokens=total_tokens - 10, output_tokens=10, total_tokens=total_tokens),
)
return response.model_copy(
update={"usage": ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response.usage)}
)
_SOCKET_TURN_TAG: Final = "ws-tag"
_SOCKET_TURN_TPM_LIMIT: Final = 1000
def _socket_with_tpm_limited_tag() -> None:
_request_stash.set(
RequestRateLimiterStash(owner_litellm_call_id="socket-open", tpm_limited_tags=frozenset({_SOCKET_TURN_TAG}))
)
def _handler_with_socket_tag_limit(
local_cache: DualCache, tag_tpm_limit: int | None
) -> PROXY_MaxParallelRequestsHandler_v3:
resolver, _ = _static_tag_limits(
{} if tag_tpm_limit is None else {_SOCKET_TURN_TAG: TagRateLimit(rpm_limit=None, tpm_limit=tag_tpm_limit)}
)
return _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache), tag_rate_limit_resolver=resolver
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("auth_limits", "tag_tpm_limit"),
[
({"tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({"team_id": "team-ws", "team_tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({"user_id": "user-ws", "user_tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({"org_id": "org-ws", "organization_tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({"team_id": "team-ws", "user_id": "user-ws", "team_member_tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({"end_user_id": "end-user-ws", "end_user_tpm_limit": _SOCKET_TURN_TPM_LIMIT}, None),
({}, _SOCKET_TURN_TPM_LIMIT),
],
ids=["key", "team", "user", "org", "team-member", "end-user", "tag"],
)
async def test_a_later_socket_turn_is_refused_once_a_tpm_window_it_counts_toward_has_no_tokens_left_v3(
auth_limits: dict[str, object], tag_tpm_limit: int | None
) -> None:
api_key: Final = hash_token("sk-socket-turn-check")
local_cache: Final = DualCache()
handler: Final = _handler_with_socket_tag_limit(local_cache, tag_tpm_limit)
_socket_with_tpm_limited_tag()
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, **auth_limits)
turn_metadata: Final = {
"user_api_key_hash": api_key,
"user_api_key_team_id": "team-ws",
"user_api_key_user_id": "user-ws",
"user_api_key_org_id": "org-ws",
"user_api_key_end_user_id": "end-user-ws",
}
async def log_turn(call_id: str) -> None:
await handler.async_log_success_event(
kwargs=_socket_turn_kwargs(call_id, turn_metadata),
response_obj=_responses_websocket_turn_events(total_tokens=_SOCKET_TURN_TPM_LIMIT // 2),
start_time=None,
end_time=None,
)
await log_turn("turn-1")
await handler.enforce_token_limits_for_frame(user_api_key_dict=user_api_key_dict, requested_model="gpt-5.4")
await log_turn("turn-2")
tokens_before_check: Final = await _tokens_counted(handler, local_cache, "api_key", api_key)
with pytest.raises(ProxyRateLimitError) as exc_info:
await handler.enforce_token_limits_for_frame(user_api_key_dict=user_api_key_dict, requested_model="gpt-5.4")
assert exc_info.value.status_code == 429
assert await _tokens_counted(handler, local_cache, "api_key", api_key) == tokens_before_check
@pytest.mark.asyncio
@pytest.mark.parametrize(
"agent_limits",
[{"tpm_limit": _SOCKET_TURN_TPM_LIMIT}, {"session_tpm_limit": _SOCKET_TURN_TPM_LIMIT}],
ids=["agent", "agent-session"],
)
async def test_a_later_socket_turn_is_refused_once_the_sockets_agent_tpm_window_has_no_tokens_left_v3(
agent_limits: dict[str, int],
) -> None:
from litellm.types.agents import AgentResponse
api_key: Final = hash_token("sk-socket-agent-turn")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key)
agent: Final = AgentResponse(
agent_id="agent-ws", agent_name="ws-agent", agent_card_params={"name": "WS Agent"}, **agent_limits
)
turn_metadata: Final = {"user_api_key_hash": api_key, "agent_id": "agent-ws", "session_id": "sess-ws"}
async def log_turn(call_id: str, total_tokens: int) -> None:
await handler.async_log_success_event(
kwargs=_socket_turn_kwargs(call_id, turn_metadata),
response_obj=_responses_websocket_turn_events(total_tokens=total_tokens),
start_time=None,
end_time=None,
)
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id",
side_effect=lambda agent_id: agent if agent_id == "agent-ws" else None,
):
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-5.4",
"litellm_call_id": "socket-open",
"metadata": {"agent_id": "agent-ws", "session_id": "sess-ws"},
},
call_type=CallTypes.aresponses_websocket.value,
)
await log_turn("turn-1", _SOCKET_TURN_TPM_LIMIT // 10)
await handler.enforce_token_limits_for_frame(user_api_key_dict=user_api_key_dict, requested_model="gpt-5.4")
await log_turn("turn-2", _SOCKET_TURN_TPM_LIMIT)
with pytest.raises(ProxyRateLimitError) as exc_info:
await handler.enforce_token_limits_for_frame(user_api_key_dict=user_api_key_dict, requested_model="gpt-5.4")
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_the_socket_turn_check_ignores_request_limits_and_charges_nothing_v3(monkeypatch) -> None:
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
api_key: Final = hash_token("sk-socket-turn-rpm")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpm_limit=_SOCKET_TURN_TPM_LIMIT)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-5.4", "litellm_call_id": "socket-open"},
call_type=CallTypes.aresponses_websocket.value,
)
tokens_after_open: Final = await _tokens_counted(handler, local_cache, "api_key", api_key)
for _ in range(3):
await handler.enforce_token_limits_for_frame(user_api_key_dict=user_api_key_dict, requested_model="gpt-5.4")
assert await _tokens_counted(handler, local_cache, "api_key", api_key) == tokens_after_open
@pytest.mark.asyncio
@pytest.mark.parametrize(
"socket_call_id, tag_tokens",
[("socket-open", 600), ("another-socket", 0), (None, 0)],
ids=["this-socket", "another-socket", "no-socket"],
)
async def test_a_managed_socket_turn_charges_the_sockets_tpm_limited_tags_v3(
socket_call_id: str | None, tag_tokens: int
) -> None:
api_key: Final = hash_token("sk-managed-socket-turn")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_socket_with_tpm_limited_tag()
turn_metadata: Final = {} if socket_call_id is None else {RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: socket_call_id}
await handler.async_log_success_event(
kwargs={
**_socket_turn_kwargs("managed-turn", {"user_api_key_hash": api_key}, CallTypes.aresponses.value),
"litellm_params": {"litellm_metadata": turn_metadata},
},
response_obj=_logged_http_responses_call("claude-opus-4-7", total_tokens=600),
start_time=None,
end_time=None,
)
assert await _tokens_counted(handler, local_cache, "api_key", api_key) == 600
assert await _tokens_counted(handler, local_cache, "tag", _SOCKET_TURN_TAG) == tag_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["socket", "http"])
async def test_a_socket_turn_reaches_the_team_tpm_window_as_an_http_call_does_v3(transport: str, monkeypatch) -> None:
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
api_key: Final = hash_token("sk-team-window")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_socket_with_tpm_limited_tag()
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, team_id="team-ws", team_tpm_limit=500)
async def pre_call() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-5.4"},
call_type=CallTypes.aresponses.value,
)
await pre_call()
await handler.async_log_success_event(
kwargs=_socket_turn_kwargs(
f"{transport}-call",
{"user_api_key_hash": api_key, "user_api_key_team_id": "team-ws"},
CallTypes.aresponses_websocket.value if transport == "socket" else CallTypes.aresponses.value,
),
response_obj=(
_responses_websocket_turn_events(total_tokens=600)
if transport == "socket"
else _logged_http_responses_call("gpt-5.4", total_tokens=600)
),
start_time=None,
end_time=None,
)
with pytest.raises(HTTPException) as exc_info:
await pre_call()
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type, completion_tokens, counted",
[
(CallTypes.aresponses_websocket.value, 7, 27),
(CallTypes.aresponses_websocket.value, 0, 0),
(CallTypes.acompletion.value, 7, 0),
],
ids=["socket-turn", "socket-turn-without-output", "unreserved-http-call"],
)
async def test_a_failed_socket_turn_counts_its_recorded_partial_usage_v3(
call_type: str, completion_tokens: int, counted: int
) -> None:
api_key: Final = hash_token("sk-failed-socket-turn")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_socket_with_tpm_limited_tag()
await handler.async_log_failure_event(
kwargs={
**_socket_turn_kwargs("failed-turn", {"user_api_key_hash": api_key}, call_type),
"combined_usage_object": Usage(
prompt_tokens=27 - completion_tokens, completion_tokens=completion_tokens, total_tokens=27
),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert await _tokens_counted(handler, local_cache, "api_key", api_key) == counted
assert await _tokens_counted(handler, local_cache, "tag", _SOCKET_TURN_TAG) == counted
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failed_usage, counted",
[({"input_tokens": 20, "output_tokens": 7, "total_tokens": 27}, 27), (None, 0)],
ids=["failed-response-reports-usage", "failed-response-without-usage"],
)
async def test_a_native_socket_turn_that_fails_counts_the_usage_its_failed_response_reports_v3(
failed_usage: dict[str, int] | None, counted: int
) -> None:
api_key: Final = hash_token("sk-native-failed-socket-turn")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
_socket_with_tpm_limited_tag()
turn_logging: Final = LiteLLMLoggingObj(
model="gpt-5.6",
messages=[],
stream=True,
call_type=CallTypes.aresponses_websocket.value,
start_time=datetime.now(),
litellm_call_id="native-failed-turn",
function_id="native-failed-turn",
)
_record_usage_for_failure(
turn_logging,
(
{"type": "response.created", "response": {"id": "resp_failed", "status": "in_progress"}},
{
"type": "response.failed",
"response": {
"id": "resp_failed",
"status": "failed",
"error": {"code": "server_error", "message": "upstream failed mid-turn"},
"usage": failed_usage,
},
},
),
)
await handler.async_log_failure_event(
kwargs={
**_socket_turn_kwargs("native-failed-turn", {"user_api_key_hash": api_key}),
"combined_usage_object": turn_logging.model_call_details.get("combined_usage_object"),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert await _tokens_counted(handler, local_cache, "api_key", api_key) == counted
assert await _tokens_counted(handler, local_cache, "tag", _SOCKET_TURN_TAG) == counted
@pytest.mark.asyncio
async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usage_v3():
"""
@ -7528,6 +7927,22 @@ async def test_fail_closed_read_only_check_rejects_with_503_when_redis_counters_
assert response["overall_code"] == "OK"
@pytest.mark.asyncio
async def test_fail_closed_socket_turn_check_refuses_with_503_when_redis_counters_are_unreachable():
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-socket-turn"), tpm_limit=1000)
with pytest.raises(ServiceUnavailableError) as exc:
await _handler_with_redis(_UnreachableRedis(), fail_closed=True).enforce_token_limits_for_frame(
user_api_key_dict=auth, requested_model="gpt-5.4"
)
assert exc.value.status_code == 503
assert "fail_closed_rate_limit_enforcement" in str(exc.value)
await _handler_with_redis(_UnreachableRedis(), fail_closed=False).enforce_token_limits_for_frame(
user_api_key_dict=auth, requested_model="gpt-5.4"
)
@pytest.mark.parametrize("stored_counter_value, expected_code", [(1, "OK"), (2, "OVER_LIMIT"), (3, "OVER_LIMIT")])
@pytest.mark.asyncio
async def test_read_only_check_reports_the_redis_counters_without_incrementing_them(
@ -8974,6 +9389,24 @@ def test_usage_that_only_reports_a_total_charges_that_total_to_the_ptu_counter()
assert _ptu_increment(handler, ops) == 150
def test_a_socket_turns_usage_charges_the_key_but_leaves_the_team_ptu_counter_alone() -> None:
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
ops: Final = handler._build_success_event_pipeline_operations(
kwargs={**_ptu_success_kwargs(), "call_type": CallTypes.aresponses_websocket.value},
response_obj=_responses_websocket_turn_events(total_tokens=150),
rate_limit_type="total",
)
key_window: Final = handler.create_rate_limit_keys("api_key", hash_token("sk-ptu"), "tokens")
assert next(op["increment_value"] for op in ops if op["key"] == key_window) == 150
assert not _ptu_increment(handler, ops)
def test_success_accounting_charges_no_ptu_counter_without_a_ceiling():
resolve, _ = _ptu_ceiling_for("u", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(

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

@ -27,6 +27,7 @@ from litellm.caching.caching import DualCache
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.proxy_server import app
from litellm.proxy.utils import ProxyLogging
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
@ -2325,6 +2326,74 @@ class TestResponsesWSFirstFrameModelAuth:
assert authenticated_models == ["gpt-4o-mini"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"open_failure, refusal, close_reason, logged_traceback",
[
(
ProxyRateLimitError(detail="Rate limit exceeded for api_key: tokens"),
(429, "rate_limit_exceeded"),
"Rate limit exceeded",
False,
),
(ValueError("bad metadata"), (None, "invalid_request_error"), "Pre-call error", True),
],
ids=["rate-limited", "other-failure"],
)
async def test_endpoint_refuses_a_rate_limited_open_with_the_turn_refusal_frame(
self,
open_failure: Exception,
refusal: tuple[int | None, str],
close_reason: str,
logged_traceback: bool,
):
import logging
from litellm._logging import verbose_proxy_logger
from litellm.proxy.response_api_endpoints.endpoints import responses_websocket_endpoint
class _Recorder(logging.Handler):
records: tuple[logging.LogRecord, ...] = ()
def emit(self, record: logging.LogRecord) -> None:
self.records = (*self.records, record)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(return_value=json.dumps({"type": "response.create", "model": "gpt-5.4"}))
ws.send_text = AsyncMock()
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(side_effect=open_failure)
recorder: Final = _Recorder(level=logging.INFO)
previous_level: Final = verbose_proxy_logger.level
verbose_proxy_logger.addHandler(recorder)
verbose_proxy_logger.setLevel(logging.INFO)
try:
with (
patch( # test-quality-ok: first-frame model auth needs a live router and has its own tests
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor's failure is the input under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
):
await responses_websocket_endpoint(websocket=ws, model="gpt-5.4", user_api_key_dict=MagicMock())
finally:
verbose_proxy_logger.removeHandler(recorder)
verbose_proxy_logger.setLevel(previous_level)
sent: Final = json.loads(ws.send_text.await_args.args[0])
assert (sent.get("status"), sent["error"]["type"]) == refusal
ws.close.assert_awaited_once_with(code=1008, reason=close_reason)
assert [record.exc_info is not None for record in recorder.records] == [logged_traceback]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
@pytest.mark.parametrize("query_model", [None, "gpt-4o-mini"])

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

@ -1607,6 +1607,39 @@ def _admission_refusing_the_next_turn():
return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), {}])
class _TokenLimits:
"""A proxy hook whose per-turn token check answers each later turn with the next outcome."""
def __init__(self, *outcomes: Exception | None) -> None:
self._outcomes: Final = iter(outcomes)
self.checked: Final[list[tuple[object, str | None]]] = []
async def enforce_token_limits_for_frame(self, user_api_key_dict: object, requested_model: str | None) -> None:
self.checked.append((user_api_key_dict, requested_model))
outcome: Final = next(self._outcomes)
if outcome is not None:
raise outcome
def _out_of_tokens() -> Exception:
from litellm.exceptions import RateLimitError
return RateLimitError(message="key is out of tokens", llm_provider="", model="")
def _counters_unverifiable() -> Exception:
from litellm.exceptions import ServiceUnavailableError
return ServiceUnavailableError(message="counters could not be verified", llm_provider="litellm_proxy", model="")
_TURN_REFUSALS: Final = pytest.mark.parametrize(
"refusal, refusal_frame",
[(_out_of_tokens(), (429, "rate_limit_exceeded")), (_counters_unverifiable(), (503, "server_error"))],
ids=["out-of-tokens", "counters-unverifiable"],
)
class _ClientSocket:
"""A client socket that, like Starlette's, stops receiving once the server closes it."""
@ -2204,6 +2237,123 @@ class TestWebSocketTurnAdmission:
assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")]
@pytest.mark.asyncio
@_TURN_REFUSALS
async def test_managed_socket_checks_token_limits_before_every_later_turn_and_marks_each_turn(
self, monkeypatch: pytest.MonkeyPatch, refusal: Exception, refusal_frame: tuple[int, str]
) -> None:
import litellm
from litellm.constants import RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.proxy._types import UserAPIKeyAuth
answered: Final[list[tuple[str, object]]] = []
async def fake_aresponses(*args: object, **kwargs: object):
turn_metadata: Final = as_str_mapping(kwargs.get("litellm_metadata"))
answered.append(
(
json.dumps(kwargs.get("input")),
None if turn_metadata is None else turn_metadata.get(RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY),
)
)
async def _empty():
return
yield
return _empty()
monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
token_limits: Final = _TokenLimits(refusal, None)
monkeypatch.setattr(litellm, "callbacks", [token_limits])
deployment_admission: Final = AsyncMock(return_value={})
user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed-key", tpm_limit=1000)
first, second, third = _turn_frames()
websocket: Final = _ClientSocket(second, third)
await BaseLLMHTTPHandler().async_responses_websocket(
model="anthropic/claude-sonnet-4-5",
websocket=websocket,
logging_obj=Logging(
model="anthropic/claude-sonnet-4-5",
messages=[],
stream=True,
call_type="aresponses",
start_time=0,
litellm_call_id="socket-open",
function_id="test-func",
),
responses_api_provider_config=None,
user_api_key_dict=user_api_key_dict,
litellm_metadata={"model_group": "claude-ws"},
first_message=first,
turn_admission=deployment_admission,
)
assert [("two" in turn_input, "three" in turn_input) for turn_input, _ in answered] == [
(False, False),
(False, True),
]
assert [socket_call_id for _, socket_call_id in answered] == ["socket-open", "socket-open"]
assert token_limits.checked == [(user_api_key_dict, "claude-ws"), (user_api_key_dict, "claude-ws")]
assert deployment_admission.await_count == 1
assert [(status, error["type"]) for status, error in _errors(websocket)] == [refusal_frame]
@pytest.mark.asyncio
@_TURN_REFUSALS
@pytest.mark.parametrize(
"keyed, forwarded, refused",
[(True, ["one", "three"], True), (False, ["one", "two", "three"], False)],
ids=["keyed", "keyless"],
)
async def test_native_socket_checks_token_limits_before_every_later_turn_of_a_keyed_socket(
self,
monkeypatch: pytest.MonkeyPatch,
keyed: bool,
forwarded: list[str],
refused: bool,
refusal: Exception,
refusal_frame: tuple[int, str],
) -> None:
from unittest.mock import patch
import litellm
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.proxy._types import UserAPIKeyAuth
token_limits: Final = _TokenLimits(refusal, None)
monkeypatch.setattr(litellm, "callbacks", [token_limits])
provider: Final = _ProviderResponsesSocket()
config: Final = MagicMock(spec=OpenAIResponsesAPIConfig)
config.supports_native_websocket.return_value = True
config.model_in_websocket_url.return_value = True
config.get_websocket_url.return_value = "wss://api.openai.com/v1/responses"
config.validate_environment.return_value = {}
logging_obj: Final = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.copy_for_new_call.return_value = logging_obj
user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed-key", tpm_limit=1000) if keyed else None
first, second, third = _turn_frames()
websocket: Final = _ClientSocket(second, third)
with patch("websockets.connect", provider.connect):
await BaseLLMHTTPHandler().async_responses_websocket(
model="gpt-5.4",
websocket=websocket,
logging_obj=logging_obj,
responses_api_provider_config=config,
api_key="sk-test",
user_api_key_dict=user_api_key_dict,
first_message=first,
)
assert [frame["input"] for frame in provider.sent] == forwarded
assert token_limits.checked == ([(user_api_key_dict, "gpt-5.4")] * 2 if keyed else [])
assert [(status, error["type"]) for status, error in _errors(websocket)] == ([refusal_frame] if refused else [])
@pytest.mark.asyncio
async def test_native_relay_runs_the_guardrail_before_the_deployment_counts_a_turn(self):
from unittest.mock import AsyncMock