From 160a0b72cff4dbbc41bd5de1130966a207f0c17d Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:34:10 -0700 Subject: [PATCH 1/3] fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899) * fix(mcp): bill both rounds of a streamed chat completions MCP auto-execute A streamed /v1/chat/completions request with an MCP tool set to require_approval "never" runs two model calls on one logging object, and the final-stream dedupe kept only the first round's spend row. The final usage chunk also showed only the second round. The MCP stream iterator now holds the first round's assembled response, prices it, and logs one success row that carries the usage and cost of both rounds, with the fixed margin charged once. The final usage chunk shows the same summed usage. A single-round stream still logs its one row when the answer finishes. When the follow-up round fails, the first round is logged from a copy of the logging object taken before the follow-up opened, and the failure row carries only the follow-up's partial usage with no second fixed margin. With proxy post-call guardrails armed, only (assembled, cache_hit) tuples are parked for the proxy, and the deferred log is armed only while the iterator reads, so the proxy's orphaned-callback cleanup never sees it after the Router's first-chunk read * fix(mcp): send the usage chunk of a streamed chat MCP request that ends in one round When a streamed chat completions MCP request ends in its first round (the model answers without a tool call, or no tool results come back), the iterator drained the inner stream to fire its spend log and dropped the usage chunk that stream produced, so a client that asked for include_usage got no usage chunk at all. The drained usage chunk is now held and sent to the client after the final content chunk and before [DONE], the same place a plain stream sends it, so the last usage chunk matches the billed row. A two-round request still drops the first round's usage chunk and sends the merged one from the follow-up, and a stream without include_usage still sends none * fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id When the follow-up call of a streamed chat MCP auto-execute fails, round one's success log now goes through the proxy's deferred guardrail release instead of straight to the handlers, so a post-call guardrail runs on it and the row carries it. Without a proxy release it logs directly as before The last usage chunk now takes its cost from the combined row's billed total, so a fixed cost margin is counted once in the chunk, as it is on the row. The follow-up's usage chunk is held until the follow-up ends to make that possible The combined row now keeps round one's id, the stream's first chunk id, which matches how a single-round stream is logged * fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop When the follow-up call fails, round 1's success row is now logged on a copy of the logging object that does not share the request's guardrail results, so a post-call guardrail that reports on the failed request counts once instead of on both rows A client hang-up during the follow-up now bills round 1 plus the follow-up's partial answer on one success row with round 1's id and the fixed margin counted once, the same way a plain stream's hang-up bills its partial answer The last usage chunk now carries round 1's usage and cost when only round 1 is billed, and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it read with instead of discarding it The deferred release's logged_on argument is typed and every release closure accepts it, and the parked dispatch constants and protocols both MCP billing paths use now live in one module. The unit tests wait on a logging signal instead of a fixed delay and check the billed row through the proxy stream hook instead of private callback state * test(mcp): drive the unassembled follow-up through an upstream audio delta that is not base64 The test replaced litellm.stream_chunk_builder with a stub that raises. Round 2 now streams an audio delta whose data is not valid base64, so the real builder raises, the stream parks only best-effort usage, and the client's last usage chunk must carry round 1's tokens and cost * fix(mcp): count the fixed margin once in the combined stream row's usage cost * fix(mcp): bill round 2's partial answer when a client hangs up on a routed stream On the proxy the Router wraps the MCP stream, and when a client hangs up during round 2 the router's own cleanup closes that stream before the proxy's disconnect cleanup runs. The close logged round 1 alone, so round 2's partial answer was never billed. The MCP stream now bills round 1 plus round 2's partial answer as soon as the hang-up's cancellation reaches its read, before anything closes it * fix(mcp): count the follow-up's own messages in round 2's partial prompt The follow-up stream shares the request's logger, so it took the client's original messages as its prompt. When round 2 ended early, from a hang-up or a broken upstream, its partial usage counted only those messages and left out the assistant tool call and the tool results the follow-up actually sent. The follow-up stream now counts the messages the follow-up call sent * fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads On a proxy the stream the disconnect cleanup holds is the router's fallback wrapper, so the cleanup never found the MCP stream's in-flight billing and released round 1 alone when the client hung up while the proxy sat at a yield. The round billing now registers its in-flight biller on the request's logging object, next to the parked release, and the cleanup looks it up there, so any wrapper chain bills round 1 plus round 2's partial once * fix(mcp): send the MCP stream usage chunk under the billed row id and count the partial off the loop A streamed chat completion that runs a LiteLLM-hosted MCP tool bills both model calls on one row under round 1's chatcmpl id, but the usage chunk it sent the client was round 2's held chunk and kept round 2's id, so a caller could not join the usage chunk to its row. The usage chunk now carries round 1's id in every path that sends it, with or without a deferred post-call guardrail, while its usage and cost and every content chunk's id stay as they were When the client hangs up during round 2, assembling round 2's partial answer counts the follow-up messages, tool result included, with the local tokenizer. That ran on the event loop and stalled every other request on the worker for about a second on a large tool result. The assembly now runs on a worker thread through asyncify from both the reader's cancellation and the proxy's disconnect cleanup, shielded from cancellation, and round 1 is checked again after the await so a release that happened meanwhile is never billed twice * test(mcp): wait on the loop with a hang guard instead of a 300 ms deadline * fix(mcp): give a proxy-rejected follow-up's failure row round 1's own cost breakdown * refactor(streaming): fold the partial response's empty guards into one check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 22 + .../litellm_core_utils/streaming_handler.py | 35 +- litellm/proxy/common_request_processing.py | 45 +- litellm/proxy/utils.py | 3 + .../responses/mcp/chat_completions_handler.py | 101 +- .../mcp/chat_stream_round_billing.py | 281 +++++ .../responses/mcp/parked_stream_dispatch.py | 72 ++ litellm/responses/mcp/stream_round_billing.py | 58 +- .../integration/_support/hook_count_probe.py | 18 + .../integration/mcp/test_mcp_llm_endpoints.py | 261 +++- .../proxy/test_common_request_processing.py | 28 + .../mcp/test_chat_completions_handler.py | 1050 ++++++++++++++++- .../mcp/test_parked_stream_dispatch.py | 17 + 13 files changed, 1894 insertions(+), 97 deletions(-) create mode 100644 litellm/responses/mcp/chat_stream_round_billing.py create mode 100644 litellm/responses/mcp/parked_stream_dispatch.py create mode 100644 tests/unit/responses/mcp/test_parked_stream_dispatch.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4dd33fa5565..8bfb0b60693 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -568,6 +568,11 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY} def _provider_response_id(source: object) -> str | None: @@ -2566,6 +2571,23 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def guardrail_scoped_copy(self) -> "Logging": + request_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **request_params, + **{ + key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value)) + for key, value in request_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + } + scoped: Final = copy.copy(self) + scoped.litellm_params = litellm_params + scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params} + return scoped + def dispatch_attempt_failure_handlers( self, exception: Exception, traceback_exception: str ) -> Coroutine[None, None, None]: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 682b2a2f538..098a404ea14 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2486,6 +2486,25 @@ class CustomStreamWrapper: def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier + def partial_response(self) -> "ModelResponse | TextCompletionResponse | None": + partial: Final = ( + litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + if self.chunks + else None + ) + usage: Final[object] = getattr(partial, "usage", None) + if partial is None or not isinstance(usage, Usage): + return None + if self.model: + partial.model = self.model + backfill_missing_cache_usage_fields(usage) + return partial + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2495,23 +2514,15 @@ class CustomStreamWrapper: moves it into terminal failure usage before opening the next attempt, so a recovered attempt logs its own usage. """ - if self.logging_obj is None or not self.chunks: + if self.logging_obj is None: return try: - partial_response: Final = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages if isinstance(self.messages, list) else None, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) + partial_response: Final = self.partial_response() if partial_response is None: return - usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) - if usage is None: + usage: Final[object] = getattr(partial_response, "usage", None) + if not isinstance(usage, Usage): return - if self.model: - partial_response.model = self.model - backfill_missing_cache_usage_fields(usage) partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 if self._earlier_rounds_spend is None: self.logging_obj.model_call_details["combined_usage_object"] = usage diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a79d8fe29db..48eb58adaf2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -140,6 +140,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_api import RoutingContext @@ -535,6 +536,16 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +async def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None: + bill_in_flight_round: Final = in_flight_round_biller(request_data.get("litellm_logging_obj")) + if bill_in_flight_round is None: + return + try: + await bill_in_flight_round() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e) + + async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( @@ -2083,6 +2094,13 @@ class _ClientStreamFailureReport: ) +def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]: + logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({}) + return MappingProxyType( + {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params} + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -4151,11 +4169,21 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + async def _on_deferred_stream_complete( + assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None + ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, + captured_data=_captured_data + if logged_on is None + else { + **_captured_data, + **{ + key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data + }, + "litellm_logging_obj": logged_on, + }, captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, + captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, cache_hit=cache_hit, ) @@ -4180,14 +4208,8 @@ class ProxyBaseLLMRequestProcessing: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - case (assembled_response, cache_hit): - await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + case _ if (success := parked_success(args)) is not None: + await success.dispatch(_captured_native_logging_obj) case _: verbose_proxy_logger.error( "Deferred stream logging dropped: unexpected stored args shape %s", @@ -4524,6 +4546,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + await _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data)) deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data) ProxyLogging.fire_deferred_stream_logging(request_data) # A disconnect-time success event (the deferred-guardrail flush diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 74e16628bc8..ddaa8d2e69d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction +from litellm.responses.mcp.parked_stream_dispatch import held_round_spend from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -4943,6 +4944,8 @@ class ProxyLogging: if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False match _args: + case (ModelResponse(), *_) if (held := held_round_spend(logging_obj)) is not None: + logging_obj.record_billed_rounds_for_failure(held) case (ModelResponse() as assembled, *_): logging_obj.record_assembled_response_for_failure(assembled) case (BilledRoundsDispatch() as billed,): diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..1e07b768f62 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,15 +1,18 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" +import asyncio import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import anyio from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -19,7 +22,7 @@ from litellm.responses.mcp.round_billing import ( first_round_billed_on_failure, summed_cost_breakdown, ) -from litellm.types.utils import Message, ModelResponse +from litellm.types.utils import Message, ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: @@ -258,10 +261,20 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self._round_billing: ChatStreamRoundBilling | None = engaged_round_billing(initial_stream.logging_obj) + self._unsent_usage_chunk: ModelResponseStream | None = None def __aiter__(self): return self + def __del__(self) -> None: + self.release_abandoned() + + def release_abandoned(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, ChatStreamRoundBilling): + round_billing.release_abandoned() + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -321,18 +334,54 @@ async def acompletion_with_mcp( return chunk - async def _drain_inner_stream(self): + async def _drain_inner_stream(self) -> tuple[ModelResponseStream, ...]: try: - while True: - await self._stream_iterator.__anext__() - except StopAsyncIteration: - pass + return tuple([chunk async for chunk in initial_stream]) except Exception: logging.getLogger("LiteLLM").exception( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) + return () - async def __anext__(self): + async def _complete_first_round(self) -> ModelResponseStream | None: + drained: Final = await self._drain_inner_stream() + if self._round_billing is not None: + self._round_billing.complete_first_round() + usage_chunks: Final = tuple( + chunk for chunk in drained if isinstance(getattr(chunk, "usage", None), Usage) + ) + return usage_chunks[-1] if usage_chunks else None + + async def _follow_up_or_release_first_round(self, usage_chunk: ModelResponseStream | None) -> None: + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + if self.follow_up_stream is not None: + return + self._unsent_usage_chunk = usage_chunk + if self._round_billing is not None: + self._round_billing.finish() + + async def __anext__(self) -> ModelResponseStream: + round_billing: Final = self._round_billing + if round_billing is None: + return await self._next_chunk() + with round_billing.armed_while_reading(): + try: + return await self._next_chunk() + except StopAsyncIteration: + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await round_billing.bill_in_flight_round() + raise + except Exception: + round_billing.fail() + raise + + async def _next_chunk(self) -> ModelResponseStream: # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: # Get the iterator from the stream wrapper @@ -362,29 +411,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True + usage_chunk: Final = await self._complete_first_round() await self._process_tool_calls() chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - # Drain inner stream so CustomStreamWrapper fires its - # end-of-stream handler (dispatch_success_handlers → - # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may - # yield one usage chunk before raising StopAsyncIteration. - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(usage_chunk) return chunk except StopAsyncIteration: self.stream_exhausted = True + drained_usage_chunk: Final = await self._complete_first_round() # Process tool calls after stream is exhausted await self._process_tool_calls() # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(drained_usage_chunk) return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -400,7 +442,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -408,6 +449,10 @@ async def acompletion_with_mcp( verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration + client_chunk: Final = ( + chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) + ) + return await self._next_chunk() if client_chunk is None else client_chunk # If we're here and follow_up_stream is None but we expected it, log a warning if ( @@ -420,7 +465,11 @@ async def acompletion_with_mcp( verbose_logger.warning("Follow-up stream was not created despite having tool results") - raise StopAsyncIteration + unsent_usage_chunk: Final = self._unsent_usage_chunk + if unsent_usage_chunk is None: + raise StopAsyncIteration + self._unsent_usage_chunk = None + return unsent_usage_chunk async def _process_tool_calls(self): """Process tool calls after streaming completes.""" @@ -488,10 +537,16 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm + earlier_rounds: Final = None if self._round_billing is None else self._round_billing.bank_first_round() follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): + if earlier_rounds is not None: + follow_up_response.bill_failure_after_earlier_rounds(earlier_rounds) + follow_up_response.messages = follow_up_messages + if self._round_billing is not None: + self._round_billing.follow_up_started(follow_up_response.partial_response) self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger @@ -548,11 +603,13 @@ async def acompletion_with_mcp( def __aiter__(self): return self._custom_iterator + async def aclose(self) -> None: + self._custom_iterator.release_abandoned() + await super().aclose() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: - import asyncio - try: self._sync_loop = asyncio.get_event_loop() except RuntimeError: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py new file mode 100644 index 00000000000..370de663606 --- /dev/null +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -0,0 +1,281 @@ +import asyncio +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.litellm_core_utils.asyncify import asyncify, is_event_loop_running +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + HELD_ROUND_SPEND, + IN_FLIGHT_ROUND, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + DeferredStreamRelease, + parked_success, +) +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import spawn_logging_task +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage + +Parked: TypeAlias = tuple[object, ...] +PartialAnswer: TypeAlias = Callable[[], object] + + +@dataclass(frozen=True, slots=True) +class _AssembledRound: + response: ModelResponse + usage: Usage + cache_hit: bool + parked: Parked + + +@dataclass(frozen=True, slots=True) +class _PricedRound: + response: ModelResponse + usage: Usage + price: UncachedPrice + cache_hit: bool + + @property + def spend(self) -> BilledRoundsSpend: + return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown) + + +@dataclass(frozen=True, slots=True) +class _HeldUsageChunk: + chunk: ModelResponseStream + usage: Usage + + +@dataclass(frozen=True, slots=True) +class _FirstRound: + priced: _PricedRound + parked: Parked + logged_alone: Parked + banked: LiteLLMLoggingObj | None = None + usage_chunk: _HeldUsageChunk | None = None + + def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: + if self.usage_chunk is None: + return None + usage: Final = ( + self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost})) + if billed is None + else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy( + update=MappingProxyType({"cost": _cost(billed)}) + ) + ) + return self.usage_chunk.chunk.model_copy( + update=MappingProxyType({"id": self.priced.response.id, "usage": usage}) + ) + + +def _cost(response: ModelResponse) -> object: + hidden_params: Final = get_hidden_params(response) + return None if hidden_params is None else hidden_params.get("response_cost") + + +def _assembled_round(parked: Parked | None) -> _AssembledRound | None: + match parked: + case (ModelResponse() as response, cache_hit) if isinstance(usage := getattr(response, "usage", None), Usage): + return _AssembledRound(response, usage, cache_hit is True, parked) + case _: + return None + + +def engaged_round_billing(logging_obj: object) -> "ChatStreamRoundBilling | None": + return ChatStreamRoundBilling(logging_obj) if isinstance(logging_obj, LiteLLMLoggingObj) else None + + +class ChatStreamRoundBilling: + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._first: _FirstRound | None = None + self._partial_follow_up: PartialAnswer | None = None + self._reading_loop: asyncio.AbstractEventLoop | None = None + + @contextmanager + def armed_while_reading(self) -> Generator[None]: + armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None + if armed: + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() + try: + yield + finally: + if armed and self._peek_parked() is None: + self._disarm() + + def complete_first_round(self) -> None: + assembled: Final = _assembled_round(self._peek_parked()) + if assembled is not None: + self._first = _FirstRound(self._priced(assembled), assembled.parked, assembled.parked) + setattr(self._logging_obj, HELD_ROUND_SPEND, self._held_round_one_spend) + + def bank_first_round(self) -> BilledRoundsSpend | None: + self._bank_apart_from_the_follow_up() + return None if self._first is None else self._first.priced.spend + + def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + usage: Final[object] = getattr(chunk, "usage", None) + if self._first is None or not isinstance(usage, Usage): + return chunk + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None + + def finish(self) -> ModelResponseStream | None: + first: Final = self._first + billed: Final = self._release_rounds() + return None if first is None else first.client_usage_chunk(billed) + + def fail(self) -> None: + self._bank_apart_from_the_follow_up() + first: Final = self._first + if first is None or first.banked is None or not self._holds_round_one_alone(first): + self.finish() + return + self._first = None + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) + + def follow_up_started(self, partial_follow_up: PartialAnswer) -> None: + self._partial_follow_up = partial_follow_up + setattr(self._logging_obj, IN_FLIGHT_ROUND, self.bill_in_flight_round) + + async def bill_in_flight_round(self) -> None: + partial_follow_up: Final = self._partial_follow_up + if partial_follow_up is None or self._first_round_held_alone() is None: + return + final_round: Final = _assembled_round((await asyncify(partial_follow_up)(), False)) + first: Final = self._first_round_held_alone() + if final_round is None or first is None: + return + self._first = None + setattr( + self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round)) + ) + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if is_event_loop_running() or self._idle_reading_loop() is not None: + self.finish() + return + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._first = None + self._disarm() + + def _bank_apart_from_the_follow_up(self) -> None: + first: Final = self._first + if first is None or first.banked is not None or self._peek_parked() is not first.parked: + return + banked: Final = self._logging_obj.attempt_scoped_copy() + if not self._owns_release(): + self._first = replace(first, banked=banked) + return + logged_alone: Final = (*first.parked, banked) + self._first = replace(first, logged_alone=logged_alone, banked=banked) + setattr(self._logging_obj, PARKED_DISPATCH, logged_alone) + + def _release_rounds(self) -> ModelResponse | None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return None + if first is None or parked is first.logged_alone: + self._release(parked) + return None + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return None + billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round)) + self._release((billed, cache_hit)) + return billed + + def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]: + first_breakdown: Final = first.price.cost_breakdown + final_breakdown: Final = final.price.cost_breakdown + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + usage: Final[object] = getattr(rounds, "usage", None) + charged_usage: Final = ( + usage.model_copy(update=MappingProxyType({"cost": _cost(rounds)})) if isinstance(usage, Usage) else usage + ) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id, "usage": charged_usage})) + cache_hit: Final = first.cache_hit and final.cache_hit + set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit}) + self._logging_obj.cost_breakdown = summed_cost_breakdown( + first.response, first_breakdown, final.response, final_breakdown + ) + return billed, cache_hit + + def _priced(self, assembled: _AssembledRound) -> _PricedRound: + served: Final = assembled.cache_hit or served_from_cache(assembled.response) + price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) + priced: Final = assembled.response.model_copy() + set_hidden_params( + priced, + {**(get_hidden_params(assembled.response) or {}), "response_cost": price.cost, "cache_hit": served}, + ) + return _PricedRound(priced, assembled.usage, price, served) + + def _peek_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) + return parked if isinstance(parked, tuple) else None + + def _take_parked(self) -> Parked | None: + parked: Final = self._peek_parked() + setattr(self._logging_obj, PARKED_DISPATCH, None) + return parked + + def _first_round_held_alone(self) -> _FirstRound | None: + first: Final = self._first + return first if first is not None and self._holds_round_one_alone(first) else None + + def _held_round_one_spend(self) -> BilledRoundsSpend | None: + first: Final = self._first_round_held_alone() + return None if first is None else first.priced.spend + + def _holds_round_one_alone(self, first: _FirstRound) -> bool: + parked: Final = self._peek_parked() + return parked is first.parked or parked is first.logged_alone + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + return release == self._log_success + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) + return + self._log_through_the_release(parked) + + def _log_through_the_release(self, parked: Parked) -> None: + armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) + logged: Final = (self._log_success if armed is None else armed)(*parked) + idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop() + if idle_reading_loop is not None: + idle_reading_loop.run_until_complete(logged) + return + spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") + + def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None: + loop: Final = self._reading_loop + return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, DISPATCH_RELEASE, None) + + async def _log_success(self, *parked: object) -> None: + success: Final = parked_success(parked) + if success is not None: + await success.dispatch(self._logging_obj) diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py new file mode 100644 index 00000000000..792363f0bab --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,72 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect" +HELD_ROUND_SPEND: Final = "_held_round_spend_for_failure" +SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +class DeferredStreamRelease(Protocol): + def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... + + +class InFlightRoundBiller(Protocol): + def __call__(self) -> Coroutine[object, object, None]: ... + + +class HeldRoundSpend(Protocol): + def __call__(self) -> BilledRoundsSpend | None: ... + + +def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +class ParkedSuccess: + assembled: object + cache_hit: object + logged_on: LiteLLMLoggingObj | None = None + + async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None: + logger: Final = request_logger if self.logged_on is None else self.logged_on + await success_dispatcher(logger).dispatch_success_handlers( + self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + + +def in_flight_round_biller(logging_obj: object) -> InFlightRoundBiller | None: + biller: Final[InFlightRoundBiller | None] = getattr(logging_obj, IN_FLIGHT_ROUND, None) + return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None + + +def held_round_spend(logging_obj: object) -> BilledRoundsSpend | None: + held: Final[HeldRoundSpend | None] = getattr(logging_obj, HELD_ROUND_SPEND, None) + return held() if isinstance(logging_obj, LiteLLMLoggingObj) and callable(held) else None + + +def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: + match parked: + case (assembled, cache_hit): + return ParkedSuccess(assembled, cache_hit) + case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on): + return ParkedSuccess(assembled, cache_hit, logged_on) + case _: + return None diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 443147230bf..5cb256a8ac0 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -2,11 +2,18 @@ import asyncio from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + parked_success, + success_dispatcher, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task from litellm.responses.utils import ResponseAPILoggingUtils @@ -16,25 +23,6 @@ from litellm.types.utils import CostBreakdown, ModelResponse, Usage Parked: TypeAlias = tuple[object, ...] RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] -_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" -_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" -_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) - - -class _DispatchesSuccessHandlers(Protocol): - async def dispatch_success_handlers( - self, - result: object = None, - start_time: object = None, - end_time: object = None, - cache_hit: object = None, - prefer_async_handlers: bool = False, - ) -> None: ... - - -def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: - return logger - def _event_loop_running() -> bool: try: @@ -198,8 +186,8 @@ class StreamedRoundBilling: self._billed: _BilledRounds | None = None self._banked: Parked | None = None self._banked_row: _RoundRow | None = None - if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: - setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + if getattr(logging_obj, DISPATCH_RELEASE, None) is None: + setattr(logging_obj, DISPATCH_RELEASE, self._run_parked) def complete_round( self, @@ -286,7 +274,7 @@ class StreamedRoundBilling: def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: if served_from_cache(response): - return _SERVED_FROM_CACHE + return SERVED_FROM_CACHE match round_parked: case (ModelResponse() as assembled, bool()): return self._logging_obj.uncached_price(assembled) @@ -297,7 +285,7 @@ class StreamedRoundBilling: self._banked = parked self._banked_row = row if parked is not None: - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + setattr(self._logging_obj, PARKED_DISPATCH, parked) def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: row: Final = self._banked_row @@ -320,7 +308,7 @@ class StreamedRoundBilling: @staticmethod async def _log_success(success: _SuccessLog) -> None: success.carry_bridged_usage() - await _success_dispatcher(success.logger).dispatch_success_handlers( + await success_dispatcher(success.logger).dispatch_success_handlers( success.row.logged, start_time=success.logger.start_time, end_time=success.row.end_time, @@ -329,10 +317,10 @@ class StreamedRoundBilling: ) def _take_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) if not isinstance(parked, tuple): return None - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked async def _run_parked(self, *parked: object) -> None: @@ -341,29 +329,27 @@ class StreamedRoundBilling: await dispatch case (dispatch,) if asyncio.iscoroutine(dispatch): await dispatch - case (assembled, cache_hit): - await _success_dispatcher(self._logging_obj).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + case _ if (success := parked_success(parked)) is not None: + await success.dispatch(self._logging_obj) case _: return def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._run_parked def _release(self, parked: Parked) -> None: - if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) spawn_logging_task( asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" ) def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) @staticmethod def _close(parked: Parked | None) -> None: diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py index 0d4056e8fb9..3cd72f1a6ba 100644 --- a/tests/integration/_support/hook_count_probe.py +++ b/tests/integration/_support/hook_count_probe.py @@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") +class UsageRecordingPostCallGuardrail(CustomGuardrail): + """Records the token usage of the response it is handed in the request's guardrail information, so the spend row + shows which response the guardrail saw.""" + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + usage: Final = getattr(response, "usage", None) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "prompt_tokens": getattr(usage, "prompt_tokens", None), + "completion_tokens": getattr(usage, "completion_tokens", None), + }, + request_data=data, + guardrail_status="success", + ) + + probe: Final = HookCountProbe() sync_only_probe: Final = SyncOnlyHookCountProbe() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 4b967802893..f6743a8751a 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -46,9 +46,10 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] -@pytest.fixture(scope="session") -def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") +@contextmanager +def _private_gateway( + directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () +) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) @@ -68,7 +69,8 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat { "guardrail_name": STREAM_BLOCK, "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, - } + }, + *guardrails, ], } ) @@ -81,15 +83,61 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield private, services -@pytest.fixture -def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: - private, services = _llm_gateway +@contextmanager +def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + private, services = served with pytest.MonkeyPatch.context() as environment: for name, value in services.items(): environment.setenv(name, value) yield private +@pytest.fixture(scope="session") +def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served: + yield served + + +@pytest.fixture +def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_llm_gateway) as private: + yield private + + +ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" +CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}" + + +@pytest.fixture(scope="module") +def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + round_usage: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + content_filter: Final[Mapping[str, JsonValue]] = { + "guardrail_name": CONTENT_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}], + }, + } + guardrails: Final = (round_usage, content_filter) + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served: + yield served + + +@pytest.fixture +def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_guarded_llm_gateway) as private: + yield private + + Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -1032,6 +1080,184 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row( ) +def _streamed_chunks(response: httpx.Response) -> tuple[dict[str, object], ...]: + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert response.status_code == 200, response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + usages: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage")) + assert usages, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == chunks[0]["id"], (row, chunks[0]) + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + final_usage: Final = usages[-1] + assert (final_usage["prompt_tokens"], final_usage["completion_tokens"], final_usage["total_tokens"]) == ( + row["prompt_tokens"], + row["completion_tokens"], + sum((*FIRST_ROUND, *FOLLOW_UP_ROUND)), + ), usages + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert float(str(succeeded[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), succeeded + failed: Final = _spend_rows(call_id, "acompletion", "failure") + assert [float(str(row["spend"])) for row in failed] == [0.0], failed + + +def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]: + entries: Final = object_value(metadata).get("guardrail_information") + return tuple(entries) if isinstance(entries, list) else () + + +def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]: + rows: Final = read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "acompletion", "success"), + ) + entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows)) + return [ + entry.get("guardrail_response") + for entry in entries + if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL + ] + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert _round_usage_seen_by_the_guardrail(call_id) == [ + {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]} + ] + + +def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, "acompletion"), + ), + lambda found: _statuses(found) == sorted(statuses), + seconds=70, + ) + + +def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + entry + for entry in _guardrail_information(row["metadata"]) + if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name + ) + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + stream=True, + stream_options={"include_usage": True}, + guardrails=[CONTENT_FILTER], + ) + assert FOLLOW_UP_OUTAGE in response.text, response.text + rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure")) + filtered: Final = tuple( + (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row) + ) + assert filtered == (("failure", "guardrail_failed_to_respond"),), rows + (guardrail_id,) = frozenset( + str(entry.get("guardrail_id") or CONTENT_FILTER) + for row in rows + for entry in _entries_of(CONTENT_FILTER, row) + ) + counted: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (guardrail_id,), + ), + lambda found: bool(found), + seconds=70, + ) + assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted + + +@pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) +def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( + gateway: Gateway, include_usage: bool +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + prompt=f"{DIRECT} {rig.alias}", + stream=True, + stream_options={"include_usage": include_usage}, + ) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1 and (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == FIRST_ROUND, rows + usage_at: Final = tuple(index for index, chunk in enumerate(chunks) if chunk.get("usage")) + if not include_usage: + assert usage_at == (), response.text + return + assert usage_at == (len(chunks) - 1,), response.text + finished_at: Final = tuple( + index + for index, chunk in enumerate(chunks) + if chunk.get("choices") and chunk["choices"][0].get("finish_reason") is not None + ) + assert finished_at and finished_at[-1] < usage_at[0], response.text + usage: Final = chunks[-1]["usage"] + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == ( + *FIRST_ROUND, + sum(FIRST_ROUND), + ), usage + + NO_CACHE: Final = {"no-cache": True} @@ -1269,6 +1495,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows +def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin( + gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert "the sum" in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(partial_answer) > 0, rows + assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows + + @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( gateway: Gateway, upstream: StreamedUpstream diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2d08609c93b..2db8941d304 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -83,6 +83,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( team_member_budget_entity_id, ) from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import ModelResponseListIterator @@ -9051,6 +9052,33 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + @pytest.mark.asyncio + async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self): + import types + + response = await self._start_partial_stream() + closed: Final = AsyncMock() + + async def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + setattr(response.logging_obj, IN_FLIGHT_ROUND, cannot_price_the_in_flight_round) + stream: Final = types.SimpleNamespace(chunks=[], messages=None, aclose=closed) + proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock()) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=stream, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + closed.assert_awaited_once() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..c0d32742e8a 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1,17 +1,25 @@ import asyncio +import gc import json +import queue import sys +import threading import types -from collections.abc import Mapping, Sequence +import warnings +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime from types import MappingProxyType from typing import Final, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from httpx import Response from mcp.types import Tool as MCPTool +from starlette.responses import StreamingResponse +from starlette.types import Message import litellm from litellm.integrations.custom_logger import CustomLogger @@ -25,7 +33,8 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.types.utils import ModelResponse, ModelResponseStream, StandardLoggingPayload, Usage +from litellm.utils import CustomStreamWrapper @pytest.mark.asyncio @@ -1556,6 +1565,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1575,16 +1585,31 @@ class _ChatRoundRecorder(CustomLogger): super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () + self._recorded: Final = asyncio.Condition() async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() + + async def _announce(self) -> None: + async with self._recorded: + self._recorded.notify_all() + + async def settled(self, logged: int = 0, failed: int = 0) -> None: + async with self._recorded: + await asyncio.wait_for( + self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed), + timeout=_NEVER_LOGGED_SECONDS, + ) + await _settled() def _chat_round_cost(body: Mapping[str, object]) -> float: @@ -1600,18 +1625,20 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _auto_execute_chat_echo( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str -) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: +def _serve_chat_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] +) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() respx_mock.post("https://mcp-chat-rounds.example.com/v1/chat/completions").mock(side_effect=list(rounds)) - recorder: Final = _ChatRoundRecorder() - logging_obj: Final = Logging( + + +def _chat_rounds_logging_obj(recorder: _ChatRoundRecorder, call_id: str, stream: bool) -> Logging: + return Logging( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], - stream=False, + stream=stream, call_type="acompletion", start_time=datetime(2026, 1, 1), litellm_call_id=call_id, @@ -1619,6 +1646,10 @@ async def _auto_execute_chat_echo( dynamic_async_success_callbacks=[recorder], dynamic_async_failure_callbacks=[recorder], ) + + +@contextmanager +def _echo_tool_served() -> Iterator[None]: echo_tool: Final = MCPTool.model_validate( {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, by_name=False, @@ -1635,6 +1666,16 @@ async def _auto_execute_chat_echo( AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]), ), ): + yield + + +async def _auto_execute_chat_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str +) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: + _serve_chat_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, call_id, stream=False) + with _echo_tool_served(): try: response: Final = await litellm.acompletion( model="openai/gpt-5.6", @@ -1712,3 +1753,994 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) assert recorder.failed[0]["response_cost"] > 0 + + +_FIXED_MARGIN: Final = 0.001 + + +def _chat_stream_events(body: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + completed: Final = json.loads(json.dumps(body, default=dict)) + choice: Final = completed["choices"][0] + message: Final = choice["message"] + calls: Final = message.get("tool_calls") or [] + delta: Final = ( + {**message, "tool_calls": [{**call, "index": index} for index, call in enumerate(calls)]} if calls else message + ) + head: Final = {"id": completed["id"], "object": "chat.completion.chunk", "created": 1, "model": completed["model"]} + return ( + {**head, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}, + {**head, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}, + {**head, "choices": [], "usage": completed["usage"]}, + ) + + +def _sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + +def _streamed(body: Mapping[str, object]) -> Response: + return Response( + 200, + content=_sse(_chat_stream_events(body)) + b"data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +class _CutOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + raise httpx.ReadError("connection reset mid-stream") + + +def _cut_once_the_answer_starts() -> Response: + return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}) + + +async def _open_chat_echo_stream( + logging_obj: Logging, metadata: dict[str, object] | None = None +) -> CustomStreamWrapper: + stream: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +async def _read_all(stream: CustomStreamWrapper) -> tuple[ModelResponseStream, ...]: + return tuple([chunk async for chunk in stream]) + + +async def _read_until(stream: CustomStreamWrapper, content: str) -> None: + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == content: + return + pytest.fail(f"the stream never sent {content!r}") + + +def _streamed_usage(chunks: Sequence[ModelResponseStream]) -> Usage: + usages: Final = tuple(usage for chunk in chunks if isinstance(usage := getattr(chunk, "usage", None), Usage)) + assert len(usages) == 1, usages + return usages[0] + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: Sequence[StandardLoggingPayload], chunks: Sequence[ModelResponseStream] +) -> None: + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert row["id"] == chunks[0].id + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert row["metadata"]["usage_object"]["cost"] == pytest.approx(row["response_cost"]) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + assert usage.cost == pytest.approx(row["response_cost"]) + + +def _prompt_tokens_the_follow_up_sent(respx_mock: respx.MockRouter) -> int: + follow_up_request: Final = json.loads(respx_mock.calls[1].request.content) + return litellm.token_counter(model="gpt-5.6", messages=follow_up_request["messages"]) + + +def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71, 19) + assert logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_read_to_the_end_logs_one_row_whose_usage_cost_counts_the_fixed_margin_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-read-to-the-end", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-deferred", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + await _settled() + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-blocked", stream=True) + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_all(stream) + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await recorder.settled(failed=1) + + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_rejected_by_the_proxy_during_the_follow_up_bills_round_one_with_its_own_breakdown( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from fastapi import HTTPException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rejected-mid-follow-up", stream=True) + rejected: Final = HTTPException(status_code=400, detail="rejected by a stream hook") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, rejected) + await logging_obj.async_failure_handler(rejected, "rejected by a stream hook") + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + breakdown: Final = failed["cost_breakdown"] + assert (failed["prompt_tokens"], failed["completion_tokens"]) == (71, 19) + assert failed["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(failed["response_cost"]) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round_once_the_answer_finishes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_TWO),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-one-round", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + async for chunk in stream: + if chunk.choices and chunk.choices[0].finish_reason is not None: + break + await recorder.settled(logged=1) + logged_once_the_answer_finished: Final = recorder.logged + after_the_answer: Final = await _read_all(stream) + await _settled() + + assert len(logged_once_the_answer_finished) == 1 + assert recorder.logged == logged_once_the_answer_finished + assert recorder.logged[0]["id"] == _CHAT_ROUND_TWO["id"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (149, 5) + assert after_the_answer[-1:] and isinstance(after_the_answer[-1].usage, Usage) + usage: Final = _streamed_usage(after_the_answer) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + recorder.logged[0]["prompt_tokens"], + recorder.logged[0]["completion_tokens"], + recorder.logged[0]["total_tokens"], + ) + assert recorder.logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_TWO)) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closes", (True, False), ids=("aclose", "dropped")) +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, closes: bool +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-caller-stops", stream=True) + + async def stop_reading_once_the_follow_up_answers() -> None: + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + if closes: + await stream.aclose() + + with _echo_tool_served(): + await stop_reading_once_the_follow_up_answers() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + with pytest.raises(litellm.ServiceUnavailableError): + await _read_all(stream) + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert recorder.failed[0]["completion_tokens"] > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging, _failure_fields_to_lift + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks", stream=True) + request_data: Final = MappingProxyType({"litellm_logging_obj": logging_obj}) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value): + ProxyLogging.fire_deferred_stream_logging(dict(request_data)) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + unmargined: Final = ( + litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["total_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + - _FIXED_MARGIN + ) + assert failed["prompt_tokens"] == _prompt_tokens_the_follow_up_sent(respx_mock) + assert failed["completion_tokens"] > 0 + assert failed["response_cost"] == pytest.approx(unmargined) + request_failure: Final = _failure_fields_to_lift(request_data) + assert cast(StandardLoggingPayload, request_failure["standard_logging_object"])["id"] != recorder.logged[0]["id"] + assert request_failure["response_cost"] == pytest.approx(unmargined) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True) + metadata: Final[dict[str, object]] = {} + request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj} + content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj, metadata) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + content_filter.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(broken.value), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value): + ProxyLogging.fire_deferred_stream_logging(request_data) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.logged[0]["guardrail_information"] is None + assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [ + "request-content-filter" + ] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True) + + with _echo_tool_served(): + served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=await _open_chat_echo_stream(logging_obj), + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj}, + ) + chunks: Final = tuple([chunk async for chunk in served]) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_the_first_round_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + nothing_streamed: Final = Response(200, content=b"data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), nothing_streamed)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-empty-follow-up", stream=True) + + with _echo_tool_served(): + await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +_ANSWER_WITH_UNREADABLE_AUDIO: Final = MappingProxyType( + { + **_CHAT_ROUND_TWO, + "choices": ( + MappingProxyType( + { + "index": 0, + "finish_reason": "stop", + "message": MappingProxyType( + { + "role": "assistant", + "content": "PING", + "audio": MappingProxyType({"id": "audio_echo", "data": "not base64", "transcript": "PING"}), + } + ), + } + ), + ), + } +) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_ANSWER_WITH_UNREADABLE_AUDIO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert any(chunk.choices and chunk.choices[0].delta.content == "PING" for chunk in chunks) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90) + assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +@pytest.mark.parametrize("answer", [_CHAT_ROUND_TWO, _ANSWER_WITH_UNREADABLE_AUDIO], ids=("assembled", "unassembled")) +async def test_stream_auto_execute_sends_its_usage_chunk_under_the_id_of_the_row_it_billed( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + deferred_for_a_post_call_guardrail: bool, + answer: Mapping[str, object], +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(answer))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-usage-chunk-id", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + usage_chunks: Final = tuple(chunk for chunk in chunks if isinstance(getattr(chunk, "usage", None), Usage)) + answer_chunks: Final = tuple(chunk for chunk in chunks if chunk.choices and chunk.choices[0].delta.content == "PING") + assert len(recorder.logged) == 1 + assert [chunk.id for chunk in answer_chunks] == [_CHAT_ROUND_TWO["id"]] + assert len(usage_chunks) == 1 + assert usage_chunks[0].id == chunks[0].id == recorder.logged[0]["id"] == _CHAT_ROUND_ONE["id"] + + +def _assert_one_row_bills_round_one_and_the_partial_answer( + recorder: _ChatRoundRecorder, respx_mock: respx.MockRouter +) -> None: + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == _CHAT_ROUND_ONE["id"] + partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) + assert partial_prompt == _prompt_tokens_the_follow_up_sent(respx_mock) + assert partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +def _chat_rounds_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mcp-chat-rounds", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-chat-rounds", + "api_base": "https://mcp-chat-rounds.example.com/v1", + }, + } + ] + ) + + +class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + await asyncio.Event().wait() + + +def _streams_content(body: object, content: str) -> bool: + lines: Final = body.decode().splitlines() if isinstance(body, bytes) else () + events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")) + return any(event["choices"] and event["choices"][0]["delta"].get("content") == content for event in events) + + +async def _serve_until_the_client_hangs_up_after(response: StreamingResponse, content: str) -> None: + read: Final = asyncio.Event() + + async def receive() -> Message: + await read.wait() + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + if message["type"] == "http.response.body" and _streams_content(message.get("body"), content): + read.set() + + uvicorn_scope: Final = {"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, "headers": []} + await response(uvicorn_scope, receive, send) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_the_client_hangs_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds( + respx_mock, + monkeypatch, + ( + _streamed(_CHAT_ROUND_ONE), + Response(200, stream=_StallsOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}), + ), + ) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = await create_response( + proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ), + "text/event-stream", + {}, + ) + assert isinstance(served, StreamingResponse) + await _serve_until_the_client_hangs_up_after(served, "PING") + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +async def _stop_reading_once_the_proxy_sent(served: AsyncIterator[object], content: str) -> None: + async for frame in served: + if _streams_content(frame, content): + return + pytest.fail(f"the proxy never sent {content!r}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_hung_up_between_reads( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up-between-reads", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS: Final = 10 + + +class _CountsTheToolResultWhileTheLoopRuns: + def __init__(self, loop: asyncio.AbstractEventLoop, count: Callable[..., int]) -> None: + self._loop: Final = loop + self._count: Final = count + self._counting: Final = asyncio.Event() + self._loop_moved: Final = threading.Event() + self.loop_moved_while_counting: tuple[bool, ...] = () + + def __call__(self, model: str | None = None, messages: list[object] | None = None, **kwargs: object) -> int: + if any(isinstance(message, Mapping) and message.get("role") == "tool" for message in messages or ()): + self._loop.call_soon_threadsafe(self._counting.set) + moved: Final = self._loop_moved.wait(timeout=_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS) + self.loop_moved_while_counting = (*self.loop_moved_while_counting, moved) + return self._count(model=model, messages=messages, **kwargs) + + async def keep_the_loop_moving(self) -> None: + await self._counting.wait() + self._loop_moved.set() + + +@contextmanager +def _partial_count_watched(monkeypatch: pytest.MonkeyPatch) -> Iterator[_CountsTheToolResultWhileTheLoopRuns]: + from litellm.litellm_core_utils import streaming_chunk_builder_utils + + counter: Final = _CountsTheToolResultWhileTheLoopRuns( + asyncio.get_running_loop(), streaming_chunk_builder_utils.token_counter + ) + monkeypatch.setattr(streaming_chunk_builder_utils, "token_counter", counter) + other_task: Final = asyncio.create_task(counter.keep_the_loop_moving()) + try: + yield counter + finally: + other_task.cancel() + + +class _StallsAfterTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.stalled: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + self.stalled.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_stream_auto_execute_cancelled_mid_answer_counts_the_partial_off_the_loop_and_bills_it_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + answer: Final = _StallsAfterTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-cancelled-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await _open_chat_echo_stream(logging_obj) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.stalled.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await stream.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_routed_through_the_proxy_and_closed_between_reads_counts_the_partial_off_the_loop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-closed-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +class _StallsBeforeTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.opened: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b": the answer has not started\n\n" + self.opened.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_before_the_answer_starts_bills_round_one_alone_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + answer: Final = _StallsBeforeTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-before-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.opened.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.failed == () + + +async def _read_through_the_usage_chunk( + chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = () +) -> tuple[ModelResponseStream, ...]: + chunk: Final = await anext(chunks) + if isinstance(getattr(chunk, "usage", None), Usage): + return (*read, chunk) + return await _read_through_the_usage_chunk(chunks, (*read, chunk)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_through_the_usage_chunk(aiter(stream)) + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + assert recorder.failed == () + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True) + + async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]: + async for chunk in stream: + yield chunk + + with _echo_tool_served(): + relayed: Final = bridged(await _open_chat_echo_stream(logging_obj)) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError): + _ = [chunk async for chunk in relayed] + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +def _read_synchronously_until_the_follow_up_answers( + logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None +) -> None: + stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj)) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == "PING": + before_dropping() + return + pytest.fail("the stream never sent 'PING'") + + +_CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success") + + +def _names_chat_stream_billing(text: str) -> bool: + return any(name in text for name in _CHAT_STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop + ) + gc.collect() + loop.run_until_complete(recorder.settled(logged=1)) + finally: + asyncio.set_event_loop(None) + loop.close() + + _assert_one_row_bills_round_one(recorder.logged) + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True), + loop, + before_dropping=loop.close, + ) + gc.collect() + finally: + asyncio.set_event_loop(None) + loop.close() + + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] diff --git a/tests/unit/responses/mcp/test_parked_stream_dispatch.py b/tests/unit/responses/mcp/test_parked_stream_dispatch.py new file mode 100644 index 00000000000..dd380c582c8 --- /dev/null +++ b/tests/unit/responses/mcp/test_parked_stream_dispatch.py @@ -0,0 +1,17 @@ +from typing import Final + +import pytest + +from litellm.responses.mcp.parked_stream_dispatch import parked_success +from litellm.types.utils import ModelResponse + +_ASSEMBLED: Final = ModelResponse(id="chatcmpl-parked", model="gpt-5.6") + + +@pytest.mark.parametrize( + "parked", + [(), (_ASSEMBLED,), (_ASSEMBLED, False, "not a logger"), (_ASSEMBLED, False, None, None)], + ids=("empty", "assembled-alone", "third-part-is-not-a-logger", "four-parts"), +) +def test_a_parked_dispatch_of_an_unknown_shape_is_never_dispatched(parked: tuple[object, ...]): + assert parked_success(parked) is None From 4786dd861b1b8c001117e0403cc2c606c193896c Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:36:01 -0700 Subject: [PATCH 2/3] fix(proxy)!: scope /health by wildcard keys and team model allowlists the way auth does (#49155) * test(proxy): cover /health scoping for wildcard keys and team model allowlists * fix(proxy): scope /health with the model access check auth runs on a request * fix(proxy): scope /health to the model group an allowed name routes to A key alias, team alias, or router model_group_alias lets auth admit a requested name by checking its target, so /health also kept the deployment under the requested name, one the key's requests never reach. Keep only the served model group each admitted name dispatches to, through the same alias layers /v1/models mirrors, and let alias_map read its own read-only output. * test(proxy): cover /health skipping a deployment that has no model name for a restricted key * fix(health): show deployments reached through wildcard and alias allowlists, and by id when an alias shadows the name * chore(router): drop should_include_deployment from the outside-use allowlist now that /health no longer calls it * fix(health): never list a default-fallback deployment for an unserved grant, and keep access-group keys to their group's deployments * fix(health): resolve ?model= through the caller's aliases the way chat does * fix(proxy): apply user and agent model limits to /health and yield between checks /health now runs the same user and agent model checks chat runs, so a personal key whose owner may call only some models, or an agent key limited by its access group, no longer sees or probes other deployments. Each per-name check yields to the event loop, and ?model= reaches the wildcard deployment serving an allowed name * fix(proxy): check owner and agent grants for wildcard routes in /health and fail on non-access auth errors * fix(proxy): answer /health auth lookup failures the way chat does and record ?model= checks under the requested name * test(proxy): pass recorded_model_name in the direct health save tests * fix(proxy): count alias names as /health candidates so aliased names a key may call get probed * fix(health): narrow wildcard-served deployments to the key's access groups * fix(health): narrow every deployment serving a name to the key's access groups, team-owned ones included * fix(health): narrow a team deployment's internal name to the key's access groups like its public name * fix(health): keep a team deployment's internal name narrow when a limit holds only that name An agent or key limit that holds one team deployment's internal name no longer reaches every deployment behind the team's public name. The internal name maps to its public name only when the public name is also admitted. * test(health): cover all-proxy-models next to * for a key that also holds an access group * refactor(auth): keep the _as_proxy_exception name so main's tests still find it * fix(health): keep an internal team name a wildcard grant admits as its own target --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 47 +- .../proxy/common_utils/model_listing_utils.py | 11 +- .../health_endpoints/_health_endpoints.py | 427 ++++++++++- ...st_auth_checks_object_access_and_lookup.py | 10 + .../health_endpoints/test_health_endpoints.py | 720 +++++++++++++++++- .../unit/proxy/test_health_check_functions.py | 2 + 6 files changed, 1164 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b8d9530d2eb..f80296f4bd7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -84,6 +84,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, CeilingResolver, resolve_agent_access_group_ceiling, ) @@ -1684,6 +1685,7 @@ def get_actual_routes(allowed_routes: list) -> list: KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" _KEY_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +_OBJECTS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: @@ -5085,6 +5087,36 @@ async def can_agent_call_model( ) +async def _agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver +) -> tuple[AgentAccessGroupCeiling, ...]: + if not valid_token.agent_id: + return () + + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + managed: Final = managed_agent_policy(valid_token) + if managed is not None: + return await resolve_managed_agent_ceilings(managed) + unmanaged: Final = await resolve_ceiling(valid_token.agent_id) + return (unmanaged,) if unmanaged is not None else () + + +def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + managed: Final = managed_agent_policy(valid_token) + granted: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(granted, (list, tuple)): + return frozenset() + return frozenset(name for name in _OBJECTS_ADAPTER.validate_python(granted) if isinstance(name, str)) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) + return _managed_agent_model_grants(valid_token).union(*(ceiling.models for ceiling in ceilings)) + + async def _check_agent_access_group_model_access( model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, @@ -5092,20 +5124,9 @@ async def _check_agent_access_group_model_access( resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" - if not model or valid_token is None or not valid_token.agent_id: + if not model or valid_token is None: return True - - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - managed: Final = managed_agent_policy(valid_token) - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if managed is None else None - ceilings: Final = ( - await resolve_managed_agent_ceilings(managed) - if managed is not None - else (unmanaged,) - if unmanaged is not None - else () - ) + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) dispatched: Final = _resolve_team_alias( model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router ) diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index ad4c3bd8770..a9424018d9a 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -219,7 +219,7 @@ def caller_alias_maps( def alias_map(aliases: object) -> Mapping[str, str]: try: - entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + entries: Final = _ALIAS_ENTRIES.validate_python(aliases) except ValidationError: return _NO_ALIASES return MappingProxyType( @@ -287,6 +287,15 @@ def _dispatch(model_id: str, layers: Sequence[_AliasLayer]) -> tuple[str, bool]: return reduce(_through_layer, layers, (model_id, False)) +def dispatched_model(model_id: str, aliases: CallerAliases) -> str: + target, _ = _dispatch(model_id, _alias_layers(aliases)) + return target + + +def requestable_alias_names(aliases: CallerAliases) -> frozenset[str]: + return frozenset(chain.from_iterable(layer.aliases for layer in _alias_layers(aliases))) + + def _lookup_id(name: str, lookup_by_response: Mapping[str, str], lookup_ids: Container[str]) -> str | None: if name in lookup_by_response: return lookup_by_response[name] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 59c63a18f0f..8545e0609db 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,13 +7,15 @@ import secrets import time import traceback from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import chain, groupby from types import MappingProxyType from typing import Final, Literal, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import ReadOnly import litellm @@ -31,7 +33,9 @@ from litellm.proxy._types import ( CallInfo, EnterpriseLicenseData, Litellm_EntityType, + LiteLLM_UserTable, LitellmUserRoles, + ModelAccessDeniedProxyException, ProxyException, SpecialModelNames, UserAPIKeyAuth, @@ -39,13 +43,28 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( _resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request + agent_model_grants, + can_agent_call_model, + can_key_call_resolved_model, + can_user_call_model, + get_user_object, + key_model_aliases_for_auth_check, + model_matches_any_wildcard_pattern_in_list, + team_model_aliases_for_auth_check, +) +from litellm.proxy.auth.auth_exception_handler import ( + _as_proxy_exception, # pyright: ignore[reportPrivateUsage] # one auth-failure mapping, shared with user_api_key_auth ) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) -from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils import http_parsing_utils +from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, + dispatched_model, + requestable_alias_names, +) from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -79,6 +98,7 @@ from litellm.router_utils.clientside_credential_handler import ( _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path clientside_credential_keys, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields from litellm.types.router import reject_server_owned_wif_fields @@ -253,6 +273,28 @@ def get_callback_identifier(callback): router: Final = APIRouter() _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) +_DEPLOYMENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +class _ModelGrants(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + models: tuple[object, ...] = () + + +class _ModelInfoAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + access_groups: tuple[object, ...] + + +class _DeploymentAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + model_info: _ModelInfoAccessGroups + + services = ( Literal[ "slack_budget_alerts", @@ -1045,42 +1087,339 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out -def _health_accessible_model_names( - user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None -) -> frozenset[str] | None: - """Model names the caller may health-check, or None when the key is unrestricted.""" +def _key_models_unrestricted(user_api_key_dict: UserAPIKeyAuth) -> bool: granted_models: Final = _resolve_key_models_for_auth_check(user_api_key_dict) - if not granted_models or SpecialModelNames.all_proxy_models.value in granted_models: + return not granted_models or SpecialModelNames.all_proxy_models.value in granted_models + + +def _own_team_names(deployment: object, team_id: str | None) -> tuple[str, str] | None: + match deployment: + case { + "model_name": str(internal_name), + "model_info": {"team_id": str(owner_team_id), "team_public_model_name": str(public_name)}, + } if owner_team_id == team_id: + return (internal_name, public_name) + case _: + return None + + +def _names_caller_may_request(deployment: object, team_id: str | None) -> frozenset[str]: + team_names: Final = _own_team_names(deployment, team_id) + if team_names is not None: + return frozenset(team_names) + match deployment: + case {"model_name": str(name)}: + return frozenset((name,)) + case _: + return frozenset() + + +def _public_names_by_internal_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, str]: + team_names: Final = (_own_team_names(deployment, team_id) for deployment in llm_model_list) + return MappingProxyType(dict(names for names in team_names if names is not None)) + + +async def _key_owner(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + if user_api_key_dict.team_id is not None or user_api_key_dict.user_id is None: return None + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + return await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: # noqa: BLE001 # auth reads an unloadable user as no user limit, so /health does too + return None + + +async def _key_may_call_model( + model_name: str, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> bool: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + await asyncio.sleep(0) + try: + await can_key_call_resolved_model( + model=model_name, llm_model_list=None, valid_token=user_api_key_dict, llm_router=llm_router + ) + await can_agent_call_model( + model=model_name, + valid_token=user_api_key_dict, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_user_call_model( + model=model_name, + llm_router=llm_router, + user_object=owner, + key_model_aliases=key_model_aliases_for_auth_check(user_api_key_dict), + ) + except ModelAccessDeniedProxyException: + return False + except ProxyException: + raise + except Exception as failure: + if isinstance(failure, HTTPException) and failure.status_code == status.HTTP_403_FORBIDDEN: + return False + raise _as_proxy_exception(failure) from failure + return True + + +async def _caller_aliases(user_api_key_dict: UserAPIKeyAuth) -> CallerAliases: + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + router_settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict, prisma_client, proxy_logging_obj + ) + return CallerAliases( + team_model_aliases_for_auth_check(user_api_key_dict), + key_model_aliases_for_auth_check(user_api_key_dict), + router_settings.get("model_group_alias") if router_settings is not None else None, + ) + + +def _model_auth_checks(model_name: str, aliases: CallerAliases) -> str: + return resolve_model_group_alias(aliases.router, model_name) or model_name + + +def _deployment_id(deployment: object) -> str | None: + match deployment: + case {"model_info": {"id": str(model_id)}}: + return model_id + case _: + return None + + +def _granted_names(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + grants: Final = ( + *_ModelGrants.model_validate(user_api_key_dict).models, + *_OBJECTS.validate_python(user_api_key_dict.team_models), + ) + return frozenset(name for name in grants if isinstance(name, str)) + + +async def _names_any_limit_grants(user_api_key_dict: UserAPIKeyAuth, owner: LiteLLM_UserTable | None) -> frozenset[str]: + owner_grants: Final = _ModelGrants.model_validate(owner).models if owner is not None else () + return ( + _granted_names(user_api_key_dict) + | frozenset(name for name in owner_grants if isinstance(name, str)) + | await agent_model_grants(user_api_key_dict) + ) + + +def _alias_names_caller_may_request(aliases: CallerAliases, llm_router: Router | None) -> frozenset[str]: + router_aliases: Final = frozenset(llm_router.registry().aliases()) if llm_router is not None else frozenset[str]() + return requestable_alias_names(aliases) | router_aliases + + +def _served_group(model_name: str, aliases: CallerAliases, llm_router: Router | None) -> str: + dispatched: Final = dispatched_model(model_name, aliases) + return llm_router.registry().resolve(dispatched) if llm_router is not None else dispatched + + +def _name_id_pairs(deployment: object, team_id: str | None) -> tuple[tuple[str, str], ...]: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return () + return tuple((name, model_id) for name in _names_caller_may_request(deployment, team_id) | {model_id}) + + +def _ids_by_requested_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, frozenset[str]]: + pairs: Final = sorted(chain.from_iterable(_name_id_pairs(deployment, team_id) for deployment in llm_model_list)) + return MappingProxyType( + {name: frozenset(model_id for _, model_id in group) for name, group in groupby(pairs, key=lambda pair: pair[0])} + ) + + +def _ids_of(deployments: Iterable[object]) -> frozenset[str]: + return frozenset(model_id for model_id in map(_deployment_id, deployments) if model_id is not None) + + +def _access_groups_of(deployment: object) -> frozenset[str]: + match deployment: + case {"model_info": {"access_groups": list()}}: + access_groups: Final = _DeploymentAccessGroups.model_validate(deployment).model_info.access_groups + return frozenset(access_group for access_group in access_groups if isinstance(access_group, str)) + case _: + return frozenset() + + +@dataclass(frozen=True, slots=True) +class _AccessGroupGrants: + granted_names: frozenset[str] + ids_in_granted_groups: frozenset[str] + + def requested_name(self, name: str, public_names: Mapping[str, str], admitted_names: frozenset[str]) -> str: + public_name: Final = public_names.get(name) + if public_name is None or public_name not in admitted_names or self._grants(name): + return name + return public_name + + def _grants(self, name: str) -> bool: + return name in self.granted_names or model_matches_any_wildcard_pattern_in_list(name, tuple(self.granted_names)) + + def narrow(self, serving_ids: frozenset[str], served_names: frozenset[str]) -> frozenset[str]: + in_granted_groups: Final = serving_ids & self.ids_in_granted_groups + granted_directly: Final = not self.granted_names.isdisjoint(served_names) + return serving_ids if granted_directly or not in_granted_groups else in_granted_groups + + +def _access_group_grants(llm_model_list: Sequence[object], user_api_key_dict: UserAPIKeyAuth) -> _AccessGroupGrants: + granted: Final = _granted_names(user_api_key_dict) + in_granted_groups: Final = ( + deployment for deployment in llm_model_list if not granted.isdisjoint(_access_groups_of(deployment)) + ) + return _AccessGroupGrants(granted, _ids_of(in_granted_groups)) + + +def _ids_serving_group( + group: str, + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + named: Final = ids_by_name.get(group) + if named is not None: + return access_group_grants.narrow(named, frozenset((group,))) if llm_router is None: - return frozenset(granted_models) - registry: Final = llm_router.registry() + return frozenset() + team_id: Final = user_api_key_dict.team_id + wildcard_served: Final = llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + served_names: Final = frozenset((group,)).union( + *(_names_caller_may_request(deployment, team_id) for deployment in wildcard_served) + ) + return access_group_grants.narrow(_ids_of(wildcard_served), served_names) + + +async def _names_key_may_request( + names: Iterable[str], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> frozenset[str]: return frozenset( - get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=registry.names(team_id=user_api_key_dict.team_id), - model_access_groups=registry.access_groups(), + [ + name + for name in names + if await _key_may_call_model(_model_auth_checks(name, aliases), user_api_key_dict, owner, llm_router) + ] + ) + + +def _ids_serving( + groups: Iterable[str], + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + return frozenset[str]().union( + *( + _ids_serving_group(group, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + for group in groups ) ) -def _caller_may_probe_deployment( - deployment: Mapping[str, object], - allowed_models: frozenset[str] | None, +def _shadowed_ids( + llm_model_list: Sequence[object], admitted_names: frozenset[str], reached_ids: frozenset[str], team_id: str | None +) -> frozenset[str]: + ids_of_admitted_names: Final = ( + _deployment_id(deployment) + for deployment in llm_model_list + if not admitted_names.isdisjoint(_names_caller_may_request(deployment, team_id)) + ) + return frozenset(model_id for model_id in ids_of_admitted_names if model_id is not None) - reached_ids + + +@dataclass(frozen=True, slots=True) +class _ReachableDeployments: + model_ids: frozenset[str] + groups: frozenset[str] + + def include(self, deployment: object, team_id: str | None) -> bool: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return not self.groups.isdisjoint(_names_caller_may_request(deployment, team_id)) + return model_id in self.model_ids + + +_NOTHING_REACHABLE: Final = _ReachableDeployments(frozenset(), frozenset()) + + +async def _deployments_key_may_reach( + llm_model_list: Sequence[object], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, llm_router: Router | None, +) -> _ReachableDeployments: + candidate_names: Final = (await _names_any_limit_grants(user_api_key_dict, owner)).union( + _alias_names_caller_may_request(aliases, llm_router), + *(_names_caller_may_request(deployment, user_api_key_dict.team_id) for deployment in llm_model_list), + ) + admitted_names: Final = await _names_key_may_request(candidate_names, aliases, user_api_key_dict, owner, llm_router) + access_group_grants: Final = _access_group_grants(llm_model_list, user_api_key_dict) + public_names: Final = _public_names_by_internal_name(llm_model_list, user_api_key_dict.team_id) + groups: Final = frozenset( + _served_group(access_group_grants.requested_name(name, public_names, admitted_names), aliases, llm_router) + for name in admitted_names + ) + ids_by_name: Final = _ids_by_requested_name(llm_model_list, user_api_key_dict.team_id) + reached_ids: Final = _ids_serving(groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + admitted_ids: Final = await _names_key_may_request( + _shadowed_ids(llm_model_list, admitted_names, reached_ids, user_api_key_dict.team_id), + aliases, + user_api_key_dict, + owner, + llm_router, + ) + id_groups: Final = frozenset(_served_group(model_id, aliases, llm_router) for model_id in admitted_ids) + return _ReachableDeployments( + reached_ids | _ids_serving(id_groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants), groups + ) + + +def _model_param_target( + model: str, + aliases: CallerAliases, + llm_model_list: Sequence[Mapping[str, object]], team_id: str | None, - caller_is_admin: bool, + llm_router: Router | None, +) -> str: + group: Final = _served_group(model, aliases, llm_router) + if llm_router is None or deployments_targeted_by_name(llm_model_list, group, team_id): + return group + patterns: Final = frozenset( + deployment.get("model_name") + for deployment in llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + ) + match tuple(patterns): + case (str(pattern),): + return pattern + case _: + return group + + +def _caller_may_probe_deployment( + deployment: Mapping[str, object], reachable: _ReachableDeployments, team_id: str | None, caller_is_admin: bool ) -> bool: """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" if not caller_is_admin and not deployment_usable_by_team(deployment, team_id): return False - if allowed_models is None: - return True - if llm_router is None: - return deployment.get("model_name") in allowed_models - model: Final = dict(deployment) - registry: Final = llm_router.registry() - return any(registry.includes(model, model_name=name, team_id=team_id) for name in allowed_models) + return reachable.include(deployment, team_id) def _resolve_targeted_model_ids( @@ -1137,6 +1476,7 @@ def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: async def _perform_health_check_and_save( model_list, target_model, + recorded_model_name: str | None, cli_model, details, prisma_client, @@ -1159,8 +1499,7 @@ async def _perform_health_check_and_save( # Optionally save health check result to database (non-blocking) if prisma_client is not None: - # For CLI model, use cli_model name; for router models, use target_model - model_name_for_db: Final = cli_model if cli_model is not None else target_model + model_name_for_db: Final = cli_model if cli_model is not None else recorded_model_name if model_name_for_db is not None: asyncio.create_task( _save_health_check_to_db( @@ -1290,6 +1629,7 @@ async def health_endpoint( cli_result: Final = await _perform_health_check_and_save( model_list=[], target_model=None, + recorded_model_name=None, cli_model=user_model, details=health_check_details, prisma_client=prisma_client, @@ -1304,15 +1644,39 @@ async def health_endpoint( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) - allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) - restrict_to_allowed_models: Final = not is_admin or allowed_models is not None + restrict_to_allowed_models: Final = not is_admin or not _key_models_unrestricted(user_api_key_dict) + aliases: Final = await _caller_aliases(user_api_key_dict) if restrict_to_allowed_models else None + reachable: Final = ( + await _deployments_key_may_reach( + _OBJECTS.validate_python(llm_model_list), + aliases, + user_api_key_dict, + await _key_owner(user_api_key_dict), + llm_router, + ) + if aliases is not None + else _NOTHING_REACHABLE + ) + requested_model: Final = ( + _model_param_target( + model, + aliases, + _DEPLOYMENTS.validate_python(llm_model_list), + user_api_key_dict.team_id, + llm_router, + ) + if model is not None and aliases is not None + else model + ) _llm_model_list: Final = [ m for m in copy.deepcopy(llm_model_list) if not restrict_to_allowed_models - or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) + or _caller_may_probe_deployment(m, reachable, user_api_key_dict.team_id, is_admin) ] - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id) + targeted_ids: Final = _resolve_targeted_model_ids( + _llm_model_list, requested_model, model_id, user_api_key_dict.team_id + ) if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1370,7 +1734,8 @@ async def health_endpoint( else: router_result: Final = await _perform_health_check_and_save( model_list=_llm_model_list, - target_model=target_model, + target_model=target_model if model_id else requested_model, + recorded_model_name=target_model, cli_model=None, details=health_check_details, prisma_client=prisma_client, diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 0bdc5f1910a..6b59f6d40dc 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( can_object_call_model, _can_object_call_vector_stores, _check_agent_access_group_model_access, + agent_model_grants, check_end_user_budget, _check_team_member_budget, personal_key_owner_teams, @@ -11294,6 +11295,15 @@ async def test_agent_access_groups_naming_no_model_deny_every_model(): assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied +@pytest.mark.asyncio +async def test_agent_model_grants_names_every_model_the_agents_access_groups_allow(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5", "openai/gpt-5.6"})) + + assert await agent_model_grants(agent_key, resolve) == frozenset({"gpt-5", "openai/gpt-5.6"}) + assert asked == ["agent-1"] + + @pytest.mark.asyncio async def test_agent_without_access_groups_adds_no_model_ceiling(): agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 9847787efce..2258eb95fb8 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -1,8 +1,9 @@ import asyncio import copy +import itertools import json import time -from collections.abc import Awaitable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from datetime import datetime, timedelta from importlib import import_module @@ -22,7 +23,7 @@ import litellm import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 from litellm.models.credentials import CredentialItem -from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, @@ -35,6 +36,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( test_model_connection as health_test_model_connection, ) from litellm.router import Router +from litellm.types.router import RouterModelGroupAliasItem from litellm.types.utils import oauth_token_exchange_litellm_params from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, @@ -925,7 +927,11 @@ async def test_test_model_connection_stored_audio_base_model_follows_the_stored_ deployment: Final = MappingProxyType( { "model_name": "custom-name-audio", - "litellm_params": {**deployment_params, "api_key": "fake-azure-key", "api_base": "https://x.openai.azure.com"}, + "litellm_params": { + **deployment_params, + "api_key": "fake-azure-key", + "api_base": "https://x.openai.azure.com", + }, "model_info": deployment_model_info, } ) @@ -1948,7 +1954,6 @@ def test_health_liveness_endpoint(proxy_client): assert duration_ms < 100, f"Health check took {duration_ms:.2f}ms, expected < 100ms for a simple endpoint" - def test_proxy_client_serves_full_proxy_routes_after_gateway_import(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy.proxy_server import app @@ -2472,6 +2477,7 @@ def _proxy_health_globals( llm_router: object, use_background_health_checks: bool = False, health_check_results: Mapping[str, object] | None = None, + prisma_client: object = None, ) -> Iterator[None]: with ( patch( # test-quality-ok: proxy module global, no injection seam @@ -2481,7 +2487,7 @@ def _proxy_health_globals( "litellm.proxy.proxy_server.llm_router", llm_router ), patch( # test-quality-ok: proxy module global, no injection seam - "litellm.proxy.proxy_server.prisma_client", None + "litellm.proxy.proxy_server.prisma_client", prisma_client ), patch( # test-quality-ok: proxy module global, no injection seam "litellm.proxy.proxy_server.use_background_health_checks", use_background_health_checks @@ -3812,7 +3818,11 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, async def _live_probed_model_ids( - model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None + model_list: Sequence[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + model: str | None = None, + llm_router: Router | None = None, + prisma_client: object = None, ) -> set[str]: from fastapi import Response @@ -3825,7 +3835,7 @@ async def _live_probed_model_ids( return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list), prisma_client=prisma_client), patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", side_effect=fake_perform, @@ -4120,6 +4130,7 @@ async def _live_narrowed_model_ids( user_api_key_dict: UserAPIKeyAuth, model: str | None = None, model_id: str | None = None, + llm_router: Router | None = None, ) -> set[str]: from fastapi import Response @@ -4130,7 +4141,7 @@ async def _live_narrowed_model_ids( return probed, [], {} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list)), patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe ), @@ -4289,6 +4300,699 @@ async def test_health_endpoint_keeps_a_team_only_public_name_off_a_team_less_key probe.assert_not_awaited() +_PROVIDER_PREFIXED_MODEL_LIST: Final = [ + { + "model_name": "openai/gpt-5.4-nano", + "litellm_params": {"model": "openai/gpt-5.4-nano"}, + "model_info": {"id": "id-prefixed-nano"}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini"}, + }, + { + "model_name": "claude-haiku-5-5", + "litellm_params": {"model": "anthropic/claude-haiku-5-5"}, + "model_info": {"id": "id-haiku"}, + }, +] +_PROVIDER_PREFIXED_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.4-nano", "model_id": "id-prefixed-nano"}, + {"model": "openai/gpt-5.4-mini", "model_id": "id-mini"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} +_PROVIDER_WILDCARD_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["openai/*"]) +_RESTRICTED_TEAM_KEY: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=[], team_id="team-a", team_models=["gpt-5.4-mini"] +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, {"id-prefixed-nano", "id-mini"}), (_RESTRICTED_TEAM_KEY, {"id-mini"})], +) +async def test_health_endpoint_probes_the_deployments_auth_lets_the_key_request( + caller: UserAPIKeyAuth, expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, ["id-prefixed-nano", "id-mini"]), (_RESTRICTED_TEAM_KEY, ["id-mini"])], +) +async def test_health_endpoint_returns_the_deployments_auth_lets_the_key_request_on_background_cache_path( + caller: UserAPIKeyAuth, expected_ids: list[str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_a_deployment_the_team_allowlist_forbids(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_RESTRICTED_TEAM_KEY, model=None, model_id="id-prefixed-nano" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +_KEY_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=["gpt-5.4-mini"], aliases={"claude-haiku-5-5": "gpt-5.4-mini"} +) +_TEAM_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["claude-haiku-5-5"], + team_id="team-a", + team_model_aliases={"claude-haiku-5-5": "gpt-5.4-mini"}, +) +_MINI_ONLY_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]) +_ALIAS_SHADOW_CASES: Final = [ + pytest.param(_KEY_ALIASED_AWAY_FROM_HAIKU, {}, id="key-alias"), + pytest.param(_TEAM_ALIASED_AWAY_FROM_HAIKU, {}, id="team-alias"), + pytest.param(_MINI_ONLY_KEY, {"claude-haiku-5-5": "gpt-5.4-mini"}, id="router-alias"), +] + + +def _router_with_aliases(model_group_alias: Mapping[str, str]) -> Router: + return Router(model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), model_group_alias=dict(model_group_alias)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_probes_the_deployment_an_alias_sends_the_key_to_not_the_one_it_shadows( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_probed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=_router_with_aliases(model_group_alias) + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_returns_the_deployment_an_alias_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_the_deployment_a_key_alias_shadows(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_KEY_ALIASED_AWAY_FROM_HAIKU, model=None, model_id="id-haiku" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_narrowed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, + caller, + model="claude-haiku-5-5", + llm_router=_router_with_aliases(model_group_alias), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_records_an_alias_check_under_the_name_the_caller_asked_for( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, _router_with_aliases(model_group_alias), prisma_client=MagicMock() + ), + patch( # test-quality-ok: the probe is the provider edge; the endpoint has no injection seam for it + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + patch( # test-quality-ok: the history row is the assertion; the endpoint has no injection seam for its writer + "litellm.proxy.health_endpoints._health_endpoints._save_health_check_to_db", new_callable=AsyncMock + ) as save, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + assert save.call_args.args[1] == "claude-haiku-5-5" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_a_deployment_without_a_model_name_for_a_restricted_key(): + nameless_deployment: Final = { + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-nameless"}, + } + + probed: Final = await _live_probed_model_ids( + [*_PROVIDER_PREFIXED_MODEL_LIST, nameless_deployment], + _MINI_ONLY_KEY, + llm_router=_router_for(_PROVIDER_PREFIXED_MODEL_LIST), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_the_deployment_a_keys_router_alias_sends_it_to_but_forbids(): + caller: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["gpt-5.4-mini", "openai/gpt-5.4-nano"], + router_settings={"model_group_alias": {"gpt-5.4-mini": "claude-haiku-5-5"}}, + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == {"id-prefixed-nano"} + + +_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild"}, +} +_CATCH_ALL_DEPLOYMENT: Final = {"model_name": "*", "litellm_params": {"model": "*"}, "model_info": {"id": "id-star"}} +_TEAM_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai-wild_team-a_7c1e", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-team-wild", "team_id": "team-a", "team_public_model_name": "openai/*"}, +} + + +def _team_key(team_models: list[str], team_model_aliases: Mapping[str, str] | None = None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-test-key", + models=[], + team_id="team-a", + team_models=team_models, + team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_list", "model_group_alias", "caller", "expected_ids"), + [ + pytest.param( + [_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-wild"}, + id="team-name-served-by-provider-wildcard", + ), + pytest.param( + [_TEAM_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-team-wild"}, + id="team-name-served-by-teams-own-wildcard", + ), + pytest.param( + [_CATCH_ALL_DEPLOYMENT], + {}, + _team_key(["gpt-5.4-mini"]), + {"id-star"}, + id="team-name-served-by-catch-all", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"gpt-4": "gpt-5.4-mini"}, + _team_key(["gpt-4"]), + {"id-mini"}, + id="team-allows-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["my-alias"], {"my-alias": "gpt-5.4-mini"}), + {"id-mini"}, + id="team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"friendly-mini": RouterModelGroupAliasItem(model="gpt-5.4-mini", hidden=True)}, + UserAPIKeyAuth(api_key="hashed-test-key", models=["friendly-*"]), + {"id-mini"}, + id="wildcard-key-allows-hidden-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["friendly-*"], {"friendly-mini": "gpt-5.4-mini"}), + {"id-mini"}, + id="wildcard-team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"claude-haiku-5-5": "gpt-5.4-mini"}, + UserAPIKeyAuth(api_key="hashed-test-key", models=[]), + {"id-prefixed-nano", "id-mini", "id-haiku"}, + id="unrestricted-key-calls-shadowed-deployment-by-id", + ), + ], +) +async def test_health_endpoint_probes_every_deployment_the_callers_allowed_names_route_to( + model_list: Sequence[Mapping[str, object]], + model_group_alias: Mapping[str, str | RouterModelGroupAliasItem], + caller: UserAPIKeyAuth, + expected_ids: set[str], +): + llm_router: Final = Router( + model_list=copy.deepcopy([dict(deployment) for deployment in model_list]), + model_group_alias=dict(model_group_alias), + ) + + probed: Final = await _live_probed_model_ids(llm_router.get_model_list() or [], caller, llm_router=llm_router) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller", + [ + pytest.param( + UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini", "gpt-4-retired"]), id="key-grant" + ), + pytest.param(_team_key(["gpt-5.4-mini", "gpt-4-retired"]), id="team-grant"), + ], +) +async def test_health_endpoint_never_probes_the_default_fallback_for_a_granted_name_no_deployment_serves( + caller: UserAPIKeyAuth, +) -> None: + llm_router: Final = Router( + model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), default_fallbacks=["claude-haiku-5-5"] + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=llm_router) + + assert probed == {"id-mini"} + + +def _personal_key_whose_owner_may_call(models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import user_api_key_cache + + owner_id: Final = "health-scope-owner-" + "-".join(models) + user_api_key_cache.set_cache(key=owner_id, value=LiteLLM_UserTable(user_id=owner_id, models=models)) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_id=owner_id) + + +def _managed_agent_key_granted(models: list[str]) -> UserAPIKeyAuth: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", agent_name="Managed", agent_card_params={}, object_permission={"models": models} + ) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[]).model_copy(update={"managed_agent_policy": policy}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("make_caller", "models", "expected_ids"), + [ + pytest.param(_personal_key_whose_owner_may_call, ["gpt-5.4-mini"], {"id-mini"}, id="key-owners-models"), + pytest.param(_managed_agent_key_granted, ["gpt-5.4-mini"], {"id-mini"}, id="managed-agents-models"), + pytest.param(_managed_agent_key_granted, [], set(), id="managed-agent-without-model-grants"), + ], +) +async def test_health_endpoint_probes_only_the_models_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], models: list[str], expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, make_caller(models)) + + assert probed == expected_ids + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_model_id_the_keys_owner_may_not_call(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals(_PROVIDER_PREFIXED_MODEL_LIST, _router_for(_PROVIDER_PREFIXED_MODEL_LIST)), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=_personal_key_whose_owner_may_call(["gpt-5.4-mini"]), + model=None, + model_id="id-prefixed-nano", + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_caller", + [ + pytest.param(_personal_key_whose_owner_may_call, id="key-owner"), + pytest.param(_managed_agent_key_granted, id="managed-agent"), + ], +) +async def test_health_endpoint_keeps_the_wildcard_serving_a_name_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], +): + model_list: Final = [_WILDCARD_DEPLOYMENT, *_PROVIDER_PREFIXED_MODEL_LIST] + + probed: Final = await _live_probed_model_ids(model_list, make_caller(["openai/gpt-5.6"])) + + assert probed == {"id-wild"} + + +@pytest.mark.asyncio +async def test_health_endpoint_answers_an_agent_key_acting_for_an_unknown_user_the_way_chat_does(): + from litellm.types.agents import AgentCaller + + users_table: Final = MagicMock() + users_table.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + agent_key: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]).model_copy( + update={"agent_caller": AgentCaller(user_id="user-deleted-before-this-request")} + ) + + with pytest.raises(ProxyException) as failed: + await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, agent_key, prisma_client=users_table) + + assert (failed.value.code, failed.value.type) == ("401", ProxyErrorTypes.auth_error) + assert "user-deleted-before-this-request" in failed.value.message + + +@pytest.mark.asyncio +async def test_health_endpoint_lets_other_requests_run_while_it_scopes_a_large_model_list(): + model_list: Final = [ + {"model_name": f"model-{n}", "litellm_params": {"model": f"openai/model-{n}"}, "model_info": {"id": f"id-{n}"}} + for n in range(50) + ] + other_request_turns: Final = itertools.count() + scoping_done: Final = asyncio.Event() + + async def scope() -> set[str]: + try: + return await _live_probed_model_ids( + model_list, UserAPIKeyAuth(api_key="hashed-test-key", models=["model-0"]) + ) + finally: + scoping_done.set() + + async def other_request() -> None: + while not scoping_done.is_set(): + next(other_request_turns) + await asyncio.sleep(0) + + probed, _ = await asyncio.gather(scope(), other_request()) + + assert probed == {"id-0"} + assert next(other_request_turns) >= len(model_list) + + +_WILDCARD_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.6", "model_id": "id-wild"}, + {"model": "openai/gpt-5.6", "model_id": "id-team-wild"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +@pytest.mark.parametrize( + ("wildcard_deployment", "expected_id"), + [ + pytest.param(_WILDCARD_DEPLOYMENT, "id-wild", id="proxy-wildcard"), + pytest.param(_TEAM_WILDCARD_DEPLOYMENT, "id-team-wild", id="teams-own-wildcard"), + ], +) +async def test_health_endpoint_model_param_reaches_the_wildcard_serving_a_team_allowed_name( + wildcard_deployment: Mapping[str, object], expected_id: str, use_background_health_checks: bool +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + model_list: Final = [wildcard_deployment, _PROVIDER_PREFIXED_MODEL_LIST[2]] + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + model_list, + _router_for(model_list), + use_background_health_checks=use_background_health_checks, + health_check_results=_WILDCARD_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=_team_key(["openai/gpt-5.6"]), model="openai/gpt-5.6", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == [expected_id] + + +_ACCESS_GROUP_SPLIT_MODEL_LIST: Final = [ + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-standard"}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_deployments_of_a_shared_name() -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium"]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct_grant", ["gpt-5.4-mini", "*", "all-proxy-models"]) +async def test_health_endpoint_probes_every_deployment_of_a_name_the_key_also_holds_directly(direct_grant: str) -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium", direct_grant]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium", "id-mini-standard"} + + +_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST: Final = [ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-standard", "access_groups": ["standard"]}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_wildcard_deployments_serving_a_name() -> None: + caller: Final = _personal_key_whose_owner_may_call(["openai/gpt-5.4-nano"]).model_copy( + update={"models": ["premium"]} + ) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST, caller) + + assert probed == {"id-wild-premium"} + + +def _team_deployment( + internal_name: str, public_name: str, model: str, model_id: str, access_groups: list[str] +) -> Mapping[str, object]: + return { + "model_name": internal_name, + "litellm_params": {"model": model}, + "model_info": { + "id": model_id, + "team_id": "team-a", + "team_public_model_name": public_name, + "access_groups": access_groups, + }, + } + + +def _team_split_model_list(public_name: str, model: str) -> list[Mapping[str, object]]: + return [ + _team_deployment("model_name_team-a_7c1e4b2a", public_name, model, "id-team-premium", ["premium"]), + _team_deployment("model_name_team-a_9d2f6e01", public_name, model, "id-team-standard", ["standard"]), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("public_name", "model"), + [ + pytest.param("openai/*", "openai/*", id="team-wildcard"), + pytest.param("gpt-5.4-mini", "openai/gpt-5.4-mini", id="team-model"), + ], +) +async def test_health_endpoint_probes_only_the_access_groups_team_deployments_serving_a_name( + public_name: str, model: str +) -> None: + probed: Final = await _live_probed_model_ids(_team_split_model_list(public_name, model), _team_key(["premium"])) + + assert probed == {"id-team-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("internal_name_grant", ["model_name_team-a_9d2f6e01", "model_name_team-a_*"]) +async def test_health_endpoint_probes_a_team_deployment_whose_internal_name_the_key_holds( + internal_name_grant: str, +) -> None: + caller: Final = _team_key(["premium", internal_name_grant]) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-premium", "id-team-standard"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_team_deployment_whose_internal_name_an_agent_holds() -> None: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", + agent_name="Managed", + agent_card_params={}, + object_permission={"models": ["model_name_team-a_9d2f6e01"]}, + ) + caller: Final = _team_key(["gpt-5.4-mini"]).model_copy(update={"managed_agent_policy": policy}) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-standard"} + + def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): resolve = _health_endpoints_module._resolve_targeted_model_ids diff --git a/tests/unit/proxy/test_health_check_functions.py b/tests/unit/proxy/test_health_check_functions.py index 2c0198cd538..79eec44958f 100644 --- a/tests/unit/proxy/test_health_check_functions.py +++ b/tests/unit/proxy/test_health_check_functions.py @@ -741,6 +741,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c result = await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, @@ -777,6 +778,7 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, From 43c07a0f1246c95da710408df4d7846bc5f848e2 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:48:32 -0700 Subject: [PATCH 3/3] fix(router)!: check and count deployment itpm and otpm on every Responses WebSocket turn (#49069) * fix(router): check and count deployment itpm and otpm on every Responses WebSocket turn The success hook only reconciled an ITPM/OTPM reservation found in top-level metadata or the standard logging object, but on the proxy the reservation lives in litellm_params.metadata, so no reservation was ever settled there. The hook now always calls the reconcile, which already returns early without a reservation and clears what it settled, so a second success event is a no-op. A Responses WebSocket turn after the first now has its deployment's current-minute ITPM and OTPM counters checked and is refused with a 429 frame once a counter is at its limit. An admitted turn carries a zero reservation on the current minute's keys, so its actual tokens land on those counters on both the native and the managed relay. * fix(responses): settle the socket open's itpm/otpm reservation on the first turn only After the first turn takes its logging copy, both websocket relays drop the open's ITPM/OTPM reservation keys from the connection metadata, the same way the budget reservation is detached. Later turns, unsolicited backend events, turns after limits are removed mid-socket, and sockets with no turn admission no longer settle or refund the open's reservation a second time. * refactor(responses): read the itpm/otpm reservation keys from core metadata helpers instead of router_utils * fix(responses): keep a socket open's hold for its first real turn and check tokens before request limits A turn that will settle the socket open's itpm/otpm reservation no longer counts that reservation against itself or swaps it for a zero one, so Codex's warmup frame on a managed provider, or a refused first frame on the native socket, no longer blocks the first real prompt or strands the hold for the minute. The token check now runs before the rpm/rpd checks, so a turn refused for tokens never uses up a request slot. Token-check tests pin the clock, and the managed relay settle test drives a real streamed call through a custom provider. * test(router): pin the router clock in the socket turn admission tests Both tests seeded this minute and the next to dodge a minute boundary. They now use the file's pinned-minute fixture and seed only that minute, so neither depends on when the test runs. * fix(router): check a responses socket's open without reserving and read shared token counters first A Responses WebSocket open now runs the same itpm/otpm check a turn runs and reserves nothing. It used to reserve an estimate that falls back to the model's max output, so a socket could not open under any otpm below that. The open still records a zero reservation for its own minute, so the first frame, which the relay forwards without a turn check, counts its actual tokens like every later turn. This drops the plumbing that only carried the open's hold into the first turn. The turn admission callable takes no arguments again, and the native relay detaches only the budget reservation, as before this change. The itpm/otpm turn check reads Redis first and falls back to this instance's in-memory count only without Redis or when the circuit breaker is open. These counters go down when a call settles, so an in-memory copy at the limit could refuse turns after the shared count dropped. The tpm check keeps its in-memory shortcut, since its counter only grows within a minute. Responses usage that reports input_tokens_details.cached_tokens now leaves cached input out of itpm, as chat usage already did. * fix(router): leave a socket turn's cached input out of itpm The native Responses WebSocket relay logs a turn with its list of events, so the deployment rate limit hook finds no usage on the response and falls back to the standard logging payload. That fallback read cached input from metadata.cache_read_input_tokens, a key the payload's cleaned metadata never carries, so every cached token of a socket turn was counted against itpm. The fallback now reads the cached count from the payload's usage_object, the same usage the logging path already built for cost, so a native socket turn counts only its uncached input. The managed relay and HTTP /v1/responses already counted correctly through the response's own usage * fix(router): mark only responses socket opens as check-only and count a turn on the minute it finishes The check-only branch for a socket open keyed on a websocket value in the request kwargs, which any HTTP body can carry and which the realtime route also sets. The router now passes an opens_responses_socket flag into set_io_token_rate_limit_request_kwargs only for its _aresponses_websocket call, and the flag lives in the same context value as the request kwargs, so no request body can set it and realtime opens reserve and get refused as before A zero reservation holds nothing to refund, so a socket turn now counts its real tokens on the minute the call finishes, read from the clock at settle time. A non-zero reservation still settles on the keys it reserved on so its refund hits the same counter When the shared Redis read returns nothing, the check now falls back to this instance's in-memory count instead of admitting the turn with no count at all * fix(router): keep io token reservation keys out of the metadata sent to providers The itpm/otpm pre-call check wrote its reservation keys into both metadata and litellm_metadata. On the routes the router serves through its generic helper (Responses, Anthropic Messages, and the rest), metadata is the client's own field and goes to the provider as is. So _litellm_itpm_cache_key and _litellm_otpm_cache_key, which carry the deployment id and the model name, reached the provider, came back in the echoed response metadata, and counted toward the provider's 16-key cap. The managed Responses socket relay merged each turn's reservation into the frame's metadata the same way, and a native socket opened with metadata sent the keys on every turn through its request defaults The router now passes the limiter the metadata field it uses for its own fields on that call (metadata for chat, litellm_metadata for the generic routes), and the reservation goes only there. metadata_channel is a required keyword, so a new caller has to name the field it writes. The managed relay puts each turn's reservation only on litellm_metadata With the reservation only on litellm_metadata, a retried call on the generic routes refunded its first attempt twice and pushed the deployment's counters below zero. The failure hook read the reservation from the logging copy in litellm_params.metadata and cleared only that copy, so the router's litellm_metadata still held it and the retry refunded it again. The hooks now also read and clear litellm_params.litellm_metadata, which is the router's own dict The proxy does not strip client keys that start with _litellm_ on the way in, so dropping that prefix at the provider boundary would also drop keys clients send today The io token tests now run Redis in process with fakeredis instead of a TCP fake server * test(responses): count socket turns through real logging callbacks instead of a Logging subclass The native relay test now hands the rate limit check to a real Logging object as its success and failure callback, so the turn's tokens go through the same dispatch path production uses. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../internal_call_metadata.py | 6 + litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/responses/main.py | 6 +- litellm/responses/streaming_iterator.py | 103 ++-- litellm/router.py | 26 +- .../io_token_rate_limit_check.py | 209 ++++++--- ...test_responses_websocket_turn_admission.py | 113 ++++- .../test_responses_websocket_all_providers.py | 441 +++++++++++++++++- .../test_router/test_io_token_rate_limits.py | 356 ++++++++++++-- tests/unit/test_router/test_router.py | 6 +- .../test_router/test_router_helper_utils.py | 374 ++++++++++++++- 11 files changed, 1492 insertions(+), 152 deletions(-) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index f3dbd612bee..453d4358b35 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -35,6 +35,12 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" +OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" +ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" +OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" +IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY}) + MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" """Where auth records the model access groups that authorized the request, for the spend writer. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 61fc23589ac..9a2b352dd86 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -249,7 +249,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -7269,7 +7269,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, + turn_admission: "ResponsesWsTurnAdmission | None" = None, frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5c8350589b8..9259103cb5c 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 Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence +from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import partial @@ -81,7 +81,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2867,7 +2867,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, + turn_admission: "ResponsesWsTurnAdmission | None" = None, responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index f0bb2201162..fb91ba6a618 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2145,18 +2145,32 @@ async def _enforce_frame_project_quota( ) -async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> Exception | None: +@dataclass(frozen=True, slots=True) +class _AdmittedTurn: + reservation: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + + +@dataclass(frozen=True, slots=True) +class _RefusedTurn: + rejection: Exception + + +ResponsesWsTurnAdmission: TypeAlias = Callable[[], Awaitable[Mapping[str, object]]] + + +async def _admit_frame_turn( + turn_admission: ResponsesWsTurnAdmission | None, raw_message: str +) -> _AdmittedTurn | _RefusedTurn: frame: Final = _json_object_or_none(raw_message) if turn_admission is None or frame is None or frame.get("type") != "response.create": - return None + return _AdmittedTurn() try: - await turn_admission() + return _AdmittedTurn(await turn_admission()) except RateLimitError as refusal: - return refusal + return _RefusedTurn(refusal) except Exception as e: verbose_logger.exception("Responses WS turn admission failed: %s", e) - return e - return None + return _RefusedTurn(e) def _turn_refusal_frame(refusal: RateLimitError) -> str: @@ -2191,6 +2205,7 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_NO_IO_TOKEN_RESERVATION: Final[Mapping[str, object]] = MappingProxyType({}) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -2318,6 +2333,17 @@ def _set_turn_guardrail_information( _replace_guardrail_information(litellm_params.get(key), guardrail_information) +def _set_turn_io_token_reservation(turn_logging_obj: LiteLLMLoggingObj, reservation: Mapping[str, object]) -> None: + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None or not reservation: + return + for bucket in (litellm_params.get(key) for key in _PER_TURN_METADATA_KEYS): + if _is_json_object(bucket): + bucket.update(reservation) + + def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: if not _is_json_object(bucket): return @@ -2458,7 +2484,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, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket @@ -2480,10 +2506,13 @@ 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 + self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission def _start_turn( - self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str | None, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( @@ -2491,6 +2520,7 @@ class ResponsesWebSocketStreaming: ) self._detach_budget_reservation() _set_turn_guardrail_information(turn_logging_obj, guardrail_information) + _set_turn_io_token_reservation(turn_logging_obj, io_token_reservation) 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)) @@ -2510,11 +2540,14 @@ class ResponsesWebSocketStreaming: } def _record_client_frame( - self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> 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, guardrail_information) + self._start_turn(client_frame, guardrail_information, io_token_reservation) async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: """Run the proxy's pre_call guardrails on one ``response.create`` frame. @@ -2533,15 +2566,18 @@ class ResponsesWebSocketStreaming: return None async def _forward_guarded_frame( - self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission): + if guarded is None: + return + admitted: Final = await self._admit_or_reject_turn(guarded, message, turn_admission) + if admitted is None: return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) masked: Final = await self._mask_response_create(guarded.frame) - self._record_client_frame(masked, guarded.guardrail_information) + self._record_client_frame(masked, guarded.guardrail_information, admitted.reservation) await self.backend_ws.send(masked) def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: @@ -2941,15 +2977,17 @@ class ResponsesWebSocketStreaming: return True async def _admit_or_reject_turn( - self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None - ) -> bool: - rejection: Final = await _admit_frame_turn(turn_admission, message) - if rejection is None: - return True - if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) - return False + self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: ResponsesWsTurnAdmission | None + ) -> _AdmittedTurn | None: + verdict: Final = await _admit_frame_turn(turn_admission, message) + match verdict: + case _AdmittedTurn(): + return verdict + case _RefusedTurn(rejection=rejection): + if self._failure is None: + self._failure = rejection + await _reject_guarded_turn(self.websocket, guarded, rejection) + return None async def _send_turn_refusal(self, refusal: RateLimitError) -> None: try: @@ -3052,7 +3090,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, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, @@ -3076,7 +3114,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 + self.turn_admission: ResponsesWsTurnAdmission | 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} @@ -3473,7 +3511,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ async def _process_response_create( - self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, raw_message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, @@ -3520,11 +3558,11 @@ class ManagedResponsesWebSocketHandler: 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: + verdict: Final = await _admit_frame_turn(turn_admission, raw_message) + if isinstance(verdict, _RefusedTurn): if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) + self._failure = verdict.rejection + await _reject_guarded_turn(self.websocket, guarded, verdict.rejection) return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} @@ -3555,6 +3593,11 @@ class ManagedResponsesWebSocketHandler: ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) + turn_litellm_metadata: Final = call_kwargs.get("litellm_metadata") + call_kwargs["litellm_metadata"] = { + **(turn_litellm_metadata if _is_json_object(turn_litellm_metadata) else {}), + **verdict.reservation, + } try: terminal_event: Final = await self._stream_and_forward(model, call_kwargs) diff --git a/litellm/router.py b/litellm/router.py index af2181cedb5..233e7581a40 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -378,6 +378,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor EncryptedContentAffinityCheck, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_admit_io_token_turn, async_io_token_refund_failure, build_io_token_rate_limit_headers, deployment_has_io_token_limits, @@ -5113,6 +5114,7 @@ class Router(DeprecatedRouterMembers): kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: """ 5 jobs: @@ -5202,7 +5204,12 @@ class Router(DeprecatedRouterMembers): # here before it's wiped below, instead of relying on that attempt's # (possibly still-pending) failure event to do it. refund_stale_reservation_before_retry(self._core_cache, kwargs) - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) + set_io_token_rate_limit_request_kwargs( + kwargs, + store_in_context=deployment_has_io_token_limits(deployment), + opens_responses_socket=opens_responses_socket, + metadata_channel=metadata_variable_name, + ) kwargs[metadata_variable_name].setdefault( ROUTING_REQUEST_TAGS_METADATA_KEY, @@ -6615,6 +6622,7 @@ class Router(DeprecatedRouterMembers): kwargs=kwargs, function_name=function_name, splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + opens_responses_socket=handler_name == "_aresponses_websocket", ) data: Final = deployment["litellm_params"].copy() @@ -10465,9 +10473,20 @@ class Router(DeprecatedRouterMembers): return deployment if current is None else current.model_dump(exclude_none=True) async def async_admit_deployment_turn( - self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> None: + self, + model_group: str, + opened_on: Mapping[str, object], + request_kwargs: Mapping[str, object], + ) -> Mapping[str, object]: deployment: Final = self._deployment_as_configured_now(opened_on) + rate_limit_check: Final = next( + (callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), None + ) + io_token_admission: Final[Mapping[str, object]] = ( + MappingProxyType({}) + if rate_limit_check is None + else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) + ) budget_limiter: Final = self._core_get_router_deployment_budget_limiter() if budget_limiter is not None: try: @@ -10488,6 +10507,7 @@ class Router(DeprecatedRouterMembers): 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 ) + return io_token_admission async def _admit_turn_through_routing_strategy( self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] 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 7580c7dadd3..80e1f3fe29c 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 @@ -14,6 +14,9 @@ from __future__ import annotations import contextlib import contextvars from collections.abc import Mapping, MutableMapping +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -23,7 +26,15 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.internal_call_metadata import ( + IO_TOKEN_RESERVATION_KEYS, + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, +) from litellm.litellm_core_utils.token_counter import deployment_openai_model, 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 @@ -39,18 +50,27 @@ else: RoutingArgsTTL: Final = 60 -_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, object] | None]] = contextvars.ContextVar( - "io_token_rate_limit_request_kwargs", + +@dataclass(frozen=True, slots=True) +class _IoTokenRateLimitRequest: + kwargs: dict[str, object] | None + opens_responses_socket: bool + metadata_channel: str + + +_io_token_rate_limit_request: Final[contextvars.ContextVar[_IoTokenRateLimitRequest | None]] = contextvars.ContextVar( + "io_token_rate_limit_request", default=None, ) -ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" -OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" -ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" -OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" - -def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, store_in_context: bool = True) -> None: +def set_io_token_rate_limit_request_kwargs( + kwargs: dict[str, object] | None, + store_in_context: bool = True, + *, + opens_responses_socket: bool = False, + metadata_channel: str, +) -> None: # The reservation sentinels are server-only, but `metadata` is caller # controlled on proxy requests. Strip any client-supplied copies here (this # runs before the router stashes its own reservation) so a forged @@ -64,11 +84,14 @@ def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, sto # every other deployment overwrite the slot with None instead of the # kwargs; overwriting (rather than skipping) also releases a previous # request's kwargs when a context is reused. - _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None) + _io_token_rate_limit_request.set( + _IoTokenRateLimitRequest(kwargs, opens_responses_socket, metadata_channel) if store_in_context else None + ) def get_io_token_rate_limit_request_kwargs() -> dict[str, object] | None: - return _io_token_rate_limit_request_kwargs.get() + request: Final = _io_token_rate_limit_request.get() + return None if request is None else request.kwargs def seconds_until_minute_reset() -> int: @@ -224,27 +247,32 @@ def _otpm_reservation( return default if from_deployment is None else min(from_deployment, default) +def _usage_field(container: object, name: str) -> object: + if is_str_keyed_mapping(container): + return container.get(name) + return getattr(container, name, None) + + +def _cached_prompt_tokens(usage: object) -> int: + reported: Final = ( + _usage_field(_usage_field(usage, "prompt_tokens_details"), "cached_tokens"), + _usage_field(_usage_field(usage, "input_tokens_details"), "cached_tokens"), + _usage_field(usage, "cache_read_input_tokens"), + ) + return next((int(tokens) for tokens in reported if isinstance(tokens, int | float) and tokens), 0) + + def _get_usage_tokens(usage: object) -> tuple[int, int, int]: if usage is None: return 0, 0, 0 if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"): prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0) completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0) - cached = 0 - details: object = getattr(usage, "prompt_tokens_details", None) - if details is not None: - cached = int(getattr(details, "cached_tokens", 0) or 0) - if not cached: - cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) if isinstance(usage, dict): prompt = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) completion = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) - details = usage.get("prompt_tokens_details") or {} - cached = int(details.get("cached_tokens", 0) or 0) if isinstance(details, dict) else 0 - if not cached: - cached = int(usage.get("cache_read_input_tokens") or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) return 0, 0, 0 @@ -298,19 +326,36 @@ def _resolve_reconcile_usage_tokens( if isinstance(standard_logging_object, dict): prompt_tokens = int(standard_logging_object.get("prompt_tokens") or 0) completion_tokens = int(standard_logging_object.get("completion_tokens") or 0) - cached_tokens = 0 - metadata: Final = standard_logging_object.get("metadata") - if isinstance(metadata, dict): - cached_tokens = int(metadata.get("cache_read_input_tokens") or 0) + logged_usage: Final = _usage_field(standard_logging_object.get("metadata"), "usage_object") + logged_cached_tokens: Final = _cached_prompt_tokens(logged_usage) # Same rationale as _usage_is_present: a bare total_tokens with no # prompt/completion breakdown can't be split, so it isn't treated # as resolved usage - the reservation is kept instead of refunded. if prompt_tokens or completion_tokens: - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + return max(0, prompt_tokens - logged_cached_tokens), completion_tokens, True return 0, 0, False +def _reservation( + itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None +) -> Mapping[str, object]: + return MappingProxyType( + { + ITPM_RESERVED_KEY: itpm_reserved, + OTPM_RESERVED_KEY: otpm_reserved, + ITPM_CACHE_KEY: itpm_cache_key, + OTPM_CACHE_KEY: otpm_cache_key, + } + ) + + +def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, reservation: Mapping[str, object]) -> None: + metadata: Final = None if request is None or not request.kwargs else request.kwargs.get(request.metadata_channel) + if isinstance(metadata, dict): + metadata.update(reservation) + + def _key_for_this_minute(reserved_key: str) -> str: counter, _, _ = reserved_key.rpartition(":") return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" @@ -326,30 +371,6 @@ def _reconcile_increments( ) -def _stash_reservation_in_metadata( - request_kwargs: dict[str, object] | None, - *, - itpm_reserved: int, - otpm_reserved: int, - itpm_cache_key: str | None, - otpm_cache_key: str | None, -) -> None: - if not request_kwargs: - return - reservation: Final = { - ITPM_RESERVED_KEY: itpm_reserved, - OTPM_RESERVED_KEY: otpm_reserved, - ITPM_CACHE_KEY: itpm_cache_key, - OTPM_CACHE_KEY: otpm_cache_key, - } - for channel in ("metadata", "litellm_metadata"): - existing = request_kwargs.get(channel) - if isinstance(existing, dict): - existing.update(reservation) - elif channel == "metadata": - request_kwargs[channel] = dict(reservation) - - def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]: itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY) otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY) @@ -399,7 +420,7 @@ def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None: for channel in _reservation_channels(kwargs): channel_dict = _as_mutable_mapping(channel) if channel_dict is not None: - for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY): + for key in IO_TOKEN_RESERVATION_KEYS: channel_dict.pop(key, None) @@ -499,7 +520,8 @@ def io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -546,11 +568,13 @@ def io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment @@ -565,7 +589,14 @@ async def async_io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs + if request is not None and request.opens_responses_socket: + _stash_reservation_in_metadata( + request, + await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span), + ) + return deployment _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -618,15 +649,65 @@ async def async_io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment +async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: + if redis_cache is None: + return None + with contextlib.suppress(RedisCircuitBreakerOpenError): + return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span) + return None + + +@with_service_target(ROUTER_USAGE_TARGET) +async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None: + shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span) + if shared is not None: + return shared + return await dual_cache.async_get_cache(key=key, local_only=True) + + +async def _raise_if_at_limit( + dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None +) -> None: + if limit is None: + return + usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span) + if usage is not None and usage >= limit: + raise _rate_limit_error(limit_label, limit, usage) + + +@with_service_target(ROUTER_USAGE_TARGET) +async def async_admit_io_token_turn( + dual_cache: DualCache, + deployment: Mapping[str, object], + now: datetime, + parent_otel_span: Span | None = None, +) -> Mapping[str, object]: + configured: Final = dict(deployment) + itpm_limit, otpm_limit = get_deployment_io_token_limits(configured) + if itpm_limit is None and otpm_limit is None: + return MappingProxyType({}) + cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M")) + if cache_keys is None: + return MappingProxyType({}) + itpm_key, otpm_key = cache_keys + await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span) + await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span) + return _reservation( + 0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None + ) + + @with_service_target(ROUTER_USAGE_TARGET) def io_token_reconcile_success( dual_cache: DualCache, diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index 87b56a92979..7b0e14d044f 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -518,9 +518,7 @@ def test_a_tpm_limit_lowered_while_the_socket_is_open_applies_to_its_next_turn( 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: +def test_a_socket_on_an_otpm_limited_deployment_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]) @@ -536,6 +534,115 @@ def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_t assert deployment.forwarded_turns() == 5 +IO_TOKEN_LIMIT: Final = 10_000 +TOKENS_PER_TURN: Final = MappingProxyType({"itpm": USAGE["input_tokens"], "otpm": USAGE["output_tokens"]}) + + +def _io_tokens_counted(identity: str, kind: str, minute: str) -> int: + with _redis() as redis: + return sum(int(redis.get(key) or 0) for key in redis.scan_iter(f"global_router:{identity}:*:{kind}:{minute}")) + + +def _add_io_tokens(deployment: Deployment, kind: str, minute: str, tokens: int) -> None: + with _redis() as redis: + redis.incrby(f"global_router:{deployment.identity}:openai/{PROVIDER_MODEL}:{kind}:{minute}", tokens) + + +async def _say_until_counted( + connection: websockets.ClientConnection, deployment: Deployment, counted: Callable[[], float], expected: int +) -> tuple[dict[str, JsonValue], ...]: + received: Final = await _say(connection, deployment.name) + await _until(counted, lambda tokens: tokens == expected) + return received + + +@pytest.mark.parametrize("kind", ("otpm", "itpm")) +def test_a_turn_is_refused_once_earlier_turns_fill_the_deployment_io_token_limit( + limited: Gateway, cert: tuple[Path, Path], kind: str +) -> None: + per_turn: Final = TOKENS_PER_TURN[kind] + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, **{kind: IO_TOKEN_LIMIT}) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + counted: Final = partial(_io_tokens_counted, deployment.identity, kind, minute) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + first: Final = await _say_until_counted(connection, deployment, counted, per_turn) + await asyncio.to_thread(_add_io_tokens, deployment, kind, minute, IO_TOKEN_LIMIT - 5 * per_turn) + answered: Final = tuple( + [ + await _say_until_counted( + connection, deployment, counted, IO_TOKEN_LIMIT - (5 - number) * per_turn + ) + for number in range(2, 6) + ] + ) + return (first, *answered, await _say(connection, deployment.name)) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns[:5]] == ["completed"] * 5 + _assert_refused(turns[5]) + assert deployment.forwarded_turns() == 5 + assert held(counted, lambda tokens: tokens == IO_TOKEN_LIMIT, holding=3, seconds=20) == IO_TOKEN_LIMIT + + +def test_a_socket_at_half_its_deployment_otpm_serves_ten_turns(limited: Gateway, cert: tuple[Path, Path]) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, otpm=IO_TOKEN_LIMIT) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + _add_io_tokens(deployment, "otpm", minute, IO_TOKEN_LIMIT // 2) + + 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(10)]) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns] == ["completed"] * 10 + assert deployment.forwarded_turns() == 10 + expected: Final = IO_TOKEN_LIMIT // 2 + 10 * USAGE["output_tokens"] + counted: Final = partial(_io_tokens_counted, deployment.identity, "otpm", minute) + assert held(counted, lambda tokens: tokens == expected, holding=3, seconds=20) == expected + + +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_every_turn_counts_its_actual_tokens_toward_the_deployment_itpm_and_otpm( + limited: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + limits: Final = {"itpm": 1_000_000, "otpm": 1_000_000} + deployment: Final = _native(scenario, peer, **limits) if relay == "native" else _managed(scenario, **limits) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _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(3)]) + + turns: Final = asyncio.run(session()) + + usages: Final = tuple(object_value(completed(received)["usage"]) for received in turns) + expected: Final = { + "itpm": sum(int(str(usage["input_tokens"])) for usage in usages), + "otpm": sum(int(str(usage["output_tokens"])) for usage in usages), + } + assert any(tokens > 0 for tokens in expected.values()), usages + counted: Final = { + kind: held( + partial(_io_tokens_counted, deployment.identity, kind, minute), + lambda tokens, kind=kind: tokens == expected[kind], + holding=3, + seconds=20, + ) + for kind in expected + } + assert counted == expected + + 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: diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index ab52e6d29c3..d0e26798090 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,15 +8,21 @@ Tests that: """ import asyncio +import itertools import json -from collections.abc import Mapping -from typing import Final -from unittest.mock import MagicMock +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone +from functools import partial +from typing import TYPE_CHECKING, Final +from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig +from litellm.llms.custom_llm import CustomLLM from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) @@ -42,6 +48,13 @@ from litellm.llms.volcengine.responses.transformation import ( VolcEngineResponsesAPIConfig, ) from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.utils import GenericStreamingChunk + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck class TestResponsesAPIWebSocketSupport: @@ -1591,7 +1604,7 @@ def _admission_refusing_the_next_turn(): from litellm.exceptions import RateLimitError - return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), {}]) class _ClientSocket: @@ -1630,7 +1643,7 @@ class _StalledClientSocket(_ClientSocket): def _admission_crashing_the_next_turn(): from unittest.mock import AsyncMock - return AsyncMock(side_effect=[ValueError("custom check failed"), None]) + return AsyncMock(side_effect=[ValueError("custom check failed"), {}]) def _guard_blocking(word: str): @@ -1655,6 +1668,229 @@ def _guard_recording(bill_refusal: MagicMock): return guard +_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34" +_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35" +_OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc) +_WS_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "otpm": 4096}, + "model_info": {"id": "ws-dep"}, +} +_OTHER_USAGE: Final = 10 +_OUTPUT_TOKENS_PER_TURN: Final = 2 +_UPSTREAM_ERROR: Final = json.dumps({"type": "error", "error": {"type": "server_error", "message": "upstream failed"}}) + + +def _turn_completed() -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_turn", + "status": "completed", + "output": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + }, + } + ) + + +async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]: + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60) + await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60) + return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT) + + +def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: finished_at + ) + + +def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission": + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT) + + +_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34" +_ITPM_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000}, + "model_info": {"id": "ws-dep"}, +} +_TURN_INPUT_TOKENS: Final = 1_000 + + +def _turn_completed_reading_cache(cached_tokens: int) -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_cached_turn", + "status": "completed", + "output": [], + "usage": { + "input_tokens": _TURN_INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": cached_tokens}, + "output_tokens": _OUTPUT_TOKENS_PER_TURN, + "total_tokens": _TURN_INPUT_TOKENS + _OUTPUT_TOKENS_PER_TURN, + }, + }, + } + ) + + +def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes + + connection: Final = Logging( + model="gpt-5-mini", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + dynamic_async_success_callbacks=[rate_limits], + dynamic_async_failure_callbacks=[rate_limits], + ) + connection.update_environment_variables( + litellm_params={"metadata": dict(opened_with), CallTypes.aresponses_websocket.value: True}, + optional_params={}, + ) + return connection + + +def _provider_turn_stream(_request: httpx.Request) -> httpx.Response: + completed: Final = { + "id": "resp_turn", + "object": "response", + "created_at": 1767271505, + "status": "completed", + "model": "gpt-5-mini", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + } + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': completed})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + + +class _ProviderResponsesSocket: + """A provider's native Responses socket that answers each ``response.create`` with a completed turn.""" + + def __init__(self) -> None: + self.sent: Final[list[dict[str, object]]] = [] + self._replies: Final[asyncio.Queue[str]] = asyncio.Queue() + + def connect(self, _url: str, **_kwargs: object) -> "_ProviderResponsesSocket": + return self + + async def __aenter__(self) -> "_ProviderResponsesSocket": + return self + + async def __aexit__(self, *_exc_info: object) -> None: + return None + + async def send(self, message: str) -> None: + self.sent.append(json.loads(message)) + await self._replies.put(_turn_completed()) + + async def recv(self, decode: bool = False) -> str: + return await self._replies.get() + + async def close(self) -> None: + return None + + +class _TwoOutputTokenTurns(CustomLLM): + def __init__(self, failing_turns: frozenset[int] = frozenset()) -> None: + super().__init__() + self._failing_turns: Final = failing_turns + self._turns: Final = itertools.count(1) + + async def astreaming(self, *args: object, **kwargs: object) -> AsyncIterator[GenericStreamingChunk]: + if next(self._turns) in self._failing_turns: + raise RuntimeError("upstream failed") + yield GenericStreamingChunk( + text="hi", + is_finished=True, + finish_reason="stop", + index=0, + tool_use=None, + usage={"prompt_tokens": 5, "completion_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + ) + + +async def _run_managed_relay( + monkeypatch: pytest.MonkeyPatch, + provider: CustomLLM, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], +) -> _ClientSocket: + import litellm + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "ws-turn-llm", "custom_handler": provider}]) + return await _relay_turns_to(monkeypatch, "ws-turn-llm/turns", rate_limits, frames, turn_admission, opened_with) + + +async def _relay_turns_to( + monkeypatch: pytest.MonkeyPatch, + model: str, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], + **connection: object, +) -> _ClientSocket: + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler, drain_logging_tasks + + monkeypatch.setattr(litellm, "callbacks", [rate_limits]) + first, *later = frames + websocket: Final = _ClientSocket(*later) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model=model, + logging_obj=Logging( + model=model, + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1", **opened_with}, + first_message=first, + turn_admission=turn_admission, + **connection, + ) + + await handler.run() + await drain_logging_tasks() + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + return websocket + + def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] @@ -1722,6 +1958,195 @@ class TestWebSocketTurnAdmission: assert isinstance(failure, RateLimitError) assert "crossed budget" in str(failure) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("first_turn_end", "first_turn_tokens"), + [ + pytest.param(_turn_completed(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(_UPSTREAM_ERROR, 0, id="first_turn_failed"), + ], + ) + async def test_native_relay_counts_every_turn_on_the_minute_it_finishes( + self, first_turn_end: str, first_turn_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + first, second, third = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[first_turn_end, _turn_completed(), _turn_completed(), Exception("stop")] + ) + handler: Final = _make_streaming( + websocket=_ClientSocket(second, third), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + turn_admission=_later_turns_admitted_through(dual_cache), + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("cached_tokens", "counted_input_tokens"), + [ + pytest.param(800, _TURN_INPUT_TOKENS - 800, id="cached_input"), + pytest.param(0, _TURN_INPUT_TOKENS, id="uncached_input"), + ], + ) + async def test_native_relay_leaves_a_turn_s_cached_input_out_of_the_deployment_s_itpm( + self, cached_tokens: int, counted_input_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _OPENED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT) + first, _, _ = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=[_turn_completed_reading_cache(cached_tokens), Exception("stop")]) + handler: Final = _make_streaming( + websocket=_ClientSocket(), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failing_turns", "first_turn_tokens"), + [ + pytest.param(frozenset(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(frozenset({1}), 0, id="first_turn_failed"), + ], + ) + async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes( + self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int + ) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + + await _run_managed_relay( + monkeypatch, + _TwoOutputTokenTurns(failing_turns), + ModelRateLimitingCheck(dual_cache), + _turn_frames(), + _later_turns_admitted_through(dual_cache), + opened_with, + ) + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + async def test_managed_relay_sends_the_provider_only_the_metadata_each_turn_carries( + self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + ) -> None: + import litellm + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + _turns_finish_at(monkeypatch, _TURNED_AT) + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream) + client_metadata: Final = {"session": "abc", "team": "search"} + dual_cache: Final = DualCache() + + await _relay_turns_to( + monkeypatch, + "openai/gpt-5-mini", + ModelRateLimitingCheck(dual_cache), + tuple( + json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) + for word in ("one", "two") + ), + _later_turns_admitted_through(dual_cache), + await _socket_opened_on(dual_cache), + api_key="sk-fake", + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [ + client_metadata, + client_metadata, + ] + assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN + + @pytest.mark.asyncio + async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from unittest.mock import patch + + from litellm import Router + + monkeypatch.setenv("LITELLM_RUST", "false") + provider: Final = _ProviderResponsesSocket() + client_metadata: Final = {"session": "abc", "team": "search"} + first, second, _ = _turn_frames() + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "itpm": 40_000, + "otpm": 4096, + }, + "model_info": {"id": "ws-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + with patch("websockets.connect", provider.connect): + await router._aresponses_websocket( + model="ws-io", + websocket=_ClientSocket(second), + first_message=first, + metadata=dict(client_metadata), + ) + + assert [frame.get("metadata") for frame in provider.sent] == [client_metadata, client_metadata] + @pytest.mark.asyncio async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock @@ -1875,7 +2300,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.bidirectional_forward() @@ -1904,7 +2329,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) relay: Final = asyncio.ensure_future(handler.bidirectional_forward()) await asyncio.wait_for(client.send_started.wait(), timeout=1) @@ -1955,7 +2380,7 @@ class TestWebSocketTurnAdmission: ), first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.run() diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index e5ff4531f84..819c6b9d4f5 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -4,14 +4,19 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). import asyncio from collections.abc import Awaitable, Callable, Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -22,6 +27,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( OTPM_RESERVED_KEY, _explicit_max_tokens, _reservation_value, + async_admit_io_token_turn, async_io_token_pre_call_check, async_io_token_reconcile_success, build_io_token_rate_limit_headers, @@ -35,6 +41,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_utc_datetime @@ -115,7 +122,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -173,7 +180,8 @@ class TestModelRateLimitingCheckIOTokens: ], "max_tokens": 10, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(litellm.RateLimitError) as exc_info: @@ -203,7 +211,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hi"}], "max_tokens": max_tokens, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -255,7 +264,8 @@ class TestModelRateLimitingCheckIOTokens: { "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -320,7 +330,7 @@ class TestModelRateLimitingCheckIOTokens: } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -363,15 +373,15 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") await check.async_pre_call_check(deployment) # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -392,7 +402,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 0, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + set_io_token_rate_limit_request_kwargs(zero_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -404,7 +414,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + set_io_token_rate_limit_request_kwargs(normal_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -432,7 +442,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "no-max-tokens", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) @@ -462,7 +472,7 @@ class TestModelRateLimitingCheckIOTokens: async def _attempt() -> dict[str, object] | None: request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") try: await check.async_pre_call_check(deployment) except litellm.RateLimitError: @@ -493,7 +503,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "roomy", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) @@ -507,7 +517,7 @@ class TestModelRateLimitingCheckIOTokens: async def _reserved(model_info: dict[str, str]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit}, "model_info": model_info, @@ -536,7 +546,7 @@ class TestModelRateLimitingCheckIOTokens: } async def _pre_call(request_kwargs: dict[str, object]) -> None: - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -563,12 +573,14 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100 - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}}) + set_io_token_rate_limit_request_kwargs( + {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, metadata_channel="metadata" + ) with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -585,7 +597,7 @@ class TestModelRateLimitingCheckIOTokens: with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}} for request_kwargs in (without_cap, with_cap): - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -602,7 +614,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int: request_kwargs: Final = {**request_shape, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check( { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap}, @@ -631,7 +643,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap}, "model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"}, @@ -667,7 +679,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") check.pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -803,7 +815,7 @@ class TestModelRateLimitingCheckIOTokens: # The retry proceeds to stash deployment B's own reservation on the # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) + set_io_token_rate_limit_request_kwargs(shared_kwargs, metadata_channel="metadata") itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b @@ -819,7 +831,7 @@ class TestModelRateLimitingCheckIOTokens: "litellm_metadata": {OTPM_RESERVED_KEY: 7}, "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, } - set_io_token_rate_limit_request_kwargs(forged) + set_io_token_rate_limit_request_kwargs(forged, metadata_channel="metadata") stored = get_io_token_rate_limit_request_kwargs() assert ITPM_RESERVED_KEY not in stored["metadata"] @@ -845,7 +857,7 @@ class TestModelRateLimitingCheckIOTokens: }, } # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) + set_io_token_rate_limit_request_kwargs(kwargs, metadata_channel="metadata") response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), @@ -880,7 +892,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hello world"}], "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(RuntimeError): @@ -945,7 +958,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") result = await async_io_token_pre_call_check(dual_cache, deployment) @@ -1147,7 +1160,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): with pytest.raises(litellm.RateLimitError): @@ -1495,15 +1508,15 @@ class TestContextSlotRetention: "messages": [{"role": "user", "content": "x" * 1000}], "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None assert ITPM_RESERVED_KEY not in kwargs["metadata"] assert ITPM_CACHE_KEY not in kwargs["metadata"] - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is kwargs - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None @pytest.mark.asyncio @@ -1516,7 +1529,7 @@ class TestContextSlotRetention: } ] ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("plain") assert deployment is not None @@ -1534,7 +1547,7 @@ class TestContextSlotRetention: ], optional_pre_call_checks=["enforce_model_rate_limits"], ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("limited") assert deployment is not None @@ -1551,7 +1564,7 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok "max_tokens": 5, "metadata": metadata, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await async_io_token_pre_call_check(DualCache(), deployment) return metadata[ITPM_RESERVED_KEY] @@ -1571,3 +1584,278 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok ) assert named == served + + +def _reply_using(prompt_tokens: int, completion_tokens: int) -> ModelResponse: + return ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +def _proxy_success_kwargs(reservation: Mapping[str, object]) -> dict[str, object]: + return { + "litellm_params": {"metadata": {**reservation, "user_api_key_hash": "hashed-key"}}, + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "openai/gpt-5-mini"}, + "metadata": {"user_api_key_hash": "hashed-key"}, + }, + } + + +async def _deliver_success( + check: ModelRateLimitingCheck, delivery: str, kwargs: dict[str, object], reply: object +) -> None: + if delivery == "async": + await check.async_log_success_event(kwargs, reply, None, None) + return + check.log_success_event(kwargs, reply, None, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deliveries", + [("async",), ("sync",), ("async", "async"), ("async", "sync"), ("sync", "async"), ("sync", "sync")], +) +async def test_success_settles_a_reservation_held_only_in_litellm_params_metadata_exactly_once( + deliveries: tuple[str, ...], +) -> None: + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" + otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=4096, ttl=60) + kwargs: Final = _proxy_success_kwargs( + {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key} + ) + + for delivery in deliveries: + await _deliver_success(check, delivery, kwargs, _reply_using(prompt_tokens=7, completion_tokens=3)) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + assert await dual_cache.async_get_cache(key=otpm_key) == 3 + + +_TURN_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_ACTUAL_TOKENS: Final = {"itpm": 5, "otpm": 2} + + +def _turn_deployment(kind: str) -> dict[str, object]: + return {"litellm_params": {"model": "openai/gpt-5-mini", kind: 100}, "model_info": {"id": f"ws-turn-{kind}"}} + + +def _turn_key(kind: str, minute: str = "12-34") -> str: + return f"global_router:ws-turn-{kind}:openai/gpt-5-mini:{kind}:{minute}" + + +@pytest.fixture +def settled_on_the_turn_s_minute(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: _TURN_AT + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [100, 101, 250]) +async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage={used}"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [0, 1, 99]) +async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + actual: Final = _ACTUAL_TOKENS + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + counter_key: Final = reservation[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] + assert isinstance(counter_key, str) + assert await dual_cache.async_get_cache(key=counter_key) == used + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=counter_key) == used + actual[kind] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + {"litellm_params": {"model": "openai/gpt-5-mini", "tpm": 1}, "model_info": {"id": "ws-turn-unlimited"}}, + {"litellm_params": {"model": "openai/gpt-5-mini", "otpm": 1}}, + ], +) +async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_settle( + deployment: dict[str, object], +) -> None: + assert await async_admit_io_token_turn(DualCache(), deployment, _TURN_AT) == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used_by_others", [0, 1, 99]) +async def test_a_socket_open_reserves_nothing_and_its_first_turn_counts_its_actual_tokens_on_its_own_minute( + kind: str, used_by_others: int +) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used_by_others, ttl=60) + socket_open: Final[dict[str, object]] = {"metadata": {}} + set_io_token_rate_limit_request_kwargs(socket_open, opens_responses_socket=True, metadata_channel="metadata") + + assert await async_io_token_pre_call_check(dual_cache, _turn_deployment(kind)) == _turn_deployment(kind) + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + + opened_with: Final = socket_open["metadata"] + assert isinstance(opened_with, dict) + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] + + +class _RedisCacheOnFakeServer(RedisCache): + def __init__(self, server: FakeServer) -> None: + super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) + self._fake_async_client: Final = FakeAsyncRedis(server=server) + + def init_async_client(self) -> FakeAsyncRedis: + return self._fake_async_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_refused_on_this_instance_s_count_when_reading_the_shared_counter_fails( + kind: str, +) -> None: + server: Final = FakeServer() + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnFakeServer(server)) + server.connected = False + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("delivery", ["async", "sync"]) +async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_that_minute( + kind: str, delivery: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", + lambda: datetime(2026, 1, 1, 12, 35, 10, tzinfo=timezone.utc), + ) + dual_cache: Final = DualCache() + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + await _deliver_success( + ModelRateLimitingCheck(dual_cache=dual_cache), + delivery, + _proxy_success_kwargs(reservation), + _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]), + ) + + assert [await dual_cache.async_get_cache(key=_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ + None, + _ACTUAL_TOKENS[kind], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_this_instance_last_read( + kind: str, +) -> None: + redis_cache: Final = _RedisCacheOnFakeServer(FakeServer()) + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + await redis_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "responses_reply", + [ + pytest.param( + ResponsesAPIResponse( + id="resp_cached", + created_at=0, + output=[], + usage=ResponseAPIUsage( + input_tokens=50, + output_tokens=5, + total_tokens=55, + input_tokens_details=InputTokensDetails(cached_tokens=30), + ), + ), + id="usage_object", + ), + pytest.param( + { + "usage": { + "input_tokens": 50, + "output_tokens": 5, + "total_tokens": 55, + "input_tokens_details": {"cached_tokens": 30}, + } + }, + id="usage_dict", + ), + ], +) +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +async def test_a_turn_counts_only_its_uncached_responses_input_toward_itpm(responses_reply: object) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key("itpm"), value=10, ttl=60) + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment("itpm"), _TURN_AT) + await async_io_token_reconcile_success(dual_cache, _proxy_success_kwargs(reservation), responses_reply) + + assert await dual_cache.async_get_cache(key=_turn_key("itpm")) == 10 + 50 - 30 + + +def _shared_counter_answering(answer: int | Exception) -> DualCache: + redis_cache: Final = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=[answer]) + return DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: str) -> None: + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(_shared_counter_answering(100), _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared_counter_is_unreachable( + kind: str, +) -> None: + dual_cache: Final = _shared_counter_answering(RedisCircuitBreakerOpenError("circuit open")) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 1ac44b5976d..a6436e47507 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2200,7 +2200,11 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ) def inject_alias_into_kwargs( - deployment: dict, kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False + deployment: dict, + kwargs: dict, + function_name: str | None = None, + splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: # Simulate the alias leaking into kwargs (as happens when # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5f1b6032ca8..5f31ded626d 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,13 +4,15 @@ import gc import json import os import weakref -from collections.abc import Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm from litellm import Router @@ -538,14 +540,14 @@ 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]: +async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> 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) + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call, **request) return routed @@ -715,6 +717,363 @@ async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): await admit_turn() +def _io_limited_router(kind: str, **litellm_params: object) -> Router: + return Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/ws-io-unpriced", + "api_key": "sk-fake", + kind: 10_000, + **litellm_params, + }, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + +def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str: + return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" + + +async def _count_minutes_tokens( + router: Router, kind: str, tokens: int, minute: datetime, model: str = "openai/ws-io-unpriced" +) -> None: + await router.cache.async_increment_cache(key=_minutes_counter_key(kind, minute, model), value=tokens, ttl=120) + + +async def _settle_turn( + router: Router, reservation: Mapping[str, object], input_tokens: int, output_tokens: int +) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + rate_limits: Final = next( + check for check in router.optional_callbacks or () if isinstance(check, ModelRateLimitingCheck) + ) + await rate_limits.async_log_success_event( + {"litellm_params": get_litellm_params(litellm_metadata=dict(reservation))}, + ResponsesWebSocketTokenUsageProcessor.create_logging_object( + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}, + } + ] + ), + None, + None, + ) + + +_LONG_OUTPUT_MODEL: Final = "openai/gpt-5.4-mini" + + +def _long_output_router() -> tuple[Router, Mapping[str, int]]: + max_output: Final = litellm.get_model_info(_LONG_OUTPUT_MODEL)["max_output_tokens"] + assert isinstance(max_output, int) + limits: Final = {"itpm": 10_000, "otpm": max_output // 2} + return _io_limited_router("itpm", model=_LONG_OUTPUT_MODEL, otpm=limits["otpm"]), limits + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_opens_without_counting_tokens_until_the_minute_reaches_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router, limits = _long_output_router() + + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [None, None] + await _count_minutes_tokens(router, kind, limits[kind], router_minute_pinned, _LONG_OUTPUT_MODEL) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit={limits[kind]}, current usage="): + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + +@pytest.mark.asyncio +async def test_responses_websocket_first_turn_counts_its_actual_tokens_toward_deployment_itpm_and_otpm( + router_minute_pinned: datetime, +): + router, _ = _long_output_router() + socket_metadata: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()))[ + "litellm_metadata" + ] + assert isinstance(socket_metadata, dict) + + await _settle_turn(router, socket_metadata, input_tokens=20, output_tokens=15) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [20, 15] + + +_ADMITTED_LATE_IN_THE_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 50, tzinfo=timezone.utc) +_FINISHED_IN_THE_NEXT_MINUTE: Final = datetime(2026, 1, 1, 12, 1, 15, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("turn", ["first", "later"]) +async def test_responses_websocket_turn_finishing_in_the_next_minute_counts_its_tokens_toward_that_minute( + kind: str, turn: str, monkeypatch: pytest.MonkeyPatch +): + _pin_router_clock(monkeypatch, _ADMITTED_LATE_IN_THE_MINUTE) + router: Final = _io_limited_router(kind) + socket: Final = await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + admit_turn: Final = socket["turn_admission"] + assert callable(admit_turn) + reservation: Final = socket["litellm_metadata"] if turn == "first" else await admit_turn() + assert isinstance(reservation, Mapping) + _pin_router_clock(monkeypatch, _FINISHED_IN_THE_NEXT_MINUTE) + + await _settle_turn(router, reservation, input_tokens=10_000, output_tokens=10_000) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage=10000"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [{}, {"websocket": True}], ids=["plain", "websocket_key"]) +async def test_messages_call_reserves_max_tokens_against_deployment_otpm_whatever_websocket_key_its_body_carries( + body: Mapping[str, object], router_minute_pinned: datetime +): + router: Final = _io_limited_router("otpm") + + for _ in range(3): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + assert await router.cache.async_get_cache(key=_minutes_counter_key("otpm", router_minute_pinned)) == 9_000 + with pytest.raises(litellm.RateLimitError, match="OTPM limit=10000, current usage=12000"): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + +async def _realtime_opens_until_refused(monkeypatch: pytest.MonkeyPatch, **socket: object) -> tuple[int, str]: + monkeypatch.setattr(litellm, "callbacks", []) + router: Final = _io_limited_router("otpm", otpm=50) + for opened in range(51): + try: + await _kwargs_routed_to(router, "ws-io", "_arealtime", **socket) + except litellm.RateLimitError as refused: + return opened, str(refused) + pytest.fail("realtime opens never reached the deployment's otpm") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_realtime_open_reserves_and_is_refused_at_deployment_otpm_like_an_open_without_a_socket( + monkeypatch: pytest.MonkeyPatch, +): + with_socket: Final = await _realtime_opens_until_refused(monkeypatch, websocket=object()) + + assert "OTPM limit=50, current usage=" in with_socket[1] + assert with_socket == await _realtime_opens_until_refused(monkeypatch) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_is_refused_once_the_minutes_tokens_reach_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage="): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_refused_for_tokens_leaves_the_daily_request_limit_unused( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind, rpd=2) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000"): + await admit_turn() + await _count_minutes_tokens(router, kind, -10_000, router_minute_pinned) + + await admit_turn() + + with pytest.raises(litellm.RateLimitError, match="RPD"): + await admit_turn() + + +_PROVIDER_RESPONSE: Final = { + "id": "resp_io", + "object": "response", + "created_at": 1767268830, + "status": "completed", + "model": "ws-io-unpriced", + "output": [ + { + "type": "message", + "id": "msg_io", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +async def test_responses_call_sends_the_provider_only_the_client_metadata_and_settles_deployment_io_tokens( + client_metadata: Mapping[str, str] | None, + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + router: Final = _io_limited_router("itpm", otpm=10_000) + + await router.aresponses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +def test_sync_responses_call_sends_the_provider_only_the_client_metadata( + client_metadata: Mapping[str, str] | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + + _io_limited_router("itpm", otpm=10_000).responses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + + +_ANTHROPIC_PROVIDER_RESPONSE: Final = { + "id": "msg_io", + "type": "message", + "role": "assistant", + "model": "ws-io-unpriced", + "content": [{"type": "text", "text": "Done."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +async def _responses_call(router: Router) -> object: + return await router.aresponses(model="ws-io", input="hi", max_output_tokens=50) + + +async def _messages_call(router: Router) -> object: + return await router.aanthropic_messages(model="ws-io", messages=[{"role": "user", "content": "hi"}], max_tokens=50) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "provider_url", "provider_response", "call"), + [ + pytest.param( + "openai/ws-io-unpriced", + "https://api.openai.com/v1/responses", + _PROVIDER_RESPONSE, + _responses_call, + id="responses", + ), + pytest.param( + "anthropic/ws-io-unpriced", + "https://api.anthropic.com/v1/messages", + _ANTHROPIC_PROVIDER_RESPONSE, + _messages_call, + id="messages", + ), + ], +) +async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attempt_s_tokens( + model: str, + provider_url: str, + provider_response: Mapping[str, object], + call: Callable[[Router], Awaitable[object]], + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post(provider_url).mock( + side_effect=[ + httpx.Response(500, json={"error": {"message": "upstream failed", "type": "server_error"}}), + httpx.Response(200, json=provider_response), + ] + ) + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": {"model": model, "api_key": "sk-fake", "itpm": 10_000, "otpm": 10_000}, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=1, + retry_after=0, + ) + + await call(router) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert provider.call_count == 2 + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ @@ -1194,10 +1553,17 @@ def _rpm_tpm_router(model_id: str) -> Router: ) +def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: + monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned + ) + + @pytest.fixture def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> datetime: pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) - monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + _pin_router_clock(monkeypatch, pinned) return pinned