merge: integrate released capability classifier

litellm/compression/compress.py

@@ -214,26 +214,33 @@ def _message_has_cache_control(message: Mapping[str, object]) -> bool:
     return False
 
 
+def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
+    last_breakpoint: Final = max(
+        (index for index, msg in enumerate(messages) if _message_has_cache_control(msg)),
+        default=-1,
+    )
+    return tuple(range(last_breakpoint + 1))
+
+
 def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
     """
     Return indices of messages that must never be compressed:
     - All system messages
     - The last user message
     - The last assistant message
-    - Any message carrying an Anthropic cache_control breakpoint
+    - Every message up to and including the last one carrying an Anthropic cache_control breakpoint
 
     The last user message is what the model is being asked to act on right now,
     so compressing it replaces the live instruction with a marker. Compression
     guardrails share this policy; see the Headroom guardrail. A cache_control
-    breakpoint pins the provider's prompt-cache prefix to that row's exact
-    bytes, so rewriting a marked row anywhere in history turns the next
-    request's cache read into a cache write.
+    breakpoint pins the provider's prompt-cache prefix to the exact bytes of every
+    row up to it, so rewriting any row inside that prefix turns the next request's
+    cache read into a cache write.
     """
     system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
     last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
     assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
-    cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
-    return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
+    return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages)))
 
 
 def _combine_scores(

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

litellm/proxy/_types.py

@@ -4716,6 +4716,7 @@ class JWTAuthBuilderResult(TypedDict):
     org_id: str | None
     team_membership: LiteLLM_TeamMembership | None
     jwt_claims: dict  # Decoded JWT token claims (avoids re-decoding)
+    agent_id: ReadOnly[str | None]
 
 
 class ClientSideFallbackModel(TypedDict, total=False):
@@ -4954,6 +4955,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     user_allowed_roles: list[str] | None = None
     user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.")
     end_user_id_jwt_field: str | None = None
+    agent_id_jwt_field: str | None = Field(
+        default=None,
+        description=(
+            "The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID "
+            "app token). Supports dot notation. The value is matched against a registered agent's agent_id, "
+            "then agent_name, and the request is rejected when it matches neither."
+        ),
+    )
     public_key_ttl: float = 600
     public_key_stale_ttl: float = Field(
         default=DEFAULT_JWKS_STALE_TTL,

litellm/proxy/auth/auth_checks.py

@@ -5773,8 +5773,7 @@ async def _organization_max_budget_check(
     if org_table.litellm_budget_table is not None:
         org_max_budget = org_table.litellm_budget_table.max_budget
 
-    # Only check if organization has a valid max_budget set
-    if org_max_budget is None or org_max_budget <= 0:
+    if org_max_budget is None:
         return
 
     # Read spend from cross-pod counter (Redis-first) or cached object (fallback)

litellm/proxy/auth/handle_jwt.py

@@ -14,7 +14,7 @@ import hashlib
 import os
 import re
 import time
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
 
 import httpx
@@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
 )
 from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.repositories.user_repository import UserRepository
+from litellm.types.agents import AgentResponse
 
 from .auth_checks import (
     _allowed_routes_check,
@@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol):
     def json(self) -> dict[str, object]: ...
 
 
+class AgentLookup(Protocol):
+    """The registered-agent lookups a JWT agent claim is matched against."""
+
+    def get_agent_by_id(self, agent_id: str) -> AgentResponse | None:
+        """The agent registered under ``agent_id``, if any."""
+
+    def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
+        """The agent registered under ``agent_name``, if any."""
+
+
+class _NoRegisteredAgents:
+    """The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
+
+    def get_agent_by_id(self, agent_id: str) -> None:
+        return None
+
+    def get_agent_by_name(self, agent_name: str) -> None:
+        return None
+
+
 def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody:
     """Decode an OIDC discovery response body."""
     return response.json()
@@ -198,6 +219,10 @@ class JWTHandler:
         self.leeway = 0
         # Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request.
         self._refresh_locks: dict[str, asyncio.Lock] = {}  # mutable-ok: lock registry, keyed by JWKS url
+        self.agent_lookup: AgentLookup = _NoRegisteredAgents()
+
+    def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None:
+        self.agent_lookup = agent_lookup
 
     def update_environment(
         self,
@@ -623,6 +648,12 @@ class JWTHandler:
             object_id = default_value
         return object_id
 
+    def get_agent_claim(self, token: Mapping[str, object]) -> str | None:
+        if self.litellm_jwtauth.agent_id_jwt_field is None:
+            return None
+        claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field)
+        return claim if isinstance(claim, str) and claim else None
+
     def get_org_id(self, token: dict, default_value: str | None) -> str | None:
         if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM):
             return token.get(self.LITELLM_ORG_ID_CLAIM)
@@ -1380,6 +1411,7 @@ class JWTAuthManager:
         api_key: str,
         jwt_valid_token: dict | None = None,
         user_email: str | None = None,
+        agent_id: str | None = None,
     ) -> JWTAuthBuilderResult | None:
         """Check admin status and route access permissions"""
         if not jwt_handler.is_admin(scopes=scopes):
@@ -1409,8 +1441,28 @@ class JWTAuthManager:
             org_id=org_id,
             team_membership=None,
             jwt_claims=jwt_valid_token or {},
+            agent_id=agent_id,
         )
 
+    @staticmethod
+    def resolve_agent_id(
+        jwt_handler: JWTHandler,
+        jwt_valid_token: Mapping[str, object],
+        agent_registry: AgentLookup,
+    ) -> str | None:
+        agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token)
+        if agent_claim is None:
+            return None
+        agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name(
+            agent_name=agent_claim
+        )
+        if agent is None:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}",
+            )
+        return agent.agent_id
+
     @staticmethod
     async def find_and_validate_specific_team_id(
         jwt_handler: JWTHandler,
@@ -2268,9 +2320,23 @@ class JWTAuthManager:
             elif rbac_role == LitellmUserRoles.INTERNAL_USER:
                 user_id = object_id
 
+        agent_id: Final = JWTAuthManager.resolve_agent_id(
+            jwt_handler=jwt_handler,
+            jwt_valid_token=jwt_valid_token,
+            agent_registry=jwt_handler.agent_lookup,
+        )
+
         # Check admin access
         admin_result: Final = await JWTAuthManager.check_admin_access(
-            jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email
+            jwt_handler,
+            scopes,
+            route,
+            user_id,
+            org_id,
+            api_key,
+            jwt_valid_token,
+            user_email=user_email,
+            agent_id=agent_id,
         )
         if admin_result:
             await JWTAuthManager._attach_team_from_header_for_admin(
@@ -2514,4 +2580,5 @@ class JWTAuthManager:
             token=api_key,
             team_membership=team_membership_object,
             jwt_claims=jwt_valid_token,
+            agent_id=agent_id,
         )

litellm/proxy/auth/user_api_key_auth.py

@@ -852,6 +852,7 @@ async def _auto_register_jwt_mapping(
     user_id: str | None = None,
     org_id: str | None = None,
     end_user_id: str | None = None,
+    agent_id: str | None = None,
 ) -> UserAPIKeyAuth | None:
     """
     Auto-register: create a new virtual key + mapping for an unrecognised JWT
@@ -884,6 +885,7 @@ async def _auto_register_jwt_mapping(
         team_id=team_id,
         user_id=user_id,
         organization_id=org_id,
+        agent_id=agent_id,
         metadata={
             "auto_registered": True,
             "jwt_claim_field": virtual_key_claim_field,
@@ -1567,6 +1569,7 @@ async def _user_api_key_auth_builder(
                     org_id: Final = result["org_id"]
                     team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
                     jwt_claims = result.get("jwt_claims", None)
+                    agent_id: Final[str | None] = result.get("agent_id")
 
                     if is_proxy_admin:
                         # Proxy admins authenticate via auth_builder (full
@@ -1592,6 +1595,7 @@ async def _user_api_key_auth_builder(
                             end_user_id=end_user_id,
                             parent_otel_span=parent_otel_span,
                             jwt_claims=jwt_claims,
+                            agent_id=agent_id,
                             **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
                         )
 
@@ -1612,6 +1616,7 @@ async def _user_api_key_auth_builder(
                         user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
                         user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
                         jwt_claims=jwt_claims,
+                        agent_id=agent_id,
                         **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
                     )
 
@@ -1635,6 +1640,7 @@ async def _user_api_key_auth_builder(
                             user_id=user_id,
                             org_id=org_id,
                             end_user_id=end_user_id,
+                            agent_id=agent_id,
                         )
                         if auto_registered is not None:
                             auto_registered.jwt_claims = jwt_claims

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(

litellm/proxy/proxy_server.py

@@ -9516,6 +9516,9 @@ class ProxyStartupEvent:
             user_api_key_cache=user_api_key_cache,
             litellm_jwtauth=litellm_jwtauth,
         )
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+        jwt_handler.bind_agent_lookup(global_agent_registry)
 
     @classmethod
     def _add_proxy_budget_to_db(cls):

litellm/proxy/utils.py

@@ -430,6 +430,14 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback:
         detail.setdefault("guardrail_mode", event_hook)
 
 
+def _is_client_error_exception(exc: Exception) -> bool:
+    if isinstance(exc, HTTPException):
+        return exc.status_code < 500
+    if isinstance(exc, ProxyException):
+        return not (exc.code.isdigit() and int(exc.code) >= 500)
+    return False
+
+
 def _exception_changes_request_flow(exc: BaseException) -> bool:
     """
     True for guardrail exceptions the proxy turns into an alternate request flow
@@ -2886,9 +2894,7 @@ class ProxyLogging:
 
         ### ALERTING ###
         await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail")
-        if AlertType.llm_exceptions in self.alert_types and not isinstance(
-            original_exception, (HTTPException, ProxyException)
-        ):
+        if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception):
             """
             Just alert on LLM API exceptions. Do not alert on user errors
 

litellm/router.py

@@ -1622,6 +1622,24 @@ class Router:
             return
         await selector.async_pre_call_check(deployment, parent_otel_span)
 
+    def _bind_override_selector_to_request(
+        self, strategy: str, selector: RouterStrategySelector | None, request_kwargs: Mapping[str, object] | None
+    ) -> None:
+        if selector is None or request_kwargs is None or strategy in self._globally_registered_strategies():
+            return
+        logging_obj: Final = request_kwargs.get("litellm_logging_obj")
+        if isinstance(logging_obj, LiteLLMLogging):
+            logging_obj.add_dynamic_callback(selector)
+
+    def _globally_registered_strategies(self) -> frozenset[str]:
+        configured: Final = (
+            self.routing_strategy,
+            *(group.routing_strategy for group in self._routing_groups.values()),
+        )
+        return frozenset(
+            normalized for normalized in map(self._normalize_strategy, configured) if normalized is not None
+        )
+
     def _get_routing_context(
         self, model: str, request_kwargs: dict | None = None
     ) -> tuple[str | None, RouterStrategySelector | None]:
@@ -1647,7 +1665,9 @@ class Router:
         override: Final = self._get_request_routing_strategy_override(request_kwargs)
         if override is not None:
             verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override)
-            return override, self._get_override_strategy_selector(override)
+            override_selector: Final = self._get_override_strategy_selector(override)
+            self._bind_override_selector_to_request(override, override_selector, request_kwargs)
+            return override, override_selector
 
         group_name: Final = model if self.get_routing_group(model) is not None else self._model_to_group.get(model)
         if group_name is None:
@@ -2461,7 +2481,7 @@ class Router:
 
             ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit)
             ## only run if model group given, not model id
-            if not self.has_model_id(model):
+            if model in self.model_names or not self.has_model_id(model):
                 self.routing_strategy_pre_call_checks(deployment=deployment)
 
             input_kwargs: Final = {
@@ -12512,7 +12532,7 @@ class Router:
         # check if aliases set on litellm model alias map
         if specific_deployment is True:
             return model, self._get_deployment_by_litellm_model(model=model)
-        elif self.has_model_id(model):
+        elif model not in self.model_names and self.has_model_id(model):
             deployment: Final = self.get_deployment(model_id=model)
             if deployment is not None:
                 deployment_model: Final = deployment.litellm_params.model

litellm/router_strategy/complexity_router/complexity_router.py

@@ -2217,8 +2217,18 @@ class ComplexityRouter(CustomLogger):
         if capability is None or classifier_system_prompt is None:
             raise ValueError("capability classifier is not configured")
 
-        asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), self._reminder_markers))
-        opening_task: Final = asks_newest_first[-1] if asks_newest_first else prompt
+        markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
+        encrypted_task: Final = _encrypted_classifier_task(request_kwargs, markers)
+        asks_newest_first: Final = (
+            () if encrypted_task is not None else tuple(_iter_human_asks_newest_first(messages or (), markers))
+        )
+        opening_task: Final = (
+            "The delegated task in the following agent_message."
+            if encrypted_task is not None
+            else asks_newest_first[-1]
+            if asks_newest_first
+            else prompt
+        )
         latest_follow_up: Final = asks_newest_first[0] if len(asks_newest_first) > 1 else None
         task_messages: list[AllMessageValues] = [  # mutable-ok: the latest message gains optional image parts below
             {"role": "user", "content": opening_task},  # mutable-ok: SDK messages are dict-shaped
@@ -2246,9 +2256,7 @@ class ComplexityRouter(CustomLogger):
             messages_for_call,
             request_kwargs,
             max_output_tokens=capability.max_output_tokens,
-            encrypted_task=_encrypted_classifier_task(
-                request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
-            ),
+            encrypted_task=encrypted_task,
         )
         verdict: Final = parse_capability_classifier_verdict(content)
         forecast: Final = capability.classify(verdict)

litellm/router_strategy/complexity_router/config.py

@@ -641,7 +641,7 @@ class CapabilityCalibrationConfig(BaseModel):
 class CapabilityClassifierConfig(BaseModel):
     """Switchyard-compatible probability threshold policy for two model tiers."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(extra="forbid", frozen=True)
     card: CapabilityCardConfig | None = None
     selective_policy: SelectivePolicy | None = None
     empirical_supplement: str | None = Field(default=None, min_length=1, max_length=4000)

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:

tests/test_litellm/compression/test_compress.py

@@ -56,11 +56,87 @@ def test_no_user_or_assistant_rows():
     assert get_protected_indices([]) == ()
 
 
+def test_rows_before_last_cache_control_breakpoint_are_protected():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "old question"},
+        {
+            "role": "assistant",
+            "content": "old answer",
+            "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
+        },
+        {"role": "tool", "tool_call_id": "t1", "content": "large file body"},
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
+        },
+        {
+            "role": "assistant",
+            "content": "ack",
+            "tool_calls": [{"id": "t2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
+        },
+        {"role": "tool", "tool_call_id": "t2", "content": "later tool output"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert protected == [0, 1, 2, 3, 4, 5, 7]
+    assert 6 not in protected
+
+
+def test_cache_control_directly_on_message_protects_prefix():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "tool", "tool_call_id": "before", "content": "large file body"},
+        {"role": "user", "content": "old question"},
+        {
+            "role": "tool",
+            "tool_call_id": "marked",
+            "content": "cached tool",
+            "cache_control": {"type": "ephemeral"},
+        },
+        {"role": "tool", "tool_call_id": "after", "content": "later tool output"},
+        {"role": "assistant", "content": "ack"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert 1 in protected
+    assert 3 in protected
+    assert 4 not in protected
+
+
+def test_no_cache_control_leaves_history_compressible():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "old question"},
+        {"role": "assistant", "content": "old answer"},
+        {"role": "tool", "tool_call_id": "t1", "content": "large file body"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    assert sorted(get_protected_indices(messages)) == [0, 2, 4]
+
+
+def test_non_mapping_content_parts_are_not_cache_control():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": ["not", "a", "dict"]},
+        {"role": "assistant", "content": "old answer"},
+        {"role": "tool", "tool_call_id": "t1", "content": "plain string"},
+        {"role": "user", "content": "live instruction"},
+    ]
+
+    protected = sorted(get_protected_indices(messages))
+
+    assert protected == [0, 2, 4]
+    assert 1 not in protected
+    assert 3 not in protected
+
+
 def test_mid_history_cache_control_part_is_protected():
-    # A large cached tool result from a few turns back, not the last user or
-    # last assistant row -- exactly the row a provider prompt-cache pins to
-    # exact bytes. Rewriting it (even leaving the marker on) changes those
-    # bytes and turns the next request's cache read into a cache write.
     messages = [
         {"role": "user", "content": "old question"},
         {"role": "assistant", "content": "old answer"},
@@ -74,9 +150,7 @@ def test_mid_history_cache_control_part_is_protected():
         {"role": "user", "content": "live instruction"},
     ]
 
-    # index 3 = last assistant, index 4 = last user (both protected by role
-    # regardless), index 2 = the cache_control-marked row itself.
-    assert sorted(get_protected_indices(messages)) == [2, 3, 4]
+    assert sorted(get_protected_indices(messages)) == [0, 1, 2, 3, 4]
 
 
 def test_cache_control_directly_on_message_is_protected():
@@ -116,8 +190,6 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(
 
 
 def test_compress_keeps_part_level_cache_control_row_verbatim():
-    # compress() scores text-only copies of the rows, where a part-level marker
-    # is gone; protection has to read the original rows or the pinned row is stubbed.
     stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
     pinned = {
         "role": "user",
@@ -126,9 +198,9 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
         ],
     }
     messages = [
-        stale_log,
-        {"role": "assistant", "content": "old answer"},
         pinned,
+        {"role": "assistant", "content": "old answer"},
+        stale_log,
         {"role": "assistant", "content": "ack"},
         {"role": "user", "content": "live instruction"},
     ]
@@ -142,6 +214,6 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
     )
 
     assert len(result["messages"]) == len(messages)
-    assert result["messages"][2] == pinned
-    assert result["messages"][0] != stale_log
+    assert result["messages"][0] == pinned
+    assert result["messages"][2] != stale_log
     assert len(result["cache"]) >= 1

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]

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -5855,6 +5855,71 @@ async def test_organization_budget_check_carries_org_state_on_the_token():
     assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0)
 
 
+@pytest.mark.parametrize(
+    "max_budget, spend, expect_blocked",
+    [
+        (0.0, 0.0, True),  # explicit zero budget blocks even a fresh org with no spend
+        (0.0, 7.4e-06, True),  # any spend at all against a zero budget blocks
+        (None, 999.0, False),  # unlimited (None) never blocks, regardless of spend
+        (5.0, 4.99, False),  # a positive budget under its cap still passes
+    ],
+)
+@pytest.mark.asyncio
+async def test_organization_zero_max_budget_is_enforced(max_budget, spend, expect_blocked):
+    """An explicit organization max_budget of 0 must mean zero allowance, matching
+    key/team/user semantics, not unlimited.
+
+    Regression for LIT-7797: `_organization_max_budget_check` returned early
+    whenever `org_max_budget <= 0`, so an org configured with max_budget=0 could
+    spend without limit.
+    """
+    from litellm.proxy._types import LiteLLM_OrganizationTable
+    from litellm.proxy.auth.auth_checks import _organization_max_budget_check
+
+    org_table = LiteLLM_OrganizationTable(
+        organization_id="o1",
+        organization_alias="zero-budget-org",
+        budget_id="b1",
+        created_by="admin",
+        updated_by="admin",
+        spend=spend,
+        litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget) if max_budget is not None else None,
+    )
+    token = UserAPIKeyAuth(token="k1", org_id="o1")
+    user_api_key_cache = UserApiKeyCache()
+    await user_api_key_cache.async_set_cache(
+        key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable
+    )
+
+    async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs):
+        return spend
+
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.budget_alerts = AsyncMock()
+
+    with patch(  # test-quality-ok: _organization_max_budget_check imports get_current_spend locally
+        "litellm.proxy.proxy_server.get_current_spend", _spend
+    ):
+        if expect_blocked:
+            with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                await _organization_max_budget_check(
+                    valid_token=token,
+                    team_object=None,
+                    prisma_client=MagicMock(),
+                    user_api_key_cache=user_api_key_cache,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
+            assert exc_info.value.max_budget == max_budget
+        else:
+            await _organization_max_budget_check(
+                valid_token=token,
+                team_object=None,
+                prisma_client=MagicMock(),
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+
+
 @pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"])
 @pytest.mark.asyncio
 async def test_spend_capable_non_llm_routes_still_enforce_budget(route):

tests/test_litellm/proxy/auth/test_handle_jwt.py

@@ -23,6 +23,7 @@ from litellm.proxy._types import (
     ProxyException,
 )
 from litellm.caching.dual_cache import DualCache
+from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
 from litellm.proxy.auth.handle_jwt import (
     JWKS_FETCH_ATTEMPTS,
     STALE_CACHE_KEY_PREFIX,
@@ -32,6 +33,7 @@ from litellm.proxy.auth.handle_jwt import (
     JWTHandler,
     NoMatchingJWTPublicKeyError,
 )
+from litellm.types.agents import AgentResponse
 
 
 @pytest.mark.asyncio
@@ -6786,3 +6788,180 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla
     }
     assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == []
     assert user.teams == []
+
+
+def _entra_agent_registry() -> AgentRegistry:
+    registry = AgentRegistry()
+    registry.register_agent(
+        AgentResponse(
+            agent_id="canonical-agent-id",
+            agent_name="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
+            agent_card_params={"name": "research-agent", "url": "http://localhost:9999/a2a", "version": "1.0.0"},
+            litellm_params={"require_trace_id_on_calls_by_agent": True},
+        )
+    )
+    return registry
+
+
+def _entra_agent_jwt_handler(agent_id_jwt_field: str | None) -> JWTHandler:
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=DualCache(),
+        litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", agent_id_jwt_field=agent_id_jwt_field),
+    )
+    return jwt_handler
+
+
+@pytest.mark.parametrize(
+    "claim_value",
+    ["canonical-agent-id", "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"],
+    ids=["matches_agent_id", "matches_agent_name"],
+)
+def test_resolve_agent_id_returns_canonical_agent_id(claim_value: str):
+    """An Entra app token's azp claim binds to the registered agent by id or by name and yields its canonical id."""
+    jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
+
+    resolved = JWTAuthManager.resolve_agent_id(
+        jwt_handler=jwt_handler,
+        jwt_valid_token={"sub": "sp-object-id-1234", "azp": claim_value},
+        agent_registry=_entra_agent_registry(),
+    )
+
+    assert resolved == "canonical-agent-id"
+
+
+def test_resolve_agent_id_reads_nested_claim():
+    jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="entra.client_id")
+
+    resolved = JWTAuthManager.resolve_agent_id(
+        jwt_handler=jwt_handler,
+        jwt_valid_token={"sub": "sp-object-id-1234", "entra": {"client_id": "canonical-agent-id"}},
+        agent_registry=_entra_agent_registry(),
+    )
+
+    assert resolved == "canonical-agent-id"
+
+
+def test_resolve_agent_id_rejects_claim_for_unregistered_agent():
+    """A configured agent claim naming no registered agent fails closed with 403 instead of falling back to an unbound identity."""
+    jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
+
+    with pytest.raises(HTTPException) as exc_info:
+        JWTAuthManager.resolve_agent_id(
+            jwt_handler=jwt_handler,
+            jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
+            agent_registry=_entra_agent_registry(),
+        )
+
+    assert exc_info.value.status_code == 403
+
+
+@pytest.mark.parametrize(
+    "token",
+    [
+        {"sub": "sp-object-id-1234"},
+        {"sub": "sp-object-id-1234", "azp": ""},
+        {"sub": "sp-object-id-1234", "azp": ["canonical-agent-id"]},
+    ],
+    ids=["claim_absent", "claim_empty", "claim_not_a_string"],
+)
+def test_resolve_agent_id_returns_none_when_claim_unusable(token: dict):
+    jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
+
+    assert (
+        JWTAuthManager.resolve_agent_id(
+            jwt_handler=jwt_handler, jwt_valid_token=token, agent_registry=_entra_agent_registry()
+        )
+        is None
+    )
+
+
+def test_resolve_agent_id_ignores_claim_when_field_not_configured():
+    """Without agent_id_jwt_field an azp claim (even an unknown one) leaves JWT auth behaviour unchanged."""
+    jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field=None)
+
+    resolved = JWTAuthManager.resolve_agent_id(
+        jwt_handler=jwt_handler,
+        jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
+        agent_registry=_entra_agent_registry(),
+    )
+
+    assert resolved is None
+
+
+def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandler, str]:
+    """A JWTHandler that verifies RS256 tokens against a pre-cached JWKS, plus a signed Entra-style app token."""
+    jwks_url = "https://login.microsoftonline.test/discovery/v2.0/keys"
+    monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
+    monkeypatch.delenv("JWT_AUDIENCE", raising=False)
+    private_key, jwk = _get_rsa_key_and_jwk(kid="entra-kid")
+    cache = DualCache()
+    cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk])
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=cache,
+        litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="azp"),
+    )
+    token = _encode_rsa_jwt(
+        private_key,
+        issuer="https://login.microsoftonline.test/lit7664-tenant/v2.0",
+        audience="api://litellm",
+        kid="entra-kid",
+        extra_claims={"sub": "sp-object-id-1234", "azp": azp, "scope": scope},
+    )
+    return jwt_handler, token
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_admin_token", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
+async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool):
+    """auth_builder carries the resolved agent id into JWTAuthBuilderResult on both the admin and standard paths."""
+    jwt_handler, token = _entra_signed_app_token(
+        monkeypatch,
+        azp="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
+        scope=LiteLLM_JWTAuth().admin_jwt_scope if is_admin_token else "",
+    )
+    jwt_handler.bind_agent_lookup(_entra_agent_registry())
+
+    result = await JWTAuthManager.auth_builder(
+        api_key=token,
+        jwt_handler=jwt_handler,
+        request_data={"model": "gpt-5.6"},
+        general_settings={"enforce_rbac": False},
+        route="/key/info" if is_admin_token else "/chat/completions",
+        prisma_client=None,
+        user_api_key_cache=None,
+        parent_otel_span=None,
+        proxy_logging_obj=None,
+    )
+
+    assert result["is_proxy_admin"] is is_admin_token
+    assert result["agent_id"] == "canonical-agent-id"
+
+
+@pytest.mark.asyncio
+async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch):
+    """An unknown agent claim is rejected even when the token would otherwise be a proxy admin."""
+    jwt_handler, token = _entra_signed_app_token(
+        monkeypatch,
+        azp="00000000-0000-0000-0000-000000000000",
+        scope=LiteLLM_JWTAuth().admin_jwt_scope,
+    )
+    jwt_handler.bind_agent_lookup(_entra_agent_registry())
+
+    with pytest.raises(HTTPException) as exc_info:
+        await JWTAuthManager.auth_builder(
+            api_key=token,
+            jwt_handler=jwt_handler,
+            request_data={"model": "gpt-5.6"},
+            general_settings={"enforce_rbac": False},
+            route="/key/info",
+            prisma_client=None,
+            user_api_key_cache=None,
+            parent_otel_span=None,
+            proxy_logging_obj=None,
+        )
+
+    assert exc_info.value.status_code == 403

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -1937,6 +1937,76 @@ async def test_standard_jwt_auth_propagates_user_email():
     assert result.api_key is None
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
+async def test_jwt_auth_propagates_agent_id_to_user_api_key_auth(is_proxy_admin: bool):
+    """The agent id resolved by auth_builder must land on UserAPIKeyAuth.agent_id so
+    agent-scoped checks (trace id requirement, MCP server/tool restrictions, spend
+    attribution) apply to JWT callers the same way they apply to agent-bound keys."""
+    jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
+    general_settings = {"enable_jwt_auth": True}
+    user_api_key_cache = DualCache()
+    jwt_handler = MagicMock()
+    jwt_handler.is_jwt.return_value = True
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(agent_id_jwt_field="azp")
+
+    user_object = LiteLLM_UserTable(user_id="sp-object-id-1234", user_role="internal_user")
+    mock_jwt_result = {
+        "is_proxy_admin": is_proxy_admin,
+        "team_object": None,
+        "user_object": user_object,
+        "end_user_object": None,
+        "org_object": None,
+        "token": jwt_token,
+        "team_id": None,
+        "user_id": "sp-object-id-1234",
+        "user_email": None,
+        "end_user_id": None,
+        "org_id": None,
+        "team_membership": None,
+        "jwt_claims": {"sub": "sp-object-id-1234", "azp": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
+        "agent_id": "canonical-agent-id",
+    }
+
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+    mock_request.method = "POST"
+    mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
+    mock_request.query_params = {}
+    mock_request.state = SimpleNamespace()
+
+    with (
+        patch.multiple(  # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
+            "litellm.proxy.proxy_server",
+            general_settings=general_settings,
+            premium_user=True,
+            master_key="sk-master",
+            prisma_client=None,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=MagicMock(),
+            jwt_handler=jwt_handler,
+        ),
+        patch(  # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
+            "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
+            new_callable=AsyncMock,
+            return_value=mock_jwt_result,
+        ),
+    ):
+        result = await _user_api_key_auth_builder(
+            request=mock_request,
+            api_key=jwt_token,
+            azure_api_key_header="",
+            anthropic_api_key_header=None,
+            google_ai_studio_api_key_header=None,
+            azure_apim_header=None,
+            request_data={"model": "gpt-5.6"},
+        )
+
+    assert result.agent_id == "canonical-agent-id"
+    assert result.user_id == "sp-object-id-1234"
+    assert result.api_key is None
+
+
 @pytest.mark.asyncio
 async def test_auto_register_binds_api_key_to_token_hash():
     """
@@ -2106,6 +2176,222 @@ async def test_auto_register_first_request_propagates_user_email():
     assert result.api_key == "hashed-auto-key"
 
 
+@pytest.mark.asyncio
+async def test_auto_register_stamps_new_key_with_jwt_agent_id():
+    """The virtual key AUTO_REGISTER creates must carry the agent id auth_builder bound
+    from the JWT claim, and the first request's principal must carry it too, or the
+    mapped-key path would drop the agent policies on that request and every later one."""
+    from litellm.proxy.auth.auth_method import AuthMethod
+    from litellm.proxy.auth.resolvers.models import CredentialRef
+    from litellm.proxy.auth.resolvers.store import IdentityStore
+    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
+    from litellm.proxy.proxy_server import hash_token
+
+    plaintext = "sk-auto-registered-agent"
+    token_hash = hash_token(plaintext)
+    persisted_principal = IdentityStore._principal_from_key(
+        UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"),
+        auth_method=AuthMethod.API_KEY,
+        credential_ref=CredentialRef(token_id=token_hash),
+    )
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
+    user_api_key_cache = MagicMock()
+    user_api_key_cache.async_set_cache = AsyncMock()
+    jwt_handler = MagicMock()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
+    generate_key = AsyncMock(return_value={"token": plaintext})
+
+    with (
+        patch(  # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
+            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
+            generate_key,
+        ),
+        patch(  # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
+            "litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
+            new_callable=AsyncMock,
+            return_value=persisted_principal,
+        ),
+    ):
+        result = await _auto_register_jwt_mapping(
+            virtual_key_claim_field="appid",
+            claim_value="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
+            jwt_handler=jwt_handler,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=MagicMock(),
+            cache_key="jwt_key_mapping:appid:2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
+            team_id="validated-team",
+            user_id="validated-user",
+            agent_id="canonical-agent-id",
+        )
+
+    assert generate_key.await_args is not None
+    assert generate_key.await_args.kwargs["agent_id"] == "canonical-agent-id"
+    assert result is not None
+    assert result.agent_id == "canonical-agent-id"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("losing_agent_id", ["other-agent", None], ids=["different_agent", "no_agent_claim"])
+async def test_auto_register_race_loser_keeps_winners_agent_id(losing_agent_id: str | None):
+    """When two requests race to AUTO_REGISTER the same mapping claim, the loser must run as
+    the persisted key, agent binding included. Every later request on that mapping uses the
+    winner's key, so stamping the loser's own (or missing) agent id on it would give one request
+    different agent policies and spend attribution than all the others."""
+    from litellm.proxy.auth.auth_method import AuthMethod
+    from litellm.proxy.auth.resolvers.models import CredentialRef
+    from litellm.proxy.auth.resolvers.store import IdentityStore
+    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
+
+    winner_hash = "winner-key-hash"
+    winner_principal = IdentityStore._principal_from_key(
+        UserAPIKeyAuth(token=winner_hash, user_id="validated-user", team_id="validated-team", agent_id="winner-agent"),
+        auth_method=AuthMethod.API_KEY,
+        credential_ref=CredentialRef(token_id=winner_hash),
+    )
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
+        side_effect=Exception("Unique constraint failed on the fields: (`jwt_claim_name`,`jwt_claim_value`)")
+    )
+    prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
+    user_api_key_cache = MagicMock()
+    user_api_key_cache.async_set_cache = AsyncMock()
+    jwt_handler = MagicMock()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
+
+    with (
+        patch(  # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
+            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
+            new_callable=AsyncMock,
+            return_value={"token": "sk-orphaned-loser-key"},
+        ),
+        patch(  # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
+            "litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object",
+            new_callable=AsyncMock,
+            return_value=winner_hash,
+        ),
+        patch(  # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
+            "litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
+            new_callable=AsyncMock,
+            return_value=winner_principal,
+        ),
+    ):
+        result = await _auto_register_jwt_mapping(
+            virtual_key_claim_field="tid",
+            claim_value="shared-tenant",
+            jwt_handler=jwt_handler,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=MagicMock(),
+            cache_key="jwt_key_mapping:tid:shared-tenant",
+            team_id="validated-team",
+            user_id="validated-user",
+            agent_id=losing_agent_id,
+        )
+
+    assert result is not None
+    assert result.token == winner_hash
+    assert result.agent_id == "winner-agent"
+
+
+@pytest.mark.asyncio
+async def test_jwt_auto_register_forwards_bound_agent_id():
+    """When a JWT under AUTO_REGISTER also carries the configured agent claim, the agent
+    id auth_builder resolved must reach the key creation, not be dropped when
+    valid_token is swapped for the freshly registered key."""
+    jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
+    user_api_key_cache = DualCache()
+    jwt_handler = MagicMock()
+    jwt_handler.is_jwt.return_value = True
+    jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"})
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        virtual_key_claim_field="sub",
+        virtual_key_mapping_cache_ttl=300,
+        agent_id_jwt_field="appid",
+    )
+    user_object = LiteLLM_UserTable(user_id="validated-user", user_role="internal_user")
+    mock_jwt_result = {
+        "is_proxy_admin": False,
+        "team_object": None,
+        "user_object": user_object,
+        "end_user_object": None,
+        "org_object": None,
+        "token": jwt_token,
+        "team_id": "validated-team",
+        "user_id": "validated-user",
+        "user_email": None,
+        "end_user_id": None,
+        "org_id": None,
+        "team_membership": None,
+        "jwt_claims": {"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
+        "agent_id": "canonical-agent-id",
+    }
+    auto_register = AsyncMock(
+        return_value=UserAPIKeyAuth(
+            token="hashed-auto-key",
+            api_key="hashed-auto-key",
+            team_id="validated-team",
+            user_id="validated-user",
+            agent_id="canonical-agent-id",
+        )
+    )
+
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+    mock_request.method = "POST"
+    mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
+    mock_request.query_params = {}
+    mock_request.state = SimpleNamespace()
+
+    with (
+        patch.multiple(  # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
+            "litellm.proxy.proxy_server",
+            general_settings={"enable_jwt_auth": True},
+            premium_user=True,
+            master_key="sk-master",
+            prisma_client=MagicMock(),
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=MagicMock(),
+            jwt_handler=jwt_handler,
+        ),
+        patch(  # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
+            "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
+            new_callable=AsyncMock,
+            return_value=_PendingAutoRegister(
+                claim_field="sub",
+                claim_value="user1",
+                cache_key="jwt_key_mapping:sub:user1",
+            ),
+        ),
+        patch(  # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
+            "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
+            new_callable=AsyncMock,
+            return_value=mock_jwt_result,
+        ),
+        patch(  # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
+            "litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping",
+            auto_register,
+        ),
+    ):
+        result = await _user_api_key_auth_builder(
+            request=mock_request,
+            api_key=jwt_token,
+            azure_api_key_header="",
+            anthropic_api_key_header=None,
+            google_ai_studio_api_key_header=None,
+            azure_apim_header=None,
+            request_data={"model": "gpt-5.6"},
+        )
+
+    assert auto_register.await_args is not None
+    assert auto_register.await_args.kwargs["agent_id"] == "canonical-agent-id"
+    assert result.agent_id == "canonical-agent-id"
+    assert result.api_key == "hashed-auto-key"
+
+
 class TestJWTOAuth2Coexistence:
     """
     Test that JWT and OAuth2 auth can coexist on the same instance.

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py

@@ -900,7 +900,7 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr
     )
 
     assert has_headroom_retrieve_tool(result.get("tools") or [])
-    (issued, _expiry), = guardrail._issued_hashes_by_call_id.values()
+    ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values()
     assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"})
 
 
@@ -953,7 +953,6 @@ async def test_anthropic_assistant_history_never_reaches_compression_service(gua
     assert result["messages"][1]["content"] == [{"type": "text", "text": table}]
 
 
-
 def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape():
     """By the time an Anthropic Messages API response reaches the agentic-loop
     gate, the OpenAI-shaped tool this guardrail injects (type: "function")
@@ -2342,9 +2341,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
     )
     assert streamed_text == final_answer
     assert not any("function_call" in str(getattr(event, "type", "")) for event in events)
-    assert not any(
-        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
-    )
+    assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
     mock_get.assert_called_once()
     assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
 
@@ -2399,9 +2396,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
         getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
     )
     assert streamed_text == final_answer
-    assert not any(
-        getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
-    )
+    assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
     mock_get.assert_called_once()
     assert len(upstream.calls) == 2
     assert not json.loads(upstream.calls[1].request.content).get("stream")
@@ -2514,6 +2509,38 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail):
     assert messages[3] == compressed_history[1]
 
 
+CACHED_PREFIX_MESSAGES = [
+    {"role": "system", "content": "You are Claude Code. " + "S" * 5000},
+    {"role": "user", "content": "old question " + "Q" * 5000},
+    {
+        "role": "assistant",
+        "content": "Reading the file now.",
+        "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
+    },
+    {"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000},
+    {
+        "role": "user",
+        "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
+    },
+    {
+        "role": "assistant",
+        "content": "Listing now.",
+        "tool_calls": [{"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
+    },
+    {"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000},
+    {"role": "assistant", "content": "Finished listing."},
+    {"role": "user", "content": "live instruction"},
+]
+
+
+@pytest.mark.asyncio
+async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail):
+    wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES)
+
+    assert [row.get("tool_call_id") for row in wire] == ["new_1"]
+    assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5]
+
+
 CACHE_MARKED_HISTORY_MESSAGES = [
     {"role": "system", "content": "You are Claude Code. " + "S" * 5000},
     {"role": "user", "content": "old question " + "Q" * 5000},
@@ -2529,6 +2556,7 @@ CACHE_MARKED_HISTORY_MESSAGES = [
         "cache_control": {"type": "ephemeral"},
     },
     {"role": "assistant", "content": "Summarized the file for you."},
+    {"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000},
     {"role": "user", "content": "live instruction"},
 ]
 

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():
     """

tests/test_litellm/proxy/proxy_server/test_proxy_config.py

@@ -3764,6 +3764,55 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
     ]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("agents_source", ["config", "db", "api"])
