diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index b80f78a50c1..bc2cb7c9fdc 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -214,26 +214,33 @@ def _message_has_cache_control(message: Mapping[str, object]) -> bool: return False +def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + last_breakpoint: Final = max( + (index for index, msg in enumerate(messages) if _message_has_cache_control(msg)), + default=-1, + ) + return tuple(range(last_breakpoint + 1)) + + def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]: """ Return indices of messages that must never be compressed: - All system messages - The last user message - The last assistant message - - Any message carrying an Anthropic cache_control breakpoint + - Every message up to and including the last one carrying an Anthropic cache_control breakpoint The last user message is what the model is being asked to act on right now, so compressing it replaces the live instruction with a marker. Compression guardrails share this policy; see the Headroom guardrail. A cache_control - breakpoint pins the provider's prompt-cache prefix to that row's exact - bytes, so rewriting a marked row anywhere in history turns the next - request's cache read into a cache write. + breakpoint pins the provider's prompt-cache prefix to the exact bytes of every + row up to it, so rewriting any row inside that prefix turns the next request's + cache read into a cache write. """ system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant") - cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg)) - return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices)) + return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages))) def _combine_scores( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9ba9fd082f3..4ddb9ce5b8e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -644,6 +644,24 @@ class Logging(LiteLLMLoggingBaseClass): """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable + def add_dynamic_callback(self, callback: CustomLogger) -> None: + self.dynamic_input_callbacks = self._with_dynamic_callback(self.dynamic_input_callbacks, callback) + self.dynamic_success_callbacks = self._with_dynamic_callback(self.dynamic_success_callbacks, callback) + self.dynamic_async_success_callbacks = self._with_dynamic_callback( + self.dynamic_async_success_callbacks, callback + ) + self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, callback) + self.dynamic_async_failure_callbacks = self._with_dynamic_callback( + self.dynamic_async_failure_callbacks, callback + ) + + @staticmethod + def _with_dynamic_callback( + callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger + ) -> list[str | Callable | CustomLogger]: + existing: Final = tuple(callbacks or ()) + return [*existing, *(() if callback in existing else (callback,))] + def process_dynamic_callbacks(self): """ Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9f85c3d1894..993f705a918 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4716,6 +4716,7 @@ class JWTAuthBuilderResult(TypedDict): org_id: str | None team_membership: LiteLLM_TeamMembership | None jwt_claims: dict # Decoded JWT token claims (avoids re-decoding) + agent_id: ReadOnly[str | None] class ClientSideFallbackModel(TypedDict, total=False): @@ -4954,6 +4955,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): user_allowed_roles: list[str] | None = None user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.") end_user_id_jwt_field: str | None = None + agent_id_jwt_field: str | None = Field( + default=None, + description=( + "The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID " + "app token). Supports dot notation. The value is matched against a registered agent's agent_id, " + "then agent_name, and the request is rejected when it matches neither." + ), + ) public_key_ttl: float = 600 public_key_stale_ttl: float = Field( default=DEFAULT_JWKS_STALE_TTL, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5afe35c0338..355fc3f6a21 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5773,8 +5773,7 @@ async def _organization_max_budget_check( if org_table.litellm_budget_table is not None: org_max_budget = org_table.litellm_budget_table.max_budget - # Only check if organization has a valid max_budget set - if org_max_budget is None or org_max_budget <= 0: + if org_max_budget is None: return # Read spend from cross-pod counter (Redis-first) or cached object (fallback) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4304542fc83..94ca3047f45 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,7 +14,7 @@ import hashlib import os import re import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast import httpx @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.user_repository import UserRepository +from litellm.types.agents import AgentResponse from .auth_checks import ( _allowed_routes_check, @@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol): def json(self) -> dict[str, object]: ... +class AgentLookup(Protocol): + """The registered-agent lookups a JWT agent claim is matched against.""" + + def get_agent_by_id(self, agent_id: str) -> AgentResponse | None: + """The agent registered under ``agent_id``, if any.""" + + def get_agent_by_name(self, agent_name: str) -> AgentResponse | None: + """The agent registered under ``agent_name``, if any.""" + + +class _NoRegisteredAgents: + """The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches.""" + + def get_agent_by_id(self, agent_id: str) -> None: + return None + + def get_agent_by_name(self, agent_name: str) -> None: + return None + + def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody: """Decode an OIDC discovery response body.""" return response.json() @@ -198,6 +219,10 @@ class JWTHandler: self.leeway = 0 # Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request. self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url + self.agent_lookup: AgentLookup = _NoRegisteredAgents() + + def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None: + self.agent_lookup = agent_lookup def update_environment( self, @@ -623,6 +648,12 @@ class JWTHandler: object_id = default_value return object_id + def get_agent_claim(self, token: Mapping[str, object]) -> str | None: + if self.litellm_jwtauth.agent_id_jwt_field is None: + return None + claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field) + return claim if isinstance(claim, str) and claim else None + def get_org_id(self, token: dict, default_value: str | None) -> str | None: if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM): return token.get(self.LITELLM_ORG_ID_CLAIM) @@ -1380,6 +1411,7 @@ class JWTAuthManager: api_key: str, jwt_valid_token: dict | None = None, user_email: str | None = None, + agent_id: str | None = None, ) -> JWTAuthBuilderResult | None: """Check admin status and route access permissions""" if not jwt_handler.is_admin(scopes=scopes): @@ -1409,8 +1441,28 @@ class JWTAuthManager: org_id=org_id, team_membership=None, jwt_claims=jwt_valid_token or {}, + agent_id=agent_id, ) + @staticmethod + def resolve_agent_id( + jwt_handler: JWTHandler, + jwt_valid_token: Mapping[str, object], + agent_registry: AgentLookup, + ) -> str | None: + agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token) + if agent_claim is None: + return None + agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name( + agent_name=agent_claim + ) + if agent is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}", + ) + return agent.agent_id + @staticmethod async def find_and_validate_specific_team_id( jwt_handler: JWTHandler, @@ -2268,9 +2320,23 @@ class JWTAuthManager: elif rbac_role == LitellmUserRoles.INTERNAL_USER: user_id = object_id + agent_id: Final = JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, + jwt_valid_token=jwt_valid_token, + agent_registry=jwt_handler.agent_lookup, + ) + # Check admin access admin_result: Final = await JWTAuthManager.check_admin_access( - jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email + jwt_handler, + scopes, + route, + user_id, + org_id, + api_key, + jwt_valid_token, + user_email=user_email, + agent_id=agent_id, ) if admin_result: await JWTAuthManager._attach_team_from_header_for_admin( @@ -2514,4 +2580,5 @@ class JWTAuthManager: token=api_key, team_membership=team_membership_object, jwt_claims=jwt_valid_token, + agent_id=agent_id, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4b53541004c..0618875be55 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -852,6 +852,7 @@ async def _auto_register_jwt_mapping( user_id: str | None = None, org_id: str | None = None, end_user_id: str | None = None, + agent_id: str | None = None, ) -> UserAPIKeyAuth | None: """ Auto-register: create a new virtual key + mapping for an unrecognised JWT @@ -884,6 +885,7 @@ async def _auto_register_jwt_mapping( team_id=team_id, user_id=user_id, organization_id=org_id, + agent_id=agent_id, metadata={ "auto_registered": True, "jwt_claim_field": virtual_key_claim_field, @@ -1567,6 +1569,7 @@ async def _user_api_key_auth_builder( org_id: Final = result["org_id"] team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None) jwt_claims = result.get("jwt_claims", None) + agent_id: Final[str | None] = result.get("agent_id") if is_proxy_admin: # Proxy admins authenticate via auth_builder (full @@ -1592,6 +1595,7 @@ async def _user_api_key_auth_builder( end_user_id=end_user_id, parent_otel_span=parent_otel_span, jwt_claims=jwt_claims, + agent_id=agent_id, **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) @@ -1612,6 +1616,7 @@ async def _user_api_key_auth_builder( user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), jwt_claims=jwt_claims, + agent_id=agent_id, **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) @@ -1635,6 +1640,7 @@ async def _user_api_key_auth_builder( user_id=user_id, org_id=org_id, end_user_id=end_user_id, + agent_id=agent_id, ) if auto_registered is not None: auto_registered.jwt_claims = jwt_claims diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 40e4c03993c..bc4b3610b1c 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1856,9 +1856,9 @@ def validate_team_org_change( # Check if the team's budget is less than the org's max_budget if ( - team.max_budget - and organization.litellm_budget_table - and organization.litellm_budget_table.max_budget + team.max_budget is not None + and organization.litellm_budget_table is not None + and organization.litellm_budget_table.max_budget is not None and team.max_budget > organization.litellm_budget_table.max_budget ): raise HTTPException( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 963fca104f8..34b742a082b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9516,6 +9516,9 @@ class ProxyStartupEvent: user_api_key_cache=user_api_key_cache, litellm_jwtauth=litellm_jwtauth, ) + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + jwt_handler.bind_agent_lookup(global_agent_registry) @classmethod def _add_proxy_budget_to_db(cls): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 48a6baa3587..e611984bf4c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -430,6 +430,14 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: detail.setdefault("guardrail_mode", event_hook) +def _is_client_error_exception(exc: Exception) -> bool: + if isinstance(exc, HTTPException): + return exc.status_code < 500 + if isinstance(exc, ProxyException): + return not (exc.code.isdigit() and int(exc.code) >= 500) + return False + + def _exception_changes_request_flow(exc: BaseException) -> bool: """ True for guardrail exceptions the proxy turns into an alternate request flow @@ -2886,9 +2894,7 @@ class ProxyLogging: ### ALERTING ### await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail") - if AlertType.llm_exceptions in self.alert_types and not isinstance( - original_exception, (HTTPException, ProxyException) - ): + if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception): """ Just alert on LLM API exceptions. Do not alert on user errors diff --git a/litellm/router.py b/litellm/router.py index fcdcf91c2cf..ef89d611075 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1622,6 +1622,24 @@ class Router: return await selector.async_pre_call_check(deployment, parent_otel_span) + def _bind_override_selector_to_request( + self, strategy: str, selector: RouterStrategySelector | None, request_kwargs: Mapping[str, object] | None + ) -> None: + if selector is None or request_kwargs is None or strategy in self._globally_registered_strategies(): + return + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLogging): + logging_obj.add_dynamic_callback(selector) + + def _globally_registered_strategies(self) -> frozenset[str]: + configured: Final = ( + self.routing_strategy, + *(group.routing_strategy for group in self._routing_groups.values()), + ) + return frozenset( + normalized for normalized in map(self._normalize_strategy, configured) if normalized is not None + ) + def _get_routing_context( self, model: str, request_kwargs: dict | None = None ) -> tuple[str | None, RouterStrategySelector | None]: @@ -1647,7 +1665,9 @@ class Router: override: Final = self._get_request_routing_strategy_override(request_kwargs) if override is not None: verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override) - return override, self._get_override_strategy_selector(override) + override_selector: Final = self._get_override_strategy_selector(override) + self._bind_override_selector_to_request(override, override_selector, request_kwargs) + return override, override_selector group_name: Final = model if self.get_routing_group(model) is not None else self._model_to_group.get(model) if group_name is None: @@ -2461,7 +2481,7 @@ class Router: ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) ## only run if model group given, not model id - if not self.has_model_id(model): + if model in self.model_names or not self.has_model_id(model): self.routing_strategy_pre_call_checks(deployment=deployment) input_kwargs: Final = { @@ -12512,7 +12532,7 @@ class Router: # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) - elif self.has_model_id(model): + elif model not in self.model_names and self.has_model_id(model): deployment: Final = self.get_deployment(model_id=model) if deployment is not None: deployment_model: Final = deployment.litellm_params.model diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index e226901dd5c..900bd65ade9 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2217,8 +2217,18 @@ class ComplexityRouter(CustomLogger): if capability is None or classifier_system_prompt is None: raise ValueError("capability classifier is not configured") - asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), self._reminder_markers)) - opening_task: Final = asks_newest_first[-1] if asks_newest_first else prompt + markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) + encrypted_task: Final = _encrypted_classifier_task(request_kwargs, markers) + asks_newest_first: Final = ( + () if encrypted_task is not None else tuple(_iter_human_asks_newest_first(messages or (), markers)) + ) + opening_task: Final = ( + "The delegated task in the following agent_message." + if encrypted_task is not None + else asks_newest_first[-1] + if asks_newest_first + else prompt + ) latest_follow_up: Final = asks_newest_first[0] if len(asks_newest_first) > 1 else None task_messages: list[AllMessageValues] = [ # mutable-ok: the latest message gains optional image parts below {"role": "user", "content": opening_task}, # mutable-ok: SDK messages are dict-shaped @@ -2246,9 +2256,7 @@ class ComplexityRouter(CustomLogger): messages_for_call, request_kwargs, max_output_tokens=capability.max_output_tokens, - encrypted_task=_encrypted_classifier_task( - request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) - ), + encrypted_task=encrypted_task, ) verdict: Final = parse_capability_classifier_verdict(content) forecast: Final = capability.classify(verdict) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 4e957f258b6..84577018e91 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -641,7 +641,7 @@ class CapabilityCalibrationConfig(BaseModel): class CapabilityClassifierConfig(BaseModel): """Switchyard-compatible probability threshold policy for two model tiers.""" - model_config = ConfigDict(frozen=True) + model_config = ConfigDict(extra="forbid", frozen=True) card: CapabilityCardConfig | None = None selective_policy: SelectivePolicy | None = None empirical_supplement: str | None = Field(default=None, min_length=1, max_length=4000) diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..af22b11224b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2202,15 +2202,20 @@ def _is_streaming_request( def _select_tokenizer(model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None = None): if custom_tokenizer is not None: - _tokenizer: Final = create_pretrained_tokenizer( + return _select_custom_tokenizer_helper( identifier=custom_tokenizer["identifier"], revision=custom_tokenizer["revision"], auth_token=custom_tokenizer["auth_token"], ) - return _tokenizer return _select_tokenizer_helper(model=model) +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _select_custom_tokenizer_helper(identifier: str, revision: str, auth_token: str | None) -> SelectTokenizerResponse: + verbose_logger.debug("Loading custom HuggingFace tokenizer %s (revision %s)", identifier, revision) + return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token) + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: if litellm.disable_hf_tokenizer_download is True: diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index 6e908bcbdcd..70281f75196 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/tests/test_litellm/compression/test_compress.py @@ -56,11 +56,87 @@ def test_no_user_or_assistant_rows(): assert get_protected_indices([]) == () +def test_rows_before_last_cache_control_breakpoint_are_protected(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "old question"}, + { + "role": "assistant", + "content": "old answer", + "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "t1", "content": "large file body"}, + { + "role": "user", + "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}], + }, + { + "role": "assistant", + "content": "ack", + "tool_calls": [{"id": "t2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "t2", "content": "later tool output"}, + {"role": "user", "content": "live instruction"}, + ] + + protected = sorted(get_protected_indices(messages)) + + assert protected == [0, 1, 2, 3, 4, 5, 7] + assert 6 not in protected + + +def test_cache_control_directly_on_message_protects_prefix(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "tool", "tool_call_id": "before", "content": "large file body"}, + {"role": "user", "content": "old question"}, + { + "role": "tool", + "tool_call_id": "marked", + "content": "cached tool", + "cache_control": {"type": "ephemeral"}, + }, + {"role": "tool", "tool_call_id": "after", "content": "later tool output"}, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "live instruction"}, + ] + + protected = sorted(get_protected_indices(messages)) + + assert 1 in protected + assert 3 in protected + assert 4 not in protected + + +def test_no_cache_control_leaves_history_compressible(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "old answer"}, + {"role": "tool", "tool_call_id": "t1", "content": "large file body"}, + {"role": "user", "content": "live instruction"}, + ] + + assert sorted(get_protected_indices(messages)) == [0, 2, 4] + + +def test_non_mapping_content_parts_are_not_cache_control(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": ["not", "a", "dict"]}, + {"role": "assistant", "content": "old answer"}, + {"role": "tool", "tool_call_id": "t1", "content": "plain string"}, + {"role": "user", "content": "live instruction"}, + ] + + protected = sorted(get_protected_indices(messages)) + + assert protected == [0, 2, 4] + assert 1 not in protected + assert 3 not in protected + + def test_mid_history_cache_control_part_is_protected(): - # A large cached tool result from a few turns back, not the last user or - # last assistant row -- exactly the row a provider prompt-cache pins to - # exact bytes. Rewriting it (even leaving the marker on) changes those - # bytes and turns the next request's cache read into a cache write. messages = [ {"role": "user", "content": "old question"}, {"role": "assistant", "content": "old answer"}, @@ -74,9 +150,7 @@ def test_mid_history_cache_control_part_is_protected(): {"role": "user", "content": "live instruction"}, ] - # index 3 = last assistant, index 4 = last user (both protected by role - # regardless), index 2 = the cache_control-marked row itself. - assert sorted(get_protected_indices(messages)) == [2, 3, 4] + assert sorted(get_protected_indices(messages)) == [0, 1, 2, 3, 4] def test_cache_control_directly_on_message_is_protected(): @@ -116,8 +190,6 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control( def test_compress_keeps_part_level_cache_control_row_verbatim(): - # compress() scores text-only copies of the rows, where a part-level marker - # is gone; protection has to read the original rows or the pinned row is stubbed. stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]} pinned = { "role": "user", @@ -126,9 +198,9 @@ def test_compress_keeps_part_level_cache_control_row_verbatim(): ], } messages = [ - stale_log, - {"role": "assistant", "content": "old answer"}, pinned, + {"role": "assistant", "content": "old answer"}, + stale_log, {"role": "assistant", "content": "ack"}, {"role": "user", "content": "live instruction"}, ] @@ -142,6 +214,6 @@ def test_compress_keeps_part_level_cache_control_row_verbatim(): ) assert len(result["messages"]) == len(messages) - assert result["messages"][2] == pinned - assert result["messages"][0] != stale_log + assert result["messages"][0] == pinned + assert result["messages"][2] != stale_log assert len(result["cache"]) >= 1 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 70f9bae283b..dd1ad9c9623 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -7155,3 +7155,21 @@ def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy(): assert copied["llm_provider-x-custom-1999"] == "1999" _run_while_a_thread_grows(headers, read, reads=300) + + +def test_add_dynamic_callback_registers_once_per_list_without_touching_the_callers_list(logging_obj: LitellmLogging): + callback: Final = CustomLogger() + caller_owned: Final = ["langfuse"] + logging_obj.dynamic_success_callbacks = caller_owned + + logging_obj.add_dynamic_callback(callback) + logging_obj.add_dynamic_callback(callback) + + assert caller_owned == ["langfuse"] + assert logging_obj.dynamic_success_callbacks == ["langfuse", callback] + assert logging_obj.dynamic_input_callbacks == [callback] + assert logging_obj.dynamic_async_success_callbacks == [callback] + assert logging_obj.dynamic_failure_callbacks == [callback] + assert logging_obj.dynamic_async_failure_callbacks == [callback] + assert LitellmLogging._with_dynamic_callback(None, callback) == [callback] + assert LitellmLogging._with_dynamic_callback((callback,), callback) == [callback] diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f93658c740..8c8b755195f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5855,6 +5855,71 @@ async def test_organization_budget_check_carries_org_state_on_the_token(): assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) +@pytest.mark.parametrize( + "max_budget, spend, expect_blocked", + [ + (0.0, 0.0, True), # explicit zero budget blocks even a fresh org with no spend + (0.0, 7.4e-06, True), # any spend at all against a zero budget blocks + (None, 999.0, False), # unlimited (None) never blocks, regardless of spend + (5.0, 4.99, False), # a positive budget under its cap still passes + ], +) +@pytest.mark.asyncio +async def test_organization_zero_max_budget_is_enforced(max_budget, spend, expect_blocked): + """An explicit organization max_budget of 0 must mean zero allowance, matching + key/team/user semantics, not unlimited. + + Regression for LIT-7797: `_organization_max_budget_check` returned early + whenever `org_max_budget <= 0`, so an org configured with max_budget=0 could + spend without limit. + """ + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import _organization_max_budget_check + + org_table = LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="zero-budget-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget) if max_budget is not None else None, + ) + token = UserAPIKeyAuth(token="k1", org_id="o1") + user_api_key_cache = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable + ) + + async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return spend + + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( # test-quality-ok: _organization_max_budget_check imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", _spend + ): + if expect_blocked: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _organization_max_budget_check( + valid_token=token, + team_object=None, + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.max_budget == max_budget + else: + await _organization_max_budget_check( + valid_token=token, + team_object=None, + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + @pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"]) @pytest.mark.asyncio async def test_spend_capable_non_llm_routes_still_enforce_budget(route): diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 94226b5404d..814e31535e0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( ProxyException, ) from litellm.caching.dual_cache import DualCache +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.handle_jwt import ( JWKS_FETCH_ATTEMPTS, STALE_CACHE_KEY_PREFIX, @@ -32,6 +33,7 @@ from litellm.proxy.auth.handle_jwt import ( JWTHandler, NoMatchingJWTPublicKeyError, ) +from litellm.types.agents import AgentResponse @pytest.mark.asyncio @@ -6786,3 +6788,180 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla } assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == [] assert user.teams == [] + + +def _entra_agent_registry() -> AgentRegistry: + registry = AgentRegistry() + registry.register_agent( + AgentResponse( + agent_id="canonical-agent-id", + agent_name="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21", + agent_card_params={"name": "research-agent", "url": "http://localhost:9999/a2a", "version": "1.0.0"}, + litellm_params={"require_trace_id_on_calls_by_agent": True}, + ) + ) + return registry + + +def _entra_agent_jwt_handler(agent_id_jwt_field: str | None) -> JWTHandler: + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", agent_id_jwt_field=agent_id_jwt_field), + ) + return jwt_handler + + +@pytest.mark.parametrize( + "claim_value", + ["canonical-agent-id", "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"], + ids=["matches_agent_id", "matches_agent_name"], +) +def test_resolve_agent_id_returns_canonical_agent_id(claim_value: str): + """An Entra app token's azp claim binds to the registered agent by id or by name and yields its canonical id.""" + jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp") + + resolved = JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, + jwt_valid_token={"sub": "sp-object-id-1234", "azp": claim_value}, + agent_registry=_entra_agent_registry(), + ) + + assert resolved == "canonical-agent-id" + + +def test_resolve_agent_id_reads_nested_claim(): + jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="entra.client_id") + + resolved = JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, + jwt_valid_token={"sub": "sp-object-id-1234", "entra": {"client_id": "canonical-agent-id"}}, + agent_registry=_entra_agent_registry(), + ) + + assert resolved == "canonical-agent-id" + + +def test_resolve_agent_id_rejects_claim_for_unregistered_agent(): + """A configured agent claim naming no registered agent fails closed with 403 instead of falling back to an unbound identity.""" + jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp") + + with pytest.raises(HTTPException) as exc_info: + JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, + jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"}, + agent_registry=_entra_agent_registry(), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.parametrize( + "token", + [ + {"sub": "sp-object-id-1234"}, + {"sub": "sp-object-id-1234", "azp": ""}, + {"sub": "sp-object-id-1234", "azp": ["canonical-agent-id"]}, + ], + ids=["claim_absent", "claim_empty", "claim_not_a_string"], +) +def test_resolve_agent_id_returns_none_when_claim_unusable(token: dict): + jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp") + + assert ( + JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, jwt_valid_token=token, agent_registry=_entra_agent_registry() + ) + is None + ) + + +def test_resolve_agent_id_ignores_claim_when_field_not_configured(): + """Without agent_id_jwt_field an azp claim (even an unknown one) leaves JWT auth behaviour unchanged.""" + jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field=None) + + resolved = JWTAuthManager.resolve_agent_id( + jwt_handler=jwt_handler, + jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"}, + agent_registry=_entra_agent_registry(), + ) + + assert resolved is None + + +def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandler, str]: + """A JWTHandler that verifies RS256 tokens against a pre-cached JWKS, plus a signed Entra-style app token.""" + jwks_url = "https://login.microsoftonline.test/discovery/v2.0/keys" + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + private_key, jwk = _get_rsa_key_and_jwk(kid="entra-kid") + cache = DualCache() + cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk]) + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="azp"), + ) + token = _encode_rsa_jwt( + private_key, + issuer="https://login.microsoftonline.test/lit7664-tenant/v2.0", + audience="api://litellm", + kid="entra-kid", + extra_claims={"sub": "sp-object-id-1234", "azp": azp, "scope": scope}, + ) + return jwt_handler, token + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_admin_token", [False, True], ids=["standard_jwt", "proxy_admin_jwt"]) +async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool): + """auth_builder carries the resolved agent id into JWTAuthBuilderResult on both the admin and standard paths.""" + jwt_handler, token = _entra_signed_app_token( + monkeypatch, + azp="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21", + scope=LiteLLM_JWTAuth().admin_jwt_scope if is_admin_token else "", + ) + jwt_handler.bind_agent_lookup(_entra_agent_registry()) + + result = await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=jwt_handler, + request_data={"model": "gpt-5.6"}, + general_settings={"enforce_rbac": False}, + route="/key/info" if is_admin_token else "/chat/completions", + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert result["is_proxy_admin"] is is_admin_token + assert result["agent_id"] == "canonical-agent-id" + + +@pytest.mark.asyncio +async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch): + """An unknown agent claim is rejected even when the token would otherwise be a proxy admin.""" + jwt_handler, token = _entra_signed_app_token( + monkeypatch, + azp="00000000-0000-0000-0000-000000000000", + scope=LiteLLM_JWTAuth().admin_jwt_scope, + ) + jwt_handler.bind_agent_lookup(_entra_agent_registry()) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=jwt_handler, + request_data={"model": "gpt-5.6"}, + general_settings={"enforce_rbac": False}, + route="/key/info", + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index c9ae105d982..866ea0b20e4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1937,6 +1937,76 @@ async def test_standard_jwt_auth_propagates_user_email(): assert result.api_key is None +@pytest.mark.asyncio +@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard_jwt", "proxy_admin_jwt"]) +async def test_jwt_auth_propagates_agent_id_to_user_api_key_auth(is_proxy_admin: bool): + """The agent id resolved by auth_builder must land on UserAPIKeyAuth.agent_id so + agent-scoped checks (trace id requirement, MCP server/tool restrictions, spend + attribution) apply to JWT callers the same way they apply to agent-bound keys.""" + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + general_settings = {"enable_jwt_auth": True} + user_api_key_cache = DualCache() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(agent_id_jwt_field="azp") + + user_object = LiteLLM_UserTable(user_id="sp-object-id-1234", user_role="internal_user") + mock_jwt_result = { + "is_proxy_admin": is_proxy_admin, + "team_object": None, + "user_object": user_object, + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": None, + "user_id": "sp-object-id-1234", + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "sp-object-id-1234", "azp": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}, + "agent_id": "canonical-agent-id", + } + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists + "litellm.proxy.proxy_server", + general_settings=general_settings, + premium_user=True, + master_key="sk-master", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=MagicMock(), + jwt_handler=jwt_handler, + ), + patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=mock_jwt_result, + ), + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-5.6"}, + ) + + assert result.agent_id == "canonical-agent-id" + assert result.user_id == "sp-object-id-1234" + assert result.api_key is None + + @pytest.mark.asyncio async def test_auto_register_binds_api_key_to_token_hash(): """ @@ -2106,6 +2176,222 @@ async def test_auto_register_first_request_propagates_user_email(): assert result.api_key == "hashed-auto-key" +@pytest.mark.asyncio +async def test_auto_register_stamps_new_key_with_jwt_agent_id(): + """The virtual key AUTO_REGISTER creates must carry the agent id auth_builder bound + from the JWT claim, and the first request's principal must carry it too, or the + mapped-key path would drop the agent policies on that request and every later one.""" + from litellm.proxy.auth.auth_method import AuthMethod + from litellm.proxy.auth.resolvers.models import CredentialRef + from litellm.proxy.auth.resolvers.store import IdentityStore + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping + from litellm.proxy.proxy_server import hash_token + + plaintext = "sk-auto-registered-agent" + token_hash = hash_token(plaintext) + persisted_principal = IdentityStore._principal_from_key( + UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"), + auth_method=AuthMethod.API_KEY, + credential_ref=CredentialRef(token_id=token_hash), + ) + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock() + user_api_key_cache = MagicMock() + user_api_key_cache.async_set_cache = AsyncMock() + jwt_handler = MagicMock() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300) + generate_key = AsyncMock(return_value={"token": plaintext}) + + with ( + patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + generate_key, + ), + patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists + "litellm.proxy.auth.resolvers.store.IdentityStore.resolve", + new_callable=AsyncMock, + return_value=persisted_principal, + ), + ): + result = await _auto_register_jwt_mapping( + virtual_key_claim_field="appid", + claim_value="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21", + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + cache_key="jwt_key_mapping:appid:2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21", + team_id="validated-team", + user_id="validated-user", + agent_id="canonical-agent-id", + ) + + assert generate_key.await_args is not None + assert generate_key.await_args.kwargs["agent_id"] == "canonical-agent-id" + assert result is not None + assert result.agent_id == "canonical-agent-id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("losing_agent_id", ["other-agent", None], ids=["different_agent", "no_agent_claim"]) +async def test_auto_register_race_loser_keeps_winners_agent_id(losing_agent_id: str | None): + """When two requests race to AUTO_REGISTER the same mapping claim, the loser must run as + the persisted key, agent binding included. Every later request on that mapping uses the + winner's key, so stamping the loser's own (or missing) agent id on it would give one request + different agent policies and spend attribution than all the others.""" + from litellm.proxy.auth.auth_method import AuthMethod + from litellm.proxy.auth.resolvers.models import CredentialRef + from litellm.proxy.auth.resolvers.store import IdentityStore + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping + + winner_hash = "winner-key-hash" + winner_principal = IdentityStore._principal_from_key( + UserAPIKeyAuth(token=winner_hash, user_id="validated-user", team_id="validated-team", agent_id="winner-agent"), + auth_method=AuthMethod.API_KEY, + credential_ref=CredentialRef(token_id=winner_hash), + ) + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock( + side_effect=Exception("Unique constraint failed on the fields: (`jwt_claim_name`,`jwt_claim_value`)") + ) + prisma_client.db.litellm_verificationtoken.delete = AsyncMock() + user_api_key_cache = MagicMock() + user_api_key_cache.async_set_cache = AsyncMock() + jwt_handler = MagicMock() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300) + + with ( + patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "sk-orphaned-loser-key"}, + ), + patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object", + new_callable=AsyncMock, + return_value=winner_hash, + ), + patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists + "litellm.proxy.auth.resolvers.store.IdentityStore.resolve", + new_callable=AsyncMock, + return_value=winner_principal, + ), + ): + result = await _auto_register_jwt_mapping( + virtual_key_claim_field="tid", + claim_value="shared-tenant", + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + cache_key="jwt_key_mapping:tid:shared-tenant", + team_id="validated-team", + user_id="validated-user", + agent_id=losing_agent_id, + ) + + assert result is not None + assert result.token == winner_hash + assert result.agent_id == "winner-agent" + + +@pytest.mark.asyncio +async def test_jwt_auto_register_forwards_bound_agent_id(): + """When a JWT under AUTO_REGISTER also carries the configured agent claim, the agent + id auth_builder resolved must reach the key creation, not be dropped when + valid_token is swapped for the freshly registered key.""" + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + user_api_key_cache = DualCache() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + virtual_key_mapping_cache_ttl=300, + agent_id_jwt_field="appid", + ) + user_object = LiteLLM_UserTable(user_id="validated-user", user_role="internal_user") + mock_jwt_result = { + "is_proxy_admin": False, + "team_object": None, + "user_object": user_object, + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": "validated-team", + "user_id": "validated-user", + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}, + "agent_id": "canonical-agent-id", + } + auto_register = AsyncMock( + return_value=UserAPIKeyAuth( + token="hashed-auto-key", + api_key="hashed-auto-key", + team_id="validated-team", + user_id="validated-user", + agent_id="canonical-agent-id", + ) + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists + "litellm.proxy.proxy_server", + general_settings={"enable_jwt_auth": True}, + premium_user=True, + master_key="sk-master", + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=MagicMock(), + jwt_handler=jwt_handler, + ), + patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=_PendingAutoRegister( + claim_field="sub", + claim_value="user1", + cache_key="jwt_key_mapping:sub:user1", + ), + ), + patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=mock_jwt_result, + ), + patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists + "litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping", + auto_register, + ), + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-5.6"}, + ) + + assert auto_register.await_args is not None + assert auto_register.await_args.kwargs["agent_id"] == "canonical-agent-id" + assert result.agent_id == "canonical-agent-id" + assert result.api_key == "hashed-auto-key" + + class TestJWTOAuth2Coexistence: """ Test that JWT and OAuth2 auth can coexist on the same instance. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index d4531398ba1..bcecb5b27db 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -900,7 +900,7 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr ) assert has_headroom_retrieve_tool(result.get("tools") or []) - (issued, _expiry), = guardrail._issued_hashes_by_call_id.values() + ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) @@ -953,7 +953,6 @@ async def test_anthropic_assistant_history_never_reaches_compression_service(gua assert result["messages"][1]["content"] == [{"type": "text", "text": table}] - def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape(): """By the time an Anthropic Messages API response reaches the agentic-loop gate, the OpenAI-shaped tool this guardrail injects (type: "function") @@ -2342,9 +2341,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( ) assert streamed_text == final_answer assert not any("function_call" in str(getattr(event, "type", "")) for event in events) - assert not any( - getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events - ) + assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) mock_get.assert_called_once() assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) @@ -2399,9 +2396,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" ) assert streamed_text == final_answer - assert not any( - getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events - ) + assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) mock_get.assert_called_once() assert len(upstream.calls) == 2 assert not json.loads(upstream.calls[1].request.content).get("stream") @@ -2514,6 +2509,38 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert messages[3] == compressed_history[1] +CACHED_PREFIX_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "old question " + "Q" * 5000}, + { + "role": "assistant", + "content": "Reading the file now.", + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000}, + { + "role": "user", + "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}], + }, + { + "role": "assistant", + "content": "Listing now.", + "tool_calls": [{"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000}, + {"role": "assistant", "content": "Finished listing."}, + {"role": "user", "content": "live instruction"}, +] + + +@pytest.mark.asyncio +async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES) + + assert [row.get("tool_call_id") for row in wire] == ["new_1"] + assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5] + + CACHE_MARKED_HISTORY_MESSAGES = [ {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, {"role": "user", "content": "old question " + "Q" * 5000}, @@ -2529,6 +2556,7 @@ CACHE_MARKED_HISTORY_MESSAGES = [ "cache_control": {"type": "ephemeral"}, }, {"role": "assistant", "content": "Summarized the file for you."}, + {"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000}, {"role": "user", "content": "live instruction"}, ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 4c9920eaa08..1f17bc29f3b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -245,6 +245,55 @@ async def test_validate_team_org_change_same_org_id(): mock_access_check.assert_not_called() # Ensure access check wasn't called +@pytest.mark.parametrize( + "org_max_budget, team_max_budget, expect_blocked", + [ + (0.0, 100.0, True), # explicit zero org budget must still cap the team's budget + (0.0, None, False), # team has no budget of its own, nothing to compare + (None, 100.0, False), # unlimited (None) org budget never blocks + (50.0, 100.0, True), # a positive org budget is still enforced normally + ], +) +@pytest.mark.asyncio +async def test_validate_team_org_change_zero_org_budget_is_enforced( + org_max_budget, team_max_budget, expect_blocked +): + """An organization with an explicit max_budget of 0 must still block moving in a + team with a larger budget, matching key/team/user zero-budget semantics. + + Regression for LIT-7797: the truthy check `organization.litellm_budget_table.max_budget` + treated an explicit 0 the same as no budget table at all, silently skipping this guard. + """ + org_id = "team-org-123" + new_org_id = "new-org-456" + + team = MagicMock(spec=LiteLLM_TeamTable) + team.organization_id = org_id + team.models = [] + team.max_budget = team_max_budget + team.tpm_limit = None + team.rpm_limit = None + team.members_with_roles = [] + + organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) + organization.organization_id = new_org_id + organization.models = [] + organization.litellm_budget_table = ( + LiteLLM_BudgetTable(max_budget=org_max_budget) if org_max_budget is not None else None + ) + organization.members = [] + + mock_router = MagicMock(spec=Router) + + if expect_blocked: + with pytest.raises(HTTPException) as exc_info: + validate_team_org_change(team=team, organization=organization, llm_router=mock_router) + assert exc_info.value.status_code == 403 + else: + result = validate_team_org_change(team=team, organization=organization, llm_router=mock_router) + assert result is None or result is True + + @pytest.mark.asyncio async def test_validate_team_org_change_members_in_org(): """ diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index d6851116324..d3578455a35 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -3764,6 +3764,55 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_ ] +@pytest.mark.asyncio +@pytest.mark.parametrize("agents_source", ["config", "db", "api"]) +async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry( + clean_agent_registry, agents_source +): + """A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup.""" + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTAuthManager + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.types.agents import AgentResponse + + original_lookup = proxy_server.jwt_handler.agent_lookup + try: + proxy_server.ProxyStartupEvent._initialize_jwt_auth( + general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}}, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + ) + if agents_source == "config": + await ProxyConfig()._init_non_llm_configs( + config={"agents": [_config_agent("loaded-agent")]}, + config_file_path=None, + ) + elif agents_source == "db": + prisma_client = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock( + return_value=[_FakeAgentRow("db-id", "loaded-agent")] + ) + await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client) + else: + clean_agent_registry.register_agent( + agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent")) + ) + + resolved = JWTAuthManager.resolve_agent_id( + jwt_handler=proxy_server.jwt_handler, + jwt_valid_token={"appid": "loaded-agent"}, + agent_registry=proxy_server.jwt_handler.agent_lookup, + ) + finally: + proxy_server.jwt_handler.bind_agent_lookup(original_lookup) + proxy_server.jwt_handler.update_environment( + prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth() + ) + + assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id + + @pytest.mark.asyncio @pytest.mark.parametrize( "config, expected_agent_names", diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 04173ced776..42af8e0af21 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13521,3 +13521,54 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp assert response.tokenizer_type == "huggingface_tokenizer" assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from litellm.types.router import DeploymentTypedDict + + claude_tokenizer: Final[Tokenizer] = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + from_pretrained: Final = MagicMock(return_value=claude_tokenizer) + + def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": { + "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": revision, "auth_token": auth_token} + }, + } + + monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained)) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + deployment("self-hosted", "main", None), + deployment("self-hosted-pinned", "v2", None), + deployment("self-hosted-private", "main", "hf_test_token"), + ] + ), + ) + litellm.utils._select_custom_tokenizer_helper.cache_clear() + try: + responses: Final = [ + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) + for _ in range(3) + ] + assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)] + assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses) + assert len({response.total_tokens for response in responses}) == 1 + assert responses[0].total_tokens > 0 + + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once")) + assert from_pretrained.call_args_list == [ + mock.call("my-org/tokenizer", revision="main", token=None), + mock.call("my-org/tokenizer", revision="v2", token=None), + mock.call("my-org/tokenizer", revision="main", token="hf_test_token"), + ] + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear() diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9def21c0573..94ccc2762c5 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -13,7 +13,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.guardrails import GuardrailEventHooks -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy.utils import get_custom_url, join_paths @@ -1303,12 +1303,10 @@ class TestPostCallFailureHookLLMExceptionAlerting: """The llm_exceptions alert is for infra / LLM-API failures, not user errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized client errors must be excluded so a guardrail content-policy block never - pages on-call. ProxyException is such an error; before LIT-3751 only - HTTPException was excluded, so AIM blocks paged as if the LLM API failed.""" + pages on-call. 5xx proxy errors still alert.""" - async def _alerted(self, exc) -> bool: + async def _alerted(self, exc: Exception) -> AsyncMock: import asyncio - from unittest.mock import AsyncMock from litellm.proxy._types import AlertType, UserAPIKeyAuth @@ -1325,7 +1323,7 @@ class TestPostCallFailureHookLLMExceptionAlerting: user_api_key_dict=UserAPIKeyAuth(), ) await asyncio.sleep(0) # let the fire-and-forget alert task run - return alerting_handler.called + return alerting_handler @pytest.mark.asyncio async def test_proxy_exception_does_not_alert(self): @@ -1338,15 +1336,49 @@ class TestPostCallFailureHookLLMExceptionAlerting: code=400, openai_code="content_policy_violation", ) - assert await self._alerted(exc) is False + assert (await self._alerted(exc)).called is False @pytest.mark.asyncio async def test_http_exception_does_not_alert(self): - assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False + assert (await self._alerted(HTTPException(status_code=400, detail="blocked"))).called is False @pytest.mark.asyncio async def test_genuine_llm_api_error_still_alerts(self): - assert await self._alerted(Exception("upstream 503")) is True + assert (await self._alerted(Exception("upstream 503"))).called is True + + @pytest.mark.asyncio + async def test_http_exception_5xx_alerts(self): + alerting_handler = await self._alerted( + HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service returned an error", + "status_code": 503, + "guardrail_name": "headroom-compression-global", + }, + ) + ) + assert alerting_handler.called is True + assert "headroom-compression-global" in alerting_handler.call_args.kwargs["message"] + + @pytest.mark.asyncio + async def test_proxy_exception_5xx_alerts(self): + from litellm.proxy._types import ProxyException + + alerting_handler = await self._alerted( + ProxyException( + message="guardrail backend down", + type="internal_server_error", + param=None, + code=503, + ) + ) + assert alerting_handler.called is True + + @pytest.mark.asyncio + async def test_http_exception_429_does_not_alert(self): + alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited")) + assert alerting_handler.called is False class TestPostCallFailureHookProxyExceptionLogging: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8df9306ba44..6e62b96625f 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2573,6 +2573,17 @@ class TestCapabilityClassifierConfig: with pytest.raises(ValidationError, match="requires classifier_type 'capability'"): ComplexityRouterConfig(**config) + def test_rejects_misspelled_optional_policy_instead_of_using_defaults(self) -> None: + with pytest.raises(ValidationError, match="threshold_steps"): + CapabilityClassifierConfig.model_validate( + { + "efficient_tier": "SIMPLE", + "capable_tier": "REASONING", + "base_threshold": 0.5, + "threshold_steps": 0.2, + } + ) + def test_threshold_defaults_match_switchyard(self): config = CapabilityClassifierConfig(efficient_tier=" SIMPLE ", capable_tier=" REASONING ", base_threshold=0.5) assert config.efficient_tier == "SIMPLE" @@ -2691,6 +2702,62 @@ class TestCapabilityClassifier: assert "Route to the Efficient model" not in sent mock_router_instance.acompletion.assert_awaited_once() + @pytest.mark.asyncio + async def test_encrypted_task_is_not_replaced_by_plaintext_envelope(self, mock_router_instance: MagicMock) -> None: + mock_router_instance.aresponses = AsyncMock( + return_value=_native_classifier_response(_capability_reply(p_solve=0.8)) + ) + router: Final = self._router(mock_router_instance) + task: Final = _encrypted_agent_task() + request: Final = {"input": [task]} + original: Final = deepcopy(request) + result: Final = await router.async_pre_routing_hook(model="capability-router", request_kwargs=request) + assert result is not None and result.model == "efficient-model" + assert result.routing_decision is not None + assert result.routing_decision["cause"] == "capability_classifier" + mock_router_instance.aresponses.assert_awaited_once() + call: Final = mock_router_instance.aresponses.call_args.kwargs + assert call["input"][-1] == task + plaintext: Final = json.dumps(call["input"][:-1]) + assert "The delegated task in the following agent_message." in plaintext + assert "Message Type: NEW_TASK" not in plaintext + assert "opaque-provider-task" not in plaintext + assert request == original + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_markers", (False, True)) + async def test_task_forecast_uses_request_scoped_codex_markers( + self, mock_router_instance: MagicMock, custom_markers: bool + ) -> None: + completion: Final = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8))) + mock_router_instance.acompletion = completion + router: Final = self._router( + mock_router_instance, + escalation_keywords=[], + **({"reminder_markers": [{"open": "", "close": ""}]} if custom_markers else {}), + ) + envelope: Final = "\n".join(_CODEX_ENVELOPES) + opening: Final = f"{envelope}\nFix nested behavior" + messages: Final = [ + {"role": "user", "content": opening}, + {"role": "user", "content": "Preserve empty inputs"}, + {"role": "user", "content": envelope}, + ] + original: Final = deepcopy(messages) + for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"): + result: Final = await router.async_pre_routing_hook( + model="capability-router", messages=messages, request_kwargs={"metadata": {"user_agent": user_agent}} + ) + assert result is not None and result.model == "efficient-model" + sent: Final = completion.call_args.kwargs["messages"] + if user_agent.startswith("codex") and not custom_markers: + assert [message["content"] for message in sent[1:]] == ["Fix nested behavior", "Preserve empty inputs"] + else: + assert [message["content"] for message in sent[1:]] == [opening, envelope] + assert result.messages == original + assert completion.await_count == 3 + assert messages == original + @pytest.mark.asyncio @pytest.mark.parametrize("p_solve,expected_model", ((0.95, "capable-model"), (0.98, "efficient-model"))) async def test_fitted_probability_controls_routing_and_preserves_raw_score( @@ -6216,10 +6283,16 @@ class TestTierModelAffinity: returned: Final = await self._route(router, metadata, "model-b") assert (first.model, repeated.model, reasoning.model, returned.model) == ( - "model-a", "model-a", "model-b", "model-a" + "model-a", + "model-a", + "model-b", + "model-a", ) assert tuple(result.routing_decision["tier"] for result in (first, repeated, reasoning, returned)) == ( - "SIMPLE", "SIMPLE", "REASONING", "SIMPLE" + "SIMPLE", + "SIMPLE", + "REASONING", + "SIMPLE", ) assert returned.litellm_params == {"temperature": 0.1} assert reasoning.litellm_params == {"temperature": 0.9} @@ -6257,9 +6330,7 @@ class TestTierModelAffinity: deployment_affinity: bool, plugins: bool, ) -> None: - router: Final = self._router( - mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins - ) + router: Final = self._router(mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins) assert (await self._route(router, metadata, "model-a")).model == "model-a" assert (await self._route(router, metadata, "model-b")).model == "model-b" @@ -6332,9 +6403,7 @@ class TestTierModelAffinity: { "role": "assistant", "content": None, - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} - ], + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], }, {"role": "tool", "tool_call_id": "call_1", "content": [IMG_PART] if gate == "image" else "done"}, ] @@ -6379,9 +6448,7 @@ class TestTierModelAffinity: { "role": "assistant", "content": None, - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} - ], + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], }, {"role": "tool", "tool_call_id": "call_1", "content": "done"}, ] @@ -6411,8 +6478,7 @@ class TestTierModelAffinity: "SIMPLE": "base", **{ tier: [ - {"model_name": model, "litellm_params": {"temperature": temperature}} - for model in models + {"model_name": model, "litellm_params": {"temperature": temperature}} for model in models ] for tier, models, temperature in ( ("MEDIUM", ("shared", "middle"), 0.4), @@ -6486,7 +6552,11 @@ class TestTierModelAffinity: model_name="affinity-router", litellm_router_instance=mock_router_instance, complexity_router_config=_custom_tier_config( - tiers={"SIMPLE": ["model-a", "model-b"], "SECURITY_REVIEW": ["model-a", "model-b"], "COMPLEX": "model-a"}, + tiers={ + "SIMPLE": ["model-a", "model-b"], + "SECURITY_REVIEW": ["model-a", "model-b"], + "COMPLEX": "model-a", + }, deployment_affinity=True, classification_mode=classification_mode, keyword_tier_rules=[ diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5f37842305d..25b657b8cd0 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -5,15 +5,20 @@ the implicit `"default"` group driven by the router's top-level `routing_strategy` / `routing_strategy_args`. """ +import asyncio +import datetime +import time +import uuid +from collections.abc import Callable from unittest.mock import patch import pytest - import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy +from litellm.utils import Rules, function_setup def _model_list(): @@ -954,6 +959,223 @@ def test_sync_pass_through_specific_deployment_runs_the_override_pre_call_check( assert plain["model_info"]["id"] == "deploy-3" +def _two_deployment_model_list(**d1_params: object) -> list[dict[str, object]]: + return [ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test-1", "mock_response": "ok", **d1_params}, + "model_info": {"id": "d1"}, + }, + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test-2", "mock_response": "ok"}, + "model_info": {"id": "d2"}, + }, + ] + + +def _proxy_shaped_request(**data: object) -> dict[str, object]: + """The proxy builds the request's `Logging` object before it hands the call to the router.""" + logging_obj, kwargs = function_setup( + "acompletion", + Rules(), + datetime.datetime.now(), + litellm_call_id=str(uuid.uuid4()), + messages=[{"role": "user", "content": "hi"}], + **data, + ) + return {**kwargs, "litellm_logging_obj": logging_obj} + + +async def _async_override_pick(router: Router, strategy: str) -> str: + deployment = await router.async_get_available_deployment( + "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) + ) + return deployment["model_info"]["id"] + + +def _sync_override_pick(router: Router, strategy: str) -> str: + deployment = router.get_available_deployment( + "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) + ) + return deployment["model_info"]["id"] + + +def _in_flight(router: Router, deployment_id: str) -> int | None: + return router.cache.get_cache(f"grp_request_count:{deployment_id}") + + +async def _async_wait_until(predicate: Callable[[], bool]) -> None: + for _ in range(100): + if predicate(): + return + await asyncio.sleep(0.02) + raise AssertionError("lifecycle callback never reached the override selector") + + +def _sync_wait_until(predicate: Callable[[], bool]) -> None: + for _ in range(100): + if predicate(): + return + time.sleep(0.02) + raise AssertionError("lifecycle callback never reached the override selector") + + +def _selector_is_not_global(selector: CustomLogger) -> bool: + global_lists = ( + litellm.callbacks, + litellm.input_callback, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ) + return not any(cb is selector for cbs in global_lists for cb in cbs) + + +@pytest.mark.asyncio +async def test_least_busy_override_sees_the_overriding_request_in_flight(): + router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) + + stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) + busy = stream._hidden_params["model_id"] + idle = "d2" if busy == "d1" else "d1" + assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] + + async for _ in stream: + pass + await _async_wait_until(lambda: _in_flight(router, busy) == 0) + assert await _async_override_pick(router, "least-busy") == "d1" + assert _selector_is_not_global(router._override_selectors["least-busy"]) + + +def test_sync_least_busy_override_sees_the_overriding_request_in_flight(): + router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) + + stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) + busy = stream._hidden_params["model_id"] + idle = "d2" if busy == "d1" else "d1" + assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] + + for _ in stream: + pass + _sync_wait_until(lambda: _in_flight(router, busy) == 0) + assert _sync_override_pick(router, "least-busy") == "d1" + assert _selector_is_not_global(router._override_selectors["least-busy"]) + + +@pytest.mark.asyncio +async def test_least_busy_override_releases_the_slot_when_the_overriding_request_fails(): + router = Router( + model_list=_two_deployment_model_list(mock_response="litellm.InternalServerError"), + routing_strategy="simple-shuffle", + num_retries=0, + ) + + with pytest.raises(litellm.InternalServerError): + await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy")) + + await _async_wait_until(lambda: _in_flight(router, "d1") == 0) + assert await _async_override_pick(router, "least-busy") == "d1" + assert _selector_is_not_global(router._override_selectors["least-busy"]) + + +@pytest.mark.asyncio +async def test_latency_based_override_learns_from_the_overriding_requests(): + router = Router( + model_list=_two_deployment_model_list(mock_delay=0.05), routing_strategy="simple-shuffle", num_retries=0 + ) + + def samples(deployment_id: str) -> list[float]: + recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency", []) + return [latency for latency in recorded if latency > 0] + + async def overriding_call() -> str: + sampled_before = {"d1": len(samples("d1")), "d2": len(samples("d2"))} + response = await router.acompletion( + **_proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") + ) + deployment_id = response._hidden_params["model_id"] + await _async_wait_until(lambda: len(samples(deployment_id)) > sampled_before[deployment_id]) + return deployment_id + + served = [await overriding_call() for _ in range(6)] + + assert "d1" in served + assert served[2:] == ["d2"] * 4 + assert _selector_is_not_global(router._override_selectors["latency-based-routing"]) + + +def test_override_selector_is_bound_only_to_the_request_that_asked_for_it(): + router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") + overriding = _proxy_shaped_request(model="grp", routing_strategy="least-busy") + plain = _proxy_shaped_request(model="grp") + + router.get_available_deployment("grp", request_kwargs=overriding) + router.get_available_deployment("grp", request_kwargs=overriding) + router.get_available_deployment("grp", request_kwargs=plain) + + selector = router._override_selectors["least-busy"] + bound = overriding["litellm_logging_obj"] + for callbacks in ( + bound.dynamic_input_callbacks, + bound.dynamic_success_callbacks, + bound.dynamic_async_success_callbacks, + bound.dynamic_failure_callbacks, + bound.dynamic_async_failure_callbacks, + ): + assert callbacks == [selector] + unbound = plain["litellm_logging_obj"] + assert unbound.dynamic_input_callbacks is None and unbound.dynamic_success_callbacks is None + assert unbound.dynamic_failure_callbacks is None and unbound.dynamic_async_failure_callbacks is None + + +def test_override_matching_the_router_strategy_is_not_bound_twice(): + router = Router(model_list=_two_deployment_model_list(), routing_strategy="least-busy") + request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") + + router.get_available_deployment("grp", request_kwargs=request) + + assert request["litellm_logging_obj"].dynamic_input_callbacks is None + + +@pytest.mark.asyncio +async def test_override_matching_a_routing_group_strategy_records_each_request_once(): + router = Router( + model_list=_two_deployment_model_list(), + routing_strategy="simple-shuffle", + routing_groups=[RoutingGroup(group_name="lat", models=["grp"], routing_strategy="latency-based-routing")], + num_retries=0, + ) + request = _proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") + assert router._globally_registered_strategies() == {"simple-shuffle", "latency-based-routing"} + + response = await router.acompletion(**request) + deployment_id = response._hidden_params["model_id"] + await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None) + + assert len(router.cache.get_cache("grp_map")[deployment_id]["latency"]) == 1 + assert request["litellm_logging_obj"].dynamic_success_callbacks is None + + +def test_bind_override_selector_to_request_binds_once_and_ignores_requests_without_logging(): + router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") + selector = router._get_override_strategy_selector("least-busy") + request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") + request["litellm_logging_obj"].dynamic_success_callbacks = ["langfuse"] + + router._bind_override_selector_to_request("least-busy", selector, request) + router._bind_override_selector_to_request("least-busy", selector, request) + router._bind_override_selector_to_request("least-busy", selector, None) + router._bind_override_selector_to_request("least-busy", selector, {"model": "grp"}) + + logging_obj = request["litellm_logging_obj"] + assert logging_obj.dynamic_success_callbacks == ["langfuse", selector] + assert logging_obj.dynamic_input_callbacks == [selector] + assert logging_obj.dynamic_async_failure_callbacks == [selector] + assert _selector_is_not_global(selector) + + def _quality_group(strategy="latency-based-routing"): return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9c29b9d829a..c46a080976c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15975,6 +15975,60 @@ async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warni assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) +@pytest.mark.asyncio +async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the_group(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"}, + "model_info": {"id": "gpt-5-nano"}, + }, + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"}, + "model_info": {"id": "gpt-5-mini-dep"}, + }, + ], + routing_strategy="simple-shuffle", + ) + + by_group = await router.acompletion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) + by_id = await router.acompletion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) + + assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" + assert by_group.choices[0].message.content == "B" + assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" + + +def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_deployment_id(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"}, + "model_info": {"id": "gpt-5-nano"}, + }, + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"}, + "model_info": {"id": "gpt-5-mini-dep"}, + }, + ], + routing_strategy="simple-shuffle", + ) + + with patch.object(router, "routing_strategy_pre_call_checks") as pre_call_checks: + by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) + assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" + pre_call_checks.assert_called_once() + assert pre_call_checks.call_args.kwargs["deployment"]["model_info"]["id"] == "gpt-5-mini-dep" + + by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) + assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" + pre_call_checks.assert_called_once() + + class TestMemberAutoRouterInference: @pytest.fixture(autouse=True) def runtime(self, monkeypatch: pytest.MonkeyPatch) -> None: