From 63483e6a88340f0053e264b2f94f4c7f31d76219 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:11:12 -0700 Subject: [PATCH 01/25] fix(responses): write one spend-log row per native WebSocket turn Each response.create turn on a native Responses WebSocket connection now logs its own call as it finishes, with its own usage and the resp_ id the client saw as request_id. A failed turn is booked alone and no longer flips the other turns. --- litellm/cost_calculator.py | 21 + litellm/litellm_core_utils/litellm_logging.py | 47 +- litellm/llms/custom_httpx/llm_http_handler.py | 14 +- .../hooks/parallel_request_limiter_v3.py | 3 +- .../proxy/hooks/proxy_track_cost_callback.py | 6 +- .../proxy/response_api_endpoints/endpoints.py | 52 ++- litellm/responses/main.py | 13 +- litellm/responses/streaming_iterator.py | 205 +++++---- litellm/router.py | 3 +- litellm/router_strategy/budget_limiter.py | 3 +- .../types/responses/streaming_websocket.py | 11 +- litellm/types/utils.py | 1 + .../test_litellm_logging.py | 86 +++- .../hooks/test_parallel_request_limiter_v3.py | 34 ++ .../response_api_endpoints/test_endpoints.py | 72 ++-- .../test_responses_websocket_all_providers.py | 403 +++++++++++------- .../router_strategy/test_budget_limiter.py | 4 +- 17 files changed, 649 insertions(+), 329 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index da420ab0578..6256bd5b8d9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2892,6 +2892,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re class _ResponsesWsEventResponse(LiteLLMBaseModel): + id: str | None = None usage: Mapping[str, object] | None = None service_tier: str | None = None @@ -2943,6 +2944,26 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): ) return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects)) + @staticmethod + def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None: + response_ids: Final = tuple( + event.response.id + for event in (_ResponsesWsEvent.model_validate(result) for result in results) + if event.response is not None and event.response.id + ) + return response_ids[-1] if response_ids else None + + @staticmethod + def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject: + tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results) + raw_events: Final = cast(OpenAIRealtimeStreamList, list(results)) # cast-ok: raw WS frames, unvalidated field + return LiteLLMRealtimeStreamLoggingObject( + id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results), + usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results), + results=raw_events, + service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None, + ) + _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 715e492b159..8c1b55b38c6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -18,7 +18,7 @@ from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from httpx import Response -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, TypeAdapter import litellm from litellm import _custom_logger_compatible_callbacks_literal @@ -536,6 +536,10 @@ def _resolve_mantle_region_for_cost( return resolve_mantle_region(litellm_params or MappingProxyType({})) +_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) +_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -760,6 +764,32 @@ class Logging(LiteLLMLoggingBaseClass): self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None + def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": + connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params) + litellm_params: Final[dict[str, object]] = { + **connection_params, + **{ + key: {**value} + for key, value in connection_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + "litellm_call_id": litellm_call_id, + } + new_call: Final = copy.copy(self) + new_call.litellm_call_id = litellm_call_id + new_call.start_time = start_time + new_call.litellm_params = litellm_params + new_call.model_call_details = { + **self.model_call_details, + "litellm_call_id": litellm_call_id, + "start_time": start_time, + "first_api_call_start_time": None, + "litellm_params": litellm_params, + } + new_call.streaming_chunks = [] + new_call.sync_streaming_chunks = [] + return new_call + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -2362,20 +2392,9 @@ class Logging(LiteLLMLoggingBaseClass): ) elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped - combined_ws_usage: Final = ( - ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - ) - ) - ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier( + logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object( results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream ) - ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None - logging_result = LiteLLMRealtimeStreamLoggingObject( - usage=combined_ws_usage, - results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - service_tier=ws_service_tier, - ) elif ( self.call_type == CallTypes.llm_passthrough_route.value @@ -6628,7 +6647,7 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id", kwargs.get("litellm_call_id")) + id = _provider_response_id(response_obj) or kwargs.get("litellm_call_id") _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8..7a766ebd904 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -155,7 +155,10 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult -from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults +from litellm.types.responses.streaming_websocket import ( + ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CallTypes, @@ -6745,8 +6748,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None, **kwargs: object, - ) -> Exception | None: + ) -> None: """ Handles Responses API WebSocket mode. @@ -6780,7 +6784,7 @@ class BaseLLMHTTPHandler: **kwargs, ) await handler.run() - return None + return import websockets from websockets.asyncio.client import ClientConnection @@ -6901,8 +6905,9 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + on_turn_failure=on_turn_failure, ) - return await streaming.bidirectional_forward() + await streaming.bidirectional_forward() except websockets.exceptions.InvalidStatus as e: verbose_logger.exception("Error connecting to responses WS backend: %s", e) @@ -6916,7 +6921,6 @@ class BaseLLMHTTPHandler: pass else: raise Exception(f"Unexpected error while closing WebSocket: {close_error}") - return None def image_edit_handler( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5b865f3b5c3..54c30073dcd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5377,7 +5377,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): that usage instead of refunding it. """ try: - stash: Final = get_request_stash() + failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failed_request)) if stash is None: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0938d3f35c3..58ec0413d3a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -469,11 +469,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 7f9fcca369b..eaa6ec4c204 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -44,9 +44,12 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.responses.streaming_websocket import ResponsesWebSocketTurnFailureHook from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.utils import ProxyLogging from litellm.router import Router router: Final = APIRouter() @@ -1436,6 +1439,40 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec return MappingProxyType({key: value for key, value in hints.items() if value is not None}) +_RESPONSES_WS_FIRST_FRAME_KEYS: Final = frozenset({"first_message", "input", "previous_response_id"}) +_RESPONSES_WS_PER_TURN_DICT_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) + + +def _responses_ws_turn_request_data( + connection_request_data: Mapping[str, object], turn_logging_obj: "LiteLLMLoggingObj" +) -> dict[str, object]: + return { + **{key: value for key, value in connection_request_data.items() if key not in _RESPONSES_WS_FIRST_FRAME_KEYS}, + **{ + key: {**value} + for key, value in connection_request_data.items() + if key in _RESPONSES_WS_PER_TURN_DICT_KEYS and isinstance(value, dict) + }, + "litellm_call_id": turn_logging_obj.litellm_call_id, + "litellm_logging_obj": turn_logging_obj, + } + + +def _responses_ws_turn_failure_hook( + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: UserAPIKeyAuth, + connection_request_data: Mapping[str, object], +) -> ResponsesWebSocketTurnFailureHook: + async def book_turn_failure(turn_logging_obj: "LiteLLMLoggingObj", exception: Exception) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exception, + request_data=_responses_ws_turn_request_data(connection_request_data, turn_logging_obj), + ) + + return book_turn_failure + + 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 @@ -1584,7 +1621,10 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_turn_failure_hook=_responses_ws_turn_failure_hook(proxy_logging_obj, user_api_key_dict, data), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: @@ -1594,13 +1634,7 @@ async def _responses_websocket_session( llm_router=llm_router, user_model=user_model, ) - failure: Final = await llm_call - if isinstance(failure, Exception): - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=failure, - request_data=routed_data, - ) + await llm_call except Exception as e: verbose_proxy_logger.exception("Responses WebSocket error") with contextlib.suppress(Exception): @@ -1611,6 +1645,8 @@ async def _responses_websocket_session( request_data=routed_data, ) await websocket.close(code=1011, reason="Internal server error") + finally: + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f6bd4bfba54..5df161fc366 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -59,7 +59,10 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * -from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults +from litellm.types.responses.streaming_websocket import ( + ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import all_litellm_params from litellm.utils import ( @@ -2422,12 +2425,13 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_turn_failure_hook: ResponsesWebSocketTurnFailureHook | None = None, **kwargs, -) -> Exception | None: +) -> None: """ Private function to handle the Responses API WebSocket mode. - For PROXY use only. Returns the provider failure that ended the connection, if any. + For PROXY use only. ``responses_ws_turn_failure_hook`` books each failed turn of a native socket. Resolves the LLM provider from ``model``, looks up the matching ``BaseResponsesAPIConfig``, and hands off to @@ -2500,7 +2504,7 @@ async def _aresponses_websocket( deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS} first_message: Final = kwargs.get("first_message") - return await base_llm_http_handler.async_responses_websocket( + await base_llm_http_handler.async_responses_websocket( model=resolved_model, websocket=websocket, logging_obj=litellm_logging_obj, @@ -2517,5 +2521,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + on_turn_failure=responses_ws_turn_failure_hook, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 32a4f4284be..e4291bd366c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,8 +7,10 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable @@ -56,6 +58,7 @@ if TYPE_CHECKING: ResponsesBackendWebSocket, ResponsesClientWebSocket, ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, ) from litellm.types.router import LiteLLM_Params @@ -1831,9 +1834,59 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) +_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES + _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) +def _json_object_or_none(message: object) -> Mapping[str, object] | None: + if _is_json_object(message): + return message + if not isinstance(message, (str, bytes)): + return None + try: + parsed: Final = _load_json_value(message) + except json.JSONDecodeError: + return None + return parsed if _is_json_object(parsed) else None + + +def _user_texts_in_input_item(item: object) -> tuple[str, ...]: + if not _is_json_object(item) or item.get("type") != "message" or item.get("role") != "user": + return () + content: Final = item.get("content", []) + if isinstance(content, str): + return (content,) + return tuple( + text + for block in _json_array_or_empty(content) + if _is_json_object(block) + and block.get("type") == "input_text" + and isinstance(text := block.get("text"), str) + and text + ) + + +def _logged_user_messages(client_frame: object) -> tuple[dict[str, object], ...]: + frame: Final = _json_object_or_none(client_frame) + if frame is None or frame.get("type") != "response.create": + return () + input_items: Final = frame.get("input", []) + if isinstance(input_items, str): + return ({"role": "user", "content": input_items},) + texts: Final = chain.from_iterable(_user_texts_in_input_item(item) for item in _json_array_or_empty(input_items)) + return tuple({"role": "user", "content": text} for text in texts) + + +@dataclass(frozen=True, slots=True) +class _ResponsesWsTurn: + logging_obj: LiteLLMLoggingObj + events: tuple[_MutableJsonObject, ...] = () + + def with_event(self, event: _MutableJsonObject) -> _ResponsesWsTurn: + return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) + + def _ws_event_error(event: Mapping[str, object]) -> object: if event.get("type") == "error": return event.get("error") @@ -1916,6 +1969,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1925,8 +1979,8 @@ class ResponsesWebSocketStreaming: litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} self.custom_llm_provider: str | None = custom_llm_provider - self.messages: list[_MutableJsonObject] = [] - self.input_messages: list[dict[str, object]] = [] + self.on_turn_failure: ResponsesWebSocketTurnFailureHook | None = on_turn_failure + self._open_turns: tuple[_ResponsesWsTurn, ...] = () self.first_message = first_message self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] @@ -1936,104 +1990,70 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: - return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES + def _start_turn(self, client_frame: str | None) -> None: + turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + turn_logging_obj: Final = self.logging_obj.copy_for_new_call( + litellm_call_id=str(uuid.uuid4()), start_time=turn_start + ) + turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) + turn_logging_obj.pre_call(input=client_frame, api_key="") + self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) - def _store_event(self, event: str | bytes | dict[str, object]) -> None: - if isinstance(event, bytes): - event = event.decode("utf-8") - if isinstance(event, str): - try: - event_obj = _load_json_object(event) - except (json.JSONDecodeError, TypeError): - return - else: - event_obj = event + def _record_client_frame(self, client_frame: str) -> None: + frame: Final = _json_object_or_none(client_frame) + if frame is not None and frame.get("type") == "response.create": + self._start_turn(client_frame) - if self._should_store_event(event_obj): - self.messages.append(event_obj) + def _record_backend_event(self, event: str) -> None: + event_obj: Final = _json_object_or_none(event) + if event_obj is None or event_obj.get("type") not in RESPONSES_WS_LOGGED_EVENT_TYPES: + return + if not self._open_turns: + self._start_turn(None) + turn: Final = self._open_turns[0].with_event(event_obj) + if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES: + self._open_turns = self._open_turns[1:] + self._log_turn(turn) + return + self._open_turns = (turn, *self._open_turns[1:]) - def _collect_input_from_client_event(self, message: object) -> None: - """Extract user input content from response.create for logging.""" - try: - if isinstance(message, str): - msg_obj = _load_json_object(message) - elif _is_json_object(message): - msg_obj = message - else: - return + def _log_open_turns(self) -> None: + open_turns: Final = self._open_turns + self._open_turns = () + for turn in open_turns: + if turn.events: + self._log_turn(turn) - if msg_obj.get("type") != "response.create": - return + def _log_turn(self, turn: _ResponsesWsTurn) -> None: + _spawn_logging_task( + asyncio.get_running_loop(), self._dispatch_turn(turn), task_name="Responses WS turn logging" + ) - input_items: Final = msg_obj.get("input", []) - if isinstance(input_items, str): - self.input_messages.append({"role": "user", "content": input_items}) - return - - if _is_json_array(input_items): - for item in input_items: - if not _is_json_object(item): - continue - if item.get("type") == "message" and item.get("role") == "user": - content = item.get("content", []) - if isinstance(content, str): - self.input_messages.append({"role": "user", "content": content}) - elif _is_json_array(content): - for c in content: - if _is_json_object(c) and c.get("type") == "input_text": - text = c.get("text", "") - if text: - self.input_messages.append({"role": "user", "content": text}) - except (json.JSONDecodeError, AttributeError, TypeError): - pass - - def _store_input(self, message: object) -> None: - self._collect_input_from_client_event(message) - if self.logging_obj: - self.logging_obj.pre_call(input=message, api_key="") - - def _failure_exception(self) -> Exception | None: + async def _dispatch_turn(self, turn: _ResponsesWsTurn) -> None: + events: Final = list(turn.events) failed_event: Final = next( - (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None + (event for event in events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None ) if failed_event is None: - return None - return _map_stream_error_to_exception( + await turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True) + return + exception: Final = _map_stream_error_to_exception( _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" ) - - async def _log_messages(self) -> None: - if not self.logging_obj: - return - if self.input_messages: - self.logging_obj.model_call_details["messages"] = self.input_messages - if not self.messages: - return - exception: Final = self._failure_exception() - if exception is None: - asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) - return - self._record_usage_for_failure() - traceback_exception: Final = "".join(traceback.format_exception(exception)) - asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + self._record_usage_for_failure(turn.logging_obj, events) + await turn.logging_obj.dispatch_failure_handlers( + exception, "".join(traceback.format_exception(exception)), prefer_async_handlers=True ) + if self.on_turn_failure is not None: + await self.on_turn_failure(turn.logging_obj, exception) - def _record_usage_for_failure(self) -> None: + @staticmethod + def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[_MutableJsonObject]) -> None: from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor - from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject - usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - self.messages - ) - tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages) - service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None - logging_result: Final = LiteLLMRealtimeStreamLoggingObject( - usage=usage, results=self.messages, service_tier=service_tier - ) - response_cost: Final = self.logging_obj.response_cost_calculator(result=logging_result) or 0.0 - self.logging_obj.record_partial_usage_for_failure(usage, response_cost) + logging_result: Final = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events) + response_cost: Final = turn_logging_obj.response_cost_calculator(result=logging_result) or 0.0 + turn_logging_obj.record_partial_usage_for_failure(logging_result.usage, response_cost) def _wrap_response_event(self, response_str: str) -> str: try: @@ -2092,7 +2112,7 @@ class ResponsesWebSocketStreaming: # Log the output-masked form so PII redacted by apply_to_output # guardrails does not appear in success logs. - self._store_event(wrapped_str) + self._record_backend_event(wrapped_str) await self.websocket.send_text(wrapped_str) @@ -2101,7 +2121,7 @@ class ResponsesWebSocketStreaming: except Exception as e: verbose_logger.exception("Error in responses WS backend_to_client: %s", e) finally: - await self._log_messages() + self._log_open_turns() def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ @@ -2446,8 +2466,7 @@ class ResponsesWebSocketStreaming: try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) - self._store_input(masked_first) - self._store_event(masked_first) + self._record_client_frame(masked_first) await self.backend_ws.send(masked_first) while True: @@ -2455,14 +2474,13 @@ class ResponsesWebSocketStreaming: if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) - self._store_input(masked) - self._store_event(masked) + self._record_client_frame(masked) await self.backend_ws.send(masked) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) - async def bidirectional_forward(self) -> Exception | None: + async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client()) try: await self.client_to_backend() @@ -2479,7 +2497,6 @@ class ResponsesWebSocketStreaming: await self.backend_ws.close() except Exception: pass - return self._failure_exception() # --------------------------------------------------------------------------- diff --git a/litellm/router.py b/litellm/router.py index f173cdb37b4..309e8c2bd1e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8406,8 +8406,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 diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 8ba1da824c6..8f81ec9414d 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -495,8 +495,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/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index f369cbcebf8..de76398b5cc 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -2,10 +2,13 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass -from typing import Protocol +from typing import TYPE_CHECKING, Protocol from litellm.types.guardrails import PresidioPerRequestConfig +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + class ResponsesClientWebSocket(Protocol): """Client-facing websocket surface used by the Responses API websocket handlers.""" @@ -43,6 +46,12 @@ class PresidioGuardrailCallback(Protocol): ) -> str: ... +class ResponsesWebSocketTurnFailureHook(Protocol): + """Books one failed ``response.create`` turn of a native Responses API socket from that turn's logging object.""" + + async def __call__(self, turn_logging_obj: Logging, exception: Exception) -> None: ... + + @dataclass(frozen=True, slots=True) class ResponsesWebSocketRequestDefaults: """Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4a962e5b3bd..90d6f86986e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4478,6 +4478,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase): results: SkipValidation[OpenAIRealtimeStreamList] usage: Usage service_tier: str | None = None + id: str | None = None _hidden_params: dict = {} @property diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 771ed5d2e15..8ea3f36068a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1998,34 +1998,34 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio -async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( +async def test_aresponses_websocket_writes_no_connection_level_success_log( monkeypatch: pytest.MonkeyPatch, ): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler - success_events = [] + logged_call_ids = [] class CaptureLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - success_events.append(response_obj) + logged_call_ids.append(kwargs.get("litellm_call_id")) monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()]) monkeypatch.setattr(litellm, "failure_callback", []) monkeypatch.setattr(litellm, "_async_failure_callback", []) monkeypatch.setattr(litellm, "success_callback", []) monkeypatch.setattr(litellm, "_async_success_callback", []) - failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai") - with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test - base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure) + with patch.object( # test-quality-ok: the provider socket is the seam; each turn logs inside it, so the wrapper must not + base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=None) ): - outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test") + await litellm._aresponses_websocket( + model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test", litellm_call_id="socket-call" + ) await asyncio.sleep(0) with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) - assert outcome is failure - assert success_events == [] + assert "socket-call" not in logged_call_ids @pytest.mark.asyncio @@ -8078,6 +8078,74 @@ def test_normalize_logging_result_prices_responses_websocket_at_returned_service assert priority_http_cost > default_http_cost +@pytest.mark.parametrize( + ("completed_response_id", "expected_request_id"), + (("resp_wrapped_turn_2", "resp_wrapped_turn_2"), (None, "turn-2-call")), +) +def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw( + completed_response_id: str | None, expected_request_id: str +): + events = [ + {"type": "response.created", "response": {"id": completed_response_id}}, + { + "type": "response.completed", + "response": { + "id": completed_response_id, + "usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}, + }, + }, + ] + logging_obj = _responses_ws_logging_obj() + now = datetime.datetime.now() + + payload = litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload( + kwargs={"litellm_call_id": "turn-2-call", "model": "gpt-4o", "messages": []}, + init_response_obj=logging_obj.normalize_logging_result(result=events), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["id"] == expected_request_id + + +def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): + reservation = {"reserved_cost": 0.5, "finalized": False} + connection = _responses_ws_logging_obj() + connection_start = connection.start_time + connection.litellm_params = { + "metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation}, + "litellm_metadata": {"model_info": {"id": "dep-1"}}, + "litellm_call_id": connection.litellm_call_id, + } + connection.model_call_details = { + **connection.model_call_details, + "litellm_params": connection.litellm_params, + "first_api_call_start_time": connection_start, + } + turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47) + + turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start) + turn.litellm_params["metadata"]["status"] = "failure" + turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o" + turn.model_call_details["standard_logging_object"] = {"id": "resp_1"} + + assert (turn.litellm_call_id, turn.start_time) == ("turn-1-call", turn_start) + assert turn.model_call_details["litellm_call_id"] == "turn-1-call" + assert turn.model_call_details["start_time"] == turn_start + assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call" + assert turn.model_call_details["first_api_call_start_time"] is None + assert turn.litellm_params["metadata"]["budget_reservation"] is reservation + assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start) + assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test" + assert "status" not in connection.litellm_params["metadata"] + assert "model_group" not in connection.litellm_params["litellm_metadata"] + assert "standard_logging_object" not in connection.model_call_details + assert connection.model_call_details["first_api_call_start_time"] == connection_start + + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d9cf8536912..28589a2f400 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4294,6 +4294,40 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): ) +@pytest.mark.asyncio +async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3(): + """ + A native Responses WebSocket holds one parallel slot for the whole socket and books each failed turn + under that turn's own call id. A failed turn must not free the slot while the socket is still open. + """ + _api_key = hash_token("sk-ws-turn-failure") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type="" + ) + + await handler.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "litellm_call_id": "turn-2-call"}, + original_exception=Exception("previous_response_not_found"), + user_api_key_dict=user_api_key_dict, + ) + in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) + await handler.async_post_call_failure_hook( + request_data=connection_data, + original_exception=Exception("socket closed"), + user_api_key_dict=user_api_key_dict, + ) + + assert in_flight_after_turn_failure == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + + @pytest.mark.asyncio async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5dac1d25572..ce967403265 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -542,29 +542,6 @@ class TestResponsesWSStreamingFirstMessage: 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.""" @@ -820,8 +797,7 @@ class TestResponsesWSFirstFrameModelAuth: ws.close.assert_not_awaited() @pytest.mark.asyncio - @pytest.mark.parametrize("provider_rejected", [True, False]) - async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): + async def test_endpoint_books_each_failed_turn_as_its_own_failed_request(self): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) @@ -833,23 +809,36 @@ class TestResponsesWSFirstFrameModelAuth: ws.url = "ws://testserver/v1/responses" ws.accept = AsyncMock() ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"}) ) ws.close = AsyncMock() + connection_metadata = {"user_api_key_alias": "codex", "model_group": "gpt-4o-mini"} processor = MagicMock() processor.common_processing_pre_call_logic = AsyncMock( - return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + return_value=( + { + "model": "gpt-4o-mini", + "metadata": connection_metadata, + "litellm_metadata": {}, + "litellm_call_id": "connection-call", + }, + MagicMock(), + ) ) + turn_logging_obj = MagicMock() + turn_logging_obj.litellm_call_id = "turn-2-call" failure = litellm.BadRequestError( - message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai" + message="previous_response_not_found", model="gpt-4o-mini", llm_provider="openai" ) - async def fake_llm_call(): - return failure if provider_rejected else None + async def relay_with_one_failed_turn(): + routed = mock_route_request.await_args.kwargs["data"] + await routed["responses_ws_turn_failure_hook"](turn_logging_obj, failure) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock() user_api_key_dict = MagicMock() with ( @@ -857,15 +846,15 @@ class TestResponsesWSFirstFrameModelAuth: "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", new_callable=AsyncMock, ), - patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test + patch( # test-quality-ok: the pre-call processor needs a live proxy; how the endpoint books a failed turn is under test "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch( # test-quality-ok: routing is the seam that hands back the relay's outcome + patch( # test-quality-ok: routing is the seam that hands the turn failure hook to the relay "litellm.proxy.route_llm_request.route_request", new_callable=AsyncMock, - return_value=fake_llm_call(), - ), + side_effect=lambda **_: relay_with_one_failed_turn(), + ) as mock_route_request, patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj, @@ -878,14 +867,19 @@ class TestResponsesWSFirstFrameModelAuth: ) ws.close.assert_not_awaited() - if not provider_rejected: - proxy_logging_obj.post_call_failure_hook.assert_not_awaited() - return proxy_logging_obj.post_call_failure_hook.assert_awaited_once() booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs assert booked["original_exception"] is failure assert booked["user_api_key_dict"] is user_api_key_dict - assert booked["request_data"]["model"] == "gpt-4o-mini" + turn_request = booked["request_data"] + assert turn_request["litellm_call_id"] == "turn-2-call" + assert turn_request["litellm_logging_obj"] is turn_logging_obj + assert turn_request["model"] == "gpt-4o-mini" + assert turn_request["metadata"] == connection_metadata + assert turn_request["metadata"] is not connection_metadata + assert "input" not in turn_request + assert "first_message" not in turn_request + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) @pytest.mark.asyncio async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self): @@ -914,6 +908,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -951,6 +946,7 @@ class TestResponsesWSFirstFrameModelAuth: assert booked["original_exception"] is rejection assert booked["user_api_key_dict"] is user_api_key_dict assert booked["request_data"]["model"] == "gpt-4o-mini" + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) @pytest.mark.asyncio async def test_reruns_model_auth_for_first_frame_model(self): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index b29e9a6b7db..5b95c472480 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1898,6 +1898,29 @@ def _make_streaming(**kwargs): return ResponsesWebSocketStreaming(**kwargs) +class _TurnLoggingRecorder: + """Stands in for a socket's connection-level Logging object and keeps every per-turn copy it hands out.""" + + def __init__(self, response_cost: float = 0.0) -> None: + self.response_cost = response_cost + self.turns: list[MagicMock] = [] + self.connection = MagicMock() + self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call) + + def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock: + from unittest.mock import AsyncMock + + turn = MagicMock() + turn.litellm_call_id = litellm_call_id + turn.start_time = start_time + turn.model_call_details = {} + turn.dispatch_success_handlers = AsyncMock() + turn.dispatch_failure_handlers = AsyncMock() + turn.response_cost_calculator = MagicMock(return_value=self.response_cost) + self.turns.append(turn) + return turn + + class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" @@ -3051,12 +3074,11 @@ class TestNativeWebSocketEncryptedContentAffinity: Exception("stop"), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() + recorder = _TurnLoggingRecorder() handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, - logging_obj=logging_obj, + logging_obj=recorder.connection, request_data={"litellm_metadata": dict(_AFFINITY_METADATA)}, custom_llm_provider="openai", ) @@ -3075,7 +3097,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content await asyncio.sleep(0) - logged = logging_obj.dispatch_success_handlers.await_args[0][0] + logged = recorder.turns[0].dispatch_success_handlers.await_args[0][0] assert logged[0]["response"]["id"] == completed["response"]["id"] @pytest.mark.asyncio @@ -3163,14 +3185,11 @@ class TestNativeWebSocketEncryptedContentAffinity: Exception("stop"), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.0) + recorder = _TurnLoggingRecorder() handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, - logging_obj=logging_obj, + logging_obj=recorder.connection, request_data={}, authorized_model="gpt-5.6", custom_llm_provider="openai", @@ -3179,163 +3198,257 @@ class TestNativeWebSocketEncryptedContentAffinity: await handler.backend_to_client() await asyncio.sleep(0) - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() - exception = logging_obj.dispatch_failure_handlers.await_args[0][0] + (turn,) = recorder.turns + turn.dispatch_success_handlers.assert_not_awaited() + turn.dispatch_failure_handlers.assert_awaited_once() + exception = turn.dispatch_failure_handlers.await_args[0][0] assert exception.status_code == expected_status assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception) - @pytest.mark.asyncio - async def test_backend_to_client_bills_completed_turns_before_a_failure(self): + +def _completed_turn(response_id: str, input_tokens: int, output_tokens: int) -> list[dict[str, object]]: + return [ + {"type": "response.created", "response": {"id": response_id, "status": "in_progress"}}, + {"type": "response.output_text.delta", "delta": "hi"}, + { + "type": "response.completed", + "response": { + "id": response_id, + "status": "completed", + "output": [], + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + }, + }, + ] + + +def _response_create(text: str) -> str: + return json.dumps( + { + "type": "response.create", + "model": "gpt-5.4-mini", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], + } + ) + + +class _ScriptedBackend: + """Answers each forwarded response.create with the next scripted turn of upstream events.""" + + def __init__(self, turns: list[list[dict[str, object]]]) -> None: import asyncio + + self.turns = list(turns) + self.events: asyncio.Queue[str] = asyncio.Queue() + self.sent: list[str] = [] + + async def send(self, frame: str) -> None: + self.sent.append(frame) + if json.loads(frame).get("type") == "response.create" and self.turns: + for event in self.turns.pop(0): + self.events.put_nowait(json.dumps(event)) + + async def recv(self, decode: bool = False) -> str: + return await self.events.get() + + async def close(self) -> None: + return None + + +async def _until(condition) -> bool: + import asyncio + + for _ in range(200): + if condition(): + return True + await asyncio.sleep(0) + return False + + +class _ScriptedClient: + """Sends each frame only after the previous turn's log was dispatched, then hangs up.""" + + def __init__(self, frames: list[str], recorder: _TurnLoggingRecorder) -> None: + self.frames = list(frames) + self.recorder = recorder + self.sent_text: list[str] = [] + self.logged_turns_before_each_frame: list[int] = [] + self.logged_turns_before_hang_up: int | None = None + + def _logged_turns(self) -> int: + return sum( + 1 + for turn in self.recorder.turns + if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count + ) + + async def receive_text(self) -> str: + sent_so_far = len(self.logged_turns_before_each_frame) + await _until(lambda: self._logged_turns() >= sent_so_far) + if self.frames: + self.logged_turns_before_each_frame.append(self._logged_turns()) + return self.frames.pop(0) + self.logged_turns_before_hang_up = self._logged_turns() + raise Exception("client hung up") + + async def send_text(self, text: str) -> None: + self.sent_text.append(text) + + +class TestNativeWebSocketPerTurnLogging: + """One native Responses socket carries many response.create turns, and each turn is its own logged call.""" + + @pytest.mark.asyncio + async def test_each_turn_is_logged_as_its_own_call_while_the_socket_stays_open(self): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [_completed_turn("resp_1", 10, 1), _completed_turn("resp_2", 20, 2), _completed_turn("resp_3", 30, 3)] + ) + handler = _make_streaming( + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + ) + + await handler.bidirectional_forward() + + assert client.logged_turns_before_each_frame == [0, 1, 2] + assert client.logged_turns_before_hang_up == 3 + assert len({turn.litellm_call_id for turn in recorder.turns}) == 3 + assert recorder.connection.litellm_call_id not in {turn.litellm_call_id for turn in recorder.turns} + logged_events = [turn.dispatch_success_handlers.await_args[0][0] for turn in recorder.turns] + sent_frames = [json.loads(text) for text in client.sent_text] + ids_the_client_saw = [frame["response"]["id"] for frame in sent_frames if frame["type"] == "response.completed"] + assert [events[-1]["response"]["id"] for events in logged_events] == ids_the_client_saw + assert [ + ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id)["response_id"] + for response_id in ids_the_client_saw + ] == ["resp_1", "resp_2", "resp_3"] + assert [len(events) for events in logged_events] == [2, 2, 2] + assert [events[-1]["response"]["usage"]["input_tokens"] for events in logged_events] == [10, 20, 30] + assert [turn.model_call_details["messages"] for turn in recorder.turns] == [ + [{"role": "user", "content": "one"}], + [{"role": "user", "content": "two"}], + [{"role": "user", "content": "three"}], + ] + for turn in recorder.turns: + turn.dispatch_failure_handlers.assert_not_awaited() + + @pytest.mark.asyncio + async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self): from unittest.mock import AsyncMock import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - websocket = MagicMock() - websocket.send_text = AsyncMock() - backend_ws = MagicMock() - backend_ws.recv = AsyncMock( - side_effect=[ - json.dumps( + recorder = _TurnLoggingRecorder() + on_turn_failure = AsyncMock() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 1), + [ { - "type": "response.completed", - "response": { - "id": "resp_1", - "status": "completed", - "output": [], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "code": "previous_response_not_found", + "message": "Previous response with id 'resp_x' not found.", }, } - ), - json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}), - Exception("stop"), + ], + _completed_turn("resp_3", 30, 3), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.01) - handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={}) + handler = _make_streaming( + websocket=client, + backend_ws=backend, + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini", + custom_llm_provider="openai", + on_turn_failure=on_turn_failure, + ) - await handler.backend_to_client() - await asyncio.sleep(0) + await handler.bidirectional_forward() - logging_obj.record_partial_usage_for_failure.assert_called_once() - usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) + first, failed, third = recorder.turns + for successful in (first, third): + successful.dispatch_success_handlers.assert_awaited_once() + successful.dispatch_failure_handlers.assert_not_awaited() + failed.dispatch_success_handlers.assert_not_awaited() + failed.dispatch_failure_handlers.assert_awaited_once() + on_turn_failure.assert_awaited_once() + booked_turn, booked_exception = on_turn_failure.await_args[0] + assert booked_turn is failed + assert booked_exception is failed.dispatch_failure_handlers.await_args[0][0] + assert booked_exception.status_code == 400 + assert "not found" in str(booked_exception) + + @pytest.mark.asyncio + async def test_a_failed_turn_records_only_its_own_partial_usage(self): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder(response_cost=0.01) + client = _ScriptedClient([_response_create(text) for text in ("one", "two")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 5), + [ + {"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}, + { + "type": "response.failed", + "response": { + "id": "resp_2", + "status": "failed", + "error": {"code": "server_error", "message": "upstream blew up"}, + "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + }, + }, + ], + ] + ) + handler = _make_streaming( + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + ) + + await handler.bidirectional_forward() + + first, failed = recorder.turns + first.record_partial_usage_for_failure.assert_not_called() + usage, response_cost = failed.record_partial_usage_for_failure.call_args[0] + assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0) assert response_cost == 0.01 - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() @pytest.mark.asyncio - async def test_bidirectional_forward_returns_the_provider_failure(self): - import asyncio - from unittest.mock import AsyncMock - + async def test_a_turn_still_in_flight_when_the_socket_closes_is_logged_with_what_it_received(self): import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - backend_drained = asyncio.Event() - backend_events = [ - json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), - json.dumps( - { - "type": "error", - "status": 400, - "error": { - "type": "invalid_request_error", - "code": "invalid_encrypted_content", - "message": "could not be verified", - }, - } - ), - ] - - async def recv(decode=False): - if backend_events: - return backend_events.pop(0) - backend_drained.set() - raise Exception("stop") - - async def receive_text(): - await backend_drained.wait() - raise Exception("client gone") - - websocket = MagicMock() - websocket.send_text = AsyncMock() - websocket.receive_text = receive_text - backend_ws = MagicMock() - backend_ws.recv = recv - backend_ws.send = AsyncMock() - backend_ws.close = AsyncMock() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.0) + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 1), + [{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}], + ] + ) handler = _make_streaming( - websocket=websocket, - backend_ws=backend_ws, - logging_obj=logging_obj, - request_data={}, - authorized_model="gpt-5.6", - custom_llm_provider="openai", + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} ) - failure = await handler.bidirectional_forward() + await handler.bidirectional_forward() + await _until(lambda: recorder.turns[1].dispatch_success_handlers.await_count == 1) - assert isinstance(failure, Exception) - assert failure.status_code == 400 - assert "could not be verified" in str(failure) - - @pytest.mark.asyncio - async def test_bidirectional_forward_returns_none_after_a_completed_turn(self): - import asyncio - from unittest.mock import AsyncMock - - import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - - backend_drained = asyncio.Event() - backend_events = [ - json.dumps( - { - "type": "response.completed", - "response": { - "id": "resp_1", - "status": "completed", - "output": [], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - }, - } - ), - ] - - async def recv(decode=False): - if backend_events: - return backend_events.pop(0) - backend_drained.set() - raise Exception("stop") - - async def receive_text(): - await backend_drained.wait() - raise Exception("client gone") - - websocket = MagicMock() - websocket.send_text = AsyncMock() - websocket.receive_text = receive_text - backend_ws = MagicMock() - backend_ws.recv = recv - backend_ws.send = AsyncMock() - backend_ws.close = AsyncMock() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - handler = _make_streaming( - websocket=websocket, - backend_ws=backend_ws, - logging_obj=logging_obj, - request_data={}, - authorized_model="gpt-5.6", - custom_llm_provider="openai", - ) - - assert await handler.bidirectional_forward() is None + first, in_flight, unanswered = recorder.turns + first.dispatch_success_handlers.assert_awaited_once() + in_flight.dispatch_success_handlers.assert_awaited_once() + assert [event["type"] for event in in_flight.dispatch_success_handlers.await_args[0][0]] == ["response.created"] + unanswered.dispatch_success_handlers.assert_not_awaited() + unanswered.dispatch_failure_handlers.assert_not_awaited() diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 90da579d910..522f4987e13 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -56,7 +56,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( From 72493425ab4b052aa3499236a5c0d17c4d61fe03 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:52:39 -0700 Subject: [PATCH 02/25] fix(responses): bill the usage a failed native WebSocket turn reports --- litellm/cost_calculator.py | 2 +- .../test_responses_websocket_all_providers.py | 110 +++++++++++++----- 2 files changed, 83 insertions(+), 29 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6256bd5b8d9..5db7dfe69db 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2888,7 +2888,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) -_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"}) +_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) class _ResponsesWsEventResponse(LiteLLMBaseModel): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 5b95c472480..1c36d993523 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,6 +8,7 @@ Tests that: """ import json +from collections.abc import Callable from typing import Final from unittest.mock import MagicMock @@ -1901,8 +1902,7 @@ def _make_streaming(**kwargs): class _TurnLoggingRecorder: """Stands in for a socket's connection-level Logging object and keeps every per-turn copy it hands out.""" - def __init__(self, response_cost: float = 0.0) -> None: - self.response_cost = response_cost + def __init__(self) -> None: self.turns: list[MagicMock] = [] self.connection = MagicMock() self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call) @@ -1916,10 +1916,15 @@ class _TurnLoggingRecorder: turn.model_call_details = {} turn.dispatch_success_handlers = AsyncMock() turn.dispatch_failure_handlers = AsyncMock() - turn.response_cost_calculator = MagicMock(return_value=self.response_cost) + turn.response_cost_calculator = MagicMock(return_value=0.0) self.turns.append(turn) return turn + def logged_turns(self) -> int: + return sum( + 1 for turn in self.turns if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count + ) + class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" @@ -3272,27 +3277,20 @@ async def _until(condition) -> bool: class _ScriptedClient: """Sends each frame only after the previous turn's log was dispatched, then hangs up.""" - def __init__(self, frames: list[str], recorder: _TurnLoggingRecorder) -> None: + def __init__(self, frames: list[str], logged_turns: Callable[[], int]) -> None: self.frames = list(frames) - self.recorder = recorder + self.logged_turns = logged_turns self.sent_text: list[str] = [] self.logged_turns_before_each_frame: list[int] = [] self.logged_turns_before_hang_up: int | None = None - def _logged_turns(self) -> int: - return sum( - 1 - for turn in self.recorder.turns - if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count - ) - async def receive_text(self) -> str: sent_so_far = len(self.logged_turns_before_each_frame) - await _until(lambda: self._logged_turns() >= sent_so_far) + await _until(lambda: self.logged_turns() >= sent_so_far) if self.frames: - self.logged_turns_before_each_frame.append(self._logged_turns()) + self.logged_turns_before_each_frame.append(self.logged_turns()) return self.frames.pop(0) - self.logged_turns_before_hang_up = self._logged_turns() + self.logged_turns_before_hang_up = self.logged_turns() raise Exception("client hung up") async def send_text(self, text: str) -> None: @@ -3309,7 +3307,7 @@ class TestNativeWebSocketPerTurnLogging: from litellm.responses.utils import ResponsesAPIRequestUtils recorder = _TurnLoggingRecorder() - client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns) backend = _ScriptedBackend( [_completed_turn("resp_1", 10, 1), _completed_turn("resp_2", 20, 2), _completed_turn("resp_3", 30, 3)] ) @@ -3349,7 +3347,7 @@ class TestNativeWebSocketPerTurnLogging: recorder = _TurnLoggingRecorder() on_turn_failure = AsyncMock() - client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns) backend = _ScriptedBackend( [ _completed_turn("resp_1", 10, 1), @@ -3393,11 +3391,46 @@ class TestNativeWebSocketPerTurnLogging: assert "not found" in str(booked_exception) @pytest.mark.asyncio - async def test_a_failed_turn_records_only_its_own_partial_usage(self): + async def test_a_failed_turn_bills_the_usage_its_failure_event_reports(self): + from datetime import datetime + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - recorder = _TurnLoggingRecorder(response_cost=0.01) - client = _ScriptedClient([_response_create(text) for text in ("one", "two")], recorder) + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + from litellm.types.utils import CallTypes + + logged: list[dict[str, object]] = [] + + class _SpendRows(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + logged.append(kwargs["standard_logging_object"]) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + logged.append(kwargs["standard_logging_object"]) + + spend_rows = _SpendRows() + connection = Logging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=datetime.now(), # noqa: DTZ005 # logging pipeline uses naive datetimes + litellm_call_id="ws-connection", + function_id="ws-connection", + dynamic_async_success_callbacks=[spend_rows], + dynamic_async_failure_callbacks=[spend_rows], + ) + connection.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {}}, + optional_params={}, + model="gpt-4o", + custom_llm_provider="openai", + ) + failed_usage = {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9} + client = _ScriptedClient([_response_create(text) for text in ("one", "two")], lambda: len(logged)) backend = _ScriptedBackend( [ _completed_turn("resp_1", 10, 5), @@ -3409,30 +3442,51 @@ class TestNativeWebSocketPerTurnLogging: "id": "resp_2", "status": "failed", "error": {"code": "server_error", "message": "upstream blew up"}, - "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + "usage": failed_usage, }, }, ], ] ) handler = _make_streaming( - websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + websocket=client, + backend_ws=backend, + logging_obj=connection, + request_data={}, + authorized_model="gpt-4o", + custom_llm_provider="openai", ) await handler.bidirectional_forward() + await _until(lambda: len(logged) == 2) - first, failed = recorder.turns - first.record_partial_usage_for_failure.assert_not_called() - usage, response_cost = failed.record_partial_usage_for_failure.call_args[0] - assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0) - assert response_cost == 0.01 + completed, failed = logged + assert (completed["status"], failed["status"]) == ("success", "failure") + assert (completed["prompt_tokens"], completed["completion_tokens"]) == (10, 5) + assert (failed["prompt_tokens"], failed["completion_tokens"]) == (7, 2) + http_failed_cost = litellm.completion_cost( + completion_response=ResponsesAPIResponse( + id="resp_2", + created_at=0, + model="gpt-4o", + object="response", + output=[], + status="failed", + usage=ResponseAPIUsage(**failed_usage), + ), + model="gpt-4o", + call_type=CallTypes.aresponses.value, + custom_llm_provider="openai", + ) + assert http_failed_cost > 0 + assert failed["response_cost"] == pytest.approx(http_failed_cost) @pytest.mark.asyncio async def test_a_turn_still_in_flight_when_the_socket_closes_is_logged_with_what_it_received(self): import websockets.exceptions # noqa: F401 (lazy submodule must be importable) recorder = _TurnLoggingRecorder() - client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns) backend = _ScriptedBackend( [ _completed_turn("resp_1", 10, 1), From d2a342f6ab835da3ba5f28161a59b0a969d9ef30 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:05:23 -0700 Subject: [PATCH 03/25] fix(proxy): pin the limiter slot owner to the proxy call id and refund unsettled reservations on disconnect The v3 limiter stashed whichever litellm_call_id the request body carried as the slot owner, so a client-supplied id on a pass-through body left the proxy's own success and failure events unable to release the slot. The owner is now the proxy's logging call id. A WebSocket connection whose turns log under their own ids never settles the TPM, ITPM, or OTPM estimate the connection reserved, so the disconnect release now refunds any reservation no callback settled, once. --- .../hooks/parallel_request_limiter_v3.py | 31 +++- .../hooks/test_parallel_request_limiter_v3.py | 141 ++++++++++++++++++ .../test_router/test_router_helper_utils.py | 46 ++++++ 3 files changed, 215 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 54c30073dcd..c3e73d7d5d4 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -740,9 +740,11 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: + from litellm.proxy.common_request_processing import request_litellm_call_id + stash: Final = get_or_create_request_stash() - owner_call_id: Final = data.get("litellm_call_id") - if isinstance(owner_call_id, str): + owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) + if owner_call_id is not None: stash.owner_litellm_call_id = owner_call_id return stash @@ -5092,6 +5094,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations=io_token_operations, parent_otel_span=litellm_parent_otel_span, ) + if stash is not None: + stash.reservation_released = True except Exception as e: verbose_proxy_logger.exception("Error in rate limit success event: %s", e) @@ -5293,8 +5297,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TTL prunes it. The stashed acquisition's presence (not the key object's current max_parallel_requests configuration, which can change mid-request) decides whether there is anything to release. + A TPM, ITPM or OTPM reservation no callback settled is refunded the + same way, so a native Responses WebSocket connection, whose turns log + under their own call ids, hands back the estimate it reserved. """ - await self._release_stashed_parallel_slot(get_request_stash(), None) + stash: Final = get_request_stash() + await self._release_stashed_parallel_slot(stash, None) + await self._refund_unsettled_reservation(stash) + + async def _refund_unsettled_reservation(self, stash: RequestRateLimiterStash | None) -> None: + if stash is None or stash.reservation_released: + return + stash.reservation_released = True + await self._refund_reserved_tokens(scopes=tuple(stash.reserved_scopes), amount=stash.reserved_tokens) + await self._refund_reserved_tokens( + scopes=tuple(stash.itpm_reserved_scopes), + amount=stash.itpm_reserved_tokens, + reservation_windows=stash.itpm_reserved_window_identities, + ) + await self._refund_reserved_tokens( + scopes=tuple(stash.otpm_reserved_scopes), + amount=stash.otpm_reserved_tokens, + reservation_windows=stash.otpm_reserved_window_identities, + ) @with_service_target("rate_limits") async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response) -> None: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 28589a2f400..906f012c34b 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -20,6 +20,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -4235,6 +4236,77 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): ) == 0 +async def _admit_with_tpm_reservation( + handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth +) -> int: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "owner-call-id", + }, + call_type="completion", + ) + stash = get_request_stash() + assert stash is not None and stash.reserved_tokens > 0 + return stash.reserved_tokens + + +@pytest.mark.asyncio +async def test_release_on_disconnect_refunds_an_unsettled_tpm_reservation_v3(monkeypatch): + """ + A native Responses WebSocket connection reserves tokens at pre-call and logs its + turns under their own call ids, so no callback ever settles the connection's + reservation. The disconnect release must hand the estimate back. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-unsettled-reservation") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict) + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 + assert get_request_stash().reservation_released is True + + +@pytest.mark.asyncio +async def test_release_on_disconnect_keeps_a_settled_tpm_reservation_v3(monkeypatch): + """ + When the owner's success event already settled the reservation to the actual + usage, the disconnect release must not refund the estimate a second time. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-settled-reservation") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict) + + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": "owner-call-id", + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10 + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10 + + @pytest.mark.asyncio async def test_post_call_failure_hook_releases_parallel_slot_v3(): """ @@ -4328,6 +4400,75 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3() assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 +def _pass_through_logging_obj(litellm_call_id: str) -> Logging: + return Logging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id=litellm_call_id, + function_id="pass-through", + ) + + +async def _admit_pass_through_with_client_call_id( + handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth +) -> Logging: + """ + A pass-through body may carry its own litellm_call_id. The proxy books the request under the id of the + logging object it put in the same data, so that id, not the client's, must own the parallel slot. + """ + logging_obj = _pass_through_logging_obj("proxy-call") + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj}, + call_type="pass_through_endpoint", + ) + return logging_obj + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(): + _api_key = hash_token("sk-passthrough-client-id-failure") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict) + + await handler.async_post_call_failure_hook( + request_data={"litellm_call_id": "proxy-call", "litellm_logging_obj": logging_obj}, + original_exception=Exception("upstream unreachable"), + user_api_key_dict=user_api_key_dict, + ) + + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + + +@pytest.mark.asyncio +async def test_success_event_releases_the_slot_the_proxy_call_id_owns_v3(): + _api_key = hash_token("sk-passthrough-client-id-success") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict) + + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": "proxy-call", + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + + @pytest.mark.asyncio async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index cbbfe654bf6..1ed063d01dc 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -482,6 +482,52 @@ 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(): + """Each native Responses WebSocket turn logs under its own call, so the deployment's + tpm and rpm counters move once per turn, while a realtime session still counts nothing.""" + 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 + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ From 3a8953462bd4d9bd5f45cd197285e8781b160a7a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:37:06 -0700 Subject: [PATCH 04/25] fix(proxy): bill failed WebSocket turns and refund only the reserved window on disconnect A failed turn's reported usage now counts toward the key's TPM and the project ITPM/OTPM scopes the same way a success does, and closing the socket hands back only the estimate still unused in the window that holds it, so a window that already rolled is left alone. An id-less error event is booked against the newest pending turn, the one the backend rejected at parse time, and a turn that nests its input under response logs its user message. A connection's one budget reservation goes to the first turn that settles it, every later turn charging its cost directly. --- .../hooks/parallel_request_limiter_v3.py | 101 ++++++++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 39 ++++-- .../spend_tracking/budget_reservation.py | 14 ++ litellm/responses/streaming_iterator.py | 42 +++++- .../hooks/test_parallel_request_limiter_v3.py | 128 +++++++++++++++++- .../hooks/test_proxy_track_cost_callback.py | 91 +++++++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 84 +++++++++++- tests/unit/proxy/test_budget_reservation.py | 61 +++++++++ .../test_responses_websocket_all_providers.py | 63 +++++++++ 9 files changed, 593 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c3e73d7d5d4..4737595a049 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -658,6 +658,7 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field(default_factory=frozenset) itpm_reserved_tokens: int = 0 itpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) itpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field( @@ -671,6 +672,7 @@ class RequestRateLimiterStash: batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False + websocket_connection: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) @@ -4043,6 +4045,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + stash.websocket_connection = call_type == CallTypes.aresponses_websocket.value ######################################################### # Check if the call type has a specific rate limiter @@ -4225,6 +4228,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) and (d.get("rate_limit") or {}).get("tokens_per_unit") is not None ) + stash.reserved_window_identities = frozenset(tpm_response.get("reservation_windows", frozenset())) tpm_reservation_scopes = tuple( # rebind-ok: record successful reservation scopes stash.reserved_scopes ) @@ -4719,7 +4723,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return () stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) if stash is None: - return () + return self._build_websocket_turn_io_ops(kwargs, response_obj) itpm_reserved: Final = stash.itpm_reserved_tokens otpm_reserved: Final = stash.otpm_reserved_tokens @@ -4784,9 +4788,68 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return tuple((*itpm_ops, *otpm_ops)) + def _websocket_turn_usage(self, kwargs: Mapping[str, object], response_obj: object) -> tuple[int, int, bool]: + response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) + if response_usage[2]: + return response_usage + return self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object")) + + def _build_websocket_turn_io_ops( + self, + kwargs: Mapping[str, object], + response_obj: object, + ) -> Sequence[RedisPipelineIncrementOperation]: + """ + A native Responses WebSocket turn logs under its own call id, so no + stash is its own. The connection's stash, inherited through the + context, names the project ITPM/OTPM scopes the turn's usage counts + toward; the connection's reservation itself is refunded at disconnect. + """ + connection: Final = get_request_stash() + if connection is None or not connection.websocket_connection: + return [] + billable_input, completion_tokens, usage_resolved = self._websocket_turn_usage(kwargs, response_obj) + if not usage_resolved: + return [] + return self._build_reservation_aware_tpm_ops( + targets=tuple(connection.itpm_reserved_scopes), + reserved_scopes=frozenset(), + actual_tokens=billable_input, + reserved_tokens=0, + ) + self._build_reservation_aware_tpm_ops( + targets=tuple(connection.otpm_reserved_scopes), + reserved_scopes=frozenset(), + actual_tokens=completion_tokens, + reserved_tokens=0, + ) + + async def _charge_websocket_turn_usage( + self, + kwargs: Mapping[str, object], + response_obj: object, + parent_otel_span: Span | None, + ) -> None: + """A failed turn still consumed the usage its failure event reports; count it the way a success does.""" + connection: Final = get_request_stash() + if connection is None or not connection.websocket_connection: + return + pipeline_operations: Final = [ + *self._build_success_event_pipeline_operations( + kwargs=dict(kwargs), + response_obj=response_obj, + rate_limit_type=self.get_rate_limit_type(), + ), + *self._build_websocket_turn_io_ops(kwargs, response_obj), + ] + if pipeline_operations: + await self.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + def _collect_tpm_scope_targets( self, - standard_logging_metadata: dict[str, Any], + standard_logging_metadata: Mapping[str, Any], kwargs: object, model_group: str | None, tpm_limited_tags: Set[str] = frozenset(), @@ -5172,7 +5235,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) @with_service_target("rate_limits") - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: """ On failure: decrement max_parallel_requests and refund the upfront TPM reservation only against the scopes the reservation actually @@ -5186,12 +5251,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) try: - litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs) + failure_kwargs: Final = dict(_REQUEST_RATE_LIMIT_DATA.validate_python(kwargs)) + litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(failure_kwargs) pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [] - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failure_kwargs)) await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True) + if stash is None: + await self._charge_websocket_turn_usage(failure_kwargs, response_obj, litellm_parent_otel_span) # Skip the reservation refund if async_post_call_failure_hook # already released it (proxy-level rejection that also bubbles up @@ -5202,7 +5270,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if stash is None or stash.reservation_released else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) ) - tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(kwargs) + tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(failure_kwargs) if stash is not None and reserved_tokens > 0: verbose_proxy_logger.debug( @@ -5297,19 +5365,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TTL prunes it. The stashed acquisition's presence (not the key object's current max_parallel_requests configuration, which can change mid-request) decides whether there is anything to release. - A TPM, ITPM or OTPM reservation no callback settled is refunded the - same way, so a native Responses WebSocket connection, whose turns log - under their own call ids, hands back the estimate it reserved. + A native Responses WebSocket connection logs its turns under their own + call ids, so no callback of its own ever settles the TPM, ITPM or OTPM + reservation it made at pre-call; its disconnect refunds that estimate. + Every other request keeps its reservation for the success or failure + event that settles it, even when the client went away mid-stream. """ stash: Final = get_request_stash() await self._release_stashed_parallel_slot(stash, None) - await self._refund_unsettled_reservation(stash) + if stash is not None and stash.websocket_connection: + await self._refund_unsettled_reservation(stash) - async def _refund_unsettled_reservation(self, stash: RequestRateLimiterStash | None) -> None: - if stash is None or stash.reservation_released: + async def _refund_unsettled_reservation(self, stash: RequestRateLimiterStash) -> None: + if stash.reservation_released: return stash.reservation_released = True - await self._refund_reserved_tokens(scopes=tuple(stash.reserved_scopes), amount=stash.reserved_tokens) + await self._refund_reserved_tokens( + scopes=tuple(stash.reserved_scopes), + amount=stash.reserved_tokens, + reservation_windows=stash.reserved_window_identities, + ) await self._refund_reserved_tokens( scopes=tuple(stash.itpm_reserved_scopes), amount=stash.itpm_reserved_tokens, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 58ec0413d3a..f7d806c2a45 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,6 +4,8 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast +from pydantic import BaseModel, ConfigDict + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -166,14 +168,18 @@ class _ProxyDBLogger(CustomLogger): user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, ) -> None: + budget_reservation: Final = _claim_budget_reservation( + budget_reservation=user_api_key_dict.budget_reservation, + litellm_call_id=_RequestCallId.model_validate(request_data).litellm_call_id, + ) try: - await _release_budget_reservation(budget_reservation=user_api_key_dict.budget_reservation) + await _release_budget_reservation(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception("Failed to release budget reservation during failure handling") try: - await _invalidate_budget_reservation_counters(budget_reservation=user_api_key_dict.budget_reservation) - if user_api_key_dict.budget_reservation is not None: - user_api_key_dict.budget_reservation["finalized"] = True + await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + if budget_reservation is not None: + budget_reservation["finalized"] = True except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after failure release failed" @@ -343,7 +349,10 @@ class _ProxyDBLogger(CustomLogger): resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) - budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + budget_reservation: Final = _claim_budget_reservation( + budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata), + litellm_call_id=_RequestCallId.model_validate(kwargs).litellm_call_id, + ) if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value @@ -615,7 +624,7 @@ class _ProxyDBLogger(CustomLogger): ProxyDBLogger: Final = _ProxyDBLogger -def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: +def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: Mapping[str, object]) -> None: patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None} if not patch: return @@ -691,13 +700,27 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: return tuple(sorted(str(key) for key in metadata)) -def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: +def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict[str, object] | None: return budget_reservation_from_metadata(metadata) +class _RequestCallId(BaseModel): + model_config = ConfigDict(frozen=True) + + litellm_call_id: str | None = None + + +def _claim_budget_reservation( + budget_reservation: dict[str, object] | None, litellm_call_id: str | None +) -> dict[str, object] | None: + from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + + return claim_budget_reservation(budget_reservation=budget_reservation, litellm_call_id=litellm_call_id) + + def _get_request_tags_for_cost_tracking( sl_object: StandardLoggingPayload | None, - metadata: dict, + metadata: Mapping[str, object], ) -> list[str] | None: if sl_object is not None: request_tags: Final = sl_object.get("request_tags", None) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4c9df95cd8d..8e8f4def110 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -378,6 +378,20 @@ async def reconcile_budget_reservation( return pending +def claim_budget_reservation( + budget_reservation: dict[str, object] | None, litellm_call_id: str | None +) -> dict[str, object] | None: + """Hand the reservation to the first call that settles it. + + A native Responses WebSocket connection shares its one reservation with every turn it logs, each under its + own call id. Two turns settling it at once would both move the counter from the estimate to their own cost, + so the first claimant keeps the reservation and every later call sees none and charges its cost directly.""" + if budget_reservation is None or litellm_call_id is None: + return budget_reservation + settling_call_id: Final = budget_reservation.setdefault("settling_call_id", litellm_call_id) + return budget_reservation if settling_call_id == litellm_call_id else None + + def stamp_budget_reservation_actual_cost(budget_reservation: dict | None, actual_cost: float | None) -> None: """Record that every reserved counter now holds ``actual_cost``, once the adjustments handed back by ``reconcile_budget_reservation(apply_consistent=False)`` have been written.""" diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index e4291bd366c..6aa88ddee75 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1871,18 +1871,31 @@ def _logged_user_messages(client_frame: object) -> tuple[dict[str, object], ...] frame: Final = _json_object_or_none(client_frame) if frame is None or frame.get("type") != "response.create": return () - input_items: Final = frame.get("input", []) + nested_response: Final = frame.get("response") + input_items: Final = frame.get( + "input", nested_response.get("input", []) if _is_json_object(nested_response) else [] + ) if isinstance(input_items, str): return ({"role": "user", "content": input_items},) texts: Final = chain.from_iterable(_user_texts_in_input_item(item) for item in _json_array_or_empty(input_items)) return tuple({"role": "user", "content": text} for text in texts) +def _ws_event_response_id(event: object) -> str | None: + response: Final = event.get("response") if _is_json_object(event) else None + response_id: Final = response.get("id") if _is_json_object(response) else None + return response_id if isinstance(response_id, str) else None + + @dataclass(frozen=True, slots=True) class _ResponsesWsTurn: logging_obj: LiteLLMLoggingObj events: tuple[_MutableJsonObject, ...] = () + @property + def response_id(self) -> str | None: + return next(filter(None, map(_ws_event_response_id, self.events)), None) + def with_event(self, event: _MutableJsonObject) -> _ResponsesWsTurn: return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) @@ -2010,12 +2023,33 @@ class ResponsesWebSocketStreaming: return if not self._open_turns: self._start_turn(None) - turn: Final = self._open_turns[0].with_event(event_obj) + index: Final = self._turn_index_for_event(event_obj) + turn: Final = self._open_turns[index].with_event(event_obj) + other_turns: Final = (*self._open_turns[:index], *self._open_turns[index + 1 :]) if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES: - self._open_turns = self._open_turns[1:] + self._open_turns = other_turns self._log_turn(turn) return - self._open_turns = (turn, *self._open_turns[1:]) + self._open_turns = (*self._open_turns[:index], turn, *self._open_turns[index + 1 :]) + + def _turn_index_for_event(self, event_obj: Mapping[str, object]) -> int: + """ + An event naming a response id belongs to the turn whose ``response.created`` + carried it. The backend answers accepted turns strictly in order, so an + id-less event goes to the oldest turn still waiting for its ``response.created``, + except an id-less ``error``: that is the backend rejecting the newest pending + ``response.create`` at parse time while an earlier turn is still streaming. + """ + response_id: Final = _ws_event_response_id(event_obj) + turn_ids: Final = tuple(turn.response_id for turn in self._open_turns) + if response_id is not None and response_id in turn_ids: + return turn_ids.index(response_id) + pending: Final = tuple(index for index, turn_id in enumerate(turn_ids) if turn_id is None) + if not pending: + return 0 + if response_id is None and event_obj.get("type") == "error": + return pending[-1] + return pending[0] def _log_open_turns(self) -> None: open_turns: Final = self._open_turns diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 906f012c34b..6c155a521c8 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -46,6 +46,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( + CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, @@ -4237,7 +4238,10 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): async def _admit_with_tpm_reservation( - handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth + handler: _PROXY_MaxParallelRequestsHandler, + local_cache: DualCache, + user_api_key_dict: UserAPIKeyAuth, + call_type: str = "completion", ) -> int: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4248,7 +4252,7 @@ async def _admit_with_tpm_reservation( "max_tokens": 50, "litellm_call_id": "owner-call-id", }, - call_type="completion", + call_type=call_type, ) stash = get_request_stash() assert stash is not None and stash.reserved_tokens > 0 @@ -4268,7 +4272,9 @@ async def test_release_on_disconnect_refunds_an_unsettled_tpm_reservation_v3(mon handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") - reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict) + reserved = await _admit_with_tpm_reservation( + handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value + ) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) @@ -4289,7 +4295,7 @@ async def test_release_on_disconnect_keeps_a_settled_tpm_reservation_v3(monkeypa handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") - await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict) + await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value) await handler.async_log_success_event( kwargs={ @@ -7938,3 +7944,117 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( await _rpm_request(limiter, cache, auth, "a2a/cheap") assert denied.value.status_code == 429 assert "expensive" in str(denied.value.detail) + + +def _websocket_turn_kwargs(call_id: str, api_key_hash: str, usage: Usage) -> dict[str, object]: + return { + "litellm_call_id": call_id, + "model": "gpt-4o-mini", + "combined_usage_object": usage, + "litellm_params": {"metadata": {"user_api_key": api_key_hash}}, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key_hash}}, + } + + +@pytest.mark.asyncio +async def test_release_on_disconnect_keeps_a_plain_request_tpm_reservation_v3(monkeypatch): + """ + A cancelled HTTP stream keeps its reservation for the success or failure event + that still settles it; only a WebSocket connection, whose turns log under their + own call ids, hands its estimate back at disconnect. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-http-cancelled-stream") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, "completion") + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + assert get_request_stash().reservation_released is False + + +@pytest.mark.asyncio +async def test_release_on_disconnect_skips_a_tpm_window_that_rolled_v3(monkeypatch): + """ + The TPM window can roll while a WebSocket connection stays open. The disconnect + refund belongs to the window the connection reserved against, so it must leave + the window that is live now untouched instead of handing it tokens it never took. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-rolled-window") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + window_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="window") + await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value) + reserved_window_start = await local_cache.async_get_cache(key=window_key) + assert reserved_window_start is not None + await local_cache.async_set_cache(key=window_key, value=int(reserved_window_start) + handler.window_size) + await local_cache.async_set_cache(key=tokens_key, value=500) + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 500 + assert get_request_stash().reservation_released is True + + +@pytest.mark.asyncio +async def test_websocket_turn_failure_counts_its_usage_toward_tpm_v3(monkeypatch): + """ + A turn that ends in ``response.failed`` still consumed the usage the event + reports. It logs under its own call id, so the failure event owns no stash; it + must still charge the connection's TPM scopes the way a successful turn does. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-failed-turn") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation( + handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value + ) + + await handler.async_log_failure_event( + kwargs=_websocket_turn_kwargs( + "turn-1", _api_key, Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17) + ), + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17 + assert get_request_stash().reservation_released is False + + +@pytest.mark.asyncio +async def test_failure_event_without_a_stash_charges_nothing_outside_a_websocket_connection_v3(monkeypatch): + """ + A nested call made inside a plain request (an LLM-judge guardrail) logs under a + fresh call id too. Its failure owns no stash and must not touch the request's + counters: only a WebSocket connection's turns are charged that way. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-nested-failed-call") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, "completion") + + await handler.async_log_failure_event( + kwargs=_websocket_turn_kwargs( + "nested-call-id", _api_key, Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17) + ), + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved 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 100e79e371d..76e398784cf 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2972,3 +2972,94 @@ async def test_track_cost_callback_enqueue_emits_no_service_span(): # test-qual assert update_database.await_count == 1, "the spend enqueue itself must still run" assert success_hook.await_count == 0, [call.kwargs for call in success_hook.await_args_list] + + +def _shared_reservation_turn_kwargs(call_id: str, reservation: dict) -> dict: + return { + "model": "gpt-5.4-mini", + "litellm_call_id": call_id, + "litellm_params": { + "metadata": {"user_api_key": "test-key-hash", "user_api_key_budget_reservation": reservation} + }, + "stream": False, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, + } + + +async def _reservation_handed_to_the_counters(kwargs) -> dict | None: + """The reservation the callback hands the spend counters for one logged call.""" + logger = ProxyDBLogger() + with ( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the arguments to this call are the boundary under test + "litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters", + new=AsyncMock(), + ) as mock_update, + patch( # test-quality-ok: llm_router is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup({}) + ), + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + return mock_update.await_args.kwargs["budget_reservation"] + + +@pytest.mark.asyncio +async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_only(): + """Every turn of a Responses WebSocket connection carries the connection's one reservation under its + own call id. The first turn settles it; a later turn reconciling the same estimate again would + under-count the key, so it must charge its cost without the reservation.""" + reservation = {"reserved_cost": 0.5, "entries": []} + + first_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-1", reservation)) + second_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-2", reservation)) + + assert first_turn is reservation + assert second_turn is None + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_releases_a_shared_reservation_for_the_first_turn_only(): + logger = ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + request_route="/custom/route", + budget_reservation=budget_reservation, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ), + ): + for call_id in ("turn-1", "turn-2"): + await logger.async_post_call_failure_hook( + request_data={"litellm_call_id": call_id}, + original_exception=Exception("Test exception"), + user_api_key_dict=user_api_key_dict, + ) + + assert mock_release_budget_reservation.await_count == 1 + assert ( + mock_release_budget_reservation.await_args.kwargs["budget_reservation"] + is user_api_key_dict.budget_reservation + ) + assert user_api_key_dict.budget_reservation["settling_call_id"] == "turn-1" diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 2d23b134433..51c8de77d0c 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -42,7 +42,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.rerank import RerankResponse -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import CallTypes, ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture @@ -3710,3 +3710,85 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +@pytest.mark.asyncio +async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_limiter): + """ + A native Responses WebSocket connection reserves project ITPM/OTPM at pre-call and + its turns log under their own call ids. Every turn's real usage must reach the + project counters, success and failure alike, and the disconnect must hand back + only the connection's estimate, never the turns' usage. + """ + handler, cache = rate_limiter + api_key = hash_token("sk-ws-project-turns") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-ws-turns", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 500}, + }, + ) + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + "litellm_call_id": "connection-call-id", + }, + call_type=CallTypes.aresponses_websocket.value, + ) + stash = get_request_stash() + assert stash is not None and stash.itpm_reserved_tokens > 0 and stash.otpm_reserved_tokens > 0 + (itpm_scope,) = stash.itpm_reserved_scopes + (otpm_scope,) = stash.otpm_reserved_scopes + itpm_key = handler.create_rate_limit_keys(itpm_scope[0], itpm_scope[1], "tokens") + otpm_key = handler.create_rate_limit_keys(otpm_scope[0], otpm_scope[1], "tokens") + + async def counters() -> tuple[int, int]: + return ( + int(await cache.async_get_cache(key=itpm_key) or 0), + int(await cache.async_get_cache(key=otpm_key) or 0), + ) + + def turn_kwargs(call_id: str, usage: Usage | None = None) -> dict[str, object]: + return { + "litellm_call_id": call_id, + "model": "bedrock_mantle/claude-opus", + "combined_usage_object": usage, + "litellm_params": {"metadata": {"user_api_key": api_key}}, + "standard_logging_object": { + "metadata": {"user_api_key_hash": api_key, "user_api_key_project_id": "proj-ws-turns"} + }, + } + + assert await counters() == (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) + + await handler.async_log_success_event( + kwargs=turn_kwargs("turn-1"), + response_obj=ResponsesAPIResponse( + id="resp_1", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=30, output_tokens=7, total_tokens=37), + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await counters() == (stash.itpm_reserved_tokens + 30, stash.otpm_reserved_tokens + 7) + + await handler.async_log_failure_event( + kwargs=turn_kwargs("turn-2", Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)), + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await counters() == (stash.itpm_reserved_tokens + 40, stash.otpm_reserved_tokens + 10) + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict=user_api_key_dict) + + assert await counters() == (40, 10) + assert stash.reservation_released is True diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 4635e145914..4591daaded6 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -40,6 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_model_access_group_budget_counters, + claim_budget_reservation, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -3720,3 +3721,63 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one assert counter_cache.in_memory_cache.get_cache( key=model_access_group_spend_counter_key("starter") ) == pytest.approx(4.2) + + +@pytest.mark.asyncio +async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at_once( + spend_counter_state, +): + """A Responses WebSocket connection's turns share its reservation under their own call ids. Only the + first claim keeps it; the other turn adds its cost on top instead of settling the same estimate again, + which would leave the key counter short by the estimate.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-shared-turns", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/responses", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + from litellm.proxy.proxy_server import increment_spend_counters + + first_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") + second_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") + assert first_turn is reservation + assert second_turn is None + assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") is reservation + + await asyncio.gather( + increment_spend_counters( + token="key-budget-shared-turns", + team_id="team-without-budget", + user_id=None, + response_cost=0.2, + budget_reservation=first_turn, + ), + increment_spend_counters( + token="key-budget-shared-turns", + team_id="team-without-budget", + user_id=None, + response_cost=0.3, + budget_reservation=second_turn, + ), + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-shared-turns") == pytest.approx(0.5) + assert counter_cache.in_memory_cache.get_cache(key="spend:team:team-without-budget") == pytest.approx(0.5) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 1c36d993523..48bb4340fba 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -3339,6 +3339,69 @@ class TestNativeWebSocketPerTurnLogging: for turn in recorder.turns: turn.dispatch_failure_handlers.assert_not_awaited() + @pytest.mark.asyncio + async def test_a_turn_whose_input_is_nested_under_response_logs_its_user_message(self): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder() + nested_frame = json.dumps( + {"type": "response.create", "response": {"model": "gpt-5.4-mini", "input": "nested"}} + ) + client = _ScriptedClient([nested_frame], recorder.logged_turns) + backend = _ScriptedBackend([_completed_turn("resp_1", 10, 1)]) + handler = _make_streaming( + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + ) + + await handler.bidirectional_forward() + + assert [turn.model_call_details["messages"] for turn in recorder.turns] == [ + [{"role": "user", "content": "nested"}] + ] + + @pytest.mark.asyncio + async def test_an_overlapped_turn_rejected_at_parse_time_fails_alone_and_the_streaming_turn_keeps_its_usage( + self, + ): + """OpenAI rejects an unparseable response.create with an id-less error while the earlier turn is + still streaming. That error is the second turn's failure, and the first turn's completed event with + its usage stays on the first turn.""" + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder() + handler = _make_streaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini", + custom_llm_provider="openai", + ) + created, _delta, completed = _completed_turn("resp_1", 300, 23) + rejection = { + "type": "error", + "status": None, + "error": {"type": "invalid_request_error", "code": None, "message": "Invalid 'max_output_tokens'"}, + } + + handler._record_client_frame(_response_create("first")) + handler._record_client_frame(_response_create("second")) + for event in (created, rejection, completed): + handler._record_backend_event(json.dumps(event)) + assert await _until(lambda: recorder.logged_turns() == 2) + + first, second = recorder.turns + first.dispatch_failure_handlers.assert_not_awaited() + first.dispatch_success_handlers.assert_awaited_once() + assert [event["type"] for event in first.dispatch_success_handlers.await_args[0][0]] == [ + "response.created", + "response.completed", + ] + second.dispatch_success_handlers.assert_not_awaited() + second.dispatch_failure_handlers.assert_awaited_once() + assert "max_output_tokens" in str(second.dispatch_failure_handlers.await_args[0][0]) + assert second.record_partial_usage_for_failure.call_args[0][0].total_tokens == 0 + @pytest.mark.asyncio async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self): from unittest.mock import AsyncMock From 2d23c8fa543a6b4cd66f4cd5a21d9e38fc77b770 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:46:39 -0700 Subject: [PATCH 05/25] test(proxy): expect the settling call id on a pass-through budget reservation --- .../pass_through_endpoints/test_pass_through_endpoints.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index eb31797d460..14040a7206d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7711,7 +7711,10 @@ async def test_passthrough_success_reconciles_budget_reservation(): increment_spend_counters.assert_awaited_once() assert increment_spend_counters.await_args.kwargs["budget_reservation"] is reservation - assert increment_spend_counters.await_args.kwargs["budget_reservation"] == budget_reservation + assert increment_spend_counters.await_args.kwargs["budget_reservation"] == { + **budget_reservation, + "settling_call_id": "lit-5425-call-id", + } @pytest.mark.asyncio From 3ce696e9a492262aa75f4f66441699f1df2d57c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:48:01 -0700 Subject: [PATCH 06/25] fix(proxy): charge successful Responses WebSocket turns to TPM and ITPM/OTPM A native Responses WebSocket turn reaches the rate limiter's success callback with its raw event list as the response and no combined_usage_object, so the turn's tokens never reached the key's TPM window or the project's input and output token limits. Read the usage from the turn's events when the call type is a Responses WebSocket turn. --- .../hooks/parallel_request_limiter_v3.py | 24 +++++++++--- .../hooks/test_parallel_request_limiter_v3.py | 39 +++++++++++++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 15 +++---- 3 files changed, 66 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 46381eebc0c..adcb467e34d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -42,6 +42,7 @@ from litellm.caching.redis_batch import ( ) from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, @@ -767,6 +768,20 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None: return call_id if isinstance(call_id, str) else None +_RESPONSES_WS_TURN_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: + combined_usage: Final = kwargs.get("combined_usage_object") + if isinstance(combined_usage, Usage): + return combined_usage + if kwargs.get("call_type") != CallTypes.aresponses_websocket.value or not isinstance(response_obj, list): + return None + return ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + _RESPONSES_WS_TURN_EVENTS.validate_python(response_obj) + ) + + def _parse_output_cap_value(raw_value: object) -> int | None: if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float, str)): return None @@ -4793,7 +4808,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) if response_usage[2]: return response_usage - return self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object")) + return self._resolve_io_token_reconcile_usage(_unmodeled_response_usage(kwargs, response_obj)) def _build_websocket_turn_io_ops( self, @@ -5036,7 +5051,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get total tokens from response. Responses LiteLLM does not model # (e.g. pass-through, whose usage is reported by the upstream rather # than parsed out of the body) carry their usage in - # ``combined_usage_object`` instead, and would otherwise never charge + # ``combined_usage_object`` instead, and a native Responses WebSocket + # turn in its raw event list; either would otherwise never charge # the TPM window. _usage: Usage | dict | None = None if isinstance( @@ -5050,9 +5066,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ): _usage = getattr(response_obj, "usage", None) else: - _combined_usage: Final = kwargs.get("combined_usage_object") - if isinstance(_combined_usage, Usage): - _usage = _combined_usage + _usage = _unmodeled_response_usage(kwargs, response_obj) total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) if total_tokens == 0: total_tokens = self._aggregate_only_total_tokens(usage=_usage) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index faf7dd7283e..49d063aa0df 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -8063,6 +8063,45 @@ async def test_websocket_turn_failure_counts_its_usage_toward_tpm_v3(monkeypatch assert get_request_stash().reservation_released is False +@pytest.mark.asyncio +async def test_websocket_turn_success_counts_its_usage_toward_tpm_v3(monkeypatch): + """ + A successful native Responses WebSocket turn is logged with the raw event list + as its response, not a modeled response object. Its usage must still reach the + key's TPM window right away, while the socket stays open. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-successful-turn") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation( + handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value + ) + turn_kwargs = { + key: value + for key, value in _websocket_turn_kwargs("turn-1", _api_key, Usage()).items() + if key != "combined_usage_object" + } + + await handler.async_log_success_event( + kwargs={**turn_kwargs, "call_type": CallTypes.aresponses_websocket.value}, + response_obj=[ + {"type": "response.created", "response": {"id": "resp_turn_1"}}, + { + "type": "response.completed", + "response": {"id": "resp_turn_1", "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}}, + }, + ], + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17 + assert get_request_stash().reservation_released is False + + @pytest.mark.asyncio async def test_failure_event_without_a_stash_charges_nothing_outside_a_websocket_connection_v3(monkeypatch): """ diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 51c8de77d0c..699b7e1cfb1 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -3768,13 +3768,14 @@ async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_lim assert await counters() == (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) await handler.async_log_success_event( - kwargs=turn_kwargs("turn-1"), - response_obj=ResponsesAPIResponse( - id="resp_1", - created_at=int(datetime.now().timestamp()), - output=[], - usage=ResponseAPIUsage(input_tokens=30, output_tokens=7, total_tokens=37), - ), + kwargs={**turn_kwargs("turn-1"), "call_type": CallTypes.aresponses_websocket.value}, + response_obj=[ + {"type": "response.created", "response": {"id": "resp_1"}}, + { + "type": "response.completed", + "response": {"id": "resp_1", "usage": {"input_tokens": 30, "output_tokens": 7, "total_tokens": 37}}, + }, + ], start_time=datetime.now(), end_time=datetime.now(), ) From 730b4f9713475f4a2653c983f20521759a64d3cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:54:51 -0700 Subject: [PATCH 07/25] fix(proxy): release a rejected pass-through's slot and count each websocket turn once toward project quotas The limiter's failure hook looked up the parallel slot by the request body's litellm_call_id, while the slot belongs to the proxy's logging call id, so a pass-through request a pre-call hook rejected held its slot for an hour when its body carried an id. Both now resolve the id the same way. Each response.create frame already charges its estimate to project ITPM/OTPM, so the turn's own usage no longer adds a second count. The cost callback's failure hook read the body's litellm_call_id through a strict model and raised on a non-string value, which dropped the failure row. It now reads a string id or none. --- .../hooks/parallel_request_limiter_v3.py | 65 +++++-------------- .../proxy/hooks/proxy_track_cost_callback.py | 20 ++++-- .../hooks/test_parallel_request_limiter_v3.py | 4 +- .../hooks/test_proxy_track_cost_callback.py | 31 +++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 30 ++++++--- 5 files changed, 85 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index adcb467e34d..70e67099832 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -742,11 +742,15 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: +def _request_owner_call_id(data: Mapping[str, object]) -> str | None: from litellm.proxy.common_request_processing import request_litellm_call_id + return request_litellm_call_id(data) + + +def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: stash: Final = get_or_create_request_stash() - owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) + owner_call_id: Final = _request_owner_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) if owner_call_id is not None: stash.owner_litellm_call_id = owner_call_id return stash @@ -4739,7 +4743,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return () stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) if stash is None: - return self._build_websocket_turn_io_ops(kwargs, response_obj) + return () itpm_reserved: Final = stash.itpm_reserved_tokens otpm_reserved: Final = stash.otpm_reserved_tokens @@ -4804,59 +4808,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return tuple((*itpm_ops, *otpm_ops)) - def _websocket_turn_usage(self, kwargs: Mapping[str, object], response_obj: object) -> tuple[int, int, bool]: - response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) - if response_usage[2]: - return response_usage - return self._resolve_io_token_reconcile_usage(_unmodeled_response_usage(kwargs, response_obj)) - - def _build_websocket_turn_io_ops( - self, - kwargs: Mapping[str, object], - response_obj: object, - ) -> Sequence[RedisPipelineIncrementOperation]: - """ - A native Responses WebSocket turn logs under its own call id, so no - stash is its own. The connection's stash, inherited through the - context, names the project ITPM/OTPM scopes the turn's usage counts - toward; the connection's reservation itself is refunded at disconnect. - """ - connection: Final = get_request_stash() - if connection is None or not connection.websocket_connection: - return [] - billable_input, completion_tokens, usage_resolved = self._websocket_turn_usage(kwargs, response_obj) - if not usage_resolved: - return [] - return self._build_reservation_aware_tpm_ops( - targets=tuple(connection.itpm_reserved_scopes), - reserved_scopes=frozenset(), - actual_tokens=billable_input, - reserved_tokens=0, - ) + self._build_reservation_aware_tpm_ops( - targets=tuple(connection.otpm_reserved_scopes), - reserved_scopes=frozenset(), - actual_tokens=completion_tokens, - reserved_tokens=0, - ) - async def _charge_websocket_turn_usage( self, kwargs: Mapping[str, object], response_obj: object, parent_otel_span: Span | None, ) -> None: - """A failed turn still consumed the usage its failure event reports; count it the way a success does.""" + """ + A failed turn still consumed the usage its failure event reports; count it toward TPM the way a + success does. Project ITPM/OTPM already charged the turn's ``response.create`` frame estimate. + """ connection: Final = get_request_stash() if connection is None or not connection.websocket_connection: return - pipeline_operations: Final = [ - *self._build_success_event_pipeline_operations( - kwargs=dict(kwargs), - response_obj=response_obj, - rate_limit_type=self.get_rate_limit_type(), - ), - *self._build_websocket_turn_io_ops(kwargs, response_obj), - ] + pipeline_operations: Final = self._build_success_event_pipeline_operations( + kwargs=dict(kwargs), + response_obj=response_obj, + rate_limit_type=self.get_rate_limit_type(), + ) if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( pipeline_operations=pipeline_operations, @@ -5493,7 +5462,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ try: failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data) - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failed_request)) + stash: Final = get_request_stash_for_call(_request_owner_call_id(failed_request)) if stash is None: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b088e81227d..3c03f13e02d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast -from pydantic import BaseModel, ConfigDict +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -170,7 +170,7 @@ class _ProxyDBLogger(CustomLogger): ) -> None: budget_reservation: Final = _claim_budget_reservation( budget_reservation=user_api_key_dict.budget_reservation, - litellm_call_id=_RequestCallId.model_validate(request_data).litellm_call_id, + litellm_call_id=_request_call_id(_CALLBACK_DATA.validate_python(request_data)), ) try: await _release_budget_reservation(budget_reservation=budget_reservation) @@ -353,7 +353,7 @@ class _ProxyDBLogger(CustomLogger): _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _claim_budget_reservation( budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata), - litellm_call_id=_RequestCallId.model_validate(kwargs).litellm_call_id, + litellm_call_id=_logged_call_id(_CALLBACK_DATA.validate_python(kwargs)), ) if ( isinstance(completion_response, LiteLLMBatch) @@ -706,10 +706,18 @@ def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dic return budget_reservation_from_metadata(metadata) -class _RequestCallId(BaseModel): - model_config = ConfigDict(frozen=True) +_CALLBACK_DATA: Final = TypeAdapter(Mapping[str, object]) - litellm_call_id: str | None = None + +def _request_call_id(request_data: Mapping[str, object]) -> str | None: + from litellm.proxy.common_request_processing import request_litellm_call_id + + return request_litellm_call_id(request_data) + + +def _logged_call_id(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None def _claim_budget_reservation( diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 49d063aa0df..d11b50243b5 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4445,8 +4445,8 @@ async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict) await handler.async_post_call_failure_hook( - request_data={"litellm_call_id": "proxy-call", "litellm_logging_obj": logging_obj}, - original_exception=Exception("upstream unreachable"), + request_data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj}, + original_exception=Exception("rejected by a pre-call hook"), user_api_key_dict=user_api_key_dict, ) 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 acf488145b2..1dd86c2c4a0 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3092,3 +3092,34 @@ async def test_async_post_call_failure_hook_releases_a_shared_reservation_for_th is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["settling_call_id"] == "turn-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_call_id", [12345, ["a", "b"], {"x": 1}]) +async def test_async_post_call_failure_hook_writes_the_failure_row_for_a_non_string_body_call_id(client_call_id): + """A rejected request's body is the client's, so its litellm_call_id can be any JSON value. The failure + row must still be written and the reservation released.""" + logger = ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", budget_reservation={"reserved_cost": 0.5, "entries": []} + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + ): + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-5.4-mini", "litellm_call_id": client_call_id}, + original_exception=Exception("Authentication Error"), + user_api_key_dict=user_api_key_dict, + ) + + released = mock_release_budget_reservation.await_args.kwargs["budget_reservation"] + assert released is user_api_key_dict.budget_reservation + assert mock_update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"]["status"] == "failure" diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 699b7e1cfb1..31767f9fe0e 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -3713,12 +3713,13 @@ if __name__ == "__main__": @pytest.mark.asyncio -async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_limiter): +async def test_websocket_turn_usage_is_not_counted_twice_toward_project_itpm_and_otpm(rate_limiter): """ - A native Responses WebSocket connection reserves project ITPM/OTPM at pre-call and - its turns log under their own call ids. Every turn's real usage must reach the - project counters, success and failure alike, and the disconnect must hand back - only the connection's estimate, never the turns' usage. + A native Responses WebSocket connection reserves project ITPM/OTPM at pre-call, and + each response.create frame is charged its own estimate before it is forwarded. The + turns that frame starts log under their own call ids, success and failure alike, and + must not count toward the project counters a second time. The disconnect hands back + only the connection's estimate. """ handler, cache = rate_limiter api_key = hash_token("sk-ws-project-turns") @@ -3765,7 +3766,17 @@ async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_lim }, } - assert await counters() == (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) + async def charge_frame() -> None: + await handler.enforce_project_io_token_quota_for_frame( + user_api_key_dict=user_api_key_dict, + requested_model="bedrock_mantle/claude-opus", + estimated_input_tokens=9, + estimated_output_tokens=40, + ) + + await charge_frame() + reserved_and_one_frame = (stash.itpm_reserved_tokens + 9, stash.otpm_reserved_tokens + 40) + assert await counters() == reserved_and_one_frame await handler.async_log_success_event( kwargs={**turn_kwargs("turn-1"), "call_type": CallTypes.aresponses_websocket.value}, @@ -3779,17 +3790,18 @@ async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_lim start_time=datetime.now(), end_time=datetime.now(), ) - assert await counters() == (stash.itpm_reserved_tokens + 30, stash.otpm_reserved_tokens + 7) + assert await counters() == reserved_and_one_frame + await charge_frame() await handler.async_log_failure_event( kwargs=turn_kwargs("turn-2", Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)), response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) - assert await counters() == (stash.itpm_reserved_tokens + 40, stash.otpm_reserved_tokens + 10) + assert await counters() == (stash.itpm_reserved_tokens + 18, stash.otpm_reserved_tokens + 80) await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict=user_api_key_dict) - assert await counters() == (40, 10) + assert await counters() == (18, 80) assert stash.reservation_released is True From d877c26517b99452697cff834f7018f9a83dcee2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:21:49 -0700 Subject: [PATCH 08/25] fix(proxy): keep per-turn WebSocket limits, request bodies and collector reservations right Free a rejected request's parallel slot through the request it runs in, and skip native Responses WebSocket connections, whose disconnect frees the slot. Count WebSocket turns toward the connection's tag TPM. In collector mode, claim the shared budget reservation before the spend event is built, so only the first turn settles it. Store each turn's own frame as its request body, and send an id-less error to the oldest waiting turn unless an earlier turn is streaming. --- litellm/litellm_core_utils/litellm_logging.py | 16 +++++- .../hooks/parallel_request_limiter_v3.py | 24 +++++---- .../proxy/hooks/proxy_track_cost_callback.py | 20 +++---- litellm/responses/streaming_iterator.py | 15 +++--- .../test_litellm_logging.py | 13 ++++- .../hooks/test_parallel_request_limiter_v3.py | 54 +++++++++++++++---- .../hooks/test_proxy_track_cost_callback.py | 38 +++++++++++++ .../test_responses_websocket_all_providers.py | 47 +++++++++++++++- 8 files changed, 188 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8d0cccd54f4..1d180d9be51 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -541,6 +541,17 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata" _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +def _proxy_server_request_with_body( + litellm_params: Mapping[str, object], request_body: Mapping[str, object] +) -> Mapping[str, object]: + proxy_server_request: Final = litellm_params.get("proxy_server_request") + return ( + {"proxy_server_request": {**_LITELLM_PARAMS.validate_python(proxy_server_request), "body": dict(request_body)}} + if isinstance(proxy_server_request, Mapping) + else {} + ) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -767,7 +778,9 @@ class Logging(LiteLLMLoggingBaseClass): self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None - def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": + def copy_for_new_call( + self, litellm_call_id: str, start_time: datetime.datetime, request_body: Mapping[str, object] + ) -> "Logging": connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params) litellm_params: Final[dict[str, object]] = { **connection_params, @@ -776,6 +789,7 @@ class Logging(LiteLLMLoggingBaseClass): for key, value in connection_params.items() if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) }, + **_proxy_server_request_with_body(connection_params, request_body), "litellm_call_id": litellm_call_id, } new_call: Final = copy.copy(self) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 70e67099832..0fb65b415b1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -742,15 +742,18 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def _request_owner_call_id(data: Mapping[str, object]) -> str | None: - from litellm.proxy.common_request_processing import request_litellm_call_id - - return request_litellm_call_id(data) +def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None) -> frozenset[str]: + if call_stash is not None: + return call_stash.tpm_limited_tags + connection: Final = get_request_stash() + return connection.tpm_limited_tags if connection is not None and connection.websocket_connection else frozenset() def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: + from litellm.proxy.common_request_processing import request_litellm_call_id + stash: Final = get_or_create_request_stash() - owner_call_id: Final = _request_owner_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) + owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) if owner_call_id is not None: stash.owner_litellm_call_id = owner_call_id return stash @@ -5067,7 +5070,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets: Final = self._collect_tpm_scope_targets( standard_logging_metadata=standard_logging_metadata, kwargs=kwargs, - tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(), + tpm_limited_tags=_tpm_limited_tags(stash), model_group=reconcile_model.group if reconcile_model is not None else None, ) targets.extend( @@ -5458,12 +5461,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): both this hook and async_log_failure_event end up running in the same flow, only the first release/refund applies. A mid-stream failure relayed here with recovered partial usage settles the reservation at - that usage instead of refunding it. + that usage instead of refunding it. A native Responses WebSocket + connection keeps its slot and reservation through a failed turn, and + its disconnect releases them. """ try: - failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data) - stash: Final = get_request_stash_for_call(_request_owner_call_id(failed_request)) - if stash is None: + stash: Final = get_request_stash() + if stash is None or stash.websocket_connection: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3c03f13e02d..d812b108403 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -139,6 +139,12 @@ class _ProxyDBLogger(CustomLogger): if self.spend_event_producer is None or not is_offloadable_success(response_obj): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) return + _claim_budget_reservation( + budget_reservation=_get_budget_reservation_from_metadata( + metadata=get_litellm_metadata_from_kwargs(kwargs=kwargs) + ), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)), + ) event: Final = build_spend_event( kwargs, response_obj, @@ -170,7 +176,7 @@ class _ProxyDBLogger(CustomLogger): ) -> None: budget_reservation: Final = _claim_budget_reservation( budget_reservation=user_api_key_dict.budget_reservation, - litellm_call_id=_request_call_id(_CALLBACK_DATA.validate_python(request_data)), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(request_data)), ) try: await _release_budget_reservation(budget_reservation=budget_reservation) @@ -353,7 +359,7 @@ class _ProxyDBLogger(CustomLogger): _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _claim_budget_reservation( budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata), - litellm_call_id=_logged_call_id(_CALLBACK_DATA.validate_python(kwargs)), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)), ) if ( isinstance(completion_response, LiteLLMBatch) @@ -709,14 +715,8 @@ def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dic _CALLBACK_DATA: Final = TypeAdapter(Mapping[str, object]) -def _request_call_id(request_data: Mapping[str, object]) -> str | None: - from litellm.proxy.common_request_processing import request_litellm_call_id - - return request_litellm_call_id(request_data) - - -def _logged_call_id(kwargs: Mapping[str, object]) -> str | None: - call_id: Final = kwargs.get("litellm_call_id") +def _string_call_id(callback_data: Mapping[str, object]) -> str | None: + call_id: Final = callback_data.get("litellm_call_id") return call_id if isinstance(call_id, str) else None diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 529959bc576..21b6fd8ee5e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2007,7 +2007,9 @@ class ResponsesWebSocketStreaming: def _start_turn(self, client_frame: str | None) -> None: turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes turn_logging_obj: Final = self.logging_obj.copy_for_new_call( - litellm_call_id=str(uuid.uuid4()), start_time=turn_start + litellm_call_id=str(uuid.uuid4()), + start_time=turn_start, + request_body=_json_object_or_none(client_frame) or {}, ) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") @@ -2036,10 +2038,11 @@ class ResponsesWebSocketStreaming: def _turn_index_for_event(self, event_obj: Mapping[str, object]) -> int: """ An event naming a response id belongs to the turn whose ``response.created`` - carried it. The backend answers accepted turns strictly in order, so an - id-less event goes to the oldest turn still waiting for its ``response.created``, - except an id-less ``error``: that is the backend rejecting the newest pending - ``response.create`` at parse time while an earlier turn is still streaming. + carried it. The backend answers turns strictly in order, so an id-less event, + a rejected turn's ``error`` included, goes to the oldest turn still waiting for + its ``response.created``. The one exception is an id-less ``error`` while a turn + is streaming: the backend rejects a ``response.create`` it cannot parse as soon + as it reads it, so that error belongs to the newest pending turn. """ response_id: Final = _ws_event_response_id(event_obj) turn_ids: Final = tuple(turn.response_id for turn in self._open_turns) @@ -2048,7 +2051,7 @@ class ResponsesWebSocketStreaming: pending: Final = tuple(index for index, turn_id in enumerate(turn_ids) if turn_id is None) if not pending: return 0 - if response_id is None and event_obj.get("type") == "error": + if response_id is None and event_obj.get("type") == "error" and len(pending) < len(turn_ids): return pending[-1] return pending[0] diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 35b1ed16aa2..ccc4f00ed47 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8267,10 +8267,12 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): reservation = {"reserved_cost": 0.5, "finalized": False} connection = _responses_ws_logging_obj() connection_start = connection.start_time + connection_request = {"url": "ws://proxy/v1/responses", "body": {"first_message": "turn 1 frame"}} connection.litellm_params = { "metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation}, "litellm_metadata": {"model_info": {"id": "dep-1"}}, "litellm_call_id": connection.litellm_call_id, + "proxy_server_request": connection_request, } connection.model_call_details = { **connection.model_call_details, @@ -8279,7 +8281,8 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): } turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47) - turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start) + turn_frame = {"type": "response.create", "input": "turn 2 frame"} + turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start, request_body=turn_frame) turn.litellm_params["metadata"]["status"] = "failure" turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o" turn.model_call_details["standard_logging_object"] = {"id": "resp_1"} @@ -8290,6 +8293,14 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call" assert turn.model_call_details["first_api_call_start_time"] is None assert turn.litellm_params["metadata"]["budget_reservation"] is reservation + assert turn.model_call_details["litellm_params"]["proxy_server_request"] == { + "url": "ws://proxy/v1/responses", + "body": turn_frame, + } + assert connection.litellm_params["proxy_server_request"] == { + "url": "ws://proxy/v1/responses", + "body": {"first_message": "turn 1 frame"}, + } assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start) assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test" assert "status" not in connection.litellm_params["metadata"] diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d11b50243b5..8c5db21cb9a 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4383,9 +4383,11 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"} await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type="" + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "litellm_call_id": "socket-call"}, + call_type=CallTypes.aresponses_websocket.value, ) await handler.async_post_call_failure_hook( @@ -4393,14 +4395,10 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3() original_exception=Exception("previous_response_not_found"), user_api_key_dict=user_api_key_dict, ) - in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value( + in_flight_after_turn_failure: Final = handler._gauge_in_flight_from_cache_value( await local_cache.async_get_cache(key=counter_key) ) - await handler.async_post_call_failure_hook( - request_data=connection_data, - original_exception=Exception("socket closed"), - user_api_key_dict=user_api_key_dict, - ) + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert in_flight_after_turn_failure == 1 assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @@ -4436,15 +4434,16 @@ async def _admit_pass_through_with_client_call_id( @pytest.mark.asyncio -async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(): +async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(monkeypatch: pytest.MonkeyPatch): _api_key = hash_token("sk-passthrough-client-id-failure") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict) + monkeypatch.setattr(litellm, "callbacks", [handler]) - await handler.async_post_call_failure_hook( + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_failure_hook( request_data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj}, original_exception=Exception("rejected by a pre-call hook"), user_api_key_dict=user_api_key_dict, @@ -5224,6 +5223,41 @@ async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch): assert "cell-blocked" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch): + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") + _request_stash.set(None) + resolver, _ = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=None, tpm_limit=1_000)}) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), tag_rate_limit_resolver=resolver + ) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-ws-tag-tpm")) + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "metadata": {"tags": ["cell-1"]}, "litellm_call_id": "socket-call"}, + call_type=CallTypes.aresponses_websocket.value, + ) + tokens_before_turns: Final = int(await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0) + + for turn_call_id in ("turn-1-call", "turn-2-call"): + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": turn_call_id, + "standard_logging_object": {"metadata": {"user_api_key_hash": user_api_key_dict.api_key}}, + "model": "gpt-4o-mini", + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert await local_cache.async_get_cache("{tag:cell-1}:tokens") == tokens_before_turns + 240 + + @pytest.mark.asyncio async def test_tag_object_tpm_limit_enforced_v3(monkeypatch): monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") 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 1dd86c2c4a0..3ae4af4a4bc 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2718,6 +2718,44 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski assert event.response_cost == 0.0125 +@pytest.mark.asyncio +async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservation(tmp_path): + from litellm.litellm_core_utils.core_helpers import ( + budget_reservation_from_metadata, + get_litellm_metadata_from_kwargs, + ) + from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args + + handler = _RecordingHandler() + consumer = SpendEventConsumer(handler) + address = UnixAddress(path=str(tmp_path / "spend.sock")) + server = await consumer.serve(address) + producer = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback + ) + logger = ProxyDBLogger(producer) + connection_kwargs = _offload_kwargs() + + for turn_call_id in ("turn-1", "turn-2", "turn-3"): + await logger.async_log_success_event( + {**connection_kwargs, "litellm_call_id": turn_call_id}, _offload_response(), datetime.now(), datetime.now() + ) + await producer.close(drain_timeout=5.0) + server.close() + assert await consumer.drain(timeout=5.0) == 0 + + sidecar_kwargs = [spend_event_callback_args(decode_spend_event(line)).kwargs for line in handler.lines] + assert [ + claim_budget_reservation( + budget_reservation=budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(kwargs)), + litellm_call_id=kwargs["litellm_call_id"], + ) + is not None + for kwargs in sidecar_kwargs + ] == [True, False, False] + + @pytest.mark.asyncio async def test_async_log_success_event_keeps_batch_retrieves_in_process(): producer = SpendEventProducer( diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 48bb4340fba..39bcf6710aa 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1907,12 +1907,13 @@ class _TurnLoggingRecorder: self.connection = MagicMock() self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call) - def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock: + def _copy_for_new_call(self, litellm_call_id: str, start_time: object, request_body: object) -> MagicMock: from unittest.mock import AsyncMock turn = MagicMock() turn.litellm_call_id = litellm_call_id turn.start_time = start_time + turn.request_body = request_body turn.model_call_details = {} turn.dispatch_success_handlers = AsyncMock() turn.dispatch_failure_handlers = AsyncMock() @@ -3336,6 +3337,9 @@ class TestNativeWebSocketPerTurnLogging: [{"role": "user", "content": "two"}], [{"role": "user", "content": "three"}], ] + assert [turn.request_body for turn in recorder.turns] == [ + json.loads(_response_create(text)) for text in ("one", "two", "three") + ] for turn in recorder.turns: turn.dispatch_failure_handlers.assert_not_awaited() @@ -3402,6 +3406,47 @@ class TestNativeWebSocketPerTurnLogging: assert "max_output_tokens" in str(second.dispatch_failure_handlers.await_args[0][0]) assert second.record_partial_usage_for_failure.call_args[0][0].total_tokens == 0 + @pytest.mark.asyncio + async def test_a_rejected_turn_answered_in_order_fails_and_the_next_turn_keeps_its_response(self): + """OpenAI answers a response.create it parses but rejects (previous_response_not_found) in order, + before it starts the next frame's response. With both turns still waiting, the id-less error is the + older turn's failure and the next response.created belongs to the newer turn.""" + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder() + handler = _make_streaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini", + custom_llm_provider="openai", + ) + created, _delta, completed = _completed_turn("resp_2", 40, 4) + rejection = { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "previous_response_not_found", + "message": "Previous response with id 'resp_missing' not found.", + }, + } + + handler._record_client_frame(_response_create("first")) + handler._record_client_frame(_response_create("second")) + for event in (rejection, created, completed): + handler._record_backend_event(json.dumps(event)) + assert await _until(lambda: recorder.logged_turns() == 2) + + first, second = recorder.turns + first.dispatch_success_handlers.assert_not_awaited() + assert "resp_missing" in str(first.dispatch_failure_handlers.await_args[0][0]) + second.dispatch_failure_handlers.assert_not_awaited() + assert [event["type"] for event in second.dispatch_success_handlers.await_args[0][0]] == [ + "response.created", + "response.completed", + ] + @pytest.mark.asyncio async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self): from unittest.mock import AsyncMock From 4b179a5815651f960a922e6916f83a41f1906f6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:20:37 -0700 Subject: [PATCH 09/25] fix(proxy): store a failed WebSocket turn's own request and count tag TPM only for calls that carry the tag --- litellm/proxy/hooks/parallel_request_limiter_v3.py | 13 +++++++++---- litellm/proxy/response_api_endpoints/endpoints.py | 2 ++ .../proxy/hooks/test_parallel_request_limiter_v3.py | 12 ++++++++---- .../proxy/response_api_endpoints/test_endpoints.py | 4 ++++ 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 0fb65b415b1..2850e203553 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -10,7 +10,7 @@ import itertools import logging import os import uuid -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable, Mapping, Sequence, Set from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass, field @@ -742,11 +742,13 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None) -> frozenset[str]: +def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None, call_tags: Iterable[object]) -> frozenset[str]: if call_stash is not None: return call_stash.tpm_limited_tags connection: Final = get_request_stash() - return connection.tpm_limited_tags if connection is not None and connection.websocket_connection else frozenset() + if connection is None or not connection.websocket_connection: + return frozenset() + return connection.tpm_limited_tags.intersection(call_tags) def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: @@ -776,6 +778,7 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None: _RESPONSES_WS_TURN_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_REQUEST_TAGS: Final = TypeAdapter(tuple[object, ...]) def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: @@ -5070,7 +5073,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets: Final = self._collect_tpm_scope_targets( standard_logging_metadata=standard_logging_metadata, kwargs=kwargs, - tpm_limited_tags=_tpm_limited_tags(stash), + tpm_limited_tags=_tpm_limited_tags( + stash, _REQUEST_TAGS.validate_python(standard_logging_object.get("request_tags") or ()) + ), model_group=reconcile_model.group if reconcile_model is not None else None, ) targets.extend( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index eaa6ec4c204..8b98f490a69 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1446,6 +1446,7 @@ _RESPONSES_WS_PER_TURN_DICT_KEYS: Final = frozenset({"metadata", "litellm_metada def _responses_ws_turn_request_data( connection_request_data: Mapping[str, object], turn_logging_obj: "LiteLLMLoggingObj" ) -> dict[str, object]: + turn_proxy_server_request: Final = turn_logging_obj.litellm_params.get("proxy_server_request") return { **{key: value for key, value in connection_request_data.items() if key not in _RESPONSES_WS_FIRST_FRAME_KEYS}, **{ @@ -1453,6 +1454,7 @@ def _responses_ws_turn_request_data( for key, value in connection_request_data.items() if key in _RESPONSES_WS_PER_TURN_DICT_KEYS and isinstance(value, dict) }, + **({"proxy_server_request": turn_proxy_server_request} if isinstance(turn_proxy_server_request, dict) else {}), "litellm_call_id": turn_logging_obj.litellm_call_id, "litellm_logging_obj": turn_logging_obj, } diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 8c5db21cb9a..46981381ebe 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5224,7 +5224,7 @@ async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch): @pytest.mark.asyncio -async def test_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch): +async def test_only_tagged_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch): monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") _request_stash.set(None) @@ -5242,11 +5242,15 @@ async def test_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatc ) tokens_before_turns: Final = int(await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0) - for turn_call_id in ("turn-1-call", "turn-2-call"): + calls_in_the_connection = (("turn-1-call", ["cell-1"]), ("nested-untagged-call", []), ("turn-2-call", ["cell-1"])) + for call_id, request_tags in calls_in_the_connection: await handler.async_log_success_event( kwargs={ - "litellm_call_id": turn_call_id, - "standard_logging_object": {"metadata": {"user_api_key_hash": user_api_key_dict.api_key}}, + "litellm_call_id": call_id, + "standard_logging_object": { + "metadata": {"user_api_key_hash": user_api_key_dict.api_key}, + "request_tags": request_tags, + }, "model": "gpt-4o-mini", }, response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)), diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index ce967403265..302d2d28843 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -822,12 +822,15 @@ class TestResponsesWSFirstFrameModelAuth: "metadata": connection_metadata, "litellm_metadata": {}, "litellm_call_id": "connection-call", + "proxy_server_request": {"url": "ws://testserver/v1/responses", "body": {"input": "first turn"}}, }, MagicMock(), ) ) + turn_proxy_server_request = {"url": "ws://testserver/v1/responses", "body": {"input": "second turn"}} turn_logging_obj = MagicMock() turn_logging_obj.litellm_call_id = "turn-2-call" + turn_logging_obj.litellm_params = {"proxy_server_request": turn_proxy_server_request} failure = litellm.BadRequestError( message="previous_response_not_found", model="gpt-4o-mini", llm_provider="openai" ) @@ -879,6 +882,7 @@ class TestResponsesWSFirstFrameModelAuth: assert turn_request["metadata"] is not connection_metadata assert "input" not in turn_request assert "first_message" not in turn_request + assert turn_request["proxy_server_request"] == turn_proxy_server_request proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) @pytest.mark.asyncio From b06da755d0cf4afeaf150367a908b60fd74ce8e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:50:37 -0700 Subject: [PATCH 10/25] fix(responses): log the client's turn frame, keep the payload id and reservation shape for non-WebSocket calls --- litellm/litellm_core_utils/litellm_logging.py | 19 +++++- .../proxy/response_api_endpoints/endpoints.py | 2 + .../spend_tracking/budget_reservation.py | 24 +++++--- litellm/responses/streaming_iterator.py | 34 ++++++++++- .../test_litellm_logging.py | 23 +++++++ .../hooks/test_proxy_track_cost_callback.py | 12 +++- .../test_pass_through_endpoints.py | 5 +- .../response_api_endpoints/test_endpoints.py | 61 ++++++++++++++++++- tests/unit/proxy/test_budget_reservation.py | 16 +++++ .../test_responses_websocket_all_providers.py | 48 +++++++++++++++ 10 files changed, 226 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1d180d9be51..07828b45eee 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -557,6 +557,18 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def _logged_response_id( + response: Mapping[str, object], litellm_call_id: object, falls_back_on_empty_id: bool +) -> str | None: + call_id: Final = litellm_call_id if isinstance(litellm_call_id, str) else None + if falls_back_on_empty_id: + return _provider_response_id(response) or call_id + if "id" not in response: + return call_id + response_id: Final = response["id"] + return None if response_id is None else str(response_id) + + def mask_api_base_credentials(api_base: str) -> str: if "key=" not in api_base: return api_base @@ -6688,7 +6700,12 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = _provider_response_id(response_obj) or kwargs.get("litellm_call_id") + litellm_call_id: Final[object] = kwargs.get("litellm_call_id") + id = _logged_response_id( + response_obj, + litellm_call_id, + falls_back_on_empty_id=isinstance(init_response_obj, LiteLLMRealtimeStreamLoggingObject), + ) _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 8b98f490a69..12e72f559d1 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_set_request_parsed_body, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.spend_tracking.budget_reservation import share_budget_reservation_across_calls from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1622,6 +1623,7 @@ async def _responses_websocket_session( await websocket.close(code=1008, reason="Pre-call error") return + share_budget_reservation_across_calls(user_api_key_dict.budget_reservation) routed_data: Final = dict( data, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8e8f4def110..cd0c4a8433d 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -378,18 +378,28 @@ async def reconcile_budget_reservation( return pending +_SETTLING_CALL_ID: Final = "settling_call_id" + + +def share_budget_reservation_across_calls(budget_reservation: dict[str, object] | None) -> None: + if budget_reservation is not None: + budget_reservation.setdefault(_SETTLING_CALL_ID, None) + + def claim_budget_reservation( budget_reservation: dict[str, object] | None, litellm_call_id: str | None ) -> dict[str, object] | None: - """Hand the reservation to the first call that settles it. + """Hand a shared reservation to the first call that settles it. - A native Responses WebSocket connection shares its one reservation with every turn it logs, each under its - own call id. Two turns settling it at once would both move the counter from the estimate to their own cost, - so the first claimant keeps the reservation and every later call sees none and charges its cost directly.""" - if budget_reservation is None or litellm_call_id is None: + A Responses WebSocket connection shares its one reservation with every turn it logs, each under its own + call id. Two turns settling it at once would both move the counter from the estimate to their own cost, + so the first claimant keeps the reservation and every later call sees none and charges its cost directly. + A reservation no connection shared belongs to its one call and is returned unchanged.""" + if budget_reservation is None or litellm_call_id is None or _SETTLING_CALL_ID not in budget_reservation: return budget_reservation - settling_call_id: Final = budget_reservation.setdefault("settling_call_id", litellm_call_id) - return budget_reservation if settling_call_id == litellm_call_id else None + if budget_reservation[_SETTLING_CALL_ID] is None: + budget_reservation[_SETTLING_CALL_ID] = litellm_call_id + return budget_reservation if budget_reservation[_SETTLING_CALL_ID] == litellm_call_id else None def stamp_budget_reservation_actual_cost(budget_reservation: dict | None, actual_cost: float | None) -> None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 21b6fd8ee5e..78e6c4bf0d5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1938,6 +1938,31 @@ def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> di return {**msg_obj, **top_fields, **restored_nested} +_RESPONSE_CREATE_MASKED_FIELDS: Final = ("input", "instructions") + + +def _with_masked_fields(client: Mapping[str, object], forwarded: Mapping[str, object]) -> Mapping[str, object]: + masked: Final = { + key: forwarded[key] for key in _RESPONSE_CREATE_MASKED_FIELDS if key in client and key in forwarded + } + return {**client, **masked} + + +def _client_frame_with_masked_text(client_frame: str, forwarded_frame: str) -> str: + client: Final = _json_object_or_none(client_frame) + forwarded: Final = _json_object_or_none(forwarded_frame) + if client is None or forwarded is None: + return client_frame + client_nested: Final = client.get("response") + forwarded_nested: Final = forwarded.get("response") + nested: Final = ( + {"response": _with_masked_fields(client_nested, forwarded_nested)} + if _is_json_object(client_nested) and _is_json_object(forwarded_nested) + else {} + ) + return json.dumps({**_with_masked_fields(client, forwarded), **nested}) + + def _wrap_output_item_encrypted_content( event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object] ) -> dict[str, object] | None: @@ -2015,6 +2040,11 @@ class ResponsesWebSocketStreaming: turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) + def _logged_client_frame(self, client_frame: str, forwarded_frame: str) -> str: + if not self.guardrail_callbacks or forwarded_frame is client_frame: + return client_frame + return _client_frame_with_masked_text(client_frame, forwarded_frame) + def _record_client_frame(self, client_frame: str) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": @@ -2504,7 +2534,7 @@ class ResponsesWebSocketStreaming: try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) - self._record_client_frame(masked_first) + self._record_client_frame(self._logged_client_frame(self.first_message, masked_first)) await self.backend_ws.send(masked_first) while True: @@ -2512,7 +2542,7 @@ class ResponsesWebSocketStreaming: if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) - self._record_client_frame(masked) + self._record_client_frame(self._logged_client_frame(message, masked)) await self.backend_ws.send(masked) except Exception as e: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index ccc4f00ed47..fb5f87c3016 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -38,6 +38,7 @@ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, _get_status_fields, + get_standard_logging_object_payload, set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -8263,6 +8264,28 @@ def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw assert payload["id"] == expected_request_id +@pytest.mark.parametrize( + ("response_obj", "expected_id"), + (({"id": ""}, ""), ({"id": "chatcmpl-1"}, "chatcmpl-1"), ({}, "http-call")), +) +def test_non_websocket_payload_id_is_the_response_id_field_as_returned( + logging_obj: LitellmLogging, response_obj: dict[str, str], expected_id: str +): + now = datetime.datetime.now() + + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "http-call", "model": "gpt-4o", "messages": []}, + init_response_obj=response_obj, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["id"] == expected_id + + def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): reservation = {"reserved_cost": 0.5, "finalized": False} connection = _responses_ws_logging_obj() 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 3ae4af4a4bc..89418a0f7a8 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2724,7 +2724,10 @@ async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservat budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, ) - from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + from litellm.proxy.spend_tracking.budget_reservation import ( + claim_budget_reservation, + share_budget_reservation_across_calls, + ) from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args handler = _RecordingHandler() @@ -2736,6 +2739,9 @@ async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservat ) logger = ProxyDBLogger(producer) connection_kwargs = _offload_kwargs() + share_budget_reservation_across_calls( + budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(connection_kwargs)) + ) for turn_call_id in ("turn-1", "turn-2", "turn-3"): await logger.async_log_success_event( @@ -3088,7 +3094,7 @@ async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_ """Every turn of a Responses WebSocket connection carries the connection's one reservation under its own call id. The first turn settles it; a later turn reconciling the same estimate again would under-count the key, so it must charge its cost without the reservation.""" - reservation = {"reserved_cost": 0.5, "entries": []} + reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None} first_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-1", reservation)) second_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-2", reservation)) @@ -3100,7 +3106,7 @@ async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_ @pytest.mark.asyncio async def test_async_post_call_failure_hook_releases_a_shared_reservation_for_the_first_turn_only(): logger = ProxyDBLogger() - budget_reservation = {"reserved_cost": 0.5, "entries": []} + budget_reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None} user_api_key_dict = UserAPIKeyAuth( api_key="test_api_key", request_route="/custom/route", diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 14040a7206d..eb31797d460 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7711,10 +7711,7 @@ async def test_passthrough_success_reconciles_budget_reservation(): increment_spend_counters.assert_awaited_once() assert increment_spend_counters.await_args.kwargs["budget_reservation"] is reservation - assert increment_spend_counters.await_args.kwargs["budget_reservation"] == { - **budget_reservation, - "settling_call_id": "lit-5425-call-id", - } + assert increment_spend_counters.await_args.kwargs["budget_reservation"] == budget_reservation @pytest.mark.asyncio diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 302d2d28843..439d4026af0 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3,7 +3,7 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from collections.abc import Mapping +from collections.abc import Coroutine, Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -885,6 +885,65 @@ class TestResponsesWSFirstFrameModelAuth: assert turn_request["proxy_server_request"] == turn_proxy_server_request proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) + @pytest.mark.asyncio + async def test_endpoint_hands_the_connection_reservation_to_its_first_settling_turn_only(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock( + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"}) + ) + ws.close = AsyncMock() + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-token", budget_reservation={"reserved_cost": 0.5, "entries": []} + ) + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock() + + async def relay() -> None: + return None + + def route(**_: object) -> Coroutine[object, object, None]: + return relay() + + with ( + patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; the reservation it leaves is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing starts the relay, which needs a live upstream socket + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + side_effect=route, + ), + patch( # test-quality-ok: the disconnect release reads proxy_logging_obj off proxy_server + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_obj, + ), + ): + await responses_websocket_endpoint(websocket=ws, model=None, user_api_key_dict=user_api_key_dict) + + reservation = user_api_key_dict.budget_reservation + assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") is reservation + assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") is None + @pytest.mark.asyncio async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self): from litellm.proxy.response_api_endpoints.endpoints import ( diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 4591daaded6..d902a0e7a47 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -1,4 +1,5 @@ import asyncio +import copy import threading import time from collections.abc import Mapping, Sequence @@ -46,6 +47,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( invalidate_budget_reservation_counters, release_budget_reservation, release_budget_reservation_on_cancel, + share_budget_reservation_across_calls, reserve_budget_for_request, ) from litellm.proxy.spend_tracking.input_tokens import ( @@ -3723,6 +3725,19 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one ) == pytest.approx(4.2) +def test_a_reservation_no_connection_shared_is_left_as_its_call_made_it(): + reservation: dict[str, object] = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:k", "reserved_cost": 0.5}], + } + made_by_its_call = copy.deepcopy(reservation) + + claims = [claim_budget_reservation(budget_reservation=reservation, litellm_call_id=call) for call in ("a", "b")] + + assert claims == [reservation, reservation] + assert reservation == made_by_its_call + + @pytest.mark.asyncio async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at_once( spend_counter_state, @@ -3756,6 +3771,7 @@ async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at from litellm.proxy.proxy_server import increment_spend_counters + share_budget_reservation_across_calls(reservation) first_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") second_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") assert first_turn is reservation diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 39bcf6710aa..f4b1b8c86d1 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -3343,6 +3343,54 @@ class TestNativeWebSocketPerTurnLogging: for turn in recorder.turns: turn.dispatch_failure_handlers.assert_not_awaited() + @pytest.mark.parametrize( + ("guardrail_callbacks", "logged_text"), + [([], "mail alice@example.com"), ([_FakeWSGuardrail()], "mail ")], + ) + @pytest.mark.asyncio + async def test_each_turn_logs_the_frame_the_client_sent_not_the_frame_sent_upstream( + self, guardrail_callbacks: list[_FakeWSGuardrail], logged_text: str + ): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream_prev" + ) + client_frame = { + "type": "response.create", + "model": "client-model-group", + "input": "mail alice@example.com", + "previous_response_id": wrapped_previous, + } + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([json.dumps(client_frame)], recorder.logged_turns) + backend = _ScriptedBackend([_completed_turn("resp_1", 10, 1)]) + handler = _make_streaming( + websocket=client, + backend_ws=backend, + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini-deployment", + request_defaults=_deployment_defaults(), + guardrail_callbacks=guardrail_callbacks, + ) + + await handler.bidirectional_forward() + + assert json.loads(backend.sent[0]) == { + "type": "response.create", + "model": "gpt-5.4-mini-deployment", + "input": logged_text, + "previous_response_id": "resp_upstream_prev", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + assert [turn.request_body for turn in recorder.turns] == [{**client_frame, "input": logged_text}] + assert recorder.turns[0].model_call_details["messages"] == [{"role": "user", "content": logged_text}] + @pytest.mark.asyncio async def test_a_turn_whose_input_is_nested_under_response_logs_its_user_message(self): import websockets.exceptions # noqa: F401 (lazy submodule must be importable) From 6ccd11643f387ef7b576224c59ed9920ee92d39d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:44:50 -0700 Subject: [PATCH 11/25] fix(proxy): keep the legacy limiter's WebSocket slot until the socket closes --- .../proxy/hooks/parallel_request_limiter.py | 77 +++++++++-- litellm/proxy/utils.py | 3 + .../hooks/test_parallel_request_limiter.py | 130 +++++++++++++++++- 3 files changed, 193 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 19251b53f13..54f4fbf562a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,6 +1,8 @@ import asyncio import sys from collections.abc import Callable +from contextvars import ContextVar +from dataclasses import dataclass from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn @@ -23,7 +25,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.types.utils import Usage +from litellm.types.utils import CallTypes, Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -48,6 +50,17 @@ def _response_total_tokens(response_obj: object) -> int: return response_usage.total_tokens if isinstance(response_usage, Usage) else 0 +@dataclass(frozen=True, slots=True) +class _WebSocketConnectionSlot: + request_count_key: str | None + holds_global_slot: bool + + +_websocket_connection_slot: Final[ContextVar[_WebSocketConnectionSlot | None]] = ContextVar( + "litellm_legacy_rate_limiter_websocket_connection_slot", default=None +) + + class CacheObject(TypedDict): current_global_requests: dict | None request_count_api_key: dict | None @@ -63,7 +76,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.internal_usage_cache = internal_usage_cache self._clock: Final = clock - def print_verbose(self, print_statement) -> None: + def print_verbose(self, print_statement: object) -> None: try: verbose_proxy_logger.debug(print_statement) if litellm.set_verbose: @@ -499,6 +512,41 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) # don't block execution for cache updates ) + if call_type == CallTypes.aresponses_websocket.value: + _websocket_connection_slot.set( + _WebSocketConnectionSlot( + request_count_key=(f"{api_key}::{precise_minute}::request_count" if api_key is not None else None), + holds_global_slot=global_max_parallel_requests is not None, + ) + ) + + @with_service_target("rate_limits") + async def async_release_websocket_connection_slot(self) -> None: + connection: Final = _websocket_connection_slot.get() + if connection is None: + return + _websocket_connection_slot.set(None) + if connection.holds_global_slot: + await self.internal_usage_cache.async_increment_cache( + key="global_max_parallel_requests", + value=-1, + local_only=True, + litellm_parent_otel_span=None, + ) + if connection.request_count_key is None: + return + current: Final = await self.internal_usage_cache.async_get_cache( + key=connection.request_count_key, + litellm_parent_otel_span=None, + ) + if current is None: + return + await self.internal_usage_cache.async_set_cache( + connection.request_count_key, + {**current, "current_requests": max(current["current_requests"] - 1, 0)}, + ttl=60, + litellm_parent_otel_span=None, + ) @with_service_target("rate_limits") async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: @@ -520,6 +568,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "user_api_key_model_max_budget", None ) user_api_key_end_user_id: Final = kwargs.get("user") + requests_released: Final = 0 if _websocket_connection_slot.get() is not None else 1 user_api_key_metadata: Final = kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) or {} user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_team_metadata", None) @@ -540,7 +589,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-requests_released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -562,13 +611,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -603,13 +652,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -629,13 +678,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -655,13 +704,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -681,13 +730,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -712,7 +761,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): global_max_parallel_requests: Final = _metadata.get("global_max_parallel_requests", None) user_api_key: Final = _metadata.get("user_api_key", None) self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") - if user_api_key is None: + if user_api_key is None or _websocket_connection_slot.get() is not None: return ## decrement call count if call failed diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9788c4635c9..bc70604a5b2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4405,6 +4405,9 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if isinstance(limiter, PROXY_MaxParallelRequestsHandler): + await limiter.async_release_websocket_connection_slot() + return if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 89cb6bd1ef2..3f5eef4593f 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -2,8 +2,9 @@ Unit Tests for the max parallel request limiter v1 for the proxy """ +import asyncio import itertools -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime from typing import Final @@ -12,11 +13,12 @@ import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.types.utils import EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token +from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) LAST_MICROSECOND_OF_JANUARY: Final = datetime(2026, 1, 31, 23, 59, 59, 999999) @@ -323,3 +325,125 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_ ) assert exc_info.value.headers["retry-after"] == "14.5" + + +async def _log_websocket_turn_success( + handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object] +) -> None: + await handler.async_log_success_event( + kwargs={"litellm_params": {"metadata": dict(metadata)}}, + response_obj=None, + start_time=FROZEN_INSTANT, + end_time=FROZEN_INSTANT, + ) + + +async def _log_websocket_turn_failure( + handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object] +) -> None: + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": dict(metadata)}, "exception": Exception("upstream turn failed")}, + response_obj=None, + start_time=FROZEN_INSTANT, + end_time=FROZEN_INSTANT, + ) + + +def _proxy_logging_with(handler: PROXY_MaxParallelRequestsHandler) -> ProxyLogging: + proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = handler + return proxy_logging + + +async def _requests_admitted( + handler: PROXY_MaxParallelRequestsHandler, session: UserAPIKeyAuth, metadata: Mapping[str, object] +) -> int: + for admitted in range(3): + try: + await handler.async_pre_call_hook( + user_api_key_dict=session, + cache=DualCache(), + data={"metadata": dict(metadata)}, + call_type=CallTypes.acompletion.value, + ) + except ProxyRateLimitError: + return admitted + await asyncio.sleep(0) + return 3 + + +@pytest.mark.parametrize("log_turn", [_log_websocket_turn_success, _log_websocket_turn_failure]) +@pytest.mark.parametrize( + "max_parallel_requests,limit_metadata,admitted_while_open", + [(1, {}, 0), (None, {"global_max_parallel_requests": 2}, 1)], + ids=["key_limit", "global_limit"], +) +@pytest.mark.asyncio +async def test_a_websocket_connection_holds_its_parallel_slot_through_its_turns_until_it_closes( + log_turn: Callable[[PROXY_MaxParallelRequestsHandler, Mapping[str, object]], Awaitable[None]], + max_parallel_requests: int | None, + limit_metadata: Mapping[str, object], + admitted_while_open: int, +): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock + ) + api_key: Final = hash_token("sk-websocket-slot") + session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=max_parallel_requests) + turn_logged: Final = asyncio.Event() + socket_closed: Final = asyncio.Event() + + async def websocket_connection() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=session, + cache=DualCache(), + data={"metadata": dict(limit_metadata)}, + call_type=CallTypes.aresponses_websocket.value, + ) + await asyncio.sleep(0) + await log_turn(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}, **limit_metadata}) + turn_logged.set() + await socket_closed.wait() + await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session) + + connection: Final = asyncio.create_task(websocket_connection()) + await turn_logged.wait() + assert await _requests_admitted(handler, session, limit_metadata) == admitted_while_open + + socket_closed.set() + await connection + assert await _requests_admitted(handler, session, limit_metadata) == 1 + + +@pytest.mark.asyncio +async def test_a_websocket_turn_logged_in_a_later_minute_leaves_no_slot_held_after_the_socket_closes(): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_clock_rolling_over_after_first_read() + ) + api_key: Final = hash_token("sk-websocket-rollover") + session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + async def websocket_connection() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=session, cache=DualCache(), data={}, call_type=CallTypes.aresponses_websocket.value + ) + await asyncio.sleep(0) + await _log_websocket_turn_success(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}}) + await handler.async_release_websocket_connection_slot() + + await asyncio.create_task(websocket_connection()) + + assert await _requests_admitted(handler, session, {}) == 1 + + +@pytest.mark.asyncio +async def test_a_disconnect_release_outside_a_websocket_connection_leaves_the_request_slot_held(): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock + ) + session: Final = UserAPIKeyAuth(api_key=hash_token("sk-http-slot"), max_parallel_requests=1) + + assert await _requests_admitted(handler, session, {}) == 1 + await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session) + + assert await _requests_admitted(handler, session, {}) == 0 From fbe9c3084cc0238760ce7b5344e35259f9e87c84 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:10:43 -0700 Subject: [PATCH 12/25] fix(router): count native Responses WebSocket turns toward deployment limits and budgets --- .../proxy/hooks/proxy_track_cost_callback.py | 6 +-- litellm/router.py | 3 +- litellm/router_strategy/budget_limiter.py | 3 +- .../response_api_endpoints/test_endpoints.py | 23 ---------- .../router_strategy/test_budget_limiter.py | 4 +- .../test_router/test_router_helper_utils.py | 44 +++++++++++++++++++ 6 files changed, 50 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6fbd0b8f1b3..f6e9928ae1d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -471,11 +471,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/router.py b/litellm/router.py index a549bb06e7d..130cdfab309 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8579,8 +8579,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 diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b20f71cd51f..2780455d604 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -502,8 +502,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/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0eaf619ba70..dd7a0ec7a2e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1642,29 +1642,6 @@ class TestResponsesWSStreamingFirstMessage: 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.""" diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 685ea46547e..335c399df8a 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -59,7 +59,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..39fa0970aa1 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -482,6 +482,50 @@ 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 + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ From 7a72a58353aaa7962ccb62a5175fce48e2bb3a64 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:51:48 -0700 Subject: [PATCH 13/25] fix(router): re-check deployment budget and RPM on each Responses WebSocket turn --- litellm/llms/custom_httpx/llm_http_handler.py | 5 +- .../proxy/response_api_endpoints/endpoints.py | 10 +- litellm/responses/main.py | 4 +- litellm/responses/streaming_iterator.py | 27 ++++- litellm/router.py | 31 +++++- .../hooks/test_proxy_track_cost_callback.py | 32 ++++++ .../response_api_endpoints/test_endpoints.py | 22 +++- .../test_responses_websocket_all_providers.py | 101 ++++++++++++++++++ .../test_router/test_router_helper_utils.py | 57 ++++++++++ 9 files changed, 276 insertions(+), 13 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 610321fe1f6..ba5947abf44 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2,7 +2,7 @@ import asyncio import inspect import json import ssl -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache, partial from types import MappingProxyType @@ -6806,6 +6806,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs: object, ) -> Exception | None: """ @@ -6838,6 +6839,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + turn_admission=turn_admission, **kwargs, ) await handler.run() @@ -6962,6 +6964,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + turn_admission=turn_admission, ) return await streaming.bidirectional_forward() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b7f57d93a9b..cdaf54e10ce 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1731,12 +1732,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec def _responses_ws_failure_frame(failure: Exception) -> str: - raw_status: Final = getattr(failure, "status_code", None) - status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 + proxy_exception: Final = handle_exception_on_proxy(failure) + status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500 error_type: Final = ( "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error" ) - return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}}) + return json.dumps( + {"type": "error", "status": status, "error": {"type": error_type, "message": proxy_exception.message}} + ) async def _enforce_responses_ws_first_frame_model_auth( @@ -1896,7 +1899,6 @@ async def _responses_websocket_session( request_data=routed_data, ) except Exception as e: - verbose_proxy_logger.exception("Responses WebSocket error") with contextlib.suppress(Exception): await websocket.send_text(_responses_ws_failure_frame(e)) await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f566c8f4e02..269b7367f34 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 @@ -2475,6 +2475,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs, ) -> Exception | None: """ @@ -2567,5 +2568,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + turn_admission=turn_admission, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cf68332c07..58f55ac73fd 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1892,6 +1892,13 @@ async def _enforce_frame_project_quota( ) +async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> None: + frame: Final = _json_object_or_none(raw_message) + if turn_admission is None or frame is None or frame.get("type") != "response.create": + return + await turn_admission() + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2062,6 +2069,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -2081,6 +2089,7 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission def _start_turn(self, client_frame: str | None) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time @@ -2522,8 +2531,10 @@ class ResponsesWebSocketStreaming: return json.dumps(evt_obj) if modified else response_str - async def _enforce_or_reject_frame(self, message: str) -> bool: - """Run the per-frame project quota check. + async def _enforce_or_reject_frame( + self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> bool: + """Run the per-frame project quota check, then the deployment's own admission for a turn after the first. On rejection, sends an ``error`` event to the client and reports that the frame must be dropped instead of forwarded, so the connection @@ -2533,6 +2544,7 @@ class ResponsesWebSocketStreaming: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) + await _admit_frame_turn(turn_admission, message) except RateLimitError as e: try: await self.websocket.send_text( @@ -2561,7 +2573,7 @@ class ResponsesWebSocketStreaming: while True: message = await self.websocket.receive_text() - if not await self._enforce_or_reject_frame(message): + if not await self._enforce_or_reject_frame(message, self.turn_admission): continue masked = await self._mask_response_create(message) self._record_client_frame(masked) @@ -2642,6 +2654,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + turn_admission: Callable[[], Awaitable[None]] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2660,6 +2673,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3027,7 +3041,9 @@ class ManagedResponsesWebSocketHandler: # Core request handler # ------------------------------------------------------------------ - async def _process_response_create(self, raw_message: str) -> None: + async def _process_response_create( + self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, and forward every streaming event to the client. @@ -3067,6 +3083,7 @@ class ManagedResponsesWebSocketHandler: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message ) + await _admit_frame_turn(turn_admission, raw_message) except RateLimitError as e: await self._send_error(str(e), error_type="rate_limit_exceeded") return @@ -3123,7 +3140,7 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc) break - await self._process_response_create(message) + await self._process_response_create(message, self.turn_admission) except Exception as exc: verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) diff --git a/litellm/router.py b/litellm/router.py index 130cdfab309..94debd5b03d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5632,6 +5632,7 @@ class Router: ) custom_llm_provider: Final = provider_for_generic_call(data) + generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") response_kwargs: Final = { **data, @@ -5639,6 +5640,11 @@ class Router: **kwargs, "model": model_name, **_with_router_resolved_session_model(kwargs.get("session"), model_name), + **( + {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)} + if generic_function_name == "_aresponses_websocket" + else {} + ), } # Only set custom_llm_provider if it's not None if custom_llm_provider is not None: @@ -5648,7 +5654,7 @@ class Router: self, deployment, response_kwargs, - surface_for_call(getattr(original_generic_function, "__name__", "")), + surface_for_call(generic_function_name), ) async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): response = await original_generic_function(**compacted_input) @@ -9247,6 +9253,29 @@ class Router: ) raise e + async def async_admit_deployment_turn( + self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> None: + budget_limiter: Final = self.get_router_deployment_budget_limiter() + if budget_limiter is not None: + try: + await budget_limiter.async_filter_deployments( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts + model=model_group, + healthy_deployments=[dict(deployment)], + messages=None, + request_kwargs=dict(request_kwargs), + ) + except ValueError as e: + raise litellm.RateLimitError( + message=str(e), + llm_provider="", + model=model_group, + category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from e + await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict + deployment=dict(deployment), parent_otel_span=None + ) + @contextlib.asynccontextmanager async def _deployment_slot( self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None 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..ba2af41bf8e 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1105,6 +1105,38 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +class _FailedTrackingAlerts: + def __init__(self) -> None: + self.messages: tuple[str, ...] = () + + async def failed_tracking_alert(self, error_message: str, failing_model: str) -> None: + self.messages = (*self.messages, error_message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "alerted"), [("_aresponses_websocket", True), ("_arealtime", False)]) +async def test_track_cost_callback_alerts_on_a_websocket_turn_with_no_logging_object(call_type: str, alerted: bool): + kwargs: Final = { + "call_type": call_type, + "model": "gpt-5-mini", + "litellm_call_id": "turn-call-id", + "litellm_params": {}, + "stream": False, + } + alerts: Final = _FailedTrackingAlerts() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", alerts): + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.sleep(0) + + assert any(f"call_type: {call_type}" in message for message in alerts.messages) is alerted + + @pytest.mark.asyncio async def test_track_cost_callback_defers_in_progress_background_interaction(): # test-quality-ok: writing no spend row and raising no alert is the whole observable contract of the deferral path """ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index dd7a0ec7a2e..f19a702dc13 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1596,7 +1596,7 @@ class TestManagedResponsesWSFirstMessage: processed: list = [] - async def fake_process(msg: str) -> None: + async def fake_process(msg: str, turn_admission: object = None) -> None: processed.append(msg) handler = ManagedResponsesWebSocketHandler( @@ -1641,6 +1641,26 @@ class TestResponsesWSStreamingFirstMessage: backend_ws.send.assert_awaited_once_with(first) +@pytest.mark.parametrize( + ("failure", "expected_status", "expected_type"), + [ + (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"), + (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"), + (RuntimeError("boom"), 500, "server_error"), + ], + ids=["crossed-budget", "bad-request", "unexpected"], +) +def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure( + failure: Exception, expected_status: int, expected_type: str +): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame + + frame: Final = json.loads(_responses_ws_failure_frame(failure)) + + assert frame["status"] == expected_status + assert frame["error"]["type"] == expected_type + + class TestWSSessionCostTracking: @pytest.mark.asyncio async def test_router_budget_limiter_skips_arealtime_call_type(self): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 2b78aa98b04..92c815aa4ab 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1344,6 +1344,107 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _turn_frames() -> tuple[str, str, str]: + return tuple(json.dumps({"type": "response.create", "input": word}) for word in ("one", "two", "three")) + + +def _admission_refusing_the_next_turn(): + from unittest.mock import AsyncMock + + from litellm.exceptions import RateLimitError + + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + + +class TestWebSocketTurnAdmission: + """A turn after the connection's first is admitted by the deployment again before it reaches the provider.""" + + @pytest.mark.asyncio + async def test_native_relay_drops_a_later_turn_the_deployment_refuses(self): + from unittest.mock import AsyncMock + + first, second, third = _turn_frames() + cancel: Final = json.dumps({"type": "response.cancel"}) + turn_admission: Final = _admission_refusing_the_next_turn() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, cancel, third, Exception("stop")]) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=first, + turn_admission=turn_admission, + ) + + await handler.client_to_backend() + + assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [ + json.loads(first), + json.loads(cancel), + json.loads(third), + ] + assert turn_admission.await_count == 2 + refusal: Final = json.loads(websocket.send_text.await_args.args[0]) + assert refusal["error"]["type"] == "rate_limit_exceeded" + assert "crossed budget" in refusal["error"]["message"] + + @pytest.mark.asyncio + async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch): + from unittest.mock import AsyncMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args, **kwargs): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + turn_admission: Final = _admission_refusing_the_next_turn() + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=[second, third, Exception("stop")]) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + turn_admission=turn_admission, + ) + + await handler.run() + + assert len(answered) == 2 + assert "two" not in json.dumps(answered) + assert turn_admission.await_count == 2 + refusals: Final = [ + json.loads(call.args[0]) + for call in websocket.send_text.await_args_list + if json.loads(call.args[0]).get("type") == "error" + ] + assert [refusal["error"]["type"] for refusal in refusals] == ["rate_limit_exceeded"] + + def _deployment_defaults(): from types import MappingProxyType diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 39fa0970aa1..0e046ec48ff 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -526,6 +526,63 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 +async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]: + routed: Final[dict[str, object]] = {} + + async def call(**kwargs: object) -> None: + routed.update(kwargs) + + call.__name__ = function_name + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call) + return routed + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget(): + router = Router( + model_list=[ + { + "model_name": "ws-budget", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "max_budget": 0.001, + "budget_duration": "1d", + }, + "model_info": {"id": "ws-budget-dep"}, + } + ] + ) + assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses") + admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + await admit_turn() + await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002) + + with pytest.raises(litellm.RateLimitError, match="crossed budget"): + await admit_turn() + + +@pytest.mark.asyncio +async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): + router = Router( + model_list=[ + { + "model_name": "ws-rpm", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1}, + "model_info": {"id": "ws-rpm-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await admit_turn() + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ From 131eb344329bbdc6c308441f14df2a6f7756049f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:23:46 -0700 Subject: [PATCH 14/25] test(router): list the WebSocket turn admission method as covered through its partial --- tests/code_coverage_tests/router_code_coverage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index ba2913bcfb7..17a10e98092 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -114,6 +114,7 @@ ignored_function_names = [ "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py + "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py ] From 4ceba6d5dd2da0e250bc8eac3c70fccbef53afd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:02:30 -0700 Subject: [PATCH 15/25] fix(responses): give a refused WebSocket turn's error frame a 429 status so Codex ends the turn --- litellm/responses/streaming_iterator.py | 27 +++++++++---------- .../test_responses_websocket_all_providers.py | 3 ++- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 58f55ac73fd..33adb663021 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1899,6 +1899,12 @@ async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None await turn_admission() +def _turn_refusal_frame(refusal: RateLimitError) -> str: + return json.dumps( + {"type": "error", "status": 429, "error": {"type": "rate_limit_exceeded", "message": str(refusal)}} + ) + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2547,17 +2553,7 @@ class ResponsesWebSocketStreaming: await _admit_frame_turn(turn_admission, message) except RateLimitError as e: try: - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "rate_limit_exceeded", - "message": str(e), - }, - } - ) - ) + await self.websocket.send_text(_turn_refusal_frame(e)) except Exception: # noqa: BLE001, S110 # client may already be gone pass return False @@ -2702,10 +2698,11 @@ class ManagedResponsesWebSocketHandler: return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: + await self._send_frame(json.dumps({"type": "error", "error": {"type": error_type, "message": message}})) + + async def _send_frame(self, frame: str) -> None: try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) + await self.websocket.send_text(frame) except Exception: pass @@ -3085,7 +3082,7 @@ class ManagedResponsesWebSocketHandler: ) await _admit_frame_turn(turn_admission, 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 call_kwargs: Final = self._build_base_call_kwargs(msg_obj) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 92c815aa4ab..1d80d87a75d 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1388,6 +1388,7 @@ class TestWebSocketTurnAdmission: ] 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"] @@ -1442,7 +1443,7 @@ class TestWebSocketTurnAdmission: for call in websocket.send_text.await_args_list if json.loads(call.args[0]).get("type") == "error" ] - assert [refusal["error"]["type"] for refusal in refusals] == ["rate_limit_exceeded"] + assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")] def _deployment_defaults(): From 65b3e9b74d6a2551e7c2066f24497f29f1cdf8dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:17:53 -0700 Subject: [PATCH 16/25] fix(router): leave ITPM/OTPM to socket open and record a refused WebSocket turn as the socket's failure --- litellm/llms/custom_httpx/llm_http_handler.py | 3 +- litellm/responses/streaming_iterator.py | 27 +++++++++++++--- litellm/router.py | 3 +- .../io_token_rate_limit_check.py | 19 +++++++++++ .../response_api_endpoints/test_endpoints.py | 10 +++++- .../test_responses_websocket_all_providers.py | 32 ++++++++++++++++++- .../test_router/test_router_helper_utils.py | 19 +++++++++++ 7 files changed, 103 insertions(+), 10 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ba5947abf44..70b4d2d0f36 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6842,8 +6842,7 @@ class BaseLLMHTTPHandler: turn_admission=turn_admission, **kwargs, ) - await handler.run() - return None + return await handler.run() import websockets from websockets.asyncio.client import ClientConnection diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index fd9a7889487..f69721b39ff 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2565,15 +2565,24 @@ class ResponsesWebSocketStreaming: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) + except RateLimitError as e: + await self._send_turn_refusal(e) + return False + try: await _admit_frame_turn(turn_admission, message) except RateLimitError as e: - try: - await self.websocket.send_text(_turn_refusal_frame(e)) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass + if self._failure is None: + self._failure = e + await self._send_turn_refusal(e) return False return True + 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 client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -2685,6 +2694,7 @@ class ManagedResponsesWebSocketHandler: 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. @@ -3095,8 +3105,14 @@ class ManagedResponsesWebSocketHandler: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message ) + except RateLimitError as e: + await self._send_frame(_turn_refusal_frame(e)) + return + try: await _admit_frame_turn(turn_admission, raw_message) except RateLimitError as e: + if self._failure is None: + self._failure = e await self._send_frame(_turn_refusal_frame(e)) return @@ -3136,7 +3152,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. @@ -3157,3 +3173,4 @@ class ManagedResponsesWebSocketHandler: 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 682fb71421d..39e5a154739 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -270,6 +270,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, @@ -9318,7 +9319,7 @@ class Router: category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT, ) from e await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict - deployment=dict(deployment), parent_otel_span=None + deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None ) @contextlib.asynccontextmanager 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/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 047d79f5b38..3fa930f9855 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 @@ -1647,18 +1650,23 @@ class TestResponsesWSStreamingFirstMessage: (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"], + 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: diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 1d80d87a75d..1f5421aa97c 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1392,11 +1392,39 @@ class TestWebSocketTurnAdmission: 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): 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 @@ -1433,11 +1461,13 @@ class TestWebSocketTurnAdmission: turn_admission=turn_admission, ) - await handler.run() + 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 diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 0e046ec48ff..030e051a380 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -583,6 +583,25 @@ async def test_responses_websocket_turn_after_the_first_counts_toward_deployment 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(): """ From 61ed4b78995b12133fe619612251f4760d035973 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 05:06:39 -0700 Subject: [PATCH 17/25] fix(proxy): log the routed model group on a WebSocket turn's failure row --- .../proxy/hooks/proxy_track_cost_callback.py | 19 ++++++++- .../hooks/test_proxy_track_cost_callback.py | 40 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f6e9928ae1d..c46b70c9b86 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -17,6 +17,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 +101,19 @@ 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 _routed_model_group_fallback(request_data: Mapping[str, object]) -> Mapping[str, str]: + caller_metadata: Final = request_data.get("metadata") + if is_str_keyed_mapping(caller_metadata) and caller_metadata.get("model_group"): + return {} + routed_metadata: Final = request_data.get("litellm_metadata") + model_group: Final = routed_metadata.get("model_group") if is_str_keyed_mapping(routed_metadata) else None + return {"model_group": model_group} if isinstance(model_group, str) and model_group else {} + + def _failure_snapshot_with_guardrails_recorded_after_it( failure_snapshot: object, recorded_guardrail_information: object ) -> Mapping[str, object] | None: @@ -220,10 +234,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_model_group_fallback(request_fields), **_metadata}) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") if ( 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 ba2af41bf8e..9174306eba0 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -136,6 +136,46 @@ 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, 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, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await ProxyDBLogger().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=mock_update_database.call_args[1]["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert row["model_group"] == expected_model_group + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_metadata(): logger = ProxyDBLogger() From 0e0ed1e2be101e750104feda2522e1c0bb7761b0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:39:14 -0700 Subject: [PATCH 18/25] fix(router): count WebSocket turn tokens toward TPM and re-run routing strategy checks on each turn --- litellm/cost_calculator.py | 9 +- litellm/litellm_core_utils/litellm_logging.py | 13 ++- litellm/router.py | 44 +++++++- .../router_code_coverage.py | 2 + .../hooks/test_proxy_track_cost_callback.py | 23 ++-- .../test_responses_websocket_all_providers.py | 15 +-- .../test_router/test_router_helper_utils.py | 103 ++++++++++++++++++ 7 files changed, 184 insertions(+), 25 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0e315e97432..0fb7363cbb5 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3068,14 +3068,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 e766a2be9c3..42d64eb88e4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -968,6 +968,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. @@ -2419,7 +2429,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/router.py b/litellm/router.py index 39e5a154739..a89361cc629 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9299,9 +9299,18 @@ 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, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] + 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: @@ -9318,10 +9327,43 @@ class Router: 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/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index c1e1ebb87b9..5594cdd1384 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -120,6 +120,8 @@ ignored_function_names = [ "_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/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 9174306eba0..a3ca604281c 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 @@ -146,7 +146,7 @@ async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_ ], ) async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_failure_row( - metadata_buckets: dict, expected_model_group: str + metadata_buckets: dict[str, dict[str, str]], expected_model_group: str ): request_data: Final = { "model": "ws-budget-sock", @@ -155,20 +155,15 @@ async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_f **metadata_buckets, } - with patch( - "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", - new_callable=AsyncMock, - ) as mock_update_database: - await ProxyDBLogger().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"), - ) + 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=mock_update_database.call_args[1]["kwargs"], + kwargs=spend_writer.update_database.call_args.kwargs["kwargs"], response_obj={}, start_time=datetime.now(), end_time=datetime.now(), diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 1f5421aa97c..228e0ad82a6 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1186,7 +1186,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 @@ -1232,11 +1232,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 @@ -1317,7 +1317,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_forwards_frame_allowed_by_quota_callback(self): @@ -1420,7 +1421,7 @@ class TestWebSocketTurnAdmission: assert "crossed budget" in str(failure) @pytest.mark.asyncio - async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch): + async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock import litellm @@ -1430,7 +1431,7 @@ class TestWebSocketTurnAdmission: answered: Final[list[object]] = [] - async def fake_aresponses(*args, **kwargs): + async def fake_aresponses(*args: object, **kwargs: object): answered.append(kwargs.get("input")) async def _empty(): diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 030e051a380..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, @@ -537,6 +539,27 @@ async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> d 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( @@ -583,6 +606,86 @@ async def test_responses_websocket_turn_after_the_first_counts_toward_deployment 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( From 084995dc8f3bfb7d718b5f2a4049136a1923ab61 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:05:42 -0700 Subject: [PATCH 19/25] test(integration): cover per-turn admission on Responses WebSockets Adds integration cells for budget, RPM, TPM, a mid-socket TPM edit, usage-based-routing-v2, the project quota frame, and the open-time failure frame on the native and managed relays, plus HTTP controls for TPM counting and the failure row model_group --- ...test_responses_websocket_turn_admission.py | 489 ++++++++++++++++++ .../spend/test_failed_dispatch_tokens.py | 52 ++ 2 files changed, 541 insertions(+) create mode 100644 tests/integration/routing/test_responses_websocket_turn_admission.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..a5fc6566fa8 --- /dev/null +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -0,0 +1,489 @@ +import asyncio +import hashlib +import os +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 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, + completed, + create, + drain, + responses_peer, + turn, +) +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 FROM "LiteLLM_SpendLogs" WHERE api_key = %s' +) +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]) -> 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, + "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] +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("responses-ws-turn-admission") + with gateway_from_environment() as base: + with owned_proxy( + base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory, router_settings) + ) 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 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_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue: + entries: Final = list_value(gateway.get("/model/info", {"litellm_model_id": deployment.identity})["data"]) + return object_value(object_value(entries[0])["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) + + +@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]) + + 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(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 + 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, + ) + failures: Final = [row for row in rows if row["status"] == "failure"] + assert len(failures) == 1, rows + assert failures[0]["model_group"] == deployment.name, 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 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..9ee1b0389c6 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.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,53 @@ 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}]} + + +@pytest.mark.parametrize( + ("route", "body"), + ( + ("/v1/chat/completions", _chat_body), + ("/v1/responses", _responses_body), + ("/v1/messages", _messages_body), + ), + ids=("chat", "responses", "messages"), +) +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: + 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 + ) + assert failed.status_code == 500, failed.text + assert wire.drain() + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, model_group 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" and rows[0]["model_group"] == model, rows From 355eae430a90e126efc4b5f70a801ad6d4482c06 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:24:23 -0700 Subject: [PATCH 20/25] fix(responses): count WebSocket turn tokens once and end a socket whose turn check fails A native Responses WebSocket turn is now logged as an async call, so a proxy with a sync success callback such as the response cache no longer runs each CustomLogger's sync success hook on top of the async one. That second run counted every turn's tokens twice toward deployment TPM. A turn whose admission check raises an unexpected error now gets a 500 error frame, the socket closes with 1011, and the error becomes the socket's failure, on both the native and the managed relay, instead of the socket dropping with no frame and no failure row. A refused turn's failure row now carries the request's tags from the router's metadata, the same way it already carries the model group. --- litellm/litellm_core_utils/litellm_logging.py | 1 + .../proxy/hooks/proxy_track_cost_callback.py | 23 ++-- litellm/responses/main.py | 4 +- litellm/responses/streaming_iterator.py | 58 ++++++---- .../types/responses/streaming_websocket.py | 2 + .../test_litellm_logging.py | 50 +++++++++ .../hooks/test_proxy_track_cost_callback.py | 35 ++++++ .../test_responses_websocket_all_providers.py | 101 ++++++++++++++++++ 8 files changed, 245 insertions(+), 29 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 42d64eb88e4..c1ebaaccf0e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2223,6 +2223,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 diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c46b70c9b86..0681b4bd892 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -2,6 +2,7 @@ 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 import litellm @@ -105,13 +106,19 @@ def _str_keyed_fields(value: object) -> Mapping[str, object]: return value if is_str_keyed_mapping(value) else {} -def _routed_model_group_fallback(request_data: Mapping[str, object]) -> Mapping[str, str]: - caller_metadata: Final = request_data.get("metadata") - if is_str_keyed_mapping(caller_metadata) and caller_metadata.get("model_group"): - return {} - routed_metadata: Final = request_data.get("litellm_metadata") - model_group: Final = routed_metadata.get("model_group") if is_str_keyed_mapping(routed_metadata) else None - return {"model_group": model_group} if isinstance(model_group, str) and model_group else {} +_ROUTED_METADATA_FIELD_TYPES: Final = MappingProxyType({"model_group": str, "tags": 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, field_type in _ROUTED_METADATA_FIELD_TYPES.items() + if not caller_metadata.get(field) + and isinstance(routed_metadata.get(field), field_type) + and routed_metadata[field] + } def _failure_snapshot_with_guardrails_recorded_after_it( @@ -238,7 +245,7 @@ class _ProxyDBLogger(CustomLogger): _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({**_routed_model_group_fallback(request_fields), **_metadata}) + existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") if ( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 269b7367f34..69b9d9834ca 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -64,7 +64,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 all_litellm_params +from litellm.types.utils import CallTypes, all_litellm_params from litellm.utils import ( ProviderConfigManager, client, @@ -2490,7 +2490,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, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index f69721b39ff..7662f205813 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 @@ -1907,11 +1908,18 @@ async def _enforce_frame_project_quota( ) -async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> None: +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 - await turn_admission() + 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: @@ -1920,6 +1928,22 @@ def _turn_refusal_frame(refusal: RateLimitError) -> str: ) +def _turn_admission_error_frame(failure: Exception) -> str: + return json.dumps({"type": "error", "status": 500, "error": {"type": "server_error", "message": str(failure)}}) + + +async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None: + match rejection: + case RateLimitError(): + with contextlib.suppress(Exception): + await websocket.send_text(_turn_refusal_frame(rejection)) + case _: + with contextlib.suppress(Exception): + await websocket.send_text(_turn_admission_error_frame(rejection)) + with contextlib.suppress(Exception): + await websocket.close(code=1011, reason="Internal server error") + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2568,14 +2592,13 @@ class ResponsesWebSocketStreaming: except RateLimitError as e: await self._send_turn_refusal(e) return False - try: - await _admit_frame_turn(turn_admission, message) - except RateLimitError as e: - if self._failure is None: - self._failure = e - await self._send_turn_refusal(e) - return False - return True + rejection: Final = await _admit_frame_turn(turn_admission, message) + if rejection is None: + return True + if self._failure is None: + self._failure = rejection + await _answer_rejected_turn(self.websocket, rejection) + return False async def _send_turn_refusal(self, refusal: RateLimitError) -> None: try: @@ -2726,10 +2749,8 @@ class ManagedResponsesWebSocketHandler: await self._send_frame(json.dumps({"type": "error", "error": {"type": error_type, "message": message}})) async def _send_frame(self, frame: str) -> None: - try: + with contextlib.suppress(Exception): await self.websocket.send_text(frame) - except Exception: - pass def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]: """ @@ -3108,12 +3129,11 @@ class ManagedResponsesWebSocketHandler: except RateLimitError as e: await self._send_frame(_turn_refusal_frame(e)) return - try: - await _admit_frame_turn(turn_admission, raw_message) - except RateLimitError as e: + rejection: Final = await _admit_frame_turn(turn_admission, raw_message) + if rejection is not None: if self._failure is None: - self._failure = e - await self._send_frame(_turn_refusal_frame(e)) + self._failure = rejection + await _answer_rejected_turn(self.websocket, rejection) return call_kwargs: Final = self._build_base_call_kwargs(msg_obj) diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index f369cbcebf8..b35f0fd7fbb 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -14,6 +14,8 @@ class ResponsesClientWebSocket(Protocol): async def receive_text(self) -> str: ... + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + class ResponsesBackendWebSocket(Protocol): """Upstream provider websocket surface used when proxying a native Responses API socket.""" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index df6593d29b1..6d51b1e4622 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2049,6 +2049,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 a3ca604281c..df4fd0d3624 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -171,6 +171,41 @@ async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_f 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"}}, []), + ], +) +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() diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 228e0ad82a6..d8ec8c0ca05 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1357,6 +1357,37 @@ def _admission_refusing_the_next_turn(): 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("budget store unreachable")) + + +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.""" @@ -1477,6 +1508,76 @@ class TestWebSocketTurnAdmission: assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")] + @pytest.mark.asyncio + async def test_native_relay_ends_the_socket_with_a_500_when_turn_admission_crashes(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=Exception("stop")) + 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)] + assert _errors(websocket) == [(500, {"type": "server_error", "message": "budget store unreachable"})] + assert websocket.close_code == 1011 + + @pytest.mark.asyncio + async def test_managed_relay_ends_the_socket_with_a_500_when_turn_admission_crashes( + 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"] + assert _errors(websocket) == [(500, {"type": "server_error", "message": "budget store unreachable"})] + assert websocket.close_code == 1011 + def _deployment_defaults(): from types import MappingProxyType From 83a78adda00a809adb7da75bc0e2f3d1432719bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:27:56 -0700 Subject: [PATCH 21/25] test(integration): cover a crashed WebSocket turn check and failure-row tags A turn whose admission check raises gets a 500 server_error frame, the socket closes with 1011 and its one failure row carries the error, on both relays. The refused turn's failure row keeps the caller's tags, and HTTP failure rows on chat, responses and messages keep them too. --- ...test_responses_websocket_turn_admission.py | 124 +++++++++++++++--- .../spend/test_failed_dispatch_tokens.py | 42 ++++-- 2 files changed, 142 insertions(+), 24 deletions(-) diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index a5fc6566fa8..d6a8651c0cd 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import hashlib import os import uuid @@ -8,6 +9,7 @@ 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 @@ -50,8 +52,23 @@ TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] TURN_TOKENS: Final = USAGE["total_tokens"] MANAGED_PROVIDER: Final = "mistral" SPEND_ROWS_FOR_KEY: Final = ( - 'SELECT request_id, spend, status, model_group, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s' + '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] @@ -70,13 +87,17 @@ 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]) -> Path: +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"], @@ -91,13 +112,17 @@ def _config(directory: Path, router_settings: Mapping[str, JsonValue]) -> Path: @contextmanager def _admitting_proxy( - tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path], router_settings: Mapping[str, JsonValue] + 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(directory, router_settings) - ) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=config) as proxy: yield proxy @@ -112,6 +137,12 @@ def limited(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) - 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: @@ -189,9 +220,13 @@ def _deployment_spend(deployment: Deployment) -> float: return float(redis.get(f"deployment_spend:{deployment.identity}:1d") or 0) -def _configured_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue: +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])["litellm_params"]).get("tpm") + 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: @@ -222,6 +257,21 @@ async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) return await asyncio.to_thread(eventually, read, satisfied, 15) +async def _close_code(connection: websockets.ClientConnection) -> int | None: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(connection.wait_closed(), timeout=20) + return connection.close_code + + +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 @@ -230,9 +280,10 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log 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) as connection: + 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)) @@ -243,19 +294,60 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log _assert_refused(second) _assert_refused(third) assert deployment.forwarded_turns() == 1 - 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, - ) - failures: Final = [row for row in rows if row["status"] == "failure"] - assert len(failures) == 1, rows + 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 +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_ends_the_socket_as_its_failure( + 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}" + + def fail_the_turn_check() -> None: + faulting.request( + "PATCH", + f"/model/{deployment.identity}/update", + {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: fault}}, + ).raise_for_status() + eventually( + lambda: _configured(faulting, deployment, "model_info").get(TURN_CHECK_FAULT), + lambda configured: configured == fault, + seconds=15, + ) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], int | None]: + async with _socket(faulting, key, deployment.name) as connection: + first: Final = await _say(connection, deployment.name) + await asyncio.to_thread(fail_the_turn_check) + second: Final = await _say(connection, deployment.name) + return first, second, await _close_code(connection) + + first, second, close_code = asyncio.run(session()) + + completed(first) + assert [event.get("type") for event in second] == ["error"], second + assert second[0]["status"] == 500, second + error: Final = object_value(second[0]["error"]) + assert error["type"] == "server_error", second + assert fault in string_value(error["message"]), second + assert close_code == 1011, close_code + assert deployment.forwarded_turns() == 1 + 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_socket_opened_after_websocket_turns_spent_the_budget_is_refused_and_so_is_http( limited: Gateway, cert: tuple[Path, Path] ) -> None: diff --git a/tests/integration/spend/test_failed_dispatch_tokens.py b/tests/integration/spend/test_failed_dispatch_tokens.py index 9ee1b0389c6..db5c49623eb 100644 --- a/tests/integration/spend/test_failed_dispatch_tokens.py +++ b/tests/integration/spend/test_failed_dispatch_tokens.py @@ -4,7 +4,7 @@ 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 @@ -64,7 +64,7 @@ def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]: return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} -@pytest.mark.parametrize( +FAILURE_ROUTES: Final = pytest.mark.parametrize( ("route", "body"), ( ("/v1/chat/completions", _chat_body), @@ -73,9 +73,14 @@ def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]: ), ids=("chat", "responses", "messages"), ) -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: + + +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( @@ -87,16 +92,37 @@ def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for( 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 + "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 FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + '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" and rows[0]["model_group"] == model, rows + 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 From 3125511b7e57c2726ba1259042410f535cbaae9c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:24:43 -0700 Subject: [PATCH 22/25] fix(responses): keep a WebSocket open after a crashed turn check and keep only string tags on failure rows A turn whose admission check raises an unexpected error still gets a 500 error frame, but the socket now stays open, like a refused turn. Closing it cut any turn still streaming on the same socket, which a client that sends its next response.create early would lose. A refused turn's failure row now takes the router's tags only when every tag is a string, so the fallback never copies a caller's unvalidated tags onto the row. --- .../proxy/hooks/proxy_track_cost_callback.py | 24 +++++++++++---- litellm/responses/streaming_iterator.py | 15 +++++----- .../types/responses/streaming_websocket.py | 2 -- ...test_responses_websocket_turn_admission.py | 30 ++++++++----------- .../hooks/test_proxy_track_cost_callback.py | 1 + .../test_responses_websocket_all_providers.py | 24 ++++++++------- 6 files changed, 54 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0681b4bd892..3195394b9aa 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -5,6 +5,8 @@ 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 @@ -106,7 +108,21 @@ def _str_keyed_fields(value: object) -> Mapping[str, object]: return value if is_str_keyed_mapping(value) else {} -_ROUTED_METADATA_FIELD_TYPES: Final = MappingProxyType({"model_group": str, "tags": list}) +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]: @@ -114,10 +130,8 @@ def _routed_metadata_fallback(request_data: Mapping[str, object]) -> Mapping[str routed_metadata: Final = _str_keyed_fields(request_data.get("litellm_metadata")) return { field: routed_metadata[field] - for field, field_type in _ROUTED_METADATA_FIELD_TYPES.items() - if not caller_metadata.get(field) - and isinstance(routed_metadata.get(field), field_type) - and 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)) } diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7662f205813..bcd3fef9415 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1932,16 +1932,17 @@ def _turn_admission_error_frame(failure: Exception) -> str: return json.dumps({"type": "error", "status": 500, "error": {"type": "server_error", "message": str(failure)}}) -async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None: +def _turn_rejection_frame(rejection: Exception) -> str: match rejection: case RateLimitError(): - with contextlib.suppress(Exception): - await websocket.send_text(_turn_refusal_frame(rejection)) + return _turn_refusal_frame(rejection) case _: - with contextlib.suppress(Exception): - await websocket.send_text(_turn_admission_error_frame(rejection)) - with contextlib.suppress(Exception): - await websocket.close(code=1011, reason="Internal server error") + return _turn_admission_error_frame(rejection) + + +async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None: + with contextlib.suppress(Exception): + await websocket.send_text(_turn_rejection_frame(rejection)) RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index b35f0fd7fbb..f369cbcebf8 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -14,8 +14,6 @@ class ResponsesClientWebSocket(Protocol): async def receive_text(self) -> str: ... - async def close(self, code: int = ..., reason: str | None = ...) -> None: ... - class ResponsesBackendWebSocket(Protocol): """Upstream provider websocket surface used when proxying a native Responses API socket.""" diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index d6a8651c0cd..9b5d4c4706e 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -1,5 +1,4 @@ import asyncio -import contextlib import hashlib import os import uuid @@ -257,12 +256,6 @@ async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) return await asyncio.to_thread(eventually, read, satisfied, 15) -async def _close_code(connection: websockets.ClientConnection) -> int | None: - with contextlib.suppress(TimeoutError): - await asyncio.wait_for(connection.wait_closed(), timeout=20) - return connection.close_code - - def _failure_rows(key: str) -> list[dict[str, JsonValue]]: rows: Final = eventually( lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), @@ -304,7 +297,7 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log @pytest.mark.parametrize("relay", ("native", "managed")) -def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_ends_the_socket_as_its_failure( +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: @@ -312,26 +305,27 @@ def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_ends_the_socket key: Final = scenario.key(models=[deployment.name]) fault: Final = f"turn check unavailable {uuid.uuid4().hex}" - def fail_the_turn_check() -> None: + def set_the_turn_check_fault(value: str) -> None: faulting.request( "PATCH", f"/model/{deployment.identity}/update", - {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: fault}}, + {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: value}}, ).raise_for_status() eventually( - lambda: _configured(faulting, deployment, "model_info").get(TURN_CHECK_FAULT), - lambda configured: configured == fault, + lambda: _configured(faulting, deployment, "model_info").get(TURN_CHECK_FAULT) or "", + lambda configured: configured == value, seconds=15, ) - async def session() -> tuple[tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], int | None]: + 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(fail_the_turn_check) + await asyncio.to_thread(set_the_turn_check_fault, fault) second: Final = await _say(connection, deployment.name) - return first, second, await _close_code(connection) + await asyncio.to_thread(set_the_turn_check_fault, "") + return first, second, await _say(connection, deployment.name) - first, second, close_code = asyncio.run(session()) + first, second, third = asyncio.run(session()) completed(first) assert [event.get("type") for event in second] == ["error"], second @@ -339,8 +333,8 @@ def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_ends_the_socket error: Final = object_value(second[0]["error"]) assert error["type"] == "server_error", second assert fault in string_value(error["message"]), second - assert close_code == 1011, close_code - assert deployment.forwarded_turns() == 1 + completed(third) + assert deployment.forwarded_turns() == 2 failures: Final = _failure_rows(key) assert len(failures) == 1, failures assert failures[0]["model_group"] == deployment.name, failures 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 df4fd0d3624..1f0dd67d86b 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -178,6 +178,7 @@ async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_f ({"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( diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index d8ec8c0ca05..e340d4b70ed 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 @@ -1381,7 +1382,7 @@ class _ClientSocket: def _admission_crashing_the_next_turn(): from unittest.mock import AsyncMock - return AsyncMock(side_effect=ValueError("budget store unreachable")) + return AsyncMock(side_effect=[ValueError("custom check failed"), None]) def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: @@ -1509,14 +1510,14 @@ class TestWebSocketTurnAdmission: @pytest.mark.asyncio - async def test_native_relay_ends_the_socket_with_a_500_when_turn_admission_crashes(self): + 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=Exception("stop")) + backend_ws.recv = AsyncMock(side_effect=asyncio.Event().wait) websocket: Final = _ClientSocket(second, third) handler: Final = _make_streaming( websocket=websocket, @@ -1529,12 +1530,15 @@ class TestWebSocketTurnAdmission: 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)] - assert _errors(websocket) == [(500, {"type": "server_error", "message": "budget store unreachable"})] - assert websocket.close_code == 1011 + 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_ends_the_socket_with_a_500_when_turn_admission_crashes( + async def test_managed_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving( self, monkeypatch: pytest.MonkeyPatch ): import litellm @@ -1574,9 +1578,9 @@ class TestWebSocketTurnAdmission: failure: Final = await handler.run() assert isinstance(failure, ValueError) - assert answered == ["one"] - assert _errors(websocket) == [(500, {"type": "server_error", "message": "budget store unreachable"})] - assert websocket.close_code == 1011 + 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 From fedc2a336fbc2b6551cacc6b134acc387c14f1f4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:58:16 -0700 Subject: [PATCH 23/25] test(integration): cover a turn still streaming when the next WebSocket turn's check crashes A native turn still streaming when the next turn's admission check raises now has its own cell. It must complete and be billed at the turn's price. The crashed-check cell also keeps the fault set for one more turn, so a relay that stopped checking turns after a crash goes red. --- ...test_responses_websocket_turn_admission.py | 98 ++++++++++++++----- 1 file changed, 76 insertions(+), 22 deletions(-) diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index 9b5d4c4706e..a246c89e7b3 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -1,6 +1,7 @@ import asyncio import hashlib import os +import threading import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager @@ -32,11 +33,14 @@ 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 @@ -296,6 +300,36 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log 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 @@ -305,35 +339,21 @@ def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_the_socket_serv key: Final = scenario.key(models=[deployment.name]) fault: Final = f"turn check unavailable {uuid.uuid4().hex}" - def set_the_turn_check_fault(value: str) -> None: - faulting.request( - "PATCH", - f"/model/{deployment.identity}/update", - {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: value}}, - ).raise_for_status() - eventually( - lambda: _configured(faulting, deployment, "model_info").get(TURN_CHECK_FAULT) or "", - lambda configured: configured == value, - seconds=15, - ) - 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_the_turn_check_fault, fault) + await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, fault) second: Final = await _say(connection, deployment.name) - await asyncio.to_thread(set_the_turn_check_fault, "") - return first, second, 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 = asyncio.run(session()) + first, second, third, fourth = asyncio.run(session()) completed(first) - assert [event.get("type") for event in second] == ["error"], second - assert second[0]["status"] == 500, second - error: Final = object_value(second[0]["error"]) - assert error["type"] == "server_error", second - assert fault in string_value(error["message"]), second - completed(third) + _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 @@ -342,6 +362,40 @@ def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_the_socket_serv 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: From f69f938391a59588000cb454a4cc25dafa2890f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:50:39 -0700 Subject: [PATCH 24/25] fix(responses): run per-turn guardrails before the deployment's turn check on WebSockets A turn after the first now runs the project quota check, then the proxy's pre_call guardrails, then the deployment's turn admission, on both the native and the managed relay, the order HTTP calls use. A turn a guardrail blocks gets the guardrail's error even past the budget and no longer counts toward the deployment's rpm or runs its pre-call checks. --- litellm/responses/streaming_iterator.py | 26 +++--- .../test_responses_websocket_all_providers.py | 88 +++++++++++++++++++ 2 files changed, 102 insertions(+), 12 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 4c83436218d..20ddbf906b1 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2306,9 +2306,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(message, turn_admission): return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) @@ -2690,10 +2692,8 @@ class ResponsesWebSocketStreaming: return json.dumps(evt_obj) if modified else response_str - async def _enforce_or_reject_frame( - self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None - ) -> bool: - """Run the per-frame project quota check, then the deployment's own admission for a turn after the first. + async def _enforce_or_reject_frame(self, message: str) -> bool: + """Run the per-frame project quota check. On rejection, sends an ``error`` event to the client and reports that the frame must be dropped instead of forwarded, so the connection @@ -2706,6 +2706,9 @@ class ResponsesWebSocketStreaming: except RateLimitError as e: await self._send_turn_refusal(e) return False + return True + + async def _admit_or_reject_turn(self, message: str, turn_admission: Callable[[], Awaitable[None]] | None) -> bool: rejection: Final = await _admit_frame_turn(turn_admission, message) if rejection is None: return True @@ -2736,9 +2739,9 @@ class ResponsesWebSocketStreaming: while True: message = await self.websocket.receive_text() - if not await self._enforce_or_reject_frame(message, self.turn_admission): + 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) @@ -3280,16 +3283,15 @@ class ManagedResponsesWebSocketHandler: except RateLimitError as e: 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 _answer_rejected_turn(self.websocket, rejection) return - - guarded: Final = await self._guard_frame(raw_message) - if guarded is None: - return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 163e9611584..030214d95e6 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1621,6 +1621,17 @@ def _admission_crashing_the_next_turn(): 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 _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] @@ -1745,6 +1756,83 @@ class TestWebSocketTurnAdmission: 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 async def test_native_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(self): from unittest.mock import AsyncMock From 4af9c787479014f99e13ba8a0462356be0342f53 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:16:32 -0700 Subject: [PATCH 25/25] fix(responses): bill a WebSocket turn's guardrail cost when the deployment refuses it A turn the per-turn guardrail passes and the deployment then refuses now writes its own failed spend row with that turn's guardrail information, on both relays, the way a guardrail-blocked frame and an HTTP refusal already do --- .../proxy/response_api_endpoints/endpoints.py | 3 +- litellm/responses/streaming_iterator.py | 31 +++--- .../response_api_endpoints/test_endpoints.py | 38 ++++++++ .../test_responses_websocket_all_providers.py | 95 +++++++++++++++++++ 4 files changed, 155 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 15888d2601a..540606b7190 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1858,7 +1858,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: @@ -1906,6 +1906,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 diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 20ddbf906b1..1e7368e15eb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1962,11 +1962,6 @@ def _turn_rejection_frame(rejection: Exception) -> str: return _turn_admission_error_frame(rejection) -async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None: - with contextlib.suppress(Exception): - await websocket.send_text(_turn_rejection_frame(rejection)) - - RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2039,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 @@ -2310,7 +2317,7 @@ class ResponsesWebSocketStreaming: self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None or not await self._admit_or_reject_turn(message, turn_admission): + 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) @@ -2708,13 +2715,15 @@ class ResponsesWebSocketStreaming: return False return True - async def _admit_or_reject_turn(self, message: str, turn_admission: Callable[[], Awaitable[None]] | None) -> bool: + 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 _answer_rejected_turn(self.websocket, rejection) + await _reject_guarded_turn(self.websocket, guarded, rejection) return False async def _send_turn_refusal(self, refusal: RateLimitError) -> None: @@ -3290,7 +3299,7 @@ class ManagedResponsesWebSocketHandler: if rejection is not None: if self._failure is None: self._failure = rejection - await _answer_rejected_turn(self.websocket, rejection) + await _reject_guarded_turn(self.websocket, guarded, rejection) return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 03be269fdef..34d7fabbe2e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2054,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 030214d95e6..405f1a5b731 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1632,6 +1632,17 @@ def _guard_blocking(word: str): 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"] @@ -1833,6 +1844,90 @@ class TestWebSocketTurnAdmission: ] + @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