[Infra] Merge internal dev branch with main (#25036)

* fix(proxy): enforce key-level model allowlist for custom auth

custom_auth_run_common_checks only runs common_checks (team/user/project model checks).
Custom auth now also enforces key-level model restrictions via can_key_call_model.

Move the custom-auth key-access regression tests to test_user_api_key_auth.py and keep test_custom_auth_end_user_budget.py focused on end-user budget behavior.

Made-with: Cursor

* fix(proxy): gate custom-auth key model checks behind opt-in

Keep key-level model allowlist enforcement in custom auth behind `custom_auth_run_common_checks` to preserve backwards compatibility, and update tests to verify default non-enforcement and opt-in enforcement behavior.

Made-with: Cursor

* test(proxy): isolate custom auth default check from shared settings state

Patch `proxy_server.general_settings` to an empty dict in the default custom-auth key-access test so it remains deterministic under shared module state.

Made-with: Cursor

* test(proxy): strengthen custom auth post-check assertions

Tighten custom auth regression tests by asserting exact can_key_call_model args and remove an unused common_checks mock from the default behavior path.

Made-with: Cursor

* fix(agentcore): parse A2A JSON-RPC responses in AgentCore provider

* fix(prompt-templates): ensure_alternating_roles handles tool-call chains

* feat(auth): add JWT claim routing overrides for OAuth2 validation

Made-with: Cursor

* docs(auth): document JWT-to-OAuth2 routing overrides

Add generic docs for running JWT and OAuth2 together, including routing_overrides YAML examples and list-based selector behavior for iss/client_id/aud.

Made-with: Cursor

---------

Co-authored-by: Milan <milan@berri.ai>
Co-authored-by: michelligabriele <gabriele.michelli@icloud.com>

docs/my-website/docs/proxy/oauth2.md

@@ -61,3 +61,24 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 
 Start the LiteLLM Proxy with [`--detailed_debug` mode and you should see more verbose logs](cli.md#detailed_debug)
 
+## Using OAuth2 + JWT Together
+
+If both `enable_oauth2_auth` and `enable_jwt_auth` are enabled, LiteLLM can split auth paths:
+- JWT validation for user tokens
+- OAuth2 introspection for machine tokens
+
+For JWT-shaped machine tokens, configure `litellm_jwtauth.routing_overrides`:
+
+```yaml title="config.yaml"
+general_settings:
+  enable_jwt_auth: true
+  enable_oauth2_auth: true
+  litellm_jwtauth:
+    routing_overrides:
+      - iss: "machine-issuer.example.com"
+        client_id: "MID_LITELLM"
+        path: "oauth2"
+```
+
+For full `routing_overrides` behavior and list-based selectors, see [`/proxy/token_auth`](./token_auth.md#route-jwt-shaped-machine-tokens-to-oauth2).
+

docs/my-website/docs/proxy/token_auth.md

@@ -790,6 +790,47 @@ litellm_jwtauth:
   user_roles_jwt_field: "resource_access.your-client.roles"
 ```
 
+## Route JWT-Shaped Machine Tokens to OAuth2
+
+Use this when both are enabled:
+- `enable_jwt_auth: true` for standard JWT validation
+- `enable_oauth2_auth: true` for OAuth2 introspection
+
+If some machine tokens are also JWT-shaped, configure `routing_overrides` to route matching tokens to OAuth2.
+
+```yaml title="config.yaml"
+general_settings:
+  enable_jwt_auth: true
+  enable_oauth2_auth: true
+  litellm_jwtauth:
+    user_id_jwt_field: "sub"
+    routing_overrides:
+      - iss: "machine-issuer.example.com"
+        client_id: "MID_LITELLM"
+        path: "oauth2"
+```
+
+### Matching behavior
+
+- A rule matches when all configured selectors match token claims
+- Supported selectors: `iss` (required), `client_id` (optional), `aud` (optional)
+- Selector values support both string and list forms
+- If no rule matches, LiteLLM continues with standard JWT validation
+
+### List-based override example
+
+```yaml title="config.yaml"
+general_settings:
+  enable_jwt_auth: true
+  enable_oauth2_auth: true
+  litellm_jwtauth:
+    routing_overrides:
+      - iss: ["machine-issuer.example.com", "backup-issuer.example.com"]
+        client_id: ["MID_LITELLM", "MID_BACKUP"]
+        aud: ["api://litellm", "api://fallback"]
+        path: "oauth2"
+```
+
 ## [BETA] Control Access with OIDC Roles
 
 Allow JWT tokens with supported roles to access the proxy.

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -337,24 +337,35 @@ def _insert_assistant_continue_message(
     """
     Add assistant continuation messages between consecutive user messages.
 
-    Only checks directly adjacent messages to preserve backward compatibility.
+    Skips tool messages and assistant messages with tool calls in the
+    alternation check, matching strict templates like llama.cpp.
     """
     if not ensure_alternating_roles or len(messages) <= 1:
         return messages
 
     continue_message = assistant_continue_message or DEFAULT_ASSISTANT_CONTINUE_MESSAGE
 
+    # Find indexes where assistant_continue should be inserted (before that index)
+    insert_before_indexes: set = set()
+
+    for i in range(len(messages)):
+        curr = messages[i]
+        if _counts_for_alternation(curr) and curr["role"] == "user":
+            # Look backwards for the previous counted message
+            j = i - 1
+            while j >= 0:
+                if _counts_for_alternation(messages[j]):
+                    if messages[j]["role"] == "user":
+                        insert_before_indexes.add(i)
+                    break
+                j -= 1
+
+    # Build the result with assistant_continue inserted at the right positions
     modified_messages: List[AllMessageValues] = []
     for i, message in enumerate(messages):
-        if (
-            i < len(messages) - 1
-            and message.get("role") == "user"
-            and messages[i + 1].get("role") == "user"
-        ):
-            modified_messages.append(message)
+        if i in insert_before_indexes:
             modified_messages.append(continue_message)
-        else:
-            modified_messages.append(message)
+        modified_messages.append(message)
 
     return modified_messages
 

litellm/llms/bedrock/chat/agentcore/transformation.py

@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
 from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.llms.bedrock_agentcore import (
     AgentCoreMessage,
@@ -343,6 +344,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         Parse direct JSON response (non-streaming).
 
         Supports multiple agent response schemas:
+        0. {"jsonrpc": "2.0", "result": {"message": {"parts": [...]}}} - A2A JSON-RPC
         1. {"result": {"role": "assistant", "content": [{"text": "..."}]}} - standard AgentCore
         2. {"response": [{"text": "..."}]} - Strands agent format
         3. {"result": "plain text"} or {"response": "plain text"} - simple string
@@ -361,6 +363,18 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
                 final_message=None,
             )
 
+        # Strategy 0: A2A JSON-RPC format
+        # {"jsonrpc": "2.0", "result": {"message": {"parts": [{"kind": "text", "text": "..."}]}}}
+        if "jsonrpc" in response_json:
+            content = extract_text_from_a2a_response(response_json)
+            if content:
+                return AgentCoreParsedResponse(
+                    content=content,
+                    usage=None,
+                    final_message=None,
+                )
+            # Fall through to other strategies if A2A extraction returned empty
+
         # Strategy 1: {"result": {"content": [{"text": "..."}]}} - standard AgentCore format
         if "result" in response_json and isinstance(response_json["result"], dict):
             result = response_json["result"]

litellm/proxy/_types.py

@@ -4098,6 +4098,24 @@ class ScopeMapping(OIDCPermissions):
     }
 
 
+class JWTRoutingOverride(BaseModel):
+    """
+    Override default auth routing for JWT-shaped bearer tokens.
+
+    A rule matches when all provided selectors match token claims.
+    If matched, request is routed to the configured auth path.
+    """
+
+    iss: Union[str, List[str]]
+    client_id: Optional[Union[str, List[str]]] = None
+    aud: Optional[Union[str, List[str]]] = None
+    path: Literal["oauth2"] = "oauth2"
+
+    model_config = {
+        "extra": "forbid",
+    }
+
+
 class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     """
     A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth.
@@ -4198,6 +4216,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
         default=300,
         description="TTL (seconds) for caching JWT-to-virtual-key mapping lookups.",
     )
+    routing_overrides: Optional[List[JWTRoutingOverride]] = Field(
+        default=None,
+        description="Optional claim-based routing overrides for JWT-shaped tokens. Matching rules route requests to oauth2 before default JWT flow.",
+    )
     #########################################################
 
     def __init__(self, **kwargs: Any) -> None:

litellm/proxy/auth/handle_jwt.py

@@ -16,6 +16,8 @@ from cryptography import x509
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from fastapi import HTTPException
+import jwt
+from jwt.api_jwk import PyJWK
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -71,6 +73,21 @@ class JWTHandler:
 
     prisma_client: Optional[PrismaClient]
     user_api_key_cache: DualCache
+    # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
+    # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
+    #   the key in different ways (e.g. HS* and RS*)."
+    SUPPORTED_JWT_ALGORITHMS = [
+        "RS256",
+        "RS384",
+        "RS512",
+        "PS256",
+        "PS384",
+        "PS512",
+        "ES256",
+        "ES384",
+        "ES512",
+        "EdDSA",
+    ]
 
     def __init__(
         self,
@@ -97,6 +114,30 @@ class JWTHandler:
         parts = token.split(".")
         return len(parts) == 3
 
+    @staticmethod
+    def get_unverified_claims(token: str) -> Optional[dict]:
+        """
+        Decode JWT claims without signature verification.
+        Used for routing decisions before selecting validation path.
+        """
+        if not JWTHandler.is_jwt(token):
+            return None
+
+        try:
+            claims = jwt.decode(
+                token,
+                options={"verify_signature": False, "verify_aud": False},
+                algorithms=JWTHandler.SUPPORTED_JWT_ALGORITHMS,
+            )
+            if isinstance(claims, dict):
+                return claims
+            return None
+        except Exception as e:
+            verbose_proxy_logger.debug(
+                "Failed to decode unverified JWT claims for routing: %s", e
+            )
+            return None
+
     def _rbac_role_from_role_mapping(self, token: dict) -> Optional[RBAC_ROLES]:
         """
         Returns the RBAC role the token 'belongs' to based on role mappings.
@@ -664,30 +705,11 @@ class JWTHandler:
             raise Exception(f"Failed to fetch OIDC UserInfo: {str(e)}")
 
     async def auth_jwt(self, token: str) -> dict:
-        # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
-        # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
-        #   the key in different ways (e.g. HS* and RS*)."
-        algorithms = [
-            "RS256",
-            "RS384",
-            "RS512",
-            "PS256",
-            "PS384",
-            "PS512",
-            "ES256",
-            "ES384",
-            "ES512",
-            "EdDSA",
-        ]
-
         audience = os.getenv("JWT_AUDIENCE")
         decode_options = None
         if audience is None:
             decode_options = {"verify_aud": False}
 
-        import jwt
-        from jwt.api_jwk import PyJWK
-
         header = jwt.get_unverified_header(token)
 
         verbose_proxy_logger.debug("header: %s", header)
@@ -721,7 +743,7 @@ class JWTHandler:
                 payload = jwt.decode(
                     token,
                     public_key_obj,  # type: ignore
-                    algorithms=algorithms,
+                    algorithms=self.SUPPORTED_JWT_ALGORITHMS,
                     options=decode_options,  # type: ignore[arg-type]
                     audience=audience,
                     leeway=self.leeway,  # allow testing of expired tokens
@@ -749,7 +771,7 @@ class JWTHandler:
                 payload = jwt.decode(
                     token,
                     key,
-                    algorithms=algorithms,
+                    algorithms=self.SUPPORTED_JWT_ALGORITHMS,
                     audience=audience,
                     options=decode_options,
                 )

litellm/proxy/auth/user_api_key_auth.py

@@ -11,7 +11,7 @@ import asyncio
 import re
 import secrets
 from datetime import datetime, timezone
-from typing import List, Optional, Tuple, cast
+from typing import Any, List, Optional, Tuple, cast
 
 import fastapi
 from fastapi import HTTPException, Request, WebSocket, status
@@ -139,6 +139,58 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str:
     return api_key
 
 
+def _routing_selector_matches_claim(
+    selector_value: Optional[Any], claim_value: Optional[Any]
+) -> bool:
+    if selector_value is None:
+        return True
+
+    selector_list = (
+        [str(v) for v in selector_value]
+        if isinstance(selector_value, list)
+        else [str(selector_value)]
+    )
+
+    if isinstance(claim_value, list):
+        claim_list = [str(v) for v in claim_value]
+        return any(v in claim_list for v in selector_list)
+
+    return str(claim_value) in selector_list if claim_value is not None else False
+
+
+def _matches_routing_override(
+    token_claims: dict, override: "JWTRoutingOverride"
+) -> bool:
+    return (
+        _routing_selector_matches_claim(override.iss, token_claims.get("iss"))
+        and _routing_selector_matches_claim(
+            override.client_id, token_claims.get("client_id")
+        )
+        and _routing_selector_matches_claim(override.aud, token_claims.get("aud"))
+    )
+
+
+def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> bool:
+    routing_overrides = jwt_handler.litellm_jwtauth.routing_overrides
+    if not routing_overrides:
+        return False
+
+    token_claims = jwt_handler.get_unverified_claims(token=token)
+    if token_claims is None:
+        return False
+
+    for override in routing_overrides:
+        if override.path == "oauth2" and _matches_routing_override(
+            token_claims=token_claims, override=override
+        ):
+            verbose_proxy_logger.debug(
+                "JWT routing override matched. Routing token to OAuth2 introspection."
+            )
+            return True
+
+    return False
+
+
 def _get_bearer_token(
     api_key: str,
 ):
@@ -649,12 +701,20 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 # - JWT tokens (3 dot-separated parts) -> skip OAuth2, fall through to JWT handler
                 # - Opaque tokens -> use OAuth2 handler
                 # This allows JWT for users and OAuth2 for M2M on the same instance
-                is_jwt_token = (
+                is_jwt = (
                     jwt_handler.is_jwt(token=api_key)
                     if general_settings.get("enable_jwt_auth", False) is True
                     else False
                 )
-                if not is_jwt_token:
+                # Routing uses unverified JWT claims only to choose auth path.
+                # Final authentication is enforced by the selected validator.
+                route_jwt_to_oauth2 = (
+                    is_jwt
+                    and _should_route_jwt_to_oauth2_override(
+                        token=api_key, jwt_handler=jwt_handler
+                    )
+                )
+                if not is_jwt or route_jwt_to_oauth2:
                     # return UserAPIKeyAuth object
                     # helper to check if the api_key is a valid oauth2 token
                     from litellm.proxy.proxy_server import premium_user
@@ -688,7 +748,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     jwt_claims: Optional[dict]
                     if (
                         jwt_handler.litellm_jwtauth.oidc_userinfo_enabled
-                        and not jwt_handler.is_jwt(token=api_key)
+                        and not is_jwt
                     ):
                         jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key)
                     else:
@@ -1193,49 +1253,13 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 raise Exception(
                     "Key is blocked. Update via `/key/unblock` if you're an admin."
                 )
-            config = valid_token.config
-
-            if config != {}:
-                model_list = config.get("model_list", [])
-                new_model_list = model_list
-                verbose_proxy_logger.debug(
-                    f"\n new llm router model list {new_model_list}"
-                )
-            elif (
-                isinstance(valid_token.models, list)
-                and "all-team-models" in valid_token.models
-            ):
-                # Do not do any validation at this step
-                # the validation will occur when checking the team has access to this model
-                pass
-            else:
-                model = get_model_from_request(request_data, route)
-                fallback_models = cast(
-                    Optional[List[ALL_FALLBACK_MODEL_VALUES]],
-                    request_data.get("fallbacks", None),
-                )
-
-                if model is not None:
-                    await can_key_call_model(
-                        model=model,
-                        llm_model_list=llm_model_list,
-                        valid_token=valid_token,
-                        llm_router=llm_router,
-                    )
-
-                if fallback_models is not None:
-                    for m in fallback_models:
-                        await can_key_call_model(
-                            model=m["model"] if isinstance(m, dict) else m,
-                            llm_model_list=llm_model_list,
-                            valid_token=valid_token,
-                            llm_router=llm_router,
-                        )
-                        await is_valid_fallback_model(
-                            model=m["model"] if isinstance(m, dict) else m,
-                            llm_router=llm_router,
-                            user_model=None,
-                        )
+            await _enforce_key_and_fallback_model_access(
+                valid_token=valid_token,
+                request_data=request_data,
+                route=route,
+                llm_model_list=llm_model_list,
+                llm_router=llm_router,
+            )
 
             # Check 2. If user_id for this token is in budget - done in common_checks()
             if valid_token.user_id is not None:
@@ -1764,6 +1788,61 @@ async def _lookup_end_user_and_apply_budget(
     return valid_token, end_user_object
 
 
+async def _enforce_key_and_fallback_model_access(
+    *,
+    valid_token: UserAPIKeyAuth,
+    request_data: dict,
+    route: str,
+    llm_model_list: Optional[list],
+    llm_router: Optional[Any],
+) -> None:
+    """
+    Key-level model allowlist and client fallbacks (same as standard auth).
+    Not included in common_checks — common_checks enforces team/user/project model access only.
+    """
+    config = valid_token.config
+
+    if config != {}:
+        model_list = config.get("model_list", [])
+        new_model_list = model_list
+        verbose_proxy_logger.debug(
+            f"\n new llm router model list {new_model_list}"
+        )
+    elif (
+        isinstance(valid_token.models, list)
+        and "all-team-models" in valid_token.models
+    ):
+        pass
+    else:
+        model = get_model_from_request(request_data, route)
+        fallback_models = cast(
+            Optional[List[ALL_FALLBACK_MODEL_VALUES]],
+            request_data.get("fallbacks", None),
+        )
+
+        if model is not None:
+            await can_key_call_model(
+                model=model,
+                llm_model_list=llm_model_list,
+                valid_token=valid_token,
+                llm_router=llm_router,
+            )
+
+        if fallback_models is not None:
+            for m in fallback_models:
+                await can_key_call_model(
+                    model=m["model"] if isinstance(m, dict) else m,
+                    llm_model_list=llm_model_list,
+                    valid_token=valid_token,
+                    llm_router=llm_router,
+                )
+                await is_valid_fallback_model(
+                    model=m["model"] if isinstance(m, dict) else m,
+                    llm_router=llm_router,
+                    user_model=None,
+                )
+
+
 async def _run_post_custom_auth_checks(
     valid_token: UserAPIKeyAuth,
     request: Request,
@@ -1773,6 +1852,7 @@ async def _run_post_custom_auth_checks(
 ) -> UserAPIKeyAuth:
     from litellm.proxy.proxy_server import (
         general_settings,
+        llm_model_list,
         llm_router,
         model_max_budget_limiter,
         prisma_client,
@@ -1816,6 +1896,15 @@ async def _run_post_custom_auth_checks(
                 ),
             )
 
+    if general_settings.get("custom_auth_run_common_checks", False):
+        await _enforce_key_and_fallback_model_access(
+            valid_token=valid_token,
+            request_data=request_data,
+            route=route,
+            llm_model_list=llm_model_list,
+            llm_router=llm_router,
+        )
+
     current_model = request_data.get("model", None)
 
     # 3. Check key-level model_max_budget

tests/llm_translation/test_prompt_factory.py

@@ -859,8 +859,8 @@ def test_ensure_alternating_roles_three_consecutive_assistants():
     ]
 
 
-def test_ensure_alternating_roles_does_not_split_tool_call_chain():
-    """Tool-call chains [user, assistant(tc), tool, user] are preserved as-is."""
+def test_ensure_alternating_roles_inserts_assistant_continue_across_tool_chain():
+    """[user, assistant(tc), tool, user] gets assistant_continue before the second user."""
     messages = [
         {"role": "user", "content": "Search for X"},
         {
@@ -899,15 +899,16 @@ def test_ensure_alternating_roles_does_not_split_tool_call_chain():
             ],
         },
         {"role": "tool", "tool_call_id": "c1", "content": "results"},
+        {"role": "assistant", "content": "Please continue."},
         {"role": "user", "content": "Thanks, now do Y"},
     ]
 
 
 def test_ensure_alternating_roles_assistant_tool_call_then_assistant():
     """
-    Preserve old behavior for malformed adjacent assistant turns:
-    [assistant(tool_calls), assistant(no-tool-calls), user] should insert
-    user_continue between assistant messages.
+    Malformed [assistant(tc), assistant(no-tc), user]:
+    user_continue inserts break between adjacents, then assistant_continue
+    fills the counted-sequence gap.
     """
     messages = [
         {
@@ -945,6 +946,7 @@ def test_ensure_alternating_roles_assistant_tool_call_then_assistant():
                 }
             ],
         },
+        {"role": "assistant", "content": "Please continue."},
         {"role": "user", "content": "Please continue."},
         {"role": "assistant", "content": "Here's what I found."},
         {"role": "user", "content": "Thanks"},
@@ -993,10 +995,184 @@ def test_ensure_alternating_roles_trailing_tool_call_assistant():
                 }
             ],
         },
+        {"role": "assistant", "content": "Please continue."},
         {"role": "user", "content": "Please continue."},
     ]
 
 
+def test_ensure_alternating_roles_multiple_tool_results():
+    """[user, assistant(tc), tool, tool, user] — multiple tool results before next user."""
+    messages = [
+        {"role": "user", "content": "Search for X and Y"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "search_x", "arguments": "{}"},
+                },
+                {
+                    "id": "c2",
+                    "type": "function",
+                    "function": {"name": "search_y", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "result X"},
+        {"role": "tool", "tool_call_id": "c2", "content": "result Y"},
+        {"role": "user", "content": "Thanks"},
+    ]
+
+    transformed_messages = get_completion_messages(
+        messages=messages,
+        assistant_continue_message=None,
+        user_continue_message=None,
+        ensure_alternating_roles=True,
+    )
+
+    assert transformed_messages == [
+        {"role": "user", "content": "Search for X and Y"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "search_x", "arguments": "{}"},
+                },
+                {
+                    "id": "c2",
+                    "type": "function",
+                    "function": {"name": "search_y", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "result X"},
+        {"role": "tool", "tool_call_id": "c2", "content": "result Y"},
+        {"role": "assistant", "content": "Please continue."},
+        {"role": "user", "content": "Thanks"},
+    ]
+
+
+def test_ensure_alternating_roles_chained_tool_calls():
+    """[user, assistant(tc), tool, assistant(tc), tool, user] — chained tool calls."""
+    messages = [
+        {"role": "user", "content": "Do multi-step task"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "step1", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "step1 done"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c2",
+                    "type": "function",
+                    "function": {"name": "step2", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c2", "content": "step2 done"},
+        {"role": "user", "content": "What happened?"},
+    ]
+
+    transformed_messages = get_completion_messages(
+        messages=messages,
+        assistant_continue_message=None,
+        user_continue_message=None,
+        ensure_alternating_roles=True,
+    )
+
+    assert transformed_messages == [
+        {"role": "user", "content": "Do multi-step task"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "step1", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "step1 done"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c2",
+                    "type": "function",
+                    "function": {"name": "step2", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c2", "content": "step2 done"},
+        {"role": "assistant", "content": "Please continue."},
+        {"role": "user", "content": "What happened?"},
+    ]
+
+
+def test_ensure_alternating_roles_system_prefix_with_tool_chain():
+    """[system, user, assistant(tc), tool, user] — system prefix doesn't interfere."""
+    messages = [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Search for X"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "search", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "results"},
+        {"role": "user", "content": "Thanks"},
+    ]
+
+    transformed_messages = get_completion_messages(
+        messages=messages,
+        assistant_continue_message=None,
+        user_continue_message=None,
+        ensure_alternating_roles=True,
+    )
+
+    assert transformed_messages == [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Search for X"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "c1",
+                    "type": "function",
+                    "function": {"name": "search", "arguments": "{}"},
+                },
+            ],
+        },
+        {"role": "tool", "tool_call_id": "c1", "content": "results"},
+        {"role": "assistant", "content": "Please continue."},
+        {"role": "user", "content": "Thanks"},
+    ]
+
+
 def test_alternating_roles_e2e():
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
     import json

tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py

@@ -160,6 +160,67 @@ class TestAgentCoreJsonResponseParsing:
         assert parsed["content"] == ""
         assert parsed["final_message"] == response_json["result"]
 
+    def test_parse_json_a2a_jsonrpc_nested_message(self, config):
+        """Strategy 0: A2A JSON-RPC with result.message.parts[] format."""
+        response_json = {
+            "jsonrpc": "2.0",
+            "id": "test_id",
+            "result": {
+                "message": {
+                    "role": "agent",
+                    "parts": [{"kind": "text", "text": "1 + 1 = 2"}],
+                    "messageId": "123",
+                }
+            },
+        }
+        parsed = config._parse_json_response(response_json)
+        assert parsed["content"] == "1 + 1 = 2"
+        assert parsed["usage"] is None
+
+    def test_parse_json_a2a_jsonrpc_direct_parts(self, config):
+        """Strategy 0: A2A JSON-RPC with result.parts[] format (direct message)."""
+        response_json = {
+            "jsonrpc": "2.0",
+            "id": "test_id",
+            "result": {
+                "kind": "message",
+                "parts": [{"kind": "text", "text": "Direct response"}],
+            },
+        }
+        parsed = config._parse_json_response(response_json)
+        assert parsed["content"] == "Direct response"
+        assert parsed["usage"] is None
+
+    def test_parse_json_a2a_jsonrpc_multi_parts(self, config):
+        """Strategy 0: A2A JSON-RPC with multiple text parts concatenated."""
+        response_json = {
+            "jsonrpc": "2.0",
+            "id": "test_id",
+            "result": {
+                "message": {
+                    "role": "agent",
+                    "parts": [
+                        {"kind": "text", "text": "First part"},
+                        {"kind": "text", "text": "Second part"},
+                    ],
+                }
+            },
+        }
+        parsed = config._parse_json_response(response_json)
+        assert parsed["content"] == "First part Second part"
+        assert parsed["usage"] is None
+
+    def test_parse_json_a2a_jsonrpc_empty_falls_through(self, config):
+        """Strategy 0: A2A JSON-RPC with empty result falls through to Strategy 3."""
+        response_json = {
+            "jsonrpc": "2.0",
+            "id": "test_id",
+            "result": "plain text fallback",
+        }
+        parsed = config._parse_json_response(response_json)
+        assert parsed["content"] == "plain text fallback"
+        assert parsed["usage"] is None
+
 
 class TestAgentCoreNonStreamingJsonFormats:
     """Tests for _get_parsed_response with different JSON formats (non-streaming path)."""

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -3,22 +3,30 @@ import json
 import os
 import sys
 from typing import Tuple
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
 sys.path.insert(
     0, os.path.abspath("../../..")
 )  # Adds the parent directory to the system path
 
-from unittest.mock import MagicMock
-
 import pytest
 
 import litellm.proxy.proxy_server
 from litellm.caching.dual_cache import DualCache
-from litellm.proxy._types import LiteLLM_JWTAuth, UserAPIKeyAuth
+from litellm.proxy._types import (
+    LiteLLM_JWTAuth,
+    ProxyErrorTypes,
+    ProxyException,
+    UserAPIKeyAuth,
+    JWTRoutingOverride,
+)
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.auth.route_checks import RouteChecks
-from litellm.proxy.auth.user_api_key_auth import get_api_key, user_api_key_auth
+from litellm.proxy.auth.user_api_key_auth import (
+    _run_post_custom_auth_checks,
+    get_api_key,
+    user_api_key_auth,
+)
 
 
 def test_get_api_key():
@@ -38,6 +46,86 @@ def test_get_api_key():
     ) == (api_key, passed_in_key)
 
 
+@pytest.mark.asyncio
+async def test_custom_auth_does_not_enforce_key_model_access_by_default():
+    valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])
+    request_data = {"model": "gpt-4o"}
+
+    with patch(
+        "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock
+    ) as mock_can_key, patch(
+        "litellm.proxy.proxy_server.general_settings",
+        {},
+    ):
+        await _run_post_custom_auth_checks(
+            valid_token=valid_token,
+            request=None,
+            request_data=request_data,
+            route="/v1/chat/completions",
+            parent_otel_span=None,
+        )
+        mock_can_key.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in():
+    valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])
+    request_data = {"model": "gpt-4o-mini"}
+
+    with patch(
+        "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock
+    ) as mock_can_key, patch(
+        "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
+    ), patch(
+        "litellm.proxy.proxy_server.general_settings",
+        {"custom_auth_run_common_checks": True},
+    ):
+        await _run_post_custom_auth_checks(
+            valid_token=valid_token,
+            request=None,
+            request_data=request_data,
+            route="/v1/chat/completions",
+            parent_otel_span=None,
+        )
+        mock_can_key.assert_awaited_once_with(
+            model="gpt-4o-mini",
+            llm_model_list=ANY,
+            valid_token=valid_token,
+            llm_router=ANY,
+        )
+
+
+@pytest.mark.asyncio
+async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in():
+    valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])
+    request_data = {"model": "gpt-4o"}
+
+    with patch(
+        "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock
+    ) as mock_can_key, patch(
+        "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
+    ), patch(
+        "litellm.proxy.proxy_server.general_settings",
+        {"custom_auth_run_common_checks": True},
+    ):
+        mock_can_key.side_effect = ProxyException(
+            message="Key not allowed to access model",
+            type=ProxyErrorTypes.key_model_access_denied,
+            param="model",
+            code=401,
+        )
+        with pytest.raises(ProxyException) as exc:
+            await _run_post_custom_auth_checks(
+                valid_token=valid_token,
+                request=None,
+                request_data=request_data,
+                route="/v1/chat/completions",
+                parent_otel_span=None,
+            )
+
+        assert exc.value.type == ProxyErrorTypes.key_model_access_denied
+
+
 @pytest.mark.parametrize(
     "custom_litellm_key_header, api_key, passed_in_key",
     [
@@ -689,6 +777,203 @@ class TestJWTOAuth2Coexistence:
             mock_jwt_auth.assert_called_once()
             assert result.user_id == "jwt-human-user"
 
+    @pytest.mark.asyncio
+    async def test_routing_override_routes_matching_jwt_to_oauth2(self):
+        """
+        When routing_overrides match JWT claims, route JWT-shaped token to OAuth2.
+        """
+        jwt_token = (
+            "eyJhbGciOiJSUzI1NiJ9."
+            "eyJpc3MiOiJtYWNoaW5lLWlzc3Vlci5leGFtcGxlLmNvbSIsImNsaWVudF9pZCI6Ik1JRF9MSVRFTExNIn0."
+            "c2ln"
+        )
+        general_settings = {
+            "enable_oauth2_auth": True,
+            "enable_jwt_auth": True,
+        }
+        mock_oauth2_response = UserAPIKeyAuth(
+            api_key=jwt_token,
+            user_id="machine-client-override",
+        )
+
+        mock_request = MagicMock()
+        mock_request.url.path = "/v1/chat/completions"
+        mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
+        mock_request.query_params = {}
+
+        with patch(
+            "litellm.proxy.proxy_server.general_settings", general_settings
+        ), patch("litellm.proxy.proxy_server.premium_user", True), patch(
+            "litellm.proxy.proxy_server.master_key", "sk-master"
+        ), patch(
+            "litellm.proxy.proxy_server.prisma_client", None
+        ), patch(
+            "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
+            new_callable=AsyncMock,
+            return_value=mock_oauth2_response,
+        ) as mock_oauth2, patch(
+            "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
+            new_callable=AsyncMock,
+        ) as mock_jwt_auth:
+            litellm.proxy.proxy_server.jwt_handler.update_environment(
+                prisma_client=None,
+                user_api_key_cache=DualCache(),
+                litellm_jwtauth=LiteLLM_JWTAuth(
+                    routing_overrides=[
+                        JWTRoutingOverride(
+                            iss="machine-issuer.example.com",
+                            client_id="MID_LITELLM",
+                            path="oauth2",
+                        )
+                    ]
+                ),
+            )
+
+            result = await user_api_key_auth(
+                request=mock_request,
+                api_key=f"Bearer {jwt_token}",
+            )
+
+            mock_oauth2.assert_called_once_with(token=jwt_token)
+            mock_jwt_auth.assert_not_called()
+            assert result.user_id == "machine-client-override"
+
+    @pytest.mark.asyncio
+    async def test_routing_override_does_not_match_client_id_falls_back_to_jwt(self):
+        """
+        If override ISS matches but client_id does not, continue default JWT flow.
+        """
+        jwt_token = (
+            "eyJhbGciOiJSUzI1NiJ9."
+            "eyJpc3MiOiJtYWNoaW5lLWlzc3Vlci5leGFtcGxlLmNvbSIsImNsaWVudF9pZCI6IlVTRVJfUE9SVEFMIn0."
+            "c2ln"
+        )
+        general_settings = {
+            "enable_oauth2_auth": True,
+            "enable_jwt_auth": True,
+        }
+        mock_jwt_result = {
+            "is_proxy_admin": True,
+            "team_object": None,
+            "user_object": None,
+            "end_user_object": None,
+            "org_object": None,
+            "token": jwt_token,
+            "team_id": "jwt-team",
+            "user_id": "jwt-user-no-override",
+            "end_user_id": None,
+            "org_id": None,
+            "team_membership": None,
+            "jwt_claims": {"sub": "user1"},
+        }
+
+        mock_request = MagicMock()
+        mock_request.url.path = "/v1/chat/completions"
+        mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
+        mock_request.query_params = {}
+
+        with patch(
+            "litellm.proxy.proxy_server.general_settings", general_settings
+        ), patch("litellm.proxy.proxy_server.premium_user", True), patch(
+            "litellm.proxy.proxy_server.master_key", "sk-master"
+        ), patch(
+            "litellm.proxy.proxy_server.prisma_client", None
+        ), patch(
+            "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
+            new_callable=AsyncMock,
+        ) as mock_oauth2, patch(
+            "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
+            new_callable=AsyncMock,
+            return_value=mock_jwt_result,
+        ) as mock_jwt_auth:
+            litellm.proxy.proxy_server.jwt_handler.update_environment(
+                prisma_client=None,
+                user_api_key_cache=DualCache(),
+                litellm_jwtauth=LiteLLM_JWTAuth(
+                    routing_overrides=[
+                        JWTRoutingOverride(
+                            iss="machine-issuer.example.com",
+                            client_id="MID_LITELLM",
+                            path="oauth2",
+                        )
+                    ]
+                ),
+            )
+
+            result = await user_api_key_auth(
+                request=mock_request,
+                api_key=f"Bearer {jwt_token}",
+            )
+
+            mock_oauth2.assert_not_called()
+            mock_jwt_auth.assert_called_once()
+            assert result.user_id == "jwt-user-no-override"
+
+    @pytest.mark.asyncio
+    async def test_routing_override_matches_aud_claim_list_and_list_selectors(self):
+        """
+        Match routing override when selectors are lists and token aud claim is a list.
+        """
+        jwt_token = (
+            "eyJhbGciOiJSUzI1NiJ9."
+            "eyJpc3MiOiJtYWNoaW5lLWlzc3Vlci5leGFtcGxlLmNvbSIsImNsaWVudF9pZCI6Ik1JRF9MSVRFTExNIiwiYXVkIjpbImFwaTovL2xpdGVsbG0iLCJhcGk6Ly9vdGhlciJdfQ."
+            "c2ln"
+        )
+        general_settings = {
+            "enable_oauth2_auth": True,
+            "enable_jwt_auth": True,
+        }
+        mock_oauth2_response = UserAPIKeyAuth(
+            api_key=jwt_token,
+            user_id="machine-client-aud-list",
+        )
+
+        mock_request = MagicMock()
+        mock_request.url.path = "/v1/chat/completions"
+        mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
+        mock_request.query_params = {}
+
+        with patch(
+            "litellm.proxy.proxy_server.general_settings", general_settings
+        ), patch("litellm.proxy.proxy_server.premium_user", True), patch(
+            "litellm.proxy.proxy_server.master_key", "sk-master"
+        ), patch(
+            "litellm.proxy.proxy_server.prisma_client", None
+        ), patch(
+            "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
+            new_callable=AsyncMock,
+            return_value=mock_oauth2_response,
+        ) as mock_oauth2, patch(
+            "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
+            new_callable=AsyncMock,
+        ) as mock_jwt_auth:
+            litellm.proxy.proxy_server.jwt_handler.update_environment(
+                prisma_client=None,
+                user_api_key_cache=DualCache(),
+                litellm_jwtauth=LiteLLM_JWTAuth(
+                    routing_overrides=[
+                        JWTRoutingOverride(
+                            iss=[
+                                "machine-issuer.example.com",
+                                "other-issuer.example.com",
+                            ],
+                            client_id=["MID_LITELLM", "MID_BACKUP"],
+                            aud=["api://litellm", "api://fallback"],
+                            path="oauth2",
+                        )
+                    ]
+                ),
+            )
+
+            result = await user_api_key_auth(
+                request=mock_request,
+                api_key=f"Bearer {jwt_token}",
+            )
+
+            mock_oauth2.assert_called_once_with(token=jwt_token)
+            mock_jwt_auth.assert_not_called()
+            assert result.user_id == "machine-client-aud-list"
+
     @pytest.mark.asyncio
     async def test_only_oauth2_enabled_handles_all_tokens(self):
         """