From 3604b600d3fa6b0217c751ba9166f3fcbc5a24c7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 2 Apr 2026 16:38:01 -0700 Subject: [PATCH] [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 Co-authored-by: michelligabriele --- docs/my-website/docs/proxy/oauth2.md | 21 ++ docs/my-website/docs/proxy/token_auth.md | 41 +++ .../prompt_templates/common_utils.py | 29 +- .../bedrock/chat/agentcore/transformation.py | 14 + litellm/proxy/_types.py | 22 ++ litellm/proxy/auth/handle_jwt.py | 64 ++-- litellm/proxy/auth/user_api_key_auth.py | 183 ++++++++--- tests/llm_translation/test_prompt_factory.py | 186 ++++++++++- .../test_agentcore_transformation.py | 61 ++++ .../proxy/auth/test_user_api_key_auth.py | 295 +++++++++++++++++- 10 files changed, 829 insertions(+), 87 deletions(-) diff --git a/docs/my-website/docs/proxy/oauth2.md b/docs/my-website/docs/proxy/oauth2.md index 204a01538c..c0597058cf 100644 --- a/docs/my-website/docs/proxy/oauth2.md +++ b/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). + diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md index c287ab0364..d37b05391b 100644 --- a/docs/my-website/docs/proxy/token_auth.md +++ b/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. diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index cec61405eb..46e60c24d3 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/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 diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index d6eb5a734c..f066322b81 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/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"] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8faf36df4c..e4702364a3 100644 --- a/litellm/proxy/_types.py +++ b/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: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6ad372f3b5..4a6856b6d1 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/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, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 53ae08aefb..046c39a910 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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 diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 64556c3f26..27b0539aa4 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/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 diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 3dede83032..43182926f9 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/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).""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 74c7f9bca5..6e1b245b3d 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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): """