From eb48850a1cf50a13a281cdaf8974195fa662371b Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 21:26:30 +0000 Subject: [PATCH 01/21] feat(proxy): bind JWT claims to registered agents via agent_id_jwt_field JWT auth validated Entra app tokens but never carried an agent identity into the authenticated principal, so agent policies (trace id requirement, per-agent MCP restrictions, agent spend attribution) only applied to virtual keys bound to an agent. A new litellm_jwtauth field, agent_id_jwt_field, names the claim (dot notation supported) that is matched against a registered agent's id, then name; the canonical agent_id flows through the standard and proxy-admin JWT paths, and a configured claim naming no registered agent fails closed with 403 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 9 + litellm/proxy/auth/handle_jwt.py | 46 ++++- litellm/proxy/auth/user_api_key_auth.py | 3 + .../proxy/auth/test_handle_jwt.py | 179 ++++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 70 +++++++ 5 files changed, 305 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b315a2beac9..de6972f5e76 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4694,6 +4694,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): @@ -4924,6 +4925,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/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4304542fc83..0e09fce268c 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 @@ -51,6 +51,7 @@ from litellm.proxy._types import ( TeamMemberAddRequest, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, global_agent_registry from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id from litellm.proxy.auth.route_checks import RouteChecks @@ -623,6 +624,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 +1387,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 +1417,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: AgentRegistry, + ) -> 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, @@ -2209,6 +2237,7 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, request_headers: dict | None = None, request_method: str | None = None, + agent_registry: AgentRegistry = global_agent_registry, ) -> JWTAuthBuilderResult: """Main authentication and authorization builder""" # Check if OIDC UserInfo endpoint is enabled, but fall back to standard @@ -2268,9 +2297,21 @@ 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=agent_registry + ) + # 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 +2555,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 687f36bbe8b..a7f5d2b914b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1559,6 +1559,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 @@ -1584,6 +1585,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), ) @@ -1604,6 +1606,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), ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 94226b5404d..2fe8729b78e 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 "", + ) + + 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, + agent_registry=_entra_agent_registry(), + ) + + 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, + ) + + 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, + agent_registry=_entra_agent_registry(), + ) + + 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 fded86d43af..0c656d7875a 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(): """ From 35cff949248feea53bd59e05c4f6dad5f86c5741 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:34:14 +0000 Subject: [PATCH 02/21] fix(proxy): resolve the agent registry lazily in JWT auth to break the import cycle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 0e09fce268c..fcbcf35dba9 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -51,7 +51,6 @@ from litellm.proxy._types import ( TeamMemberAddRequest, UserAPIKeyAuth, ) -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, global_agent_registry from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id from litellm.proxy.auth.route_checks import RouteChecks @@ -62,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, @@ -128,6 +128,20 @@ 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: ... + + def get_agent_by_name(self, agent_name: str) -> AgentResponse | None: ... + + +def _global_agent_lookup() -> AgentLookup: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + return global_agent_registry + + def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody: """Decode an OIDC discovery response body.""" return response.json() @@ -1424,7 +1438,7 @@ class JWTAuthManager: def resolve_agent_id( jwt_handler: JWTHandler, jwt_valid_token: Mapping[str, object], - agent_registry: AgentRegistry, + agent_registry: AgentLookup, ) -> str | None: agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token) if agent_claim is None: @@ -2237,7 +2251,7 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, request_headers: dict | None = None, request_method: str | None = None, - agent_registry: AgentRegistry = global_agent_registry, + agent_registry: AgentLookup | None = None, ) -> JWTAuthBuilderResult: """Main authentication and authorization builder""" # Check if OIDC UserInfo endpoint is enabled, but fall back to standard @@ -2298,7 +2312,9 @@ class JWTAuthManager: user_id = object_id agent_id: Final = JWTAuthManager.resolve_agent_id( - jwt_handler=jwt_handler, jwt_valid_token=jwt_valid_token, agent_registry=agent_registry + jwt_handler=jwt_handler, + jwt_valid_token=jwt_valid_token, + agent_registry=agent_registry if agent_registry is not None else _global_agent_lookup(), ) # Check admin access From 4435aa601dbcfe264fe2fb74d7695e4c1f9e4319 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:28:02 +0000 Subject: [PATCH 03/21] fix(proxy): keep JWT agent binding through AUTO_REGISTER key creation The virtual key created by AUTO_REGISTER replaced the JWT principal without the agent_id auth_builder had resolved from agent_id_jwt_field, so agent policies were skipped on that request and every later mapped-key request. Pass the bound agent_id into generate_key_helper_fn and onto the returned principal. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 4 + .../proxy/auth/test_user_api_key_auth.py | 152 ++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ef32438893e..4d428fc6eb8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -850,6 +850,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 @@ -881,6 +882,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, @@ -969,6 +971,7 @@ async def _auto_register_jwt_mapping( if auto_registered_key is not None: auto_registered_key.org_id = org_id auto_registered_key.end_user_id = end_user_id + auto_registered_key.agent_id = agent_id auto_registered_key.api_key = auto_registered_key.token return auto_registered_key @@ -1635,6 +1638,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/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index a6017f1ca35..92c87df5060 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 @@ -2176,6 +2176,158 @@ 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) + principal = IdentityStore._principal_from_key( + UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team"), + 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=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 +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. From efad8deb713ebd84200b50b7424bdf76131e4cb7 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 14 Sep 2026 20:55:02 +0000 Subject: [PATCH 04/21] fix(alerting): send llm_exceptions Slack alert for 5xx HTTPException/ProxyException Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 12 +++- tests/test_litellm/proxy/test_proxy_utils.py | 75 +++++++++++++------- 2 files changed, 58 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c095586b6c9..5bfd7f5d05f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -429,6 +429,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 @@ -2885,9 +2893,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/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9def21c0573..9506275be51 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1,6 +1,7 @@ import datetime as real_datetime import smtplib from typing import Final +from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException @@ -9,15 +10,10 @@ from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.proxy.utils import PrismaClient, ProxyLogging, get_custom_url, join_paths from litellm.types.guardrails import GuardrailEventHooks -from unittest.mock import MagicMock, patch - -from litellm.proxy.utils import get_custom_url, join_paths - - def test_get_custom_url(monkeypatch): monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") custom_url = get_custom_url(request_base_url="http://0.0.0.0:4000", route="ui/") @@ -1303,10 +1299,9 @@ 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): import asyncio from unittest.mock import AsyncMock @@ -1325,7 +1320,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 +1333,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: @@ -2110,9 +2139,7 @@ def test_create_model_info_response_resolves_alias_to_deployment_model(): ] ) - response = create_model_info_response( - model_id="bedrock-claude-opus-5", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="bedrock-claude-opus-5", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2141,9 +2168,7 @@ def test_create_model_info_response_keeps_exact_alias_over_generalized_deploymen ] ) - response = create_model_info_response( - model_id="claude-opus-5", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="claude-opus-5", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2167,9 +2192,7 @@ def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_na ] ) - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2194,9 +2217,7 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): ] ) - response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2274,7 +2295,9 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): await proxy_logging_obj.post_call_failure_hook( request_data={"metadata": {}}, - original_exception=HTTPException(status_code=400, detail="Upstream passthrough request failed with status 400"), + original_exception=HTTPException( + status_code=400, detail="Upstream passthrough request failed with status 400" + ), user_api_key_dict=UserAPIKeyAuth(), traceback_str=upstream_traceback, ) From f9298d897908ad7b82c674fbc18d8f2b32737e4f Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 14 Sep 2026 20:56:10 +0000 Subject: [PATCH 05/21] style(tests): drop unrelated formatting churn Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_utils.py | 28 +++++++++++++------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9506275be51..bc86d3311af 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1,7 +1,6 @@ import datetime as real_datetime import smtplib from typing import Final -from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException @@ -10,10 +9,15 @@ from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import PrismaClient, ProxyLogging, get_custom_url, join_paths +from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from unittest.mock import MagicMock, patch + +from litellm.proxy.utils import get_custom_url, join_paths + + def test_get_custom_url(monkeypatch): monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") custom_url = get_custom_url(request_base_url="http://0.0.0.0:4000", route="ui/") @@ -2139,7 +2143,9 @@ def test_create_model_info_response_resolves_alias_to_deployment_model(): ] ) - response = create_model_info_response(model_id="bedrock-claude-opus-5", provider="openai", llm_router=router) + response = create_model_info_response( + model_id="bedrock-claude-opus-5", provider="openai", llm_router=router + ) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2168,7 +2174,9 @@ def test_create_model_info_response_keeps_exact_alias_over_generalized_deploymen ] ) - response = create_model_info_response(model_id="claude-opus-5", provider="openai", llm_router=router) + response = create_model_info_response( + model_id="claude-opus-5", provider="openai", llm_router=router + ) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2192,7 +2200,9 @@ def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_na ] ) - response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=router) + response = create_model_info_response( + model_id="gpt-4o", provider="openai", llm_router=router + ) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2217,7 +2227,9 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): ] ) - response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) + response = create_model_info_response( + model_id="my-embeddings", provider="openai", llm_router=router + ) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2295,9 +2307,7 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): await proxy_logging_obj.post_call_failure_hook( request_data={"metadata": {}}, - original_exception=HTTPException( - status_code=400, detail="Upstream passthrough request failed with status 400" - ), + original_exception=HTTPException(status_code=400, detail="Upstream passthrough request failed with status 400"), user_api_key_dict=UserAPIKeyAuth(), traceback_str=upstream_traceback, ) From 4d4d3fb18a28bb071089b163835551f90cbfa360 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:09:14 +0000 Subject: [PATCH 06/21] fix(proxy): bind agent registry into JWTHandler and keep persisted agent id on AUTO_REGISTER race Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 23 ++++-- litellm/proxy/auth/user_api_key_auth.py | 1 - litellm/proxy/proxy_server.py | 2 + .../proxy/auth/test_handle_jwt.py | 4 +- .../proxy/auth/test_user_api_key_auth.py | 70 ++++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 41 +++++++++++ 6 files changed, 128 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index fcbcf35dba9..94ca3047f45 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -131,15 +131,21 @@ class _UserInfoResponse(Protocol): 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: ... + 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: ... + def get_agent_by_name(self, agent_name: str) -> AgentResponse | None: + """The agent registered under ``agent_name``, if any.""" -def _global_agent_lookup() -> AgentLookup: - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry +class _NoRegisteredAgents: + """The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches.""" - return global_agent_registry + 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: @@ -213,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, @@ -2251,7 +2261,6 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, request_headers: dict | None = None, request_method: str | None = None, - agent_registry: AgentLookup | None = None, ) -> JWTAuthBuilderResult: """Main authentication and authorization builder""" # Check if OIDC UserInfo endpoint is enabled, but fall back to standard @@ -2314,7 +2323,7 @@ class JWTAuthManager: agent_id: Final = JWTAuthManager.resolve_agent_id( jwt_handler=jwt_handler, jwt_valid_token=jwt_valid_token, - agent_registry=agent_registry if agent_registry is not None else _global_agent_lookup(), + agent_registry=jwt_handler.agent_lookup, ) # Check admin access diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4d428fc6eb8..1ef0c7abd80 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -971,7 +971,6 @@ async def _auto_register_jwt_mapping( if auto_registered_key is not None: auto_registered_key.org_id = org_id auto_registered_key.end_user_id = end_user_id - auto_registered_key.agent_id = agent_id auto_registered_key.api_key = auto_registered_key.token return auto_registered_key diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..919357498af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6217,6 +6217,7 @@ class ProxyConfig: ) global_agent_registry.load_agents_from_config(agent_config) + jwt_handler.bind_agent_lookup(global_agent_registry) mcp_servers_config: Final = config.get("mcp_servers", None) if mcp_servers_config: @@ -8182,6 +8183,7 @@ class ProxyConfig: global_agent_registry as AGENT_REGISTRY, ) + jwt_handler.bind_agent_lookup(AGENT_REGISTRY) try: async with AGENT_RECONCILE_LOCK: db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 2fe8729b78e..814e31535e0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -6923,6 +6923,7 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a 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, @@ -6934,7 +6935,6 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, - agent_registry=_entra_agent_registry(), ) assert result["is_proxy_admin"] is is_admin_token @@ -6949,6 +6949,7 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch 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( @@ -6961,7 +6962,6 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, - agent_registry=_entra_agent_registry(), ) 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 92c87df5060..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 @@ -2189,8 +2189,8 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id(): plaintext = "sk-auto-registered-agent" token_hash = hash_token(plaintext) - principal = IdentityStore._principal_from_key( - UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team"), + 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), ) @@ -2210,7 +2210,7 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id(): 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=principal, + return_value=persisted_principal, ), ): result = await _auto_register_jwt_mapping( @@ -2233,6 +2233,70 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id(): 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 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 e109b650da7..0a93e607313 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -3740,6 +3740,47 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_ ] +@pytest.mark.asyncio +@pytest.mark.parametrize("agents_source", ["config", "db"]) +async def test_ProxyConfig_agent_loading_binds_registry_to_jwt_agent_claims(clean_agent_registry, agents_source): + """A JWT agent claim must resolve against the agents the proxy loaded, whichever source registered them.""" + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="appid"), + ) + original_lookup = proxy_server.jwt_handler.agent_lookup + proxy_server.jwt_handler.bind_agent_lookup(jwt_handler.agent_lookup) + try: + if agents_source == "config": + await ProxyConfig()._init_non_llm_configs( + config={"agents": [_config_agent("loaded-agent")]}, + config_file_path=None, + ) + else: + 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) + + resolved = JWTAuthManager.resolve_agent_id( + jwt_handler=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) + + 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", From f8e26deb54fb46aca3df5cc60060ce0c8143e05b Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:38:00 +0000 Subject: [PATCH 07/21] fix(proxy): bind JWT agent lookup at startup regardless of agent source Move jwt_handler.bind_agent_lookup out of the YAML and DB agent loading paths and into ProxyStartupEvent._initialize_jwt_auth so agents created via the API or UI after startup, with no agents in config and no DB agent reload, still resolve for agent_id_jwt_field Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +-- .../proxy/proxy_server/test_proxy_config.py | 34 ++++++++++++------- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 919357498af..a5869be1e48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6217,7 +6217,6 @@ class ProxyConfig: ) global_agent_registry.load_agents_from_config(agent_config) - jwt_handler.bind_agent_lookup(global_agent_registry) mcp_servers_config: Final = config.get("mcp_servers", None) if mcp_servers_config: @@ -8183,7 +8182,6 @@ class ProxyConfig: global_agent_registry as AGENT_REGISTRY, ) - jwt_handler.bind_agent_lookup(AGENT_REGISTRY) try: async with AGENT_RECONCILE_LOCK: db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) @@ -9515,6 +9513,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/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 0a93e607313..805627487dd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -3741,42 +3741,50 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_ @pytest.mark.asyncio -@pytest.mark.parametrize("agents_source", ["config", "db"]) -async def test_ProxyConfig_agent_loading_binds_registry_to_jwt_agent_claims(clean_agent_registry, agents_source): - """A JWT agent claim must resolve against the agents the proxy loaded, whichever source registered them.""" +@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, JWTHandler + 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 - jwt_handler = JWTHandler() - jwt_handler.update_environment( - prisma_client=None, - user_api_key_cache=UserApiKeyCache(), - litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="appid"), - ) original_lookup = proxy_server.jwt_handler.agent_lookup - proxy_server.jwt_handler.bind_agent_lookup(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, ) - else: + 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=jwt_handler, + 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 From bd9fdd77e99bf994984ca34b98b7b66b8e5f9dee Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 14 Sep 2026 22:47:04 +0000 Subject: [PATCH 08/21] test(alerting): type the _alerted helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index bc86d3311af..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 @@ -1305,9 +1305,8 @@ class TestPostCallFailureHookLLMExceptionAlerting: client errors must be excluded so a guardrail content-policy block never pages on-call. 5xx proxy errors still alert.""" - async def _alerted(self, exc): + async def _alerted(self, exc: Exception) -> AsyncMock: import asyncio - from unittest.mock import AsyncMock from litellm.proxy._types import AlertType, UserAPIKeyAuth From bdc63d590ecf3e3a396c09a7e400b6ff59fed607 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 00:42:18 +0000 Subject: [PATCH 09/21] fix(router): keep weighted routing when a deployment id equals a model_name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 2 +- tests/test_litellm/test_router.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 0657c1e05ba..ca0f685b6df 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12502,7 +12502,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/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f1b445fb1bd..1977fe715cd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15806,3 +15806,29 @@ async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warni assert binding is None assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] 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" From 1943667fef6036437acab574d5a024ac10d563db Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 00:51:13 +0000 Subject: [PATCH 10/21] fix(router): run sync pre-call checks when a model_name collides with a deployment id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 2 +- tests/test_litellm/test_router.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index ca0f685b6df..e56530a19e3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2461,7 +2461,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 = { diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1977fe715cd..f412af4564b 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15832,3 +15832,31 @@ async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the 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() From 9c59feee7cca3de0cb9727e9463cc5a3f66bf27b Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 15 Sep 2026 01:07:41 +0000 Subject: [PATCH 11/21] fix(headroom): protect the cached prefix through the last cache_control breakpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/compression/compress.py | 24 +++++- .../test_litellm/compression/test_compress.py | 84 +++++++++++++++++++ .../guardrail_hooks/test_headroom.py | 49 ++++++++--- 3 files changed, 144 insertions(+), 13 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index c646baf9d9e..99410a533f9 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -205,21 +205,41 @@ def _extract_anthropic_tool_exchange_spans( return spans, None +def _has_cache_control(message: Mapping[str, object]) -> bool: + if message.get("cache_control") is not None: + return True + content: Final = message.get("content") + return isinstance(content, list) and any( + isinstance(part, Mapping) and part.get("cache_control") is not None for part in content + ) + + +def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + last_breakpoint: Final = max( + (index for index, msg in enumerate(messages) if _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 + - 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. + The provider caches the exact bytes of that prefix, so rewriting any row inside + it 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:] - last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] - return system_indices + last_user + last_assistant + last_assistant: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] + return tuple(dict.fromkeys(system_indices + last_user + last_assistant + _cached_prefix_indices(messages))) def _combine_scores( diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index 6827c37dfd5..57992ffcc55 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/tests/test_litellm/compression/test_compress.py @@ -53,3 +53,87 @@ def test_every_system_row_is_protected(): def test_no_user_or_assistant_rows(): assert sorted(get_protected_indices([{"role": "system", "content": "sys"}])) == [0] 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 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 d8eeb8d2b8a..315c85936f3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1797,12 +1797,8 @@ PARTS_MESSAGES = [ { "role": "user", "content": [ - {"type": "text", "text": "Earlier turn.", "cache_control": {"type": "ephemeral"}}, - { - "type": "text", - "text": "Second block. " + "B" * 5000, - "cache_control": {"type": "ephemeral", "ttl": "1h"}, - }, + {"type": "text", "text": "Earlier turn."}, + {"type": "text", "text": "Second block. " + "B" * 5000}, ], }, { @@ -1891,14 +1887,9 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( messages = result["structured_messages"] history_content = messages[1]["content"] - # Rewritten all-text row collapses to one part carrying the LAST declared - # breakpoint: an Anthropic breakpoint caches the prefix ending at its - # part, so after the merge the last one (and its TTL) still describes the - # row. assert isinstance(history_content, list) assert len(history_content) == 1 assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" - assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} # Mixed row passes through byte-identical. assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] # The service-declared hash still drives retrieve-tool injection on a restored row. @@ -2523,6 +2514,42 @@ 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] + + # --------------------------------------------------------------------------- # #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP # gateway) executes headroom_retrieve and echoes the recovered original content From e390dfbb64160e3aa6a32a47ab597e8002a7d437 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 15 Sep 2026 01:13:17 +0000 Subject: [PATCH 12/21] fix(headroom): format protected index assignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/compression/compress.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 99410a533f9..0af9618382f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -238,7 +238,9 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int """ 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:] - last_assistant: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] + last_assistant: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[ + -1: + ] return tuple(dict.fromkeys(system_indices + last_user + last_assistant + _cached_prefix_indices(messages))) From f808c6899ff7508622a4f8981cf7f5c59a9c1535 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 04:18:08 +0000 Subject: [PATCH 13/21] fix(router): bind per-request routing_strategy override selectors to the request's callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 18 ++ litellm/router.py | 22 +- .../test_litellm_logging.py | 18 ++ .../test_router_routing_groups.py | 224 +++++++++++++++++- 4 files changed, 280 insertions(+), 2 deletions(-) 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/router.py b/litellm/router.py index 1665583386f..7a852b3ef5f 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: 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/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}] From 0c611e63c86bad89568e60b6a82c9bc866c82471 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 09:12:49 +0000 Subject: [PATCH 14/21] fix(utils): cache custom HuggingFace tokenizers across /utils/token_counter requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 ++- tests/test_litellm/proxy/test_proxy_server.py | 57 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) 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/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 04173ced776..552044e2598 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13521,3 +13521,60 @@ 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 = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + loads: Final[list[tuple[str, str, str | None]]] = [] + + class CountingHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + loads.append((identifier, revision, token)) + return 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", CountingHubTokenizer) + 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 loads == [("my-org/tokenizer", "main", 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 loads == [ + ("my-org/tokenizer", "main", None), + ("my-org/tokenizer", "v2", None), + ("my-org/tokenizer", "main", "hf_test_token"), + ] + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear() From b64e430e93c7fb5a197845b4f5f8f10654d4c42d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 09:37:53 +0000 Subject: [PATCH 15/21] test(proxy): record custom tokenizer loads with a mock instead of a mutable list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 552044e2598..42af8e0af21 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13529,14 +13529,8 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi from litellm import Router from litellm.types.router import DeploymentTypedDict - claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] - loads: Final[list[tuple[str, str, str | None]]] = [] - - class CountingHubTokenizer: - @staticmethod - def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: - loads.append((identifier, revision, token)) - return claude_tokenizer + 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 { @@ -13547,7 +13541,7 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi }, } - monkeypatch.setattr(litellm.utils, "Tokenizer", CountingHubTokenizer) + monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained)) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( @@ -13564,17 +13558,17 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) for _ in range(3) ] - assert loads == [("my-org/tokenizer", "main", None)] + 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 loads == [ - ("my-org/tokenizer", "main", None), - ("my-org/tokenizer", "v2", None), - ("my-org/tokenizer", "main", "hf_test_token"), + 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() From 914ae9b248e0b3e5c0dd0c770a5a19dcb9f4bcfd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 11:38:28 -0700 Subject: [PATCH 16/21] test(ui): use the current deployment affinity label --- .../edit_auto_router_modal.integration.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 646e83a773b..0bb3340ac09 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -105,7 +105,7 @@ describe("EditAutoRouterModal keyword matching", () => { expect(screen.queryByText("Advanced: Compression")).not.toBeInTheDocument(); expect(screen.queryByText("Model Access Groups")).not.toBeInTheDocument(); await user.click(screen.getByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" })); + await user.click(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); expect(modelPatchUpdateCall).toHaveBeenLastCalledWith( From e62f0d037600a7825ecd8e360e556f75a322d3d0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 11:43:04 -0700 Subject: [PATCH 17/21] fix(router): reject unknown capability policy fields --- litellm/router_strategy/complexity_router/config.py | 2 +- .../router_strategy/test_complexity_router.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fc1ad739903..7c47bac68da 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -616,7 +616,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) efficient_tier: str = Field( description="Tier used when the efficient model's forecasted solve probability meets the adjusted threshold", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8dbd36087b5..38411ef52ea 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2500,6 +2500,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" From 501be3143d0c46144887527e514497d81c27ce5b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 15 Sep 2026 11:30:45 -0700 Subject: [PATCH 18/21] fix(proxy): enforce organization budgets when max_budget is 0 _organization_max_budget_check returned early whenever org_max_budget was <= 0, so an organization with an explicit max_budget of 0 was treated as unlimited instead of zero allowance. Key, team, and user budget checks already skip only on None; align organization budgets with that convention. validate_team_org_change had the same defect in a different shape: it used a truthy check on the org's max_budget when validating a team move, so an explicit 0 there silently skipped the guard too. Co-Authored-By: Claude Sonnet 5 --- litellm/proxy/auth/auth_checks.py | 3 +- .../management_endpoints/team_endpoints.py | 6 +- .../proxy/auth/test_auth_checks.py | 65 +++++++++++++++++++ .../test_team_endpoints.py | 49 ++++++++++++++ 4 files changed, 118 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6d61ad4d3e8..35d34f9d6de 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5767,8 +5767,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/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2d04a4d1e04..5da024e136e 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/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f87f2def93..66e8b26b957 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5829,6 +5829,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/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a2f534fbe4d..91a1d30325b 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(): """ From 896f35c7513c689469942326c3f25c75d0c1fca7 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 11:52:20 -0700 Subject: [PATCH 19/21] fix(router): extract capability tasks with request scoped markers --- .../complexity_router/complexity_router.py | 7 ++-- .../router_strategy/test_complexity_router.py | 34 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 72f22e0bab0..68dfde394e7 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2210,7 +2210,8 @@ 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)) + markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) + asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), markers)) opening_task: Final = 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 @@ -2239,9 +2240,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_classifier_task(request_kwargs, markers), ) verdict: Final = parse_capability_classifier_verdict(content) threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 38411ef52ea..c8916f10965 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2588,6 +2588,40 @@ class TestCapabilityClassifier: complexity_router_config=_capability_router_config(**overrides), ) + @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( From cadb7ee44dd5b5b6902fb40c9e7048d494642811 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 12:05:30 -0700 Subject: [PATCH 20/21] fix(router): preserve native encrypted capability tasks --- .../complexity_router/complexity_router.py | 13 ++++++++--- .../router_strategy/test_complexity_router.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 68dfde394e7..c8de22e91fb 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2211,8 +2211,15 @@ class ComplexityRouter(CustomLogger): raise ValueError("capability classifier is not configured") markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) - asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), markers)) - opening_task: Final = asks_newest_first[-1] if asks_newest_first else prompt + 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 @@ -2240,7 +2247,7 @@ class ComplexityRouter(CustomLogger): messages_for_call, request_kwargs, max_output_tokens=capability.max_output_tokens, - encrypted_task=_encrypted_classifier_task(request_kwargs, markers), + encrypted_task=encrypted_task, ) verdict: Final = parse_capability_classifier_verdict(content) threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index c8916f10965..0931b9d01a7 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2588,6 +2588,28 @@ class TestCapabilityClassifier: complexity_router_config=_capability_router_config(**overrides), ) + @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( From 55fc0deabc882be978713f7752983141ef031bb4 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 12:09:05 -0700 Subject: [PATCH 21/21] style(router): format encrypted task selection --- .../router_strategy/complexity_router/complexity_router.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index c8de22e91fb..d20abefbb2a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2218,7 +2218,9 @@ class ComplexityRouter(CustomLogger): 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 + 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