diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3976612c6e1..4fcd3bde843 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3113,14 +3113,19 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): return response_ids[-1] if response_ids else None @staticmethod - def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject: + def create_logging_object( + results: Sequence[Mapping[str, object]], litellm_model_name: str | None = None + ) -> LiteLLMRealtimeStreamLoggingObject: tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results) - return LiteLLMRealtimeStreamLoggingObject( + logging_object: Final = LiteLLMRealtimeStreamLoggingObject( id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results), usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results), results=cast(OpenAIRealtimeStreamList, list(results)), # cast-ok: raw WS frames, unvalidated field service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None, ) + if litellm_model_name is not None: + logging_object.hidden_params = {"litellm_model_name": litellm_model_name} + return logging_object _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b6bdc21bc05..e7eb5259702 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -986,6 +986,16 @@ class Logging(LiteLLMLoggingBaseClass): return model_id return None + def _routed_deployment_model(self) -> str | None: + litellm_params: Final = self._model_call_details_view.get("litellm_params") + metadata: Final = ( + StandardLoggingPayloadSetup.merge_litellm_metadata(_LITELLM_PARAMS.validate_python(litellm_params)) + if isinstance(litellm_params, dict) + else {} + ) + deployment: Final = metadata.get("deployment") + return deployment if isinstance(deployment, str) else None + def get_deployment_model_for_cost(self) -> str | None: """The provider-qualified model to price against. @@ -2253,6 +2263,7 @@ class Logging(LiteLLMLoggingBaseClass): and litellm_params.get(CallTypes.agenerate_content.value, False) is not True and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True and litellm_params.get(CallTypes.arealtime.value, False) is not True + and litellm_params.get(CallTypes.aresponses_websocket.value, False) is not True ) _is_sync_litellm_request = is_sync_litellm_request @@ -2508,7 +2519,8 @@ class Logging(LiteLLMLoggingBaseClass): elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object( - results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + litellm_model_name=self._routed_deployment_model(), ) elif self.relays_llm_passthrough and isinstance(result, Response): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index bbda4f57ec8..40e7f546bdc 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6910,6 +6910,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: @@ -6972,12 +6973,12 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + turn_admission=turn_admission, guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks frame_guard=frame_guard, **kwargs, ) - await handler.run() - return None + return await handler.run() import websockets from websockets.asyncio.client import ClientConnection @@ -7068,6 +7069,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + turn_admission=turn_admission, frame_guard=frame_guard, ) return await streaming.bidirectional_forward() diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6fbd0b8f1b3..3195394b9aa 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -2,8 +2,11 @@ import asyncio import traceback from collections.abc import Callable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, cast +from pydantic import TypeAdapter, ValidationError + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -17,6 +20,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, @@ -100,6 +104,37 @@ def _proxy_stamped_used_client_oauth_token( return stamped if isinstance(stamped, bool) else None +def _str_keyed_fields(value: object) -> Mapping[str, object]: + return value if is_str_keyed_mapping(value) else {} + + +def _is_model_group(value: object) -> bool: + return isinstance(value, str) and bool(value) + + +_TAG_LIST: Final = TypeAdapter(list[str]) + + +def _is_tag_list(value: object) -> bool: + try: + return bool(_TAG_LIST.validate_python(value, strict=True)) + except ValidationError: + return False + + +_ROUTED_METADATA_FIELD_CHECKS: Final = MappingProxyType({"model_group": _is_model_group, "tags": _is_tag_list}) + + +def _routed_metadata_fallback(request_data: Mapping[str, object]) -> Mapping[str, object]: + caller_metadata: Final = _str_keyed_fields(request_data.get("metadata")) + routed_metadata: Final = _str_keyed_fields(request_data.get("litellm_metadata")) + return { + field: routed_metadata[field] + for field, is_valid in _ROUTED_METADATA_FIELD_CHECKS.items() + if not caller_metadata.get(field) and is_valid(routed_metadata.get(field)) + } + + def _failure_snapshot_with_guardrails_recorded_after_it( failure_snapshot: object, recorded_guardrail_information: object ) -> Mapping[str, object] | None: @@ -220,10 +255,11 @@ class _ProxyDBLogger(CustomLogger): metadata=_metadata, original_exception=original_exception ) - _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_data, request_route) + request_fields: Final = _str_keyed_fields(request_data) + _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} - existing_metadata.update(_metadata) + existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") if ( @@ -471,11 +507,7 @@ class _ProxyDBLogger(CustomLogger): await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. - # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with - # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls. - if sl_object is None and ( - not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime") - ): + if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object for call_type=%s", kwargs.get("call_type", "unknown"), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index f032cd79783..e7c4aee275c 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -78,6 +78,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.utils import handle_exception_on_proxy from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel @@ -1861,12 +1862,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec def _responses_ws_failure_frame(failure: Exception) -> str: - raw_status: Final = getattr(failure, "status_code", None) - status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 + proxy_exception: Final = handle_exception_on_proxy(failure) + status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500 error_type: Final = ( "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error" ) - return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}}) + return json.dumps( + {"type": "error", "status": status, "error": {"type": error_type, "message": proxy_exception.message}} + ) async def _enforce_responses_ws_first_frame_model_auth( @@ -1949,7 +1952,7 @@ def _responses_ws_frame_guard( } model: Final = admitted.get("model") - # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation + # blocked and refused frames bill as their own failed requests; only the first turn settles the reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) async def guard(frame: str) -> GuardedResponsesWsFrame: @@ -1997,6 +2000,7 @@ def _responses_ws_frame_guard( frame=json.dumps(rewritten), guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)), pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + bill_refusal=partial(_bill_rejected_frame, rejected_frame_auth, frame_data=guarded_data), ) return guard @@ -2167,7 +2171,6 @@ async def _responses_websocket_session( request_data=routed_data, ) except Exception as e: - verbose_proxy_logger.exception("Responses WebSocket error") with contextlib.suppress(Exception): await websocket.send_text(_responses_ws_failure_frame(e)) await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e775e155dea..727a5ccfedf 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import partial @@ -73,7 +73,7 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, all_litellm_params +from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypes, all_litellm_params from litellm.utils import ( ProviderConfigManager, client, @@ -2609,6 +2609,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: @@ -2624,7 +2625,7 @@ async def _aresponses_websocket( litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") user: Final = kwargs.get("user", None) litellm_params: Final = GenericLiteLLMParams(**kwargs) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.aresponses_websocket.value: True} ( provider_model, @@ -2702,6 +2703,7 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + turn_admission=turn_admission, frame_guard=responses_ws_frame_guard, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 145189616af..1e7368e15eb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import copy import json import time @@ -1929,6 +1930,38 @@ async def _enforce_frame_project_quota( ) +async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> Exception | None: + frame: Final = _json_object_or_none(raw_message) + if turn_admission is None or frame is None or frame.get("type") != "response.create": + return None + try: + await turn_admission() + except RateLimitError as refusal: + return refusal + except Exception as e: + verbose_logger.exception("Responses WS turn admission failed: %s", e) + return e + return None + + +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_rejection_frame(rejection: Exception) -> str: + match rejection: + case RateLimitError(): + return _turn_refusal_frame(rejection) + case _: + return _turn_admission_error_frame(rejection) + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2001,16 +2034,28 @@ def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], un @dataclass(frozen=True, slots=True) class GuardedResponsesWsFrame: """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails - may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the - Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply.""" + may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, the + Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply, and how + to bill those guardrail records when the deployment then refuses the turn.""" frame: str guardrail_information: tuple[Mapping[str, object], ...] = () pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + bill_refusal: Callable[[Exception], Awaitable[None]] | None = None ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] + +async def _reject_guarded_turn( + websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception +) -> None: + with contextlib.suppress(Exception): + await websocket.send_text(_turn_rejection_frame(rejection)) + if guarded.guardrail_information and guarded.bill_refusal is not None: + await guarded.bill_refusal(rejection) + + _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES @@ -2194,6 +2239,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket @@ -2215,6 +2261,7 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission def _start_turn( self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () @@ -2266,9 +2313,11 @@ class ResponsesWebSocketStreaming: await self._send_error_event("invalid_request_error", str(e)) return None - async def _forward_guarded_frame(self, message: str) -> None: + async def _forward_guarded_frame( + self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None: + if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission): return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) @@ -2662,10 +2711,27 @@ class ResponsesWebSocketStreaming: self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) except RateLimitError as e: - await self._send_error_event("rate_limit_exceeded", str(e)) + await self._send_turn_refusal(e) return False return True + async def _admit_or_reject_turn( + self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None + ) -> bool: + rejection: Final = await _admit_frame_turn(turn_admission, message) + if rejection is None: + return True + if self._failure is None: + self._failure = rejection + await _reject_guarded_turn(self.websocket, guarded, rejection) + return False + + async def _send_turn_refusal(self, refusal: RateLimitError) -> None: + try: + await self.websocket.send_text(_turn_refusal_frame(refusal)) + except Exception: # noqa: BLE001, S110 # client may already be gone + pass + async def _send_error_event(self, error_type: str, message: str) -> None: try: await self.websocket.send_text( @@ -2684,7 +2750,7 @@ class ResponsesWebSocketStreaming: message = await self.websocket.receive_text() if not await self._enforce_or_reject_frame(message): continue - await self._forward_guarded_frame(message) + await self._forward_guarded_frame(message, self.turn_admission) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) @@ -2761,6 +2827,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, frame_guard: ResponsesWsFrameGuard | None = None, guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, @@ -2784,6 +2851,8 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission + self._failure: Exception | None = None # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -2812,12 +2881,11 @@ class ManagedResponsesWebSocketHandler: return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: - pass + await self._send_frame(json.dumps({"type": "error", "error": {"type": error_type, "message": message}})) + + async def _send_frame(self, frame: str) -> None: + with contextlib.suppress(Exception): + await self.websocket.send_text(frame) def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]: """ @@ -3179,7 +3247,9 @@ class ManagedResponsesWebSocketHandler: # Core request handler # ------------------------------------------------------------------ - async def _process_response_create(self, raw_message: str) -> None: + async def _process_response_create( + self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, and forward every streaming event to the client. @@ -3220,12 +3290,17 @@ 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_error(str(e), error_type="rate_limit_exceeded") + await self._send_frame(_turn_refusal_frame(e)) return - guarded: Final = await self._guard_frame(raw_message) if guarded is None: return + rejection: Final = await _admit_frame_turn(turn_admission, raw_message) + if rejection is not None: + if self._failure is None: + self._failure = rejection + await _reject_guarded_turn(self.websocket, guarded, rejection) + return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} @@ -3269,7 +3344,7 @@ class ManagedResponsesWebSocketHandler: # Main entry point # ------------------------------------------------------------------ - async def run(self) -> None: + async def run(self) -> Exception | None: """ Main loop: accept ``response.create`` events sequentially and handle each one before waiting for the next message. @@ -3285,8 +3360,9 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc) break - await self._process_response_create(message) + await self._process_response_create(message, self.turn_admission) except Exception as exc: verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) await self._send_error(f"Internal server error: {exc}") + return self._failure diff --git a/litellm/router.py b/litellm/router.py index c30108b4e38..54ad973edf8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -289,6 +289,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( deployment_has_io_token_limits, refund_stale_reservation_before_retry, set_io_token_rate_limit_request_kwargs, + without_io_token_limits, ) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, @@ -5962,6 +5963,7 @@ class Router: ) custom_llm_provider: Final = provider_for_generic_call(data) + generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") response_kwargs: Final = { **data, @@ -5969,6 +5971,11 @@ class Router: **kwargs, "model": model_name, **_with_router_resolved_session_model(kwargs.get("session"), model_name), + **( + {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)} + if generic_function_name == "_aresponses_websocket" + else {} + ), } # Only set custom_llm_provider if it's not None if custom_llm_provider is not None: @@ -5978,7 +5985,7 @@ class Router: self, deployment, response_kwargs, - surface_for_call(getattr(original_generic_function, "__name__", "")), + surface_for_call(generic_function_name), ) async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): response = await original_generic_function(**compacted_input) @@ -8933,8 +8940,7 @@ class Router: Track remaining tpm/rpm quota for model in model_list """ try: - # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. - if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): + if kwargs.get("call_type") == "_arealtime": return if is_batch_retrieve_call_type(kwargs.get("call_type")): return @@ -9602,6 +9608,71 @@ class Router: ) raise e + def _deployment_as_configured_now(self, deployment: Mapping[str, object]) -> Mapping[str, object]: + model_info: Final = deployment.get("model_info") + model_id: Final = ( + _MODEL_INFO_ADAPTER.validate_python(model_info).get("id") if isinstance(model_info, dict) else None + ) + current: Final = self.get_deployment(model_id) if isinstance(model_id, str) else None + return deployment if current is None else current.model_dump(exclude_none=True) + + async def async_admit_deployment_turn( + self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> None: + deployment: Final = self._deployment_as_configured_now(opened_on) + budget_limiter: Final = self.get_router_deployment_budget_limiter() + if budget_limiter is not None: + try: + await budget_limiter.async_filter_deployments( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts + model=model_group, + healthy_deployments=[dict(deployment)], + messages=None, + request_kwargs=dict(request_kwargs), + ) + except ValueError as e: + raise litellm.RateLimitError( + message=str(e), + llm_provider="", + model=model_group, + category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from e + await self._admit_turn_through_routing_strategy(model_group, deployment, request_kwargs) + await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict + deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None + ) + + async def _admit_turn_through_routing_strategy( + self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> None: + routing_kwargs: Final = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"} + strategy, selector = self._get_routing_context( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped request dict + model=model_group, request_kwargs=routing_kwargs + ) + if strategy == "simple-shuffle" or selector is None: + return + selected: Final = await self._select_deployment_async( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts + strategy=strategy, + selector=selector, + model=model_group, + healthy_deployments=[dict(deployment)], + messages=None, + input=None, + request_kwargs=routing_kwargs, + ) + if selected is None: + no_deployment: Final = await async_raise_no_deployment_exception( + litellm_router_instance=self, model=model_group, parent_otel_span=None + ) + raise litellm.RateLimitError( + message=str(no_deployment), + llm_provider="", + model=model_group, + category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from no_deployment + await self._async_override_selector_pre_call_check( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict + strategy, selector, dict(deployment), None + ) + @contextlib.asynccontextmanager async def _deployment_slot( self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1622a30b95f..7573bb6716f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -505,8 +505,7 @@ class RouterBudgetLimiting(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") - # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. - if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): + if kwargs.get("call_type") == "_arealtime": return standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 0d701411c94..086f5517c29 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -23,6 +23,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -96,6 +97,24 @@ def deployment_has_io_token_limits(deployment: dict) -> bool: return itpm is not None or otpm is not None +_IO_TOKEN_LIMIT_KEYS: Final = frozenset({"itpm", "otpm"}) +_IO_TOKEN_LIMIT_SECTIONS: Final = frozenset({"litellm_params", "model_info"}) + + +def without_io_token_limits(deployment: Mapping[str, object]) -> Mapping[str, object]: + return { + key: _without_io_token_limit_keys(value) if key in _IO_TOKEN_LIMIT_SECTIONS else value + for key, value in deployment.items() + if key not in _IO_TOKEN_LIMIT_KEYS + } + + +def _without_io_token_limit_keys(section: object) -> object: + if not is_str_keyed_mapping(section): + return section + return {key: value for key, value in section.items() if key not in _IO_TOKEN_LIMIT_KEYS} + + def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | None: model_id: Final = deployment.get("model_info", {}).get("id") deployment_name: Final = deployment.get("litellm_params", {}).get("model") diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index e669107900d..681476456c2 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -134,6 +134,9 @@ ignored_function_names = [ "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py + "_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py + "_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py ] diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py new file mode 100644 index 00000000000..a246c89e7b3 --- /dev/null +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -0,0 +1,629 @@ +import asyncio +import hashlib +import os +import threading +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import asynccontextmanager, contextmanager +from dataclasses import dataclass +from functools import partial +from datetime import UTC, datetime +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + list_value, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import ( + PROVIDER_MODEL, + USAGE, + Peer, + Script, + completed, + create, + drain, + events, + responses_peer, + turn, + until_terminal, +) +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] * OUTPUT_RATE +TURN_TOKENS: Final = USAGE["total_tokens"] +MANAGED_PROVIDER: Final = "mistral" +SPEND_ROWS_FOR_KEY: Final = ( + 'SELECT request_id, spend, status, model_group, metadata, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s' +) +TURN_CHECK_FAULT: Final = "turn_check_fault" +FAULTING_TURN_CHECK: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class FaultingTurnCheck(CustomLogger): + async def async_pre_call_check(self, deployment, parent_otel_span): + fault = (deployment.get("model_info") or {{}}).get("{TURN_CHECK_FAULT}") + if fault: + raise RuntimeError(fault) + return None + + +instance = FaultingTurnCheck() +""" +FrameReader = Callable[[], int] + + +@dataclass(frozen=True, slots=True) +class Deployment: + name: str + identity: str + forwarded_turns: FrameReader + + +def _key_hash(key: str) -> str: + return hashlib.sha256(key.encode()).hexdigest() + + +def _redis() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True) + + +def _config(directory: Path, router_settings: Mapping[str, JsonValue], callbacks: tuple[str, ...]) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(PROXY_CONFIG.read_text())) + path: Final = directory / "turn-admission.yaml" + path.write_text( + yaml.safe_dump( + { + **base, + "litellm_settings": { + **object_value(base["litellm_settings"]), + **({"callbacks": list(callbacks)} if callbacks else {}), + }, + "router_settings": { + **object_value(base["router_settings"]), + "redis_host": os.environ["REDIS_HOST"], + "redis_port": int(os.environ["REDIS_PORT"]), + **router_settings, + }, + } + ) + ) + return path + + +@contextmanager +def _admitting_proxy( + tmp_path_factory: pytest.TempPathFactory, + cert: tuple[Path, Path], + router_settings: Mapping[str, JsonValue], + callback_sources: Mapping[str, str] = MappingProxyType({}), +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("responses-ws-turn-admission") + for module, source in callback_sources.items(): + (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: + yield proxy + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-turn-admission-cert")) + + +@pytest.fixture(scope="module") +def limited(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with _admitting_proxy(tmp_path_factory, cert, {"optional_pre_call_checks": ["enforce_model_rate_limits"]}) as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def faulting(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with _admitting_proxy(tmp_path_factory, cert, {}, {"faulting_turn_check": FAULTING_TURN_CHECK}) as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def usage_routed(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with _admitting_proxy(tmp_path_factory, cert, {"routing_strategy": "usage-based-routing-v2"}) as proxy: + yield proxy + + +def _native(scenario: Scenario, peer: Peer, **limits: JsonValue) -> Deployment: + identity: Final = f"ws-admission-{uuid.uuid4().hex}" + name: Final = scenario.model( + model_info={"id": identity}, + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + **limits, + ) + return Deployment( + name, identity, lambda: sum(frame.get("type") == "response.create" for frame in drain(peer.frames)) + ) + + +def _upstream_chats(gateway: Gateway, provider_model: str) -> int: + with httpx.Client(base_url=gateway.upstream_url, timeout=15, trust_env=False) as upstream: + observed: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"] + assert isinstance(observed, list) + return sum(object_value(object_value(entry)["body"]).get("model") == provider_model for entry in observed) + + +def _managed(scenario: Scenario, **limits: JsonValue) -> Deployment: + identity: Final = f"ws-admission-{uuid.uuid4().hex}" + provider_model: Final = f"ws-managed-{uuid.uuid4().hex}" + name: Final = scenario.model( + model_info={"id": identity}, + model=f"{MANAGED_PROVIDER}/{provider_model}", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + **limits, + ) + return Deployment(name, identity, lambda: _upstream_chats(scenario.gateway, provider_model)) + + +def _socket_url(gateway: Gateway, model: str) -> str: + return f"{str(gateway.client.base_url).rstrip('/').replace('http://', 'ws://')}/v1/responses?model={model}" + + +@asynccontextmanager +async def _socket( + gateway: Gateway, key: str, model: str, headers: Mapping[str, str] | None = None +) -> AsyncIterator[websockets.ClientConnection]: + async with websockets.connect( + _socket_url(gateway, model), + additional_headers={"Authorization": f"Bearer {key}", **(headers or {})}, + open_timeout=10, + ) as connection: + yield connection + + +async def _say(connection: websockets.ClientConnection, model: str, text: str = "") -> tuple[dict[str, JsonValue], ...]: + return await turn(connection, create(model, text or f"turn {uuid.uuid4().hex}")) + + +def _assert_refused(received: tuple[dict[str, JsonValue], ...]) -> None: + assert [event.get("type") for event in received] == ["error"], received + assert received[0]["status"] == 429, received + assert object_value(received[0]["error"])["type"] == "rate_limit_exceeded", received + + +def _fresh_minute() -> str: + eventually(lambda: datetime.now(UTC), lambda now: now.second < 30, seconds=65) + return datetime.now(UTC).strftime("%H-%M") + + +def _deployment_spend(deployment: Deployment) -> float: + with _redis() as redis: + return float(redis.get(f"deployment_spend:{deployment.identity}:1d") or 0) + + +def _configured(gateway: Gateway, deployment: Deployment, section: str) -> dict[str, JsonValue]: + entries: Final = list_value(gateway.get("/model/info", {"litellm_model_id": deployment.identity})["data"]) + return object_value(object_value(entries[0])[section]) + + +def _configured_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue: + return _configured(gateway, deployment, "litellm_params").get("tpm") + + +def _tokens_counted(identity: str, minute: str) -> int: + with _redis() as redis: + return sum(int(redis.get(key) or 0) for key in redis.scan_iter(f"{identity}:*:tpm:{minute}")) + + +def _settled_count(read: Callable[[], int]) -> int: + reached: Final = eventually(read, lambda tokens: tokens > 0, seconds=15) + return held(read, lambda tokens: tokens == reached, holding=3, seconds=20) + + +async def _turns_against_tpm( + gateway: Gateway, key: str, deployment: Deployment +) -> tuple[tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], int]: + minute: Final = await asyncio.to_thread(_fresh_minute) + tokens: Final = partial(_tokens_counted, deployment.identity, minute) + async with _socket(gateway, key, deployment.name) as connection: + first: Final = await _say(connection, deployment.name) + counted_after_first: Final = await asyncio.to_thread(_settled_count, tokens) + second: Final = await _say(connection, deployment.name) + if second[-1].get("type") == "response.completed": + await _until(tokens, lambda counted: counted > counted_after_first) + return first, second, await _say(connection, deployment.name), counted_after_first + + +async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) -> float: + return await asyncio.to_thread(eventually, read, satisfied, 15) + + +def _failure_rows(key: str) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: any(row["status"] == "failure" for row in found), + seconds=70, + ) + return [row for row in rows if row["status"] == "failure"] + + +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_logged_once_as_its_failure( + limited: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + budget: Final = {"max_budget": TURN_SPEND / 2, "budget_duration": "1d"} + deployment: Final = _native(scenario, peer, **budget) if relay == "native" else _managed(scenario, **budget) + key: Final = scenario.key(models=[deployment.name]) + tag: Final = f"ws-turn-{uuid.uuid4().hex}" + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name, {"x-litellm-tags": tag}) as connection: + first: Final = await _say(connection, deployment.name) + await _until(lambda: _deployment_spend(deployment), lambda spent: spent > TURN_SPEND / 2) + return (first, await _say(connection, deployment.name), await _say(connection, deployment.name)) + + first, second, third = asyncio.run(session()) + + completed(first) + _assert_refused(second) + _assert_refused(third) + assert deployment.forwarded_turns() == 1 + failures: Final = _failure_rows(key) + assert len(failures) == 1, failures + assert failures[0]["model_group"] == deployment.name, failures + assert tag in list_value(failures[0]["request_tags"]), failures + assert float(failures[0]["spend"]) == 0.0, failures + error: Final = object_value(object_value(failures[0]["metadata"])["error_information"]) + assert str(error["error_code"]) == "429", error + + +def _set_turn_check_fault(gateway: Gateway, deployment: Deployment, fault: str) -> None: + gateway.request( + "PATCH", + f"/model/{deployment.identity}/update", + {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: fault}}, + ).raise_for_status() + eventually( + lambda: _configured(gateway, deployment, "model_info").get(TURN_CHECK_FAULT) or "", + lambda configured: configured == fault, + seconds=15, + ) + + +def _assert_turn_check_failed(received: tuple[dict[str, JsonValue], ...], fault: str) -> None: + assert [event.get("type") for event in received] == ["error"], received + assert received[0]["status"] == 500, received + error: Final = object_value(received[0]["error"]) + assert error["type"] == "server_error", received + assert fault in string_value(error["message"]), received + + +def _first_turn_held_until(release: threading.Event) -> Script: + def script(response_id: str) -> tuple[dict[str, JsonValue], ...]: + if response_id == "resp_peer_1": + release.wait(timeout=60) + return events(f"resp_{uuid.uuid4().hex}", "seven") + + return script + + +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_the_socket_serves_the_next_turn( + faulting: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, faulting.scenario() as scenario: + deployment: Final = _native(scenario, peer) if relay == "native" else _managed(scenario) + key: Final = scenario.key(models=[deployment.name]) + fault: Final = f"turn check unavailable {uuid.uuid4().hex}" + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(faulting, key, deployment.name) as connection: + first: Final = await _say(connection, deployment.name) + await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, fault) + second: Final = await _say(connection, deployment.name) + third: Final = await _say(connection, deployment.name) + await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, "") + return first, second, third, await _say(connection, deployment.name) + + first, second, third, fourth = asyncio.run(session()) + + completed(first) + _assert_turn_check_failed(second, fault) + _assert_turn_check_failed(third, fault) + completed(fourth) + assert deployment.forwarded_turns() == 2 + failures: Final = _failure_rows(key) + assert len(failures) == 1, failures + assert failures[0]["model_group"] == deployment.name, failures + recorded: Final = object_value(object_value(failures[0]["metadata"])["error_information"]) + assert fault in string_value(recorded["error_message"]), recorded + + +def test_a_turn_still_streaming_when_the_next_turns_check_fails_completes_and_is_billed( + faulting: Gateway, cert: tuple[Path, Path] +) -> None: + release: Final = threading.Event() + with responses_peer(cert, _first_turn_held_until(release)) as peer, faulting.scenario() as scenario: + deployment: Final = _native(scenario, peer) + key: Final = scenario.key(models=[deployment.name]) + fault: Final = f"turn check unavailable {uuid.uuid4().hex}" + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(faulting, key, deployment.name) as connection: + try: + await connection.send(create(deployment.name, f"streaming {uuid.uuid4().hex}")) + await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, fault) + rejected: Final = await _say(connection, deployment.name) + finally: + release.set() + return rejected, await until_terminal(connection, ()) + + rejected, streamed = asyncio.run(session()) + + _assert_turn_check_failed(rejected, fault) + billed: Final = completed(streamed) + assert deployment.forwarded_turns() == 1 + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: any(row["request_id"] == billed["id"] for row in found), + seconds=70, + ) + row: Final = next(row for row in rows if row["request_id"] == billed["id"]) + assert row["status"] == "success", row + assert float(row["spend"]) == pytest.approx(TURN_SPEND), row + + +def test_a_socket_opened_after_websocket_turns_spent_the_budget_is_refused_and_so_is_http( + limited: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, max_budget=TURN_SPEND / 2, budget_duration="1d") + key: Final = scenario.key(models=[deployment.name]) + + 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)) + await _until(lambda: _deployment_spend(deployment), lambda spent: spent > TURN_SPEND / 2) + 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 == 1011, close_code + assert peer.paths.qsize() == 1 + assert deployment.forwarded_turns() == 1 + http: Final = limited.request( + "POST", "/v1/responses", {"model": deployment.name, "input": f"after {uuid.uuid4().hex}"}, key=key + ) + assert http.status_code == 429, http.text + assert peer.paths.qsize() == 1 + + +def test_a_socket_the_router_refuses_to_open_gets_the_status_and_message_http_returns( + limited: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + name: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + tags=[f"granted-{uuid.uuid4().hex}"], + model_info={"enable_tag_filtering": True}, + ) + key: Final = scenario.key(models=[name]) + foreign: Final = {"x-litellm-tags": f"foreign-{uuid.uuid4().hex}"} + http: Final = limited.request( + "POST", "/v1/responses", {"model": name, "input": "tag routing"}, key=key, headers=foreign + ) + assert http.status_code != 200, http.text + http_error: Final = object_value(JSON_OBJECT.validate_json(http.content)["error"]) + + async def refused_open() -> tuple[dict[str, JsonValue], ...]: + async with _socket(limited, key, name, foreign) as connection: + return await _say(connection, name) + + received: Final = asyncio.run(refused_open()) + + assert [event.get("type") for event in received] == ["error"], received + assert received[0]["status"] == http.status_code, (received, http.text) + assert object_value(received[0]["error"])["message"] == string_value(http_error["message"]), ( + received, + http.text, + ) + assert peer.paths.qsize() == 0 + + +def test_each_turn_counts_as_one_request_toward_the_deployment_rpm(limited: Gateway, cert: tuple[Path, Path]) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, rpm=1) + key: Final = scenario.key(models=[deployment.name]) + _fresh_minute() + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return (await _say(connection, deployment.name), await _say(connection, deployment.name)) + + first, second = asyncio.run(session()) + + completed(first) + _assert_refused(second) + assert deployment.forwarded_turns() == 1 + + +def test_a_deployment_tpm_refuses_a_turn_only_once_earlier_turns_used_it_up( + limited: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, tpm=TURN_TOKENS + 2) + key: Final = scenario.key(models=[deployment.name]) + + first, second, third, counted_after_first = asyncio.run(_turns_against_tpm(limited, key, deployment)) + + completed(first) + completed(second) + _assert_refused(third) + assert counted_after_first == TURN_TOKENS + assert deployment.forwarded_turns() == 2 + + +def test_a_tpm_limit_lowered_while_the_socket_is_open_applies_to_its_next_turn( + limited: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, tpm=100 * TURN_TOKENS) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + + def lower_tpm() -> None: + limited.request( + "PATCH", f"/model/{deployment.identity}/update", {"litellm_params": {"tpm": TURN_TOKENS - 2}} + ).raise_for_status() + eventually(lambda: _configured_tpm(limited, deployment), lambda tpm: tpm == TURN_TOKENS - 2, seconds=15) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + first: Final = await _say(connection, deployment.name) + await _until(partial(_tokens_counted, deployment.identity, minute), lambda tokens: tokens > 0) + await asyncio.to_thread(lower_tpm) + return (first, await _say(connection, deployment.name)) + + first, second = asyncio.run(session()) + + completed(first) + _assert_refused(second) + assert deployment.forwarded_turns() == 1 + + +def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_turn( + limited: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, otpm=10_000) + key: Final = scenario.key(models=[deployment.name]) + _fresh_minute() + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(5)]) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns] == ["completed"] * 5 + assert deployment.forwarded_turns() == 5 + + +def test_usage_based_routing_refuses_a_turn_only_once_earlier_turns_used_up_the_deployment_tpm( + usage_routed: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, usage_routed.scenario() as scenario: + deployment: Final = _native(scenario, peer, tpm=TURN_TOKENS + 2) + key: Final = scenario.key(models=[deployment.name]) + + first, second, third, counted_after_first = asyncio.run(_turns_against_tpm(usage_routed, key, deployment)) + + completed(first) + completed(second) + _assert_refused(third) + assert counted_after_first == TURN_TOKENS + 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}" + name: Final = scenario.model(model_info={"id": identity}, tpm=1_000_000) + key: Final = scenario.key(models=[name]) + minute: Final = _fresh_minute() + answered: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": name, "messages": [{"role": "user", "content": f"http turn {uuid.uuid4().hex}"}]}, + key=key, + ) + assert answered.status_code == 200, answered.text + tokens: Final = object_value(JSON_OBJECT.validate_json(answered.content)["usage"])["total_tokens"] + assert isinstance(tokens, int) and tokens > 0, answered.text + + assert _settled_count(partial(_tokens_counted, identity, minute)) == tokens + + +def test_an_http_call_counts_its_tokens_once_toward_the_deployment_tpm_under_model_rate_limits( + limited: Gateway, +) -> None: + _assert_http_chat_counted_once(limited) + + +def test_an_http_call_counts_its_tokens_once_toward_the_deployment_tpm_under_usage_based_routing( + usage_routed: Gateway, +) -> None: + _assert_http_chat_counted_once(usage_routed) + + +def test_usage_based_routing_serves_every_turn_while_the_deployment_has_headroom( + usage_routed: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, usage_routed.scenario() as scenario: + deployment: Final = _native(scenario, peer, tpm=1_000_000, rpm=1_000) + key: Final = scenario.key(models=[deployment.name]) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(usage_routed, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(3)]) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns] == ["completed"] * 3 + assert deployment.forwarded_turns() == 3 + + +def test_a_frame_over_the_project_input_quota_gets_a_429_error_frame_and_never_reaches_the_provider( + limited: Gateway, +) -> None: + with limited.scenario() as scenario: + deployment: Final = _managed(scenario) + team: Final = scenario.team() + key: Final = scenario.key( + models=[deployment.name], + team_id=team, + project_id=scenario.project(team, model_itpm_limit={deployment.name: 100}), + ) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return ( + await _say(connection, deployment.name, "short"), + await _say(connection, deployment.name, "long " * 400), + await _say(connection, deployment.name, "long " * 400), + ) + + first, second, third = asyncio.run(session()) + + completed(first) + _assert_refused(third) + answered: Final = sum(received[-1].get("type") == "response.completed" for received in (first, second, third)) + assert deployment.forwarded_turns() == answered diff --git a/tests/integration/spend/test_failed_dispatch_tokens.py b/tests/integration/spend/test_failed_dispatch_tokens.py index 5778ff9654d..db5c49623eb 100644 --- a/tests/integration/spend/test_failed_dispatch_tokens.py +++ b/tests/integration/spend/test_failed_dispatch_tokens.py @@ -1,11 +1,13 @@ import json import uuid +from collections.abc import Callable, Mapping from typing import Final import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, list_value from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue @pytest.mark.covers("spend.failed_dispatch.failure_row_records_estimated_input_tokens") @@ -48,3 +50,79 @@ def test_provider_500_after_dispatch_records_estimated_prompt_tokens_on_failure_ assert row["status"] == "failure" and float(row["spend"]) == 0 and row["completion_tokens"] == 0, row assert row["prompt_tokens"] > 0, f"failure row lost the dispatched input tokens: {row}" assert row["total_tokens"] == row["prompt_tokens"], row + + +def _chat_body(model: str, text: str) -> Mapping[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": text}]} + + +def _responses_body(model: str, text: str) -> Mapping[str, JsonValue]: + return {"model": model, "input": text} + + +def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]: + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + + +FAILURE_ROUTES: Final = pytest.mark.parametrize( + ("route", "body"), + ( + ("/v1/chat/completions", _chat_body), + ("/v1/responses", _responses_body), + ("/v1/messages", _messages_body), + ), + ids=("chat", "responses", "messages"), +) + + +def _provider_500_failure_row( + gateway: Gateway, + route: str, + body: Callable[[str, str], Mapping[str, JsonValue]], + headers: Mapping[str, str] | None = None, +) -> tuple[str, dict[str, JsonValue]]: + def provider(request: Request) -> Reply: + assert request.method == "POST" + return Reply( + status=500, + body=b'{"error":{"message":"synthetic provider outage","type":"server_error","code":"500"}}', + ) + + with wire_server(provider) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + key: Final = scenario.key(models=[model]) + failed: Final = gateway.request( + "POST", route, body(model, f"failed dispatch group {uuid.uuid4().hex}"), key=key, headers=headers + ) + assert failed.status_code == 500, failed.text + assert wire.drain() + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, model_group, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (failed.headers["x-litellm-call-id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "failure", rows + return model, rows[0] + + +@FAILURE_ROUTES +def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for( + gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]] +) -> None: + model, row = _provider_500_failure_row(gateway, route, body) + + assert row["model_group"] == model, row + + +@FAILURE_ROUTES +def test_a_provider_500_failure_row_keeps_the_tags_the_caller_sent( + gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]] +) -> None: + tag: Final = f"failed-dispatch-{uuid.uuid4().hex}" + + _, row = _provider_500_failure_row(gateway, route, body, {"x-litellm-tags": tag}) + + assert tag in list_value(row["request_tags"]), row diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index c7450c852c0..0a888c41387 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2066,6 +2066,56 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): assert LitellmLogging.is_sync_litellm_request(captured["litellm_params"]) is False +@pytest.mark.asyncio +async def test_aresponses_websocket_turn_reaches_a_custom_logger_through_the_async_hook_only( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.litellm_core_utils import thread_pool_executor + from litellm.responses.main import base_llm_http_handler + + class InlineExecutor: + def submit(self, fn, *args, **kwargs): + fn(*args, **kwargs) + + sessions: list[LitellmLogging] = [] + + async def open_session(**kwargs): + sessions.append(kwargs["logging_obj"]) + + async_logged = asyncio.Event() + sync_call_ids: list[str] = [] + async_call_ids: list[str] = [] + + class CaptureLogger(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + sync_call_ids.append(kwargs["litellm_call_id"]) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async_call_ids.append(kwargs["litellm_call_id"]) + async_logged.set() + + logger = CaptureLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "success_callback", ["cache", logger]) + monkeypatch.setattr(litellm, "_async_success_callback", [logger]) + monkeypatch.setattr(thread_pool_executor, "executor", InlineExecutor()) + with patch.object( # test-quality-ok: the provider socket is the seam; how a turn of the session logs is under test + base_llm_http_handler, "async_responses_websocket", AsyncMock(side_effect=open_session) + ): + await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test") + turn = sessions[0].copy_for_new_call("ws-turn-2", dt_object.now()) + + await turn.dispatch_success_handlers( + litellm.ResponsesAPIResponse( + id="resp_turn_2", created_at=0, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[] + ), + prefer_async_handlers=True, + ) + await asyncio.wait_for(async_logged.wait(), timeout=10) + + assert (async_call_ids, sync_call_ids) == (["ws-turn-2"], []) + + @pytest.mark.asyncio async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 2f0e8a268c2..1f0dd67d86b 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4,7 +4,7 @@ import logging from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, create_autospec, patch import httpx import pytest @@ -136,6 +136,77 @@ async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_ assert metadata["original_key"] == "original_value" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_buckets", "expected_model_group"), + [ + ({"metadata": {}, "litellm_metadata": {"model_group": "ws-budget-sock"}}, "ws-budget-sock"), + ({"metadata": {"model_group": "caller-group"}, "litellm_metadata": {"model_group": "routed"}}, "caller-group"), + ({"metadata": {}, "litellm_metadata": {}}, ""), + ], +) +async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_failure_row( + metadata_buckets: dict[str, dict[str, str]], expected_model_group: str +): + request_data: Final = { + "model": "ws-budget-sock", + "input": "Reply with exactly the word two", + "proxy_server_request": {"request_id": "test_request_id"}, + **metadata_buckets, + } + + spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) + await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook( + request_data=request_data, + original_exception=litellm.RateLimitError(message="crossed budget", llm_provider="", model="ws-budget-sock"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + row: Final = get_logging_payload( + kwargs=spend_writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert row["model_group"] == expected_model_group + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_buckets", "expected_tags"), + [ + ({"metadata": {}, "litellm_metadata": {"tags": ["ws-tag"]}}, ["ws-tag"]), + ({"metadata": {"tags": ["caller-tag"]}, "litellm_metadata": {"tags": ["routed-tag"]}}, ["caller-tag"]), + ({"metadata": {}, "litellm_metadata": {"tags": "not-a-list"}}, []), + ({"metadata": {}, "litellm_metadata": {"tags": ["ws-tag", 1, {"a": 1}]}}, []), + ], +) +async def test_async_post_call_failure_hook_logs_the_routed_tags_on_the_failure_row( + metadata_buckets: dict[str, dict[str, object]], expected_tags: list[str] +): + request_data: Final = { + "model": "ws-budget-sock", + "input": "Reply with exactly the word two", + "proxy_server_request": {"request_id": "test_request_id"}, + **metadata_buckets, + } + + spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) + await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook( + request_data=request_data, + original_exception=litellm.RateLimitError(message="crossed budget", llm_provider="", model="ws-budget-sock"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + row: Final = get_logging_payload( + kwargs=spend_writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert json.loads(row["request_tags"]) == expected_tags + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_metadata(): logger = ProxyDBLogger() @@ -1105,6 +1176,38 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +class _FailedTrackingAlerts: + def __init__(self) -> None: + self.messages: tuple[str, ...] = () + + async def failed_tracking_alert(self, error_message: str, failing_model: str) -> None: + self.messages = (*self.messages, error_message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "alerted"), [("_aresponses_websocket", True), ("_arealtime", False)]) +async def test_track_cost_callback_alerts_on_a_websocket_turn_with_no_logging_object(call_type: str, alerted: bool): + kwargs: Final = { + "call_type": call_type, + "model": "gpt-5-mini", + "litellm_call_id": "turn-call-id", + "litellm_params": {}, + "stream": False, + } + alerts: Final = _FailedTrackingAlerts() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", alerts): + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.sleep(0) + + assert any(f"call_type: {call_type}" in message for message in alerts.messages) is alerted + + @pytest.mark.asyncio async def test_track_cost_callback_defers_in_progress_background_interaction(): # test-quality-ok: writing no spend row and raising no alert is the whole observable contract of the deferral path """ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index fc329eaec5f..fee4d7fab33 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -14,12 +14,15 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from fastapi import HTTPException from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.router import RouterErrors from tests._master_key import MASTER_KEY @@ -1648,7 +1651,7 @@ class TestManagedResponsesWSFirstMessage: processed: list = [] - async def fake_process(msg: str) -> None: + async def fake_process(msg: str, turn_admission: object = None) -> None: processed.append(msg) handler = ManagedResponsesWebSocketHandler( @@ -1693,30 +1696,32 @@ class TestResponsesWSStreamingFirstMessage: backend_ws.send.assert_awaited_once_with(first) +@pytest.mark.parametrize( + ("failure", "expected_status", "expected_type"), + [ + (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"), + (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"), + (RuntimeError("boom"), 500, "server_error"), + (ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['free']"), 401, "invalid_request_error"), + (ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), 401, "invalid_request_error"), + (HTTPException(status_code=403, detail="key not allowed"), 403, "invalid_request_error"), + ], + ids=["crossed-budget", "bad-request", "unexpected", "tag-routing", "proxy-auth", "http-exception"], +) +def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure( + failure: Exception, expected_status: int, expected_type: str +): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame + from litellm.proxy.utils import handle_exception_on_proxy + + frame: Final = json.loads(_responses_ws_failure_frame(failure)) + + assert frame["status"] == expected_status + assert frame["error"]["type"] == expected_type + assert frame["error"]["message"] == handle_exception_on_proxy(failure).message + + class TestWSSessionCostTracking: - @pytest.mark.asyncio - async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self): - """ - RouterBudgetLimiting.async_log_success_event must not raise when - call_type='_aresponses_websocket', even when standard_logging_object is None. - Per-turn costs are tracked by individual aresponses calls inside the session; - the outer session wrapper fires with result=None. - """ - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting - - limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) - kwargs = { - "call_type": "_aresponses_websocket", - "standard_logging_object": None, - "litellm_params": {"custom_llm_provider": "vertex_ai"}, - } - await limiter.async_log_success_event( - kwargs=kwargs, - response_obj=None, - start_time=None, - end_time=None, - ) - @pytest.mark.asyncio async def test_router_budget_limiter_skips_arealtime_call_type(self): """Same guard applies to _arealtime WS session wrappers.""" @@ -2049,6 +2054,44 @@ class TestResponsesWSFirstFrameModelAuth: assert failure["request_data"]["input"] == "hi" assert len(failure["request_data"]["litellm_call_id"]) == 36 + @pytest.mark.asyncio + async def test_frame_guard_bills_a_passed_frame_the_deployment_refuses_with_its_guardrail_records(self): + from litellm.exceptions import RateLimitError + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + return {**data, "input": data["input"] + " [guarded]"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + rate_limiter = MagicMock() + rate_limiter.async_post_call_failure_hook = AsyncMock() + user_api_key_dict = MagicMock() + refusal = RateLimitError(message="crossed budget", llm_provider="", model="") + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj + ) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi"})) + assert guarded.bill_refusal is not None + with patch.object(litellm, "callbacks", [rate_limiter, spend_logger]): + await guarded.bill_refusal(refusal) + + rate_limiter.async_post_call_failure_hook.assert_not_awaited() + spend_logger.async_post_call_failure_hook.assert_awaited_once() + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None}) + assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value + assert failure["original_exception"] is refusal + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + assert failure["request_data"]["input"] == "hi [guarded]" + @pytest.mark.asyncio async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index b38808fb57c..405f1a5b731 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -7,6 +7,7 @@ Tests that: 3. Providers without native websocket support use ManagedResponsesWebSocketHandler """ +import asyncio import json from collections.abc import Mapping from typing import Final @@ -1331,7 +1332,7 @@ class TestWebSocketProjectQuotaEnforcement: callback, not just the first one.""" @pytest.mark.asyncio - async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch): + async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock, MagicMock import litellm @@ -1377,11 +1378,11 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() assert aresponses_called is False mock_websocket.send_text.assert_called_once() - error_event = mock_websocket.send_text.call_args[0][0] - assert "rate_limit_exceeded" in error_event + error_event = json.loads(mock_websocket.send_text.call_args[0][0]) + assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded") @pytest.mark.asyncio - async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch): + async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock, MagicMock import litellm @@ -1462,7 +1463,8 @@ class TestWebSocketProjectQuotaEnforcement: assert allowed is False mock_backend_ws.send.assert_not_called() mock_websocket.send_text.assert_called_once() - assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0] + error_event = json.loads(mock_websocket.send_text.call_args[0][0]) + assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded") @pytest.mark.asyncio async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records( @@ -1580,6 +1582,425 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _turn_frames() -> tuple[str, str, str]: + return tuple(json.dumps({"type": "response.create", "input": word}) for word in ("one", "two", "three")) + + +def _admission_refusing_the_next_turn(): + from unittest.mock import AsyncMock + + from litellm.exceptions import RateLimitError + + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + + +class _ClientSocket: + """A client socket that, like Starlette's, stops receiving once the server closes it.""" + + def __init__(self, *frames: str) -> None: + self._frames: Final = iter(frames) + self.sent: Final[list[dict[str, object]]] = [] + self.close_code: int | None = None + + async def send_text(self, data: str) -> None: + self.sent.append(json.loads(data)) + + async def receive_text(self) -> str: + frame: Final = next(self._frames, None) + if self.close_code is not None or frame is None: + raise RuntimeError("WebSocket is not connected") + return frame + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.close_code = code + + +def _admission_crashing_the_next_turn(): + from unittest.mock import AsyncMock + + return AsyncMock(side_effect=[ValueError("custom check failed"), None]) + + +def _guard_blocking(word: str): + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + async def guard(frame: str) -> GuardedResponsesWsFrame: + if word in json.loads(frame)["input"]: + raise ValueError("blocked by guardrail") + return GuardedResponsesWsFrame(frame=frame) + + return guard + + +def _guard_recording(guardrail_information: tuple[Mapping[str, object], ...], bill_refusal: object): + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + async def guard(frame: str) -> GuardedResponsesWsFrame: + return GuardedResponsesWsFrame( + frame=frame, guardrail_information=guardrail_information, bill_refusal=bill_refusal + ) + + return guard + + +def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: + return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] + + +class TestWebSocketTurnAdmission: + """A turn after the connection's first is admitted by the deployment again before it reaches the provider.""" + + @pytest.mark.asyncio + async def test_native_relay_drops_a_later_turn_the_deployment_refuses(self): + from unittest.mock import AsyncMock + + first, second, third = _turn_frames() + cancel: Final = json.dumps({"type": "response.cancel"}) + turn_admission: Final = _admission_refusing_the_next_turn() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, cancel, third, Exception("stop")]) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + turn_admission=turn_admission, + ) + + await handler.client_to_backend() + + assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [ + json.loads(first), + json.loads(cancel), + json.loads(third), + ] + assert turn_admission.await_count == 2 + refusal: Final = json.loads(websocket.send_text.await_args.args[0]) + assert refusal["status"] == 429 + assert refusal["error"]["type"] == "rate_limit_exceeded" + assert "crossed budget" in refusal["error"]["message"] + + @pytest.mark.asyncio + async def test_native_relay_reports_a_refused_turn_as_the_socket_failure(self): + from unittest.mock import AsyncMock + + from litellm.exceptions import RateLimitError + + first, second, _ = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=Exception("stop")) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, Exception("client gone")]) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + turn_admission=_admission_refusing_the_next_turn(), + ) + + failure: Final = await handler.bidirectional_forward() + + assert isinstance(failure, RateLimitError) + assert "crossed budget" in str(failure) + + @pytest.mark.asyncio + async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch): + from unittest.mock import AsyncMock + + import litellm + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args: object, **kwargs: object): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + turn_admission: Final = _admission_refusing_the_next_turn() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, third, Exception("stop")]) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + turn_admission=turn_admission, + ) + + failure: Final = await handler.run() + + assert len(answered) == 2 + assert "two" not in json.dumps(answered) + assert turn_admission.await_count == 2 + assert isinstance(failure, RateLimitError) + assert "crossed budget" in str(failure) + refusals: Final = [ + json.loads(call.args[0]) + for call in websocket.send_text.await_args_list + if json.loads(call.args[0]).get("type") == "error" + ] + assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")] + + + @pytest.mark.asyncio + async def test_native_relay_runs_the_guardrail_before_the_deployment_counts_a_turn(self): + from unittest.mock import AsyncMock + + first, second, third = _turn_frames() + turn_admission: Final = _admission_refusing_the_next_turn() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + websocket: Final = _ClientSocket(second, third) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + frame_guard=_guard_blocking("two"), + turn_admission=turn_admission, + ) + + await handler.client_to_backend() + + assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [json.loads(first)] + assert turn_admission.await_count == 1 + assert [(status, error["type"]) for status, error in _errors(websocket)] == [ + (None, "invalid_request_error"), + (429, "rate_limit_exceeded"), + ] + + @pytest.mark.asyncio + async def test_managed_relay_runs_the_guardrail_before_the_deployment_counts_a_turn( + self, monkeypatch: pytest.MonkeyPatch + ): + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args: object, **kwargs: object): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + turn_admission: Final = _admission_refusing_the_next_turn() + websocket: Final = _ClientSocket(second, third) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + frame_guard=_guard_blocking("two"), + turn_admission=turn_admission, + ) + + await handler.run() + + assert answered == ["one"] + assert turn_admission.await_count == 1 + assert [(status, error["type"]) for status, error in _errors(websocket)] == [ + (None, "invalid_request_error"), + (429, "rate_limit_exceeded"), + ] + + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)] + ) + async def test_native_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses( + self, guardrail_information: tuple[Mapping[str, object], ...], billed: bool + ): + from unittest.mock import AsyncMock, call + + from litellm.exceptions import RateLimitError + + first, second, third = _turn_frames() + refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="") + bill_refusal: Final = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + handler: Final = _make_streaming( + websocket=_ClientSocket(second, third), + backend_ws=backend_ws, + request_data={}, + first_message=first, + frame_guard=_guard_recording(guardrail_information, bill_refusal), + turn_admission=AsyncMock(side_effect=[refusal, None]), + ) + + await handler.client_to_backend() + + assert [json.loads(sent.args[0])["input"] for sent in backend_ws.send.await_args_list] == ["one", "three"] + assert bill_refusal.await_args_list == ([call(refusal)] if billed else []) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)] + ) + async def test_managed_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses( + self, + monkeypatch: pytest.MonkeyPatch, + guardrail_information: tuple[Mapping[str, object], ...], + billed: bool, + ): + from unittest.mock import AsyncMock, call + + import litellm + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args: object, **kwargs: object): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="") + bill_refusal: Final = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=_ClientSocket(second, third), + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + frame_guard=_guard_recording(guardrail_information, bill_refusal), + turn_admission=AsyncMock(side_effect=[refusal, None]), + ) + + await handler.run() + + assert answered == ["one", "three"] + assert bill_refusal.await_args_list == ([call(refusal)] if billed else []) + + @pytest.mark.asyncio + async def test_native_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(self): + from unittest.mock import AsyncMock + + first, second, third = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=asyncio.Event().wait) + websocket: Final = _ClientSocket(second, third) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + turn_admission=_admission_crashing_the_next_turn(), + ) + + failure: Final = await handler.bidirectional_forward() + + assert isinstance(failure, ValueError) + assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [ + json.loads(first), + json.loads(third), + ] + assert _errors(websocket) == [(500, {"type": "server_error", "message": "custom check failed"})] + assert websocket.close_code is None + + @pytest.mark.asyncio + async def test_managed_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving( + self, monkeypatch: pytest.MonkeyPatch + ): + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args: object, **kwargs: object): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + websocket: Final = _ClientSocket(second, third) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + turn_admission=_admission_crashing_the_next_turn(), + ) + + failure: Final = await handler.run() + + assert isinstance(failure, ValueError) + assert answered == ["one", "three"] + assert _errors(websocket) == [(500, {"type": "server_error", "message": "custom check failed"})] + assert websocket.close_code is None + def _deployment_defaults(): from types import MappingProxyType diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 380aa573b3c..3c825bd560f 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -60,7 +60,9 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object]) @pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"]) +@pytest.mark.parametrize( + "call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"] +) async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type): """Non-chat surfaces carry the provider only on the standard logging payload.""" limiter = RouterBudgetLimiting( diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4ac917d706d..ce60cce8408 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,6 +1,7 @@ import json import os from datetime import datetime, timezone +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -13,6 +14,7 @@ from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( + CallTypes, ModelResponse, StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -482,6 +484,227 @@ async def test_deployment_callback_on_success(sync_mode): assert tpm_key is not None +@pytest.mark.asyncio +async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(): + import time + + model_list = [ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "rpm": 100}, + "model_info": {"id": "ws-deployment"}, + } + ] + router = Router(model_list=model_list) + standard_logging_payload = create_standard_logging_payload() + standard_logging_payload["total_tokens"] = 17 + + def turn_kwargs(call_type: str) -> dict: + return { + "call_type": call_type, + "litellm_params": { + "metadata": {"model_group": "gpt-5-mini"}, + "model_info": {"id": "ws-deployment"}, + }, + "standard_logging_object": standard_logging_payload, + } + + tpm_key = await router.deployment_callback_on_success( + kwargs=turn_kwargs("_aresponses_websocket"), + completion_response=None, + start_time=time.time(), + end_time=time.time(), + ) + assert tpm_key is not None + assert await router.cache.async_get_cache(key=tpm_key) == 17 + + realtime_key = await router.deployment_callback_on_success( + kwargs=turn_kwargs("_arealtime"), + completion_response=None, + start_time=time.time(), + end_time=time.time(), + ) + assert realtime_key is None + assert await router.cache.async_get_cache(key=tpm_key) == 17 + + +async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]: + routed: Final[dict[str, object]] = {} + + async def call(**kwargs: object) -> None: + routed.update(kwargs) + + call.__name__ = function_name + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call) + return routed + + +async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: str) -> None: + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + + turn: Final = Logging( + model="gpt-5-mini", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id=call_id, + function_id=call_id, + ) + turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}) + await turn.dispatch_success_handlers( + [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}], + prefer_async_handlers=True, + ) + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget(): + router = Router( + model_list=[ + { + "model_name": "ws-budget", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "max_budget": 0.001, + "budget_duration": "1d", + }, + "model_info": {"id": "ws-budget-dep"}, + } + ] + ) + assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses") + admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + await admit_turn() + await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002) + + with pytest.raises(litellm.RateLimitError, match="crossed budget"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): + router = Router( + model_list=[ + { + "model_name": "ws-rpm", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1}, + "model_info": {"id": "ws-rpm-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_is_checked_against_limits_changed_after_the_socket_opened(): + router = Router( + model_list=[ + { + "model_name": "ws-rpm-lowered", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 100}, + "model_info": {"id": "ws-rpm-lowered-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm-lowered", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await admit_turn() + + router.upsert_deployment( + Deployment( + model_name="ws-rpm-lowered", + litellm_params=LiteLLM_Params(model="openai/gpt-5-mini", api_key="sk-fake", rpm=1), + model_info=ModelInfo(id="ws-rpm-lowered-dep"), + ) + ) + + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_counts_toward_deployment_tpm(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "callbacks", []) + router = Router( + model_list=[ + { + "model_name": "ws-tpm", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, + "model_info": {"id": "ws-tpm-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) + routed: Final = await _kwargs_routed_to(router, "ws-tpm", "_aresponses_websocket") + admit_turn: Final = routed["turn_admission"] + assert callable(admit_turn) + await admit_turn() + + await _finish_35_token_websocket_turn(routed, "ws-tpm-turn") + + with pytest.raises(litellm.RateLimitError, match="TPM limit=30, current usage=35"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_is_refused_by_usage_based_routing_once_deployment_tpm_is_used( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", []) + router = Router( + model_list=[ + { + "model_name": "ws-tpm-v2", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, + "model_info": {"id": "ws-tpm-v2-dep"}, + } + ], + routing_strategy="usage-based-routing-v2", + ) + monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) + routed: Final = await _kwargs_routed_to(router, "ws-tpm-v2", "_aresponses_websocket") + admit_turn: Final = routed["turn_admission"] + assert callable(admit_turn) + await admit_turn() + + await _finish_35_token_websocket_turn(routed, "ws-tpm-v2-turn") + + with pytest.raises(litellm.RateLimitError): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): + router = Router( + model_list=[ + { + "model_name": "ws-otpm", + "litellm_params": {"model": "openai/ws-otpm-unpriced", "api_key": "sk-fake", "otpm": 5000}, + "model_info": {"id": "ws-otpm-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-otpm", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + for _ in range(3): + await admit_turn() + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """