diff --git a/litellm/constants.py b/litellm/constants.py index 1a1e66ecf86..e21751c7fd1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f5b9014d87c..6251bb26276 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9a2b352dd86..8a094cc6424 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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() diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 05c64317a44..cd01d15d834 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -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, } diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a515c051bbd..c83551a815c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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) diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index 7b9635dddce..308cb40b5db 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -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( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 089cb75a795..3f76bcf69ad 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index a8e55a5e94d..803b85ab3a5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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: diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index 7b0e14d044f..2de556e90e4 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -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}" diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 23d7feb2557..a8f6772a76e 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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 diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 08840978169..0f5d672b778 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -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()) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index a50aa577a79..bed185ba8f5 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 07f85cb9628..ef43fe56146 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -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"]) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index d0e26798090..7e044b55f5a 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -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