+async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry(
+    clean_agent_registry, agents_source
+):
+    """A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup."""
+    from litellm.proxy import proxy_server
+    from litellm.proxy._types import LiteLLM_JWTAuth
+    from litellm.proxy.auth.handle_jwt import JWTAuthManager
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.types.agents import AgentResponse
+
+    original_lookup = proxy_server.jwt_handler.agent_lookup
+    try:
+        proxy_server.ProxyStartupEvent._initialize_jwt_auth(
+            general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}},
+            prisma_client=None,
+            user_api_key_cache=UserApiKeyCache(),
+        )
+        if agents_source == "config":
+            await ProxyConfig()._init_non_llm_configs(
+                config={"agents": [_config_agent("loaded-agent")]},
+                config_file_path=None,
+            )
+        elif agents_source == "db":
+            prisma_client = MagicMock()
+            prisma_client.db.litellm_agentstable.find_many = AsyncMock(
+                return_value=[_FakeAgentRow("db-id", "loaded-agent")]
+            )
+            await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
+        else:
+            clean_agent_registry.register_agent(
+                agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent"))
+            )
+
+        resolved = JWTAuthManager.resolve_agent_id(
+            jwt_handler=proxy_server.jwt_handler,
+            jwt_valid_token={"appid": "loaded-agent"},
+            agent_registry=proxy_server.jwt_handler.agent_lookup,
+        )
+    finally:
+        proxy_server.jwt_handler.bind_agent_lookup(original_lookup)
+        proxy_server.jwt_handler.update_environment(
+            prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth()
+        )
+
+    assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "config, expected_agent_names",

