scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge branch 'main' into litellm_model_info_wildcard_routed_price

Browse source
This commit is contained in:
mateo-berri 2026-10-30 08:17:44 -07:00
commit a59e6ff331
32 changed files with 4550 additions and 307 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -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.

22
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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]:

35
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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:

47
litellm/proxy/auth/auth_checks.py
View file

@ -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
)

4
litellm/proxy/auth/model_checks.py
View file

@ -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(

45
litellm/proxy/common_request_processing.py
View file

@ -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

8
litellm/proxy/common_utils/model_listing_utils.py
View file

@ -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]

427
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -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,

4
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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,

3
litellm/proxy/utils.py
View file

@ -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,):

6
litellm/responses/main.py
View file

@ -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:

101
litellm/responses/mcp/chat_completions_handler.py
View file

@ -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:

281
litellm/responses/mcp/chat_stream_round_billing.py Normal file
View file

@ -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)

72
litellm/responses/mcp/parked_stream_dispatch.py Normal file
View file

@ -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

58
litellm/responses/mcp/stream_round_billing.py
View file

@ -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:

103
litellm/responses/streaming_iterator.py
View file

@ -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)

26
litellm/router.py
View file

@ -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]

209
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

@ -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,

18
tests/integration/_support/hook_count_probe.py
View file

@ -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()

261
tests/integration/mcp/test_mcp_llm_endpoints.py
View file

@ -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

113
tests/integration/routing/test_responses_websocket_turn_admission.py
View file

@ -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:

10
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -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"])

720
tests/unit/proxy/health_endpoints/test_health_endpoints.py
View file

@ -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

28
tests/unit/proxy/test_common_request_processing.py
View file

@ -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

2
tests/unit/proxy/test_health_check_functions.py
View file

@ -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,

1050
tests/unit/responses/mcp/test_chat_completions_handler.py
View file

File diff suppressed because it is too large Load diff

17
tests/unit/responses/mcp/test_parked_stream_dispatch.py Normal file
View file

@ -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

441
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -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()

356
tests/unit/test_router/test_io_token_rate_limits.py
View file

@ -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)

6
tests/unit/test_router/test_router.py
View file

@ -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

374
tests/unit/test_router/test_router_helper_utils.py
View file

@ -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