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