tests/test_litellm/proxy/test_proxy_server.py

@@ -13521,3 +13521,54 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
     assert response.tokenizer_type == "huggingface_tokenizer"
     assert response.total_tokens > 0
     assert_loop_stayed_free(took, lags)
+
+
+async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch):
+    from tokenizers import Tokenizer
+
+    from litellm import Router
+    from litellm.types.router import DeploymentTypedDict
+
+    claude_tokenizer: Final[Tokenizer] = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"]
+    from_pretrained: Final = MagicMock(return_value=claude_tokenizer)
+
+    def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict:
+        return {
+            "model_name": model_name,
+            "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
+            "model_info": {
+                "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": revision, "auth_token": auth_token}
+            },
+        }
+
+    monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained))
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.llm_router",
+        Router(
+            model_list=[
+                deployment("self-hosted", "main", None),
+                deployment("self-hosted-pinned", "v2", None),
+                deployment("self-hosted-private", "main", "hf_test_token"),
+            ]
+        ),
+    )
+    litellm.utils._select_custom_tokenizer_helper.cache_clear()
+    try:
+        responses: Final = [
+            await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once"))
+            for _ in range(3)
+        ]
+        assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)]
+        assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses)
+        assert len({response.total_tokens for response in responses}) == 1
+        assert responses[0].total_tokens > 0
+
+        await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once"))
+        await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once"))
+        assert from_pretrained.call_args_list == [
+            mock.call("my-org/tokenizer", revision="main", token=None),
+            mock.call("my-org/tokenizer", revision="v2", token=None),
+            mock.call("my-org/tokenizer", revision="main", token="hf_test_token"),
+        ]
+    finally:
+        litellm.utils._select_custom_tokenizer_helper.cache_clear()

tests/test_litellm/proxy/test_proxy_utils.py

@@ -13,7 +13,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.types.guardrails import GuardrailEventHooks
 
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from litellm.proxy.utils import get_custom_url, join_paths
 
@@ -1303,12 +1303,10 @@ class TestPostCallFailureHookLLMExceptionAlerting:
     """The llm_exceptions alert is for infra / LLM-API failures, not user
     errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized
     client errors must be excluded so a guardrail content-policy block never
-    pages on-call. ProxyException is such an error; before LIT-3751 only
-    HTTPException was excluded, so AIM blocks paged as if the LLM API failed."""
+    pages on-call. 5xx proxy errors still alert."""
 
-    async def _alerted(self, exc) -> bool:
+    async def _alerted(self, exc: Exception) -> AsyncMock:
         import asyncio
-        from unittest.mock import AsyncMock
 
         from litellm.proxy._types import AlertType, UserAPIKeyAuth
 
@@ -1325,7 +1323,7 @@ class TestPostCallFailureHookLLMExceptionAlerting:
                 user_api_key_dict=UserAPIKeyAuth(),
             )
         await asyncio.sleep(0)  # let the fire-and-forget alert task run
-        return alerting_handler.called
+        return alerting_handler
 
     @pytest.mark.asyncio
     async def test_proxy_exception_does_not_alert(self):
@@ -1338,15 +1336,49 @@ class TestPostCallFailureHookLLMExceptionAlerting:
             code=400,
             openai_code="content_policy_violation",
         )
-        assert await self._alerted(exc) is False
+        assert (await self._alerted(exc)).called is False
 
     @pytest.mark.asyncio
     async def test_http_exception_does_not_alert(self):
-        assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False
+        assert (await self._alerted(HTTPException(status_code=400, detail="blocked"))).called is False
 
     @pytest.mark.asyncio
     async def test_genuine_llm_api_error_still_alerts(self):
-        assert await self._alerted(Exception("upstream 503")) is True
+        assert (await self._alerted(Exception("upstream 503"))).called is True
+
+    @pytest.mark.asyncio
+    async def test_http_exception_5xx_alerts(self):
+        alerting_handler = await self._alerted(
+            HTTPException(
+                status_code=502,
+                detail={
+                    "error": "Headroom compression service returned an error",
+                    "status_code": 503,
+                    "guardrail_name": "headroom-compression-global",
+                },
+            )
+        )
+        assert alerting_handler.called is True
+        assert "headroom-compression-global" in alerting_handler.call_args.kwargs["message"]
+
+    @pytest.mark.asyncio
+    async def test_proxy_exception_5xx_alerts(self):
+        from litellm.proxy._types import ProxyException
+
+        alerting_handler = await self._alerted(
+            ProxyException(
+                message="guardrail backend down",
+                type="internal_server_error",
+                param=None,
+                code=503,
+            )
+        )
+        assert alerting_handler.called is True
+
+    @pytest.mark.asyncio
+    async def test_http_exception_429_does_not_alert(self):
+        alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited"))
+        assert alerting_handler.called is False
 
 
 class TestPostCallFailureHookProxyExceptionLogging:

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -2573,6 +2573,17 @@ class TestCapabilityClassifierConfig:
         with pytest.raises(ValidationError, match="requires classifier_type 'capability'"):
             ComplexityRouterConfig(**config)
 
+    def test_rejects_misspelled_optional_policy_instead_of_using_defaults(self) -> None:
+        with pytest.raises(ValidationError, match="threshold_steps"):
+            CapabilityClassifierConfig.model_validate(
+                {
+                    "efficient_tier": "SIMPLE",
+                    "capable_tier": "REASONING",
+                    "base_threshold": 0.5,
+                    "threshold_steps": 0.2,
+                }
+            )
+
     def test_threshold_defaults_match_switchyard(self):
         config = CapabilityClassifierConfig(efficient_tier=" SIMPLE ", capable_tier=" REASONING ", base_threshold=0.5)
         assert config.efficient_tier == "SIMPLE"
@@ -2691,6 +2702,62 @@ class TestCapabilityClassifier:
         assert "Route to the Efficient model" not in sent
         mock_router_instance.acompletion.assert_awaited_once()
 
+    @pytest.mark.asyncio
+    async def test_encrypted_task_is_not_replaced_by_plaintext_envelope(self, mock_router_instance: MagicMock) -> None:
+        mock_router_instance.aresponses = AsyncMock(
+            return_value=_native_classifier_response(_capability_reply(p_solve=0.8))
+        )
+        router: Final = self._router(mock_router_instance)
+        task: Final = _encrypted_agent_task()
+        request: Final = {"input": [task]}
+        original: Final = deepcopy(request)
+        result: Final = await router.async_pre_routing_hook(model="capability-router", request_kwargs=request)
+        assert result is not None and result.model == "efficient-model"
+        assert result.routing_decision is not None
+        assert result.routing_decision["cause"] == "capability_classifier"
+        mock_router_instance.aresponses.assert_awaited_once()
+        call: Final = mock_router_instance.aresponses.call_args.kwargs
+        assert call["input"][-1] == task
+        plaintext: Final = json.dumps(call["input"][:-1])
+        assert "The delegated task in the following agent_message." in plaintext
+        assert "Message Type: NEW_TASK" not in plaintext
+        assert "opaque-provider-task" not in plaintext
+        assert request == original
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("custom_markers", (False, True))
+    async def test_task_forecast_uses_request_scoped_codex_markers(
+        self, mock_router_instance: MagicMock, custom_markers: bool
+    ) -> None:
+        completion: Final = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
+        mock_router_instance.acompletion = completion
+        router: Final = self._router(
+            mock_router_instance,
+            escalation_keywords=[],
+            **({"reminder_markers": [{"open": "<custom>", "close": "</custom>"}]} if custom_markers else {}),
+        )
+        envelope: Final = "\n".join(_CODEX_ENVELOPES)
+        opening: Final = f"{envelope}\nFix nested behavior"
+        messages: Final = [
+            {"role": "user", "content": opening},
+            {"role": "user", "content": "Preserve empty inputs"},
+            {"role": "user", "content": envelope},
+        ]
+        original: Final = deepcopy(messages)
+        for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"):
+            result: Final = await router.async_pre_routing_hook(
+                model="capability-router", messages=messages, request_kwargs={"metadata": {"user_agent": user_agent}}
+            )
+            assert result is not None and result.model == "efficient-model"
+            sent: Final = completion.call_args.kwargs["messages"]
+            if user_agent.startswith("codex") and not custom_markers:
+                assert [message["content"] for message in sent[1:]] == ["Fix nested behavior", "Preserve empty inputs"]
+            else:
+                assert [message["content"] for message in sent[1:]] == [opening, envelope]
+            assert result.messages == original
+        assert completion.await_count == 3
+        assert messages == original
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("p_solve,expected_model", ((0.95, "capable-model"), (0.98, "efficient-model")))
     async def test_fitted_probability_controls_routing_and_preserves_raw_score(
@@ -6216,10 +6283,16 @@ class TestTierModelAffinity:
         returned: Final = await self._route(router, metadata, "model-b")
 
         assert (first.model, repeated.model, reasoning.model, returned.model) == (
-            "model-a", "model-a", "model-b", "model-a"
+            "model-a",
+            "model-a",
+            "model-b",
+            "model-a",
         )
         assert tuple(result.routing_decision["tier"] for result in (first, repeated, reasoning, returned)) == (
-            "SIMPLE", "SIMPLE", "REASONING", "SIMPLE"
+            "SIMPLE",
+            "SIMPLE",
+            "REASONING",
+            "SIMPLE",
         )
         assert returned.litellm_params == {"temperature": 0.1}
         assert reasoning.litellm_params == {"temperature": 0.9}
@@ -6257,9 +6330,7 @@ class TestTierModelAffinity:
         deployment_affinity: bool,
         plugins: bool,
     ) -> None:
-        router: Final = self._router(
-            mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins
-        )
+        router: Final = self._router(mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins)
         assert (await self._route(router, metadata, "model-a")).model == "model-a"
         assert (await self._route(router, metadata, "model-b")).model == "model-b"
 
@@ -6332,9 +6403,7 @@ class TestTierModelAffinity:
             {
                 "role": "assistant",
                 "content": None,
-                "tool_calls": [
-                    {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
-                ],
+                "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
             },
             {"role": "tool", "tool_call_id": "call_1", "content": [IMG_PART] if gate == "image" else "done"},
         ]
@@ -6379,9 +6448,7 @@ class TestTierModelAffinity:
             {
                 "role": "assistant",
                 "content": None,
-                "tool_calls": [
-                    {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
-                ],
+                "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
             },
             {"role": "tool", "tool_call_id": "call_1", "content": "done"},
         ]
@@ -6411,8 +6478,7 @@ class TestTierModelAffinity:
                     "SIMPLE": "base",
                     **{
                         tier: [
-                            {"model_name": model, "litellm_params": {"temperature": temperature}}
-                            for model in models
+                            {"model_name": model, "litellm_params": {"temperature": temperature}} for model in models
                         ]
                         for tier, models, temperature in (
                             ("MEDIUM", ("shared", "middle"), 0.4),
@@ -6486,7 +6552,11 @@ class TestTierModelAffinity:
             model_name="affinity-router",
             litellm_router_instance=mock_router_instance,
             complexity_router_config=_custom_tier_config(
-                tiers={"SIMPLE": ["model-a", "model-b"], "SECURITY_REVIEW": ["model-a", "model-b"], "COMPLEX": "model-a"},
+                tiers={
+                    "SIMPLE": ["model-a", "model-b"],
+                    "SECURITY_REVIEW": ["model-a", "model-b"],
+                    "COMPLEX": "model-a",
+                },
                 deployment_affinity=True,
                 classification_mode=classification_mode,
                 keyword_tier_rules=[

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

tests/test_litellm/test_router.py

@@ -15975,6 +15975,60 @@ async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warni
     assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)
 
 
+@pytest.mark.asyncio
+async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the_group():
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-5-nano",
+                "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"},
+                "model_info": {"id": "gpt-5-nano"},
+            },
+            {
+                "model_name": "gpt-5-nano",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"},
+                "model_info": {"id": "gpt-5-mini-dep"},
+            },
+        ],
+        routing_strategy="simple-shuffle",
+    )
+
+    by_group = await router.acompletion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}])
+    by_id = await router.acompletion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}])
+
+    assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep"
+    assert by_group.choices[0].message.content == "B"
+    assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep"
+
+
+def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_deployment_id():
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-5-nano",
+                "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"},
+                "model_info": {"id": "gpt-5-nano"},
+            },
+            {
+                "model_name": "gpt-5-nano",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"},
+                "model_info": {"id": "gpt-5-mini-dep"},
+            },
+        ],
+        routing_strategy="simple-shuffle",
+    )
+
+    with patch.object(router, "routing_strategy_pre_call_checks") as pre_call_checks:
+        by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}])
+        assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep"
+        pre_call_checks.assert_called_once()
+        assert pre_call_checks.call_args.kwargs["deployment"]["model_info"]["id"] == "gpt-5-mini-dep"
+
+        by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}])
+        assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep"
+        pre_call_checks.assert_called_once()
+
+
 class TestMemberAutoRouterInference:
     @pytest.fixture(autouse=True)
     def runtime(self, monkeypatch: pytest.MonkeyPatch) -> None: