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/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/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/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/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 49f1e37ce06..4a182528cce 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -14,7 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth -from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, served_model +from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, dispatched_model from litellm.proxy.common_utils.token_count_settings import cost_map_entry from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.responses.main import strip_responses_routing_prefix @@ -642,7 +642,7 @@ def wildcard_deployment_priced_model( llm_router: Router, deployment: WildcardDeployment, listed_name: str, caller_aliases: CallerAliases ) -> str | RoutedElsewhere: registry: Final = llm_router.registry() - served_name: Final = served_model(listed_name, _without_deleted_team_targets(caller_aliases, registry)) + served_name: Final = dispatched_model(listed_name, _without_deleted_team_targets(caller_aliases, registry)) routed_name: Final = routed_model_name(llm_router, served_name) priced: Final = wildcard_listed_priced_model(deployment.model_name_pattern, deployment.litellm_model, routed_name) candidates: Final = registry.names( 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/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 44ee6cd7a0d..a5d2588b35c 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -220,7 +220,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( @@ -288,11 +288,15 @@ def _dispatch(model_id: str, layers: Sequence[_AliasLayer]) -> tuple[str, bool]: return reduce(_through_layer, layers, (model_id, False)) -def served_model(model_id: str, aliases: CallerAliases) -> str: +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/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index a161ea53b98..49ad6dfd48d 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -84,7 +84,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.proxy.common_utils.model_listing_utils import CallerAliases, served_model +from litellm.proxy.common_utils.model_listing_utils import CallerAliases, dispatched_model from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient @@ -1286,7 +1286,7 @@ def _live_cost_map_model( model_id: Final = db_model.model_info.id if llm_router is None or model_id is None: return wildcard_listed_priced_model( - db_model.model_name, stored_model, served_model(listed_name, caller_aliases) + db_model.model_name, stored_model, dispatched_model(listed_name, caller_aliases) ) priced: Final = wildcard_deployment_priced_model( llm_router, 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/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/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/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/_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/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/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_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/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, 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 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