diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0c520142fb3..8e182755656 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -169,6 +169,7 @@ from litellm.proxy._types import ( is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.jwt_algorithms import APPROVED_JWT_ALGORITHMS, ApprovedJwtAlgorithm from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( @@ -386,7 +387,7 @@ class MCPServerConfig(TypedDict, total=False): id_jag_resource: str client_private_key: str client_private_key_id: str - client_assertion_signing_alg: str + client_assertion_signing_alg: ApprovedJwtAlgorithm timeout: float max_concurrent_requests: int @@ -1482,6 +1483,18 @@ def _warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: ) +def _stored_client_assertion_signing_alg(value: object, server_name: str) -> ApprovedJwtAlgorithm: + if value in APPROVED_JWT_ALGORITHMS: + return cast(ApprovedJwtAlgorithm, value) + if value is not None: + verbose_logger.warning( + "MCP server %s: client_assertion_signing_alg %r is not an approved algorithm, using RS256", + server_name, + value, + ) + return "RS256" + + def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None: """Direct legacy delegated OAuth configurations to the admitted replacement.""" if server.auth_type != MCPAuth.oauth2: @@ -2604,7 +2617,10 @@ class MCPServerManager: id_jag_resource=server_config.get("id_jag_resource", None), client_private_key=server_config.get("client_private_key", None), client_private_key_id=server_config.get("client_private_key_id", None), - client_assertion_signing_alg=server_config.get("client_assertion_signing_alg", "RS256"), + client_assertion_signing_alg=_stored_client_assertion_signing_alg( + server_config.get("client_assertion_signing_alg"), + str(server_config.get("alias") or server_config.get("server_name") or server_id), + ), token_exchange_profile=server_config.get("token_exchange_profile", "rfc8693"), allow_sampling=bool(server_config.get("allow_sampling", False)), allow_elicitation=bool(server_config.get("allow_elicitation", False)), @@ -3178,10 +3194,10 @@ class MCPServerManager: credentials_are_encrypted, ), client_private_key_id=(credentials_dict.get("client_private_key_id") if credentials_dict else None), - client_assertion_signing_alg=( - credentials_dict.get("client_assertion_signing_alg") if credentials_dict else None - ) - or "RS256", + client_assertion_signing_alg=_stored_client_assertion_signing_alg( + credentials_dict.get("client_assertion_signing_alg") if credentials_dict else None, + name_for_prefix, + ), token_exchange_profile=mcp_server.token_exchange_profile or (credentials_dict.get("token_exchange_profile") if credentials_dict else None) or "rfc8693", diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index f02e6c85d9b..1f2a492beab 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -23,20 +23,11 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.auth.jwt_algorithms import APPROVED_JWT_ALGORITHMS, jwks_keys_for from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer -_ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( - "RS256", - "RS384", - "RS512", - "ES256", - "ES384", - "ES512", - "PS256", - "PS384", - "PS512", -) +_ALLOWED_ID_TOKEN_ALGORITHMS: Final = APPROVED_JWT_ALGORITHMS _JWKS_CACHE_TTL_SECONDS: Final = 3600 _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) @@ -126,7 +117,7 @@ async def _discover_jwks_url(issuer: str) -> str: def _select_signing_key(id_token: str, keys: Sequence[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection": header: Final = jwt.get_unverified_header(id_token) kid: Final = header.get("kid") - for key in keys: + for key in jwks_keys_for(keys, _ALLOWED_ID_TOKEN_ALGORITHMS): if kid is None or key.get("kid") == kid: return jwt.PyJWK(dict(key)) # mutable-ok: PyJWT requires a concrete JWK dictionary return _BindingRejection( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 33c3a854058..6c310784370 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -40,6 +40,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy.auth.jwt_algorithms import ApprovedJwtAlgorithm from litellm.types.mcp import ( DEFAULT_CREDENTIAL_HEADER, DEFAULT_SUBJECT_TOKEN_TYPE, @@ -326,7 +327,7 @@ class PrivateKeyJwtAuth(BaseModel): source: Literal["private_key_jwt"] = "private_key_jwt" private_key: SecretStr key_id: str | None = None - signing_alg: str = "RS256" + signing_alg: ApprovedJwtAlgorithm = "RS256" class ClientSecretAuth(BaseModel): diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0b43c3864ab..858a993c4ec 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -29465,6 +29465,17 @@ "client_assertion_signing_alg": { "anyOf": [ { + "enum": [ + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512" + ], "type": "string" }, { @@ -32166,6 +32177,17 @@ "client_assertion_signing_alg": { "anyOf": [ { + "enum": [ + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512" + ], "type": "string" }, { @@ -35439,6 +35461,17 @@ "client_assertion_signing_alg": { "anyOf": [ { + "enum": [ + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512" + ], "type": "string" }, { @@ -39289,6 +39322,17 @@ "client_assertion_signing_alg": { "anyOf": [ { + "enum": [ + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512" + ], "type": "string" }, { diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index e07b20fd5d5..6f4a10c6f69 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -16,6 +16,7 @@ import re import time from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass +from functools import lru_cache from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast import httpx @@ -53,6 +54,12 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model +from litellm.proxy.auth.jwt_algorithms import ( + APPROVED_JWT_ALGORITHMS, + LEGACY_JWT_ALGORITHMS, + allowed_jwt_algorithms, + jwks_keys_for, +) from litellm.proxy.auth.model_access_denied import ( ModelAccessDeniedHTTPException, model_access_denied_client_message, @@ -60,6 +67,7 @@ from litellm.proxy.auth.model_access_denied import ( from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.team_grants import team_grants, team_model_aliases +from litellm.proxy.common_utils.fips import is_fips_mode from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -104,6 +112,14 @@ UNREACHABLE_CACHE_KEY_PREFIX: Final = "litellm_jwks_unreachable_" _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) +@lru_cache(maxsize=1) +def _log_eddsa_deprecation() -> None: + verbose_proxy_logger.warning( + "JWT Auth: accepted a token signed with EdDSA, which is deprecated and not FIPS 140-3 approved; " + "it is rejected when LITELLM_FIPS_MODE=true. Move the IdP signing key to RS256, PS256 or ES256" + ) + + class _JWTAuthSettings(Protocol): """The JWT auth settings block this handler reads back through ``getattr``, when one is configured.""" @@ -209,18 +225,7 @@ class JWTHandler: # 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", - ] + SUPPORTED_JWT_ALGORITHMS = [*APPROVED_JWT_ALGORITHMS, *LEGACY_JWT_ALGORITHMS] LITELLM_JWT_ISSUER_CLAIM = "_litellm_jwt_issuer" LITELLM_USER_ID_CLAIM = "_litellm_user_id" LITELLM_USER_EMAIL_CLAIM = "_litellm_user_email" @@ -240,13 +245,24 @@ class JWTHandler: def __init__( self, + fips_mode: Callable[[], bool] = is_fips_mode, ) -> None: + self._fips_mode: Final = fips_mode self.http_handler = HTTPHandler() 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 allowed_algorithms(self) -> tuple[str, ...]: + return allowed_jwt_algorithms(self._fips_mode()) + + def _warn_deprecated_signing_algorithm(self, token: str) -> None: + if self._fips_mode(): + return + if jwt.get_unverified_header(token).get("alg") == "EdDSA": + _log_eddsa_deprecation() + def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None: self.agent_lookup = agent_lookup @@ -943,7 +959,12 @@ class JWTHandler: log_context=f"kid={kid}", ) - public_key: Final = self.parse_keys(keys=keys, kid=kid) + usable_keys: Final[JWKKeyValue] = ( + list(jwks_keys_for(keys, self.allowed_algorithms())) + if isinstance(keys, list) + else next(iter(jwks_keys_for((keys,), self.allowed_algorithms())), {}) + ) + public_key: Final = self.parse_keys(keys=usable_keys, kid=kid) if public_key is not None: return cast(dict, public_key) @@ -1216,30 +1237,32 @@ class JWTHandler: if isinstance(public_key, dict): public_key_obj: Final = PyJWK.from_dict(self._get_jwk_from_public_key(public_key=public_key)).key - return jwt.decode( + payload: Final = jwt.decode( token, public_key_obj, - algorithms=self.SUPPORTED_JWT_ALGORITHMS, + algorithms=self.allowed_algorithms(), options=decode_options, audience=audience, issuer=issuer, leeway=self.leeway, ) - - cert: Final = x509.load_pem_x509_certificate(public_key.encode(), default_backend()) - key: Final = cert.public_key().public_bytes( - serialization.Encoding.PEM, - serialization.PublicFormat.SubjectPublicKeyInfo, - ) - return jwt.decode( - token, - key, - algorithms=self.SUPPORTED_JWT_ALGORITHMS, - audience=audience, - issuer=issuer, - options=decode_options, - leeway=self.leeway, - ) + else: + cert: Final = x509.load_pem_x509_certificate(public_key.encode(), default_backend()) + key: Final = cert.public_key().public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + payload = jwt.decode( + token, + key, + algorithms=self.allowed_algorithms(), + audience=audience, + issuer=issuer, + options=decode_options, + leeway=self.leeway, + ) + self._warn_deprecated_signing_algorithm(token) + return payload async def _auth_jwt_with_issuer(self, token: str, issuer_config: JWTIssuerConfig, kid: str | None) -> dict: try: diff --git a/litellm/proxy/auth/jwt_algorithms.py b/litellm/proxy/auth/jwt_algorithms.py new file mode 100644 index 00000000000..48e03ba9981 --- /dev/null +++ b/litellm/proxy/auth/jwt_algorithms.py @@ -0,0 +1,48 @@ +from collections.abc import Collection, Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal + +ApprovedJwtAlgorithm = Literal["RS256", "RS384", "RS512", "PS256", "PS384", "PS512", "ES256", "ES384", "ES512"] + +APPROVED_JWT_ALGORITHMS: Final[tuple[ApprovedJwtAlgorithm, ...]] = ( + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512", +) + +LEGACY_JWT_ALGORITHMS: Final = ("EdDSA",) + +_KEY_TYPE_ALGORITHMS: Final = MappingProxyType( + { + "RSA": frozenset(APPROVED_JWT_ALGORITHMS[:6]), + "EC": frozenset(APPROVED_JWT_ALGORITHMS[6:]), + "OKP": frozenset(LEGACY_JWT_ALGORITHMS), + } +) + + +def allowed_jwt_algorithms(fips_mode: bool) -> tuple[str, ...]: + return APPROVED_JWT_ALGORITHMS if fips_mode else (*APPROVED_JWT_ALGORITHMS, *LEGACY_JWT_ALGORITHMS) + + +def jwks_keys_for( + keys: Sequence[Mapping[str, object]], algorithms: Collection[str] +) -> tuple[Mapping[str, object], ...]: + """Keep keys whose declared alg is allowed; keys without alg are kept only when their kty can sign with an allowed algorithm.""" + return tuple(key for key in keys if _key_allowed(key, frozenset(algorithms))) + + +def _key_allowed(key: Mapping[str, object], algorithms: frozenset[str]) -> bool: + alg: Final = key.get("alg") + if isinstance(alg, str): + return alg in algorithms + key_type: Final = key.get("kty") + if not isinstance(key_type, str): + return False + return bool(_KEY_TYPE_ALGORITHMS.get(key_type, frozenset()) & algorithms) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index a269ad31a6b..c1a88f3aa67 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -89,6 +89,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.jwt_algorithms import APPROVED_JWT_ALGORITHMS, jwks_keys_for from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -449,12 +450,16 @@ class MCPJWTSigner(CustomGuardrail): unverified_header: Final = jwt.get_unverified_header(raw_token) kid: Final = unverified_header.get("kid") + approved_keys: Final = jwks_keys_for(jwks_keys, APPROVED_JWT_ALGORITHMS) + if not approved_keys: + raise jwt.exceptions.PyJWKSetError(f"No JWKS key at {jwks_uri!r} uses an approved signing algorithm") + # Build a JWKS object and pick the matching key. # PyJWT's PyJWKSet handles key-type parsing and kid matching correctly. from jwt import PyJWKSet try: - jwks_set: Final = PyJWKSet.from_dict({"keys": jwks_keys}) + jwks_set: Final = PyJWKSet.from_dict({"keys": list(approved_keys)}) except Exception as exc: raise jwt.exceptions.PyJWKSetError(f"Failed to parse JWKS from {jwks_uri!r}: {exc}") from exc diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 83e719810d5..df2d2b979cc 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -11,6 +11,7 @@ import httpx from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict +from litellm.proxy.auth.jwt_algorithms import ApprovedJwtAlgorithm from litellm.types.llms.base import HiddenParams if TYPE_CHECKING: @@ -251,7 +252,7 @@ class MCPCredentials(TypedDict, total=False): Key id (kid) advertised in the client_assertion JWT header """ - client_assertion_signing_alg: str | None + client_assertion_signing_alg: ApprovedJwtAlgorithm | None """ Signing algorithm for the client_assertion JWT. Default: RS256 """ diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cb32299b143..dc610724485 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -4,6 +4,7 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Self +from litellm.proxy.auth.jwt_algorithms import ApprovedJwtAlgorithm from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, @@ -145,7 +146,7 @@ class MCPServer(BaseModel): id_jag_resource: str | None = None client_private_key: str | None = None client_private_key_id: str | None = None - client_assertion_signing_alg: str = "RS256" + client_assertion_signing_alg: ApprovedJwtAlgorithm = "RS256" # Wire dialect: "rfc8693" (standard token-exchange grant) or "entra_obo" (Microsoft Entra # On-Behalf-Of, the RFC 7523 jwt-bearer grant + requested_token_use extension) token_exchange_profile: str = "rfc8693" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index cd5dae1269a..23aa8e4f4cd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -111,7 +111,6 @@ async def test_manager_sampling_preserves_explicit_headers_without_ambient_conte assert sampling.await_args.kwargs["raw_headers"] == {"x-test-caller": "sampling-caller"} - @pytest.mark.asyncio async def test_sampling_callback_keeps_creation_context_after_caller_switch(): from mcp.server.auth.middleware.auth_context import auth_context_var @@ -219,8 +218,6 @@ def _reload_mcp_manager_module(): return reloaded - - @pytest.fixture(autouse=True) def enable_eager_mcp_oauth_discovery(monkeypatch): monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1") @@ -1394,7 +1391,9 @@ class TestMCPServerManager: assert not any("oauth2_id_jag" in message for message in caplog.messages) @pytest.mark.asyncio - async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, config_only_mcp_manager_factory, monkeypatch, caplog): + async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso( + self, config_only_mcp_manager_factory, monkeypatch, caplog + ): self._clear_sso_env(monkeypatch) monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") manager = config_only_mcp_manager_factory() @@ -4680,7 +4679,9 @@ class TestMCPServerManager: @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) @pytest.mark.parametrize("is_byok", [False, True]) @pytest.mark.parametrize("scheme", ["http", "https"]) - async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): + async def test_openapi_health_loads_spec_without_mcp_handshake( + self, respx_mock, monkeypatch, auth_type, is_byok, scheme + ): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -4730,14 +4731,28 @@ class TestMCPServerManager: @pytest.mark.parametrize( ("failure", "expected_status", "expected_error"), [ - (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), + ( + httpx.Response(401, text="secret response content"), + "unhealthy", + "OpenAPI specification request failed (HTTP 401)", + ), (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), - (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), - (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), + ( + httpx.ConnectError("secret network details"), + "unhealthy", + "OpenAPI specification could not be loaded (ConnectError)", + ), + ( + httpx.Response(200, text="secret invalid JSON body"), + "unhealthy", + "OpenAPI specification could not be loaded (JSONDecodeError)", + ), ], ) - async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): + async def test_openapi_health_reports_safe_failures( + self, respx_mock, monkeypatch, failure, expected_status, expected_error + ): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -5272,8 +5287,15 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, - auth_type=None, upstream_token_header=None, + path, + method, + operation, + base_url, + headers=None, + server_label=None, + relays_upstream_auth=False, + auth_type=None, + upstream_token_header=None, ): captured["headers"] = headers captured["server_label"] = server_label @@ -5358,8 +5380,15 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, - auth_type=None, upstream_token_header=None, + path, + method, + operation, + base_url, + headers=None, + server_label=None, + relays_upstream_auth=False, + auth_type=None, + upstream_token_header=None, ): captured["headers"] = headers @@ -12485,7 +12514,9 @@ class TestConfigServerIdPinning: @pytest.mark.asyncio @pytest.mark.parametrize("aliasing_entry_first", [True, False]) - async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, config_only_mcp_manager_factory, aliasing_entry_first: bool): + async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected( + self, config_only_mcp_manager_factory, aliasing_entry_first: bool + ): """A grant naming 'docs_server' reaches both servers unpinned; the pin would narrow it to one.""" manager = config_only_mcp_manager_factory() wiki = ( @@ -12501,7 +12532,9 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(dict((wiki, docs) if aliasing_entry_first else (docs, wiki))) @pytest.mark.asyncio - async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self, config_only_mcp_manager_factory): + async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected( + self, config_only_mcp_manager_factory + ): manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): @@ -12589,7 +12622,9 @@ class TestConfigServerIdPinning: assert second_round == first_round @pytest.mark.asyncio - async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, config_only_mcp_manager_factory, caplog): + async def test_shadow_warning_fires_again_when_the_shadowed_set_changes( + self, config_only_mcp_manager_factory, caplog + ): manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) @@ -12760,7 +12795,9 @@ class TestConfigServerIdPinning: assert manager.config_mcp_servers["wiki"].url == "https://example.com/mcp" @pytest.mark.asyncio - async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, config_only_mcp_manager_factory, caplog): + async def test_a_row_that_shadows_one_id_still_reports_capturing_another( + self, config_only_mcp_manager_factory, caplog + ): """Skipping is per identifier, not per row, so the second collision is not lost.""" manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( @@ -13132,7 +13169,8 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), ], ) -async def test_debug_resolution_matches_final_header_conflict_winner(_mcp_request_ctx, +async def test_debug_resolution_matches_final_header_conflict_winner( + _mcp_request_ctx, config: Literal["stored", "static", "none"], extra_headers: dict[str, str] | None, expected_source: str, @@ -13203,7 +13241,9 @@ async def test_debug_resolution_matches_final_header_conflict_winner(_mcp_reques @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "stdio"]) -async def test_debug_reports_legacy_signing_and_non_http_transport(_mcp_request_ctx, transport: Literal["http", "stdio"]) -> None: +async def test_debug_reports_legacy_signing_and_non_http_transport( + _mcp_request_ctx, transport: Literal["http", "stdio"] +) -> None: from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from starlette.requests import Request @@ -13247,12 +13287,16 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(_mcp_request_ async def test_temporary_server_discovery_reuses_resolved_metadata_without_publishing() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="temporary-oauth-discovery", name="temporary", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, + server_id="temporary-oauth-discovery", + name="temporary", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, ) manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", ) with patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery: @@ -13272,13 +13316,18 @@ async def test_temporary_server_discovery_reuses_resolved_metadata_without_publi async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="repeated-stale", name="stale", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=auth_type, oauth2_flow="authorization_code", + server_id="repeated-stale", + name="stale", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + oauth2_flow="authorization_code", ) manager.registry[server.server_id] = server manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", ) with ( patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery, @@ -13298,13 +13347,20 @@ async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> async def test_stale_discovery_falls_back_to_resolved_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="resolved-replacement", name="replacement", url="https://old.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + server_id="resolved-replacement", + name="replacement", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + replacement: Final = original.model_copy( + update={ + "url": "https://new.example.com/mcp", + "authorization_url": "https://new.example.com/authorize", + "token_url": "https://new.example.com/token", + } ) - replacement: Final = original.model_copy(update={ - "url": "https://new.example.com/mcp", "authorization_url": "https://new.example.com/authorize", - "token_url": "https://new.example.com/token", - }) manager.registry[original.server_id] = replacement assert await manager._rejoin_oauth_metadata_discovery(original, retry_stale=False) is replacement @@ -13312,8 +13368,11 @@ async def test_stale_discovery_falls_back_to_resolved_registered_server() -> Non def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="stale-publication", name="publication", url="https://old.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + server_id="stale-publication", + name="publication", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, ) manager._set_oauth_discovery_deferred(original.server_id, True) original_slot: Final = manager._oauth_discovery_slot(original.server_id) @@ -13329,9 +13388,13 @@ def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: async def test_temporary_oauth_discovery_expires_without_more_requests() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="expiring-session", name="temporary", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + server_id="expiring-session", + name="temporary", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", ) manager._set_oauth_discovery_deferred(server.server_id, True) resolved: Final = await manager.ensure_oauth_metadata_discovered(server) @@ -13432,7 +13495,9 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r result = await manager.health_check_server(server.server_id) cached = await manager.health_check_server(server.server_id) assert result.status == "unknown" - assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + assert ( + result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + ) assert cached.health_check_error == result.health_check_error assert cached.last_health_check == result.last_health_check assert route.call_count == 1 @@ -13444,8 +13509,11 @@ async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, mon monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( - server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, - spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, + server_id="cancelled-cache", + name="cancelled-cache", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", + auth_type=MCPAuth.none, ) manager.registry = {server.server_id: server} started = asyncio.Event() @@ -13574,7 +13642,9 @@ class _DiscoveryUpstream: def _discovery_server() -> MCPServer: - return MCPServer(server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http) + return MCPServer( + server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http + ) @pytest.mark.asyncio @@ -13742,7 +13812,9 @@ async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch) assert upstream.initializes == 2 -@pytest.mark.parametrize("value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5))) +@pytest.mark.parametrize( + "value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5)) +) def test_discovery_cache_ttl_validation(value: str, expected: float, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _mcp_discovery_cache_ttl @@ -14005,26 +14077,45 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( class TestProtectedCredentialPreparation: @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,credential", [ - (MCPAuth.bearer_token, None), - (MCPAuth.bearer_token, "Bearer"), - (MCPAuth.api_key, None), - (MCPAuth.basic, "Basic"), - ]) + @pytest.mark.parametrize( + "auth_type,credential", + [ + (MCPAuth.bearer_token, None), + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.api_key, None), + (MCPAuth.basic, "Basic"), + ], + ) @pytest.mark.parametrize("dispatch", ["managed", "local"]) async def test_openapi_dispatch_rejects_unusable_effective_credentials( - self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, - auth_type: MCPAuthType, credential: str | None, dispatch: str, + self, + tmp_path: Path, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, + credential: str | None, + dispatch: str, ) -> None: from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix spec_path: Final = tmp_path / "openapi.json" - spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"}, - "paths": {"/echo": {"get": {"operationId": "echo"}}}})) + spec_path.write_text( + json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Auth", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}, + } + ) + ) server: Final = MCPServer( - server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example", - transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential, + server_id="dispatch-auth", + name="dispatch-auth", + url="https://upstream.example", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=credential, ) manager: Final = MCPServerManager() await manager._register_openapi_tools(str(spec_path), server, server.url) @@ -14047,14 +14138,21 @@ class TestProtectedCredentialPreparation: self, transport: MCPTransport, client_secret: str | None, subject: str | None ) -> None: server = MCPServer( - server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp", - transport=transport, auth_type=MCPAuth.oauth2_token_exchange, - client_id="gateway", client_secret=client_secret, - token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback", + server_id="incomplete-obo", + name="incomplete-obo", + url="https://upstream.example/mcp", + transport=transport, + auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway", + client_secret=client_secret, + token_exchange_endpoint="https://idp.example/token", + authentication_token="static-fallback", ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client( - server, mcp_auth_header="Bearer override", subject_token=subject, + server, + mcp_auth_header="Bearer override", + subject_token=subject, ) assert exc.value.status_code == (401 if subject is None else 500) assert "static-fallback" not in str(exc.value.detail) @@ -14067,8 +14165,11 @@ class TestProtectedCredentialPreparation: self, auth_type: MCPAuthType, credential: str | dict[str, str] | None ) -> None: server = MCPServer( - server_id="empty-static", name="empty-static", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, + server_id="empty-static", + name="empty-static", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential) @@ -14076,16 +14177,22 @@ class TestProtectedCredentialPreparation: assert "credential" in str(exc.value.detail).lower() @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,headers", [ - (MCPAuth.api_key, {"X-API-Key": "key"}), - (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), - ]) + @pytest.mark.parametrize( + "auth_type,headers", + [ + (MCPAuth.api_key, {"X-API-Key": "key"}), + (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), + ], + ) async def test_static_auth_accepts_actual_forwarded_credential( self, auth_type: MCPAuthType, headers: dict[str, str] ) -> None: server = MCPServer( - server_id="header-static", name="header-static", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, + server_id="header-static", + name="header-static", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, ) client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers) assert client._get_auth_headers() == headers @@ -14094,29 +14201,48 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange]) async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: server = MCPServer( - server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, + server_id="openapi-empty", + name="openapi-empty", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, token_exchange_endpoint="https://idp.example/token", ) with pytest.raises(HTTPException) as exc: await MCPServerManager().resolve_openapi_upstream_auth( - mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, - user_api_key_auth=None, forwarded_headers=None, + mcp_server=server, + oauth2_headers=None, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=None, + forwarded_headers=None, ) assert exc.value.status_code in (401, 500) @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,slot,value", [ - (MCPAuth.api_key, "X-API-Key", "token"), - (MCPAuth.authorization, "Authorization", "opaque-secret-value"), - (MCPAuth.authorization, "Authorization", "Bearer abc"), - (MCPAuth.authorization, "Authorization", "Custom abc"), - ]) + @pytest.mark.parametrize( + "auth_type,slot,value", + [ + (MCPAuth.api_key, "X-API-Key", "token"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), + (MCPAuth.authorization, "Authorization", "Bearer abc"), + (MCPAuth.authorization, "Authorization", "Custom abc"), + ], + ) async def test_raw_static_credentials_are_forwarded_unchanged( - self, auth_type: MCPAuthType, slot: str, value: str, + self, + auth_type: MCPAuthType, + slot: str, + value: str, ) -> None: - server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, authentication_token=value) + server = MCPServer( + server_id="raw-key", + name="raw-key", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=value, + ) client = await MCPServerManager()._create_mcp_client(server) assert client._resolved_auth is not None request = httpx.Request("GET", server.url) @@ -14130,17 +14256,24 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("value", ["Bearer", "basic", "token", "ApiKey", " bEaReR ", "\tTOKEN\t"]) @pytest.mark.parametrize("source", ["configured", "caller", "forwarded"]) async def test_raw_authorization_rejects_bare_schemes_before_dispatch( - self, respx_mock: MockRouter, value: str, source: str, + self, + respx_mock: MockRouter, + value: str, + source: str, ) -> None: server: Final = MCPServer( - server_id="raw-empty", name="raw-empty", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.authorization, + server_id="raw-empty", + name="raw-empty", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.authorization, authentication_token=value if source == "configured" else None, ) destination: Final = respx_mock.route().respond(200) with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: await MCPServerManager()._create_mcp_client( - server, mcp_auth_header=value if source == "caller" else None, + server, + mcp_auth_header=value if source == "caller" else None, extra_headers={"Authorization": value} if source == "forwarded" else None, ) assert exc.value.status_code == 500 @@ -14148,9 +14281,15 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: - server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True, - token_exchange_endpoint="https://idp.example/token") + server = MCPServer( + server_id="obo-byok", + name="obo-byok", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + is_byok=True, + token_exchange_endpoint="https://idp.example/token", + ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override") assert exc.value.status_code == 401 @@ -14158,41 +14297,66 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")]) async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None: - server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured) + server = MCPServer( + server_id="override", + name="override", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=configured, + ) client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override) assert client._get_auth_headers()["Authorization"] == override @pytest.mark.asyncio @pytest.mark.parametrize("token", [None, "shared"]) async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None: - server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token) + server = MCPServer( + server_id="empty-header", + name="empty-header", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=token, + ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "}) assert exc.value.status_code == 500 @pytest.mark.asyncio async def test_custom_slot_uses_its_actual_credential(self) -> None: - server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, - upstream_token_header="X-Custom", authentication_token="key") + server = MCPServer( + server_id="custom", + name="custom", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", + authentication_token="key", + ) client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"}) assert client._credential_slot == "X-Custom" assert await client.discovery_auth_fingerprint() @pytest.mark.asyncio - @pytest.mark.parametrize("static_headers,accepted", [ - ({"apikey": "static-key"}, True), - ({"apikey": ""}, False), - ({"X-Tenant": "tenant"}, True), - ]) + @pytest.mark.parametrize( + "static_headers,accepted", + [ + ({"apikey": "static-key"}, True), + ({"apikey": ""}, False), + ({"X-Tenant": "tenant"}, True), + ], + ) async def test_api_key_carried_by_static_header_passes_fail_closed_check( self, static_headers: dict[str, str], accepted: bool ) -> None: server: Final = MCPServer( - server_id="static-slot", name="static-slot", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static_headers, + server_id="static-slot", + name="static-slot", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + static_headers=static_headers, ) if not accepted: with pytest.raises(HTTPException) as exc: @@ -14204,21 +14368,36 @@ class TestProtectedCredentialPreparation: assert all(request.headers[name] == value for name, value in static_headers.items()) @pytest.mark.asyncio - @pytest.mark.parametrize("static,forwarded,caller", [ - ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), - ({}, {"X-API-Key": "forwarded"}, None), - ({}, None, "ApiKey caller"), - ({"X-API-Key": "static"}, {"Authorization": ""}, None), - ]) + @pytest.mark.parametrize( + "static,forwarded,caller", + [ + ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), + ({}, {"X-API-Key": "forwarded"}, None), + ({}, None, "ApiKey caller"), + ({"X-API-Key": "static"}, {"Authorization": ""}, None), + ], + ) async def test_openapi_static_credentials_remain_supported( - self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, - static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + self, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], + forwarded: dict[str, str] | None, + caller: str | None, ) -> None: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - _request_auth_header, _request_extra_headers, create_tool_function, + _request_auth_header, + _request_extra_headers, + create_tool_function, ) + tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key, + "/echo", + "get", + {}, + "https://upstream.example", + headers=static, + auth_type=MCPAuth.api_key, ) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") @@ -14252,8 +14431,13 @@ class TestProtectedCredentialPreparation: self.closed = True auth = CancelledAuth() - server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key) + server = MCPServer( + server_id="cancel", + name="cancel", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + ) client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth) with pytest.raises(asyncio.CancelledError): await prepare_mcp_client(server, client) @@ -14262,8 +14446,14 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization]) async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None: - server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ") + server = MCPServer( + server_id="blank-static", + name="blank-static", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=" ", + ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server) assert exc.value.status_code == 500 @@ -14271,8 +14461,13 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM=", "Basic bm8tY29sb24="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: - server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.basic) + server = MCPServer( + server_id="bad-basic", + name="bad-basic", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) assert exc.value.status_code == 500 @@ -14281,34 +14476,48 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("value", ["Basic", "Basic ", "basic"]) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_basic_scheme_alone_is_not_a_credential(self, value: str, source: str) -> None: - server = MCPServer(server_id="basic-scheme", name="basic-scheme", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.basic, - authentication_token=value if source == "configured" else None) + server = MCPServer( + server_id="basic-scheme", + name="basic-scheme", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None, + ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,value,default_slot", [ - (MCPAuth.api_key, "fixture-key", "X-API-Key"), - (MCPAuth.bearer_token, "fixture-key", "Authorization"), - (MCPAuth.basic, "user:pass", "Authorization"), - (MCPAuth.token, "fixture-key", "Authorization"), - (MCPAuth.authorization, "fixture-key", "Authorization"), - ]) + @pytest.mark.parametrize( + "auth_type,value,default_slot", + [ + (MCPAuth.api_key, "fixture-key", "X-API-Key"), + (MCPAuth.bearer_token, "fixture-key", "Authorization"), + (MCPAuth.basic, "user:pass", "Authorization"), + (MCPAuth.token, "fixture-key", "Authorization"), + (MCPAuth.authorization, "fixture-key", "Authorization"), + ], + ) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_usable_credential_survives_an_empty_alternate_header( self, auth_type: MCPAuthType, value: str, default_slot: str, source: str ) -> None: server: Final = MCPServer( - server_id="alternate", name="alternate", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, upstream_token_header="X-Custom", + server_id="alternate", + name="alternate", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + upstream_token_header="X-Custom", authentication_token=value if source == "configured" else None, ) empty_slot: Final = default_slot if source == "configured" else "X-Custom" selected_slot: Final = "X-Custom" if source == "configured" else default_slot client: Final = await MCPServerManager()._create_mcp_client( - server, mcp_auth_header=value if source == "caller" else None, extra_headers={empty_slot: ""}, + server, + mcp_auth_header=value if source == "caller" else None, + extra_headers={empty_slot: ""}, ) request: Final = await client.prepare_request_auth() assert request.headers[selected_slot] @@ -14317,8 +14526,12 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio async def test_empty_custom_and_default_headers_do_not_satisfy_auth(self) -> None: server: Final = MCPServer( - server_id="both-empty", name="both-empty", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + server_id="both-empty", + name="both-empty", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) @@ -14331,12 +14544,17 @@ class TestProtectedCredentialPreparation: self, custom_slot: str | None, source: str ) -> None: server: Final = MCPServer( - server_id="caller-auth", name="caller-auth", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header=custom_slot, + server_id="caller-auth", + name="caller-auth", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + upstream_token_header=custom_slot, ) headers: Final = {"Authorization": "Bearer caller-credential", "X-API-Key": ""} client: Final = await MCPServerManager()._create_mcp_client( - server, mcp_auth_header=headers if source == "caller" else None, + server, + mcp_auth_header=headers if source == "caller" else None, extra_headers=headers if source == "forwarded" else None, ) request: Final = await client.prepare_request_auth() @@ -14345,14 +14563,29 @@ class TestProtectedCredentialPreparation: assert custom_slot is None or custom_slot not in request.headers @pytest.mark.asyncio - @pytest.mark.parametrize("value", [ - "", " ", "Bearer", "Basic", "token", "ApiKey", - "Bearer Bearer", "ApiKey ApiKey", "token token", "bEaReR BEARER", "aPiKeY\tAPIKEY", - ]) + @pytest.mark.parametrize( + "value", + [ + "", + " ", + "Bearer", + "Basic", + "token", + "ApiKey", + "Bearer Bearer", + "ApiKey ApiKey", + "token token", + "bEaReR BEARER", + "aPiKeY\tAPIKEY", + ], + ) async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: server: Final = MCPServer( - server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, + server_id="caller-empty", + name="caller-empty", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) @@ -14363,8 +14596,11 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_basic_requires_a_username_password_separator(self, value: str, source: str) -> None: server: Final = MCPServer( - server_id="basic-pair", name="basic-pair", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.basic, + server_id="basic-pair", + name="basic-pair", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, authentication_token=value if source == "configured" else None, ) with pytest.raises(HTTPException) as exc: @@ -14377,8 +14613,12 @@ class TestProtectedCredentialPreparation: import base64 server: Final = MCPServer( - server_id="basic-valid", name="basic-valid", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value, + server_id="basic-valid", + name="basic-valid", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token=value, ) client: Final = await MCPServerManager()._create_mcp_client(server) request: Final = await client.prepare_request_auth() @@ -14387,17 +14627,27 @@ class TestProtectedCredentialPreparation: assert base64.b64decode(encoded) == value.encode() @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,value", [ - (MCPAuth.bearer_token, "Bearer"), (MCPAuth.bearer_token, "Bearer "), (MCPAuth.bearer_token, "bearer"), - (MCPAuth.token, "token"), (MCPAuth.token, "token "), (MCPAuth.token, "TOKEN"), - ]) + @pytest.mark.parametrize( + "auth_type,value", + [ + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.bearer_token, "Bearer "), + (MCPAuth.bearer_token, "bearer"), + (MCPAuth.token, "token"), + (MCPAuth.token, "token "), + (MCPAuth.token, "TOKEN"), + ], + ) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_static_scheme_only_input_cannot_hide_behind_rendered_prefix( self, auth_type: MCPAuthType, value: str, source: str ) -> None: server: Final = MCPServer( - server_id="empty-scheme", name="empty-scheme", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, + server_id="empty-scheme", + name="empty-scheme", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, authentication_token=value if source == "configured" else None, ) with pytest.raises(HTTPException) as exc: @@ -14405,17 +14655,24 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,value,expected", [ - (MCPAuth.bearer_token, "token", "Bearer token"), - (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), - (MCPAuth.token, "tokenish", "token tokenish"), - ]) + @pytest.mark.parametrize( + "auth_type,value,expected", + [ + (MCPAuth.bearer_token, "token", "Bearer token"), + (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), + (MCPAuth.token, "tokenish", "token tokenish"), + ], + ) async def test_static_credentials_that_resemble_schemes_remain_usable( self, auth_type: MCPAuthType, value: str, expected: str ) -> None: server: Final = MCPServer( - server_id="real-token", name="real-token", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, authentication_token=value, + server_id="real-token", + name="real-token", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=value, ) client: Final = await MCPServerManager()._create_mcp_client(server) request: Final = await client.prepare_request_auth() @@ -14454,16 +14711,31 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon registry.register_tool("observer-execute", "Execute", {"type": "object"}, upstream) monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) manager = MCPServerManager() - manager.registry = {"observer": MCPServer( - server_id="observer", name="observer", server_name="observer", transport="http", - url="https://observer.example/mcp", spec_path="observer.json", auth_type="none", - )} + manager.registry = { + "observer": MCPServer( + server_id="observer", + name="observer", + server_name="observer", + transport="http", + url="https://observer.example/mcp", + spec_path="observer.json", + auth_type="none", + ) + } manager.tool_name_to_mcp_server_name_mapping = {"observer-execute": "observer"} - result = await asyncio.wait_for(manager.call_tool( - server_name="observer", name="execute", arguments={"text": "hello"}, - user_api_key_auth=UserAPIKeyAuth(), proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), - guardrail_context=MCPRequestContext.resolve_guardrail_context({"metadata": {"guardrails": ["observe"] if selected else []}}), - ), timeout=5) + result = await asyncio.wait_for( + manager.call_tool( + server_name="observer", + name="execute", + arguments={"text": "hello"}, + user_api_key_auth=UserAPIKeyAuth(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + guardrail_context=MCPRequestContext.resolve_guardrail_context( + {"metadata": {"guardrails": ["observe"] if selected else []}} + ), + ), + timeout=5, + ) assert tool_started.is_set() assert guardrail_started.is_set() is selected assert result.is_error is False @@ -14492,11 +14764,21 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie from litellm.proxy._experimental.mcp_server import server as legacy_server from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback - upstream = MCPServer(server_id="explicit-empty", name="explicit_empty", url="https://example.invalid/mcp", transport=MCPTransport.http, allow_sampling=True) + upstream = MCPServer( + server_id="explicit-empty", + name="explicit_empty", + url="https://example.invalid/mcp", + transport=MCPTransport.http, + allow_sampling=True, + ) token = auth_context_var.set(None) sampling = AsyncMock() try: - legacy_server.set_auth_context(UserAPIKeyAuth(user_id="unrelated"), raw_headers={"authorization": "unrelated-credential"}, client_ip="192.0.2.99") + legacy_server.set_auth_context( + UserAPIKeyAuth(user_id="unrelated"), + raw_headers={"authorization": "unrelated-credential"}, + client_ip="192.0.2.99", + ) with ( patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as factory, patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling), @@ -14504,7 +14786,9 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie if legacy_factory: callback = _create_sampling_callback(user_api_key_auth=UserAPIKeyAuth(user_id="explicit")) else: - await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) + await MCPServerManager()._create_mcp_client( + upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None + ) callback = factory.call_args.kwargs["sampling_callback"] await callback(None, None) captured = sampling.await_args.kwargs @@ -14527,16 +14811,28 @@ class TestSharedIdentifierPrefixWarning: manager = MCPServerManager() rows = [ LiteLLM_MCPServerTable( - server_id="srv-a", server_name="alpha", alias="shared", url="https://a.example.com/mcp", - transport=MCPTransport.http, updated_at=datetime.now(), + server_id="srv-a", + server_name="alpha", + alias="shared", + url="https://a.example.com/mcp", + transport=MCPTransport.http, + updated_at=datetime.now(), ), LiteLLM_MCPServerTable( - server_id="srv-b", server_name="beta", alias="Shared", url="https://b.example.com/mcp", - transport=MCPTransport.http, updated_at=datetime.now(), + server_id="srv-b", + server_name="beta", + alias="Shared", + url="https://b.example.com/mcp", + transport=MCPTransport.http, + updated_at=datetime.now(), ), LiteLLM_MCPServerTable( - server_id="srv-c", server_name="gamma", alias="lonely", url="https://c.example.com/mcp", - transport=MCPTransport.http, updated_at=datetime.now(), + server_id="srv-c", + server_name="gamma", + alias="lonely", + url="https://c.example.com/mcp", + transport=MCPTransport.http, + updated_at=datetime.now(), ), ] raw_rows = [MagicMock(model_dump=lambda row=row: row.model_dump()) for row in rows] @@ -14575,3 +14871,37 @@ class TestSharedIdentifierPrefixWarning: assert "srv-b" in shared_warnings[0] assert "srv-c" not in shared_warnings[0] assert "'shared'" in shared_warnings[0] + + +def test_stored_client_assertion_signing_alg_falls_back_to_rs256_for_non_approved(caplog): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _stored_client_assertion_signing_alg, + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert _stored_client_assertion_signing_alg("HS256", "srv") == "RS256" + assert "client_assertion_signing_alg" in caplog.text and "HS256" in caplog.text and "srv" in caplog.text + + +def test_stored_client_assertion_signing_alg_passes_through_approved(caplog): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _stored_client_assertion_signing_alg, + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert _stored_client_assertion_signing_alg("PS384", "srv") == "PS384" + assert _stored_client_assertion_signing_alg(None, "srv") == "RS256" + assert "approved algorithm" not in caplog.text + + +def test_mcp_server_model_rejects_non_approved_client_assertion_signing_alg(): + from pydantic import ValidationError + + with pytest.raises(ValidationError) as exc: + MCPServer( + server_id="srv", + name="srv", + transport="http", + client_assertion_signing_alg="EdDSA", + ) + assert "client_assertion_signing_alg" in str(exc.value) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index 0036035f448..ee1623fe22a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -665,3 +665,35 @@ async def test_audit_matching_login_without_nonce_does_not_report_failure(caplog ) assert result is None assert "oauth_identity_binding audit" not in caplog.text + + +def test_select_signing_key_rejects_kid_matching_key_with_non_approved_alg() -> None: + import json + + from cryptography.hazmat.primitives.asymmetric import ed25519 + + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import _BindingRejection + + okp_jwk: Final = { + **json.loads(jwt.algorithms.OKPAlgorithm.to_jwk(ed25519.Ed25519PrivateKey.generate().public_key())), + "kid": KID, + "alg": "EdDSA", + } + result: Final = _select_signing_key(_sign_id_token({}), [okp_jwk]) + assert isinstance(result, _BindingRejection) + + +def test_select_signing_key_picks_approved_key_when_kid_is_shared() -> None: + import json + + from cryptography.hazmat.primitives.asymmetric import ed25519 + + okp_jwk: Final = { + **json.loads(jwt.algorithms.OKPAlgorithm.to_jwk(ed25519.Ed25519PrivateKey.generate().public_key())), + "kid": KID, + "alg": "EdDSA", + } + hs_jwk: Final = {"kty": "oct", "kid": KID, "alg": "HS256", "k": "c2VjcmV0"} + result: Final = _select_signing_key(_sign_id_token({}), [okp_jwk, hs_jwk, _PUBLIC_JWK]) + assert isinstance(result, jwt.PyJWK) + assert result.key_id == KID diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index f8b9043a23f..585f4e0bca1 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -106,9 +106,7 @@ async def test_map_user_to_teams_handles_already_in_team_exception(): ) as mock_add: with patch("litellm.proxy.auth.handle_jwt.verbose_proxy_logger") as mock_logger: # This should not raise an exception - result = await JWTAuthManager.map_user_to_teams( - user_object=user, team_object=team - ) + result = await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team) # Verify the method completed successfully assert result is None @@ -145,14 +143,10 @@ async def test_map_user_to_teams_reraises_other_proxy_exceptions(): async def test_map_user_to_teams_null_inputs(): """Test that method handles null inputs gracefully""" # Test with null user - await JWTAuthManager.map_user_to_teams( - user_object=None, team_object=LiteLLM_TeamTable(team_id="test_team_1") - ) + await JWTAuthManager.map_user_to_teams(user_object=None, team_object=LiteLLM_TeamTable(team_id="test_team_1")) # Test with null team - await JWTAuthManager.map_user_to_teams( - user_object=LiteLLM_UserTable(user_id="test_user_1"), team_object=None - ) + await JWTAuthManager.map_user_to_teams(user_object=LiteLLM_UserTable(user_id="test_user_1"), team_object=None) # Test with both null await JWTAuthManager.map_user_to_teams(user_object=None, team_object=None) @@ -208,9 +202,7 @@ async def test_find_team_with_model_access_reports_passthrough_allowlist_denial( assert exc_info.value.status_code == 403 assert "allowed_passthrough_routes" in exc_info.value.detail assert "requested model" not in exc_info.value.detail - mock_is_auth_enforced_pass_through_route.assert_called_once_with( - route="/my-pass-through", method="POST" - ) + mock_is_auth_enforced_pass_through_route.assert_called_once_with(route="/my-pass-through", method="POST") user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"] assert user_api_key_dict.metadata == {} @@ -403,9 +395,7 @@ async def test_auth_builder_proxy_admin_user_role(): route = "/chat/completions" # Create user object with PROXY_ADMIN role - user_object = LiteLLM_UserTable( - user_id="test_user_1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_object = LiteLLM_UserTable(user_id="test_user_1", user_role=LitellmUserRoles.PROXY_ADMIN) # Create mock JWT handler jwt_handler = JWTHandler() @@ -414,14 +404,10 @@ async def test_auth_builder_proxy_admin_user_role(): # Mock all the dependencies and method calls with ( patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, - patch.object( - JWTAuthManager, "check_rbac_role", new_callable=AsyncMock - ) as mock_check_rbac, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock) as mock_check_rbac, patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac, patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes, - patch.object( - jwt_handler, "get_object_id", return_value=None - ) as mock_get_object_id, + patch.object(jwt_handler, "get_object_id", return_value=None) as mock_get_object_id, patch.object( JWTAuthManager, "get_user_info", @@ -429,9 +415,7 @@ async def test_auth_builder_proxy_admin_user_role(): return_value=("test_user_1", "test@example.com", True), ) as mock_get_user_info, patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id, - patch.object( - jwt_handler, "get_end_user_id", return_value=None - ) as mock_get_end_user_id, + patch.object(jwt_handler, "get_end_user_id", return_value=None) as mock_get_end_user_id, patch.object( JWTAuthManager, "check_admin_access", @@ -444,9 +428,7 @@ async def test_auth_builder_proxy_admin_user_role(): new_callable=AsyncMock, return_value=(None, None), ) as mock_find_team, - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=set() - ) as mock_get_all_team_ids, + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()) as mock_get_all_team_ids, patch.object( JWTAuthManager, "find_team_with_model_access", @@ -459,12 +441,8 @@ async def test_auth_builder_proxy_admin_user_role(): new_callable=AsyncMock, return_value=(user_object, None, None, None, user_object.user_id), ) as mock_get_objects, - patch.object( - JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock - ) as mock_map_user, - patch.object( - JWTAuthManager, "validate_object_id", return_value=True - ) as mock_validate_object, + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock) as mock_map_user, + patch.object(JWTAuthManager, "validate_object_id", return_value=True) as mock_validate_object, ): # Set up the mock return values mock_auth_jwt.return_value = {"sub": "test_user_1", "scope": ""} @@ -498,9 +476,7 @@ async def test_auth_builder_non_proxy_admin_user_role(): route = "/chat/completions" # Create user object with regular USER role - user_object = LiteLLM_UserTable( - user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER) # Create mock JWT handler jwt_handler = JWTHandler() @@ -509,14 +485,10 @@ async def test_auth_builder_non_proxy_admin_user_role(): # Mock all the dependencies and method calls with ( patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, - patch.object( - JWTAuthManager, "check_rbac_role", new_callable=AsyncMock - ) as mock_check_rbac, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock) as mock_check_rbac, patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac, patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes, - patch.object( - jwt_handler, "get_object_id", return_value=None - ) as mock_get_object_id, + patch.object(jwt_handler, "get_object_id", return_value=None) as mock_get_object_id, patch.object( JWTAuthManager, "get_user_info", @@ -524,9 +496,7 @@ async def test_auth_builder_non_proxy_admin_user_role(): return_value=("test_user_1", "test@example.com", True), ) as mock_get_user_info, patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id, - patch.object( - jwt_handler, "get_end_user_id", return_value=None - ) as mock_get_end_user_id, + patch.object(jwt_handler, "get_end_user_id", return_value=None) as mock_get_end_user_id, patch.object( JWTAuthManager, "check_admin_access", @@ -539,9 +509,7 @@ async def test_auth_builder_non_proxy_admin_user_role(): new_callable=AsyncMock, return_value=(None, None), ) as mock_find_team, - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=set() - ) as mock_get_all_team_ids, + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()) as mock_get_all_team_ids, patch.object( JWTAuthManager, "find_team_with_model_access", @@ -554,12 +522,8 @@ async def test_auth_builder_non_proxy_admin_user_role(): new_callable=AsyncMock, return_value=(user_object, None, None, None, user_object.user_id), ) as mock_get_objects, - patch.object( - JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock - ) as mock_map_user, - patch.object( - JWTAuthManager, "validate_object_id", return_value=True - ) as mock_validate_object, + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock) as mock_map_user, + patch.object(JWTAuthManager, "validate_object_id", return_value=True) as mock_validate_object, ): # Set up the mock return values mock_auth_jwt.return_value = {"sub": "test_user_1", "scope": ""} @@ -708,11 +672,7 @@ async def test_sync_user_role_and_teams(): prisma_client=None, user_api_key_cache=mock_user_api_key_cache, litellm_jwtauth=LiteLLM_JWTAuth( - jwt_litellm_role_map=[ - JWTLiteLLMRoleMap( - jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN - ) - ], + jwt_litellm_role_map=[JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN)], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, @@ -721,9 +681,7 @@ async def test_sync_user_role_and_teams(): token = {"roles": ["ADMIN"], "my_id_teams": ["team1", "team2"]} - user = LiteLLM_UserTable( - user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"] - ) + user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"]) prisma = AsyncMock() prisma.db.litellm_usertable.update = AsyncMock() @@ -750,11 +708,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): prisma_client=None, user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( - jwt_litellm_role_map=[ - JWTLiteLLMRoleMap( - jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN - ) - ], + jwt_litellm_role_map=[JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN)], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, @@ -771,9 +725,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): prisma = AsyncMock() prisma.db.litellm_usertable.update = AsyncMock() - await JWTAuthManager.sync_user_role_and_teams( - jwt_handler, token, user, prisma, user_api_key_cache=mock_cache - ) + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma, user_api_key_cache=mock_cache) mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args @@ -793,11 +745,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): prisma_client=None, user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( - jwt_litellm_role_map=[ - JWTLiteLLMRoleMap( - jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN - ) - ], + jwt_litellm_role_map=[JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN)], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, @@ -818,9 +766,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock, ): - await JWTAuthManager.sync_user_role_and_teams( - jwt_handler, token, user, prisma, user_api_key_cache=mock_cache - ) + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma, user_api_key_cache=mock_cache) mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args @@ -840,11 +786,7 @@ async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): prisma_client=None, user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( - jwt_litellm_role_map=[ - JWTLiteLLMRoleMap( - jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN - ) - ], + jwt_litellm_role_map=[JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN)], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, @@ -860,9 +802,7 @@ async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): prisma = AsyncMock() - await JWTAuthManager.sync_user_role_and_teams( - jwt_handler, token, user, prisma, user_api_key_cache=mock_cache - ) + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma, user_api_key_cache=mock_cache) mock_cache.async_set_cache.assert_not_called() @@ -887,9 +827,7 @@ def test_get_all_jwt_team_ids_unions_singular_and_plural(): assert jwt_handler.get_all_jwt_team_ids({"teams": ["a", "b"]}) == ["a", "b"] # both populated, no overlap - assert jwt_handler.get_all_jwt_team_ids( - {"team_id": "primary", "teams": ["a", "b"]} - ) == ["a", "b", "primary"] + assert jwt_handler.get_all_jwt_team_ids({"team_id": "primary", "teams": ["a", "b"]}) == ["a", "b", "primary"] # both populated with overlap — singular dedup'd assert jwt_handler.get_all_jwt_team_ids({"team_id": "a", "teams": ["a", "b"]}) == [ @@ -898,9 +836,7 @@ def test_get_all_jwt_team_ids_unions_singular_and_plural(): ] # singular field as multi-element list (some IdPs) — merge all, preserve plural-first order - assert jwt_handler.get_all_jwt_team_ids( - {"team_id": ["primary", "secondary"], "teams": ["a"]} - ) == [ + assert jwt_handler.get_all_jwt_team_ids({"team_id": ["primary", "secondary"], "teams": ["a"]}) == [ "a", "primary", "secondary", @@ -960,19 +896,11 @@ async def test_map_jwt_role_to_litellm_role(): litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ # Exact match - JWTLiteLLMRoleMap( - jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN - ), + JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN), # Wildcard patterns - JWTLiteLLMRoleMap( - jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER - ), - JWTLiteLLMRoleMap( - jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM - ), - JWTLiteLLMRoleMap( - jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER - ), + JWTLiteLLMRoleMap(jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER), + JWTLiteLLMRoleMap(jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM), + JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER), ], roles_jwt_field="roles", ), @@ -1044,9 +972,7 @@ async def test_map_jwt_role_to_litellm_role(): # Test patterns that don't match character classes jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [ - JWTLiteLLMRoleMap( - jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER - ), + JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER), ] token = {"roles": ["dev_4"]} # 4 is not in [123] result = jwt_handler.map_jwt_role_to_litellm_role(token) @@ -1141,25 +1067,19 @@ async def test_nested_jwt_field_access(): jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", - role_mappings=[ - RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) - ], + role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)], ) assert jwt_handler.get_object_id(nested_token, None) == "obj789" # Test 5b: object_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="object_id", - role_mappings=[ - RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) - ], + role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)], ) assert jwt_handler.get_object_id(flat_token, None) == "obj789" # Test 6: end_user_id_jwt_field with nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - end_user_id_jwt_field="customer.end_user_id" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="customer.end_user_id") assert jwt_handler.get_end_user_id(nested_token, None) == "customer123" # Test 6b: end_user_id_jwt_field with flat access (backward compatibility) @@ -1175,9 +1095,7 @@ async def test_nested_jwt_field_access(): assert jwt_handler.get_team_id(flat_token, None) == "team456" # Test 8: roles_jwt_field with deeply nested access (already supported, but testing) - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - roles_jwt_field="resource_access.my-client.roles" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(roles_jwt_field="resource_access.my-client.roles") assert jwt_handler.get_jwt_role(nested_token, []) == ["admin", "user"] # Test 9: user_roles_jwt_field with nested access (already supported, but testing) @@ -1223,10 +1141,7 @@ async def test_nested_jwt_field_missing_paths(): # Test 2: Missing user.email should return default jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="user.email") - assert ( - jwt_handler.get_user_email(incomplete_token, "default@example.com") - == "default@example.com" - ) + assert jwt_handler.get_user_email(incomplete_token, "default@example.com") == "default@example.com" # Test 3: Missing groups should return empty list jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") @@ -1241,43 +1156,28 @@ async def test_nested_jwt_field_missing_paths(): jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", - role_mappings=[ - RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) - ], + role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)], ) assert jwt_handler.get_object_id(incomplete_token, "default_obj") == "default_obj" # Test 6: Missing customer.end_user_id should return default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - end_user_id_jwt_field="customer.end_user_id" - ) - assert ( - jwt_handler.get_end_user_id(incomplete_token, "default_customer") - == "default_customer" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="customer.end_user_id") + assert jwt_handler.get_end_user_id(incomplete_token, "default_customer") == "default_customer" # Test 7: Missing tenant.team_id should use team_id_default fallback - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_id_jwt_field="tenant.team_id", team_id_default="fallback_team" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="tenant.team_id", team_id_default="fallback_team") assert jwt_handler.get_team_id(incomplete_token, "default_team") == "fallback_team" # Test 8: Missing resource_access.my-client.roles should return default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - roles_jwt_field="resource_access.my-client.roles" - ) - assert jwt_handler.get_jwt_role(incomplete_token, ["default_role"]) == [ - "default_role" - ] + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(roles_jwt_field="resource_access.my-client.roles") + assert jwt_handler.get_jwt_role(incomplete_token, ["default_role"]) == ["default_role"] # Test 9: Missing nested user roles should return default jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( user_roles_jwt_field="resource_access.my-client.roles", user_allowed_roles=["admin", "user"], ) - assert jwt_handler.get_user_roles(incomplete_token, ["default_user_role"]) == [ - "default_user_role" - ] + assert jwt_handler.get_user_roles(incomplete_token, ["default_user_role"]) == ["default_user_role"] @pytest.mark.asyncio @@ -1302,9 +1202,7 @@ async def test_metadata_prefix_handling_in_nested_fields(): } # Test 1: metadata.user.email should access user.email after prefix removal - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - user_email_jwt_field="metadata.user.email" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="metadata.user.email") # The get_nested_value function removes "metadata." prefix, so "metadata.user.email" becomes "user.email" assert jwt_handler.get_user_email(token, None) == "user@example.com" @@ -1340,9 +1238,7 @@ async def test_find_team_with_model_access_model_group(monkeypatch): async def mock_get_team_object(*args, **kwargs): # type: ignore return team - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object - ) + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() @@ -1450,9 +1346,7 @@ async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatc async def mock_get_team_object(*args, **kwargs): # type: ignore return team - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object - ) + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() @@ -1518,14 +1412,10 @@ async def test_auth_builder_returns_team_membership_object(): team_id=_team_id, budget_id="budget_123", spend=10.5, - litellm_budget_table=LiteLLM_BudgetTable( - budget_id="budget_123", rpm_limit=100, tpm_limit=5000 - ), + litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget_123", rpm_limit=100, tpm_limit=5000), ) - user_object = LiteLLM_UserTable( - user_id=_user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id=_user_id, user_role=LitellmUserRoles.INTERNAL_USER) team_object = LiteLLM_TeamTable(team_id=_team_id) @@ -1536,14 +1426,10 @@ async def test_auth_builder_returns_team_membership_object(): # Mock all the dependencies and method calls with ( patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, - patch.object( - JWTAuthManager, "check_rbac_role", new_callable=AsyncMock - ) as mock_check_rbac, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock) as mock_check_rbac, patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac, patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes, - patch.object( - jwt_handler, "get_object_id", return_value=None - ) as mock_get_object_id, + patch.object(jwt_handler, "get_object_id", return_value=None) as mock_get_object_id, patch.object( JWTAuthManager, "get_user_info", @@ -1551,9 +1437,7 @@ async def test_auth_builder_returns_team_membership_object(): return_value=(_user_id, "test@example.com", True), ) as mock_get_user_info, patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id, - patch.object( - jwt_handler, "get_end_user_id", return_value=None - ) as mock_get_end_user_id, + patch.object(jwt_handler, "get_end_user_id", return_value=None) as mock_get_end_user_id, patch.object( JWTAuthManager, "check_admin_access", @@ -1566,9 +1450,7 @@ async def test_auth_builder_returns_team_membership_object(): new_callable=AsyncMock, return_value=(_team_id, team_object), ) as mock_find_team, - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=set() - ) as mock_get_all_team_ids, + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()) as mock_get_all_team_ids, patch.object( JWTAuthManager, "find_team_with_model_access", @@ -1587,15 +1469,9 @@ async def test_auth_builder_returns_team_membership_object(): user_object.user_id, ), ) as mock_get_objects, - patch.object( - JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock - ) as mock_map_user, - patch.object( - JWTAuthManager, "validate_object_id", return_value=True - ) as mock_validate_object, - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ) as mock_sync_user, + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock) as mock_map_user, + patch.object(JWTAuthManager, "validate_object_id", return_value=True) as mock_validate_object, + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock) as mock_sync_user, ): # Set up the mock return values mock_auth_jwt.return_value = {"sub": _user_id, "scope": ""} @@ -1614,24 +1490,12 @@ async def test_auth_builder_returns_team_membership_object(): ) # Verify that team_membership_object is returned - assert result["team_membership"] is not None, ( - "team_membership should be present" - ) - assert result["team_membership"] == mock_team_membership, ( - "team_membership should match the mock object" - ) - assert result["team_membership"].user_id == _user_id, ( - "team_membership user_id should match" - ) - assert result["team_membership"].team_id == _team_id, ( - "team_membership team_id should match" - ) - assert result["team_membership"].budget_id == "budget_123", ( - "team_membership budget_id should match" - ) - assert result["team_membership"].spend == 10.5, ( - "team_membership spend should match" - ) + assert result["team_membership"] is not None, "team_membership should be present" + assert result["team_membership"] == mock_team_membership, "team_membership should match the mock object" + assert result["team_membership"].user_id == _user_id, "team_membership user_id should match" + assert result["team_membership"].team_id == _team_id, "team_membership team_id should match" + assert result["team_membership"].budget_id == "budget_123", "team_membership budget_id should match" + assert result["team_membership"].spend == 10.5, "team_membership spend should match" @pytest.mark.asyncio @@ -1648,9 +1512,7 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): general_settings = {"enforce_rbac": False} route = "/chat/completions" - user_object = LiteLLM_UserTable( - user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER) # Create JWT handler with OIDC UserInfo enabled jwt_handler = JWTHandler() @@ -1677,18 +1539,12 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): # Mock all the dependencies with ( - patch.object( - jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock - ) as mock_get_userinfo, + patch.object(jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock) as mock_get_userinfo, patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, - patch.object( - JWTAuthManager, "check_rbac_role", new_callable=AsyncMock - ) as mock_check_rbac, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock) as mock_check_rbac, patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac, patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes, - patch.object( - jwt_handler, "get_object_id", return_value=None - ) as mock_get_object_id, + patch.object(jwt_handler, "get_object_id", return_value=None) as mock_get_object_id, patch.object( JWTAuthManager, "get_user_info", @@ -1696,9 +1552,7 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): return_value=("test_user_1", "test@example.com", True), ) as mock_get_user_info, patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id, - patch.object( - jwt_handler, "get_end_user_id", return_value=None - ) as mock_get_end_user_id, + patch.object(jwt_handler, "get_end_user_id", return_value=None) as mock_get_end_user_id, patch.object( JWTAuthManager, "check_admin_access", @@ -1711,9 +1565,7 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): new_callable=AsyncMock, return_value=(None, None), ) as mock_find_team, - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=set() - ) as mock_get_all_team_ids, + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()) as mock_get_all_team_ids, patch.object( JWTAuthManager, "find_team_with_model_access", @@ -1726,15 +1578,9 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): new_callable=AsyncMock, return_value=(user_object, None, None, None, user_object.user_id), ) as mock_get_objects, - patch.object( - JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock - ) as mock_map_user, - patch.object( - JWTAuthManager, "validate_object_id", return_value=True - ) as mock_validate_object, - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ) as mock_sync_user, + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock) as mock_map_user, + patch.object(JWTAuthManager, "validate_object_id", return_value=True) as mock_validate_object, + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock) as mock_sync_user, ): # Set up mock return values mock_get_userinfo.return_value = userinfo_response @@ -1775,9 +1621,7 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): general_settings = {"enforce_rbac": False} route = "/chat/completions" - user_object = LiteLLM_UserTable( - user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER) # Create JWT handler with OIDC UserInfo disabled jwt_handler = JWTHandler() @@ -1801,18 +1645,12 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): # Mock all the dependencies with ( - patch.object( - jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock - ) as mock_get_userinfo, + patch.object(jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock) as mock_get_userinfo, patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, - patch.object( - JWTAuthManager, "check_rbac_role", new_callable=AsyncMock - ) as mock_check_rbac, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock) as mock_check_rbac, patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac, patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes, - patch.object( - jwt_handler, "get_object_id", return_value=None - ) as mock_get_object_id, + patch.object(jwt_handler, "get_object_id", return_value=None) as mock_get_object_id, patch.object( JWTAuthManager, "get_user_info", @@ -1820,9 +1658,7 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): return_value=("test_user_1", None, None), ) as mock_get_user_info, patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id, - patch.object( - jwt_handler, "get_end_user_id", return_value=None - ) as mock_get_end_user_id, + patch.object(jwt_handler, "get_end_user_id", return_value=None) as mock_get_end_user_id, patch.object( JWTAuthManager, "check_admin_access", @@ -1835,9 +1671,7 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): new_callable=AsyncMock, return_value=(None, None), ) as mock_find_team, - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=set() - ) as mock_get_all_team_ids, + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()) as mock_get_all_team_ids, patch.object( JWTAuthManager, "find_team_with_model_access", @@ -1850,15 +1684,9 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): new_callable=AsyncMock, return_value=(user_object, None, None, None, user_object.user_id), ) as mock_get_objects, - patch.object( - JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock - ) as mock_map_user, - patch.object( - JWTAuthManager, "validate_object_id", return_value=True - ) as mock_validate_object, - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ) as mock_sync_user, + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock) as mock_map_user, + patch.object(JWTAuthManager, "validate_object_id", return_value=True) as mock_validate_object, + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock) as mock_sync_user, ): # Set up mock return values mock_auth_jwt.return_value = jwt_response @@ -1904,9 +1732,7 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens() general_settings = {"enforce_rbac": False} route = "/chat/completions" - user_object = LiteLLM_UserTable( - user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER) jwt_handler = JWTHandler() user_api_key_cache = DualCache() @@ -1925,9 +1751,7 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens() jwt_response = {"sub": "test_user_1", "scope": ""} with ( - patch.object( - jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock - ) as mock_get_userinfo, + patch.object(jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock) as mock_get_userinfo, patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "get_rbac_role", return_value=None), @@ -1968,9 +1792,7 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens() ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), ): mock_auth_jwt.return_value = jwt_response @@ -2052,9 +1874,7 @@ async def test_auth_builder_uses_team_from_header_e2e(): ) team_object = LiteLLM_TeamTable(team_id="team-2") - user_object = LiteLLM_UserTable( - user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_object = LiteLLM_UserTable(user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER) with ( patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, @@ -2065,9 +1885,7 @@ async def test_auth_builder_uses_team_from_header_e2e(): new_callable=AsyncMock, return_value=None, ), - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_team, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_team, patch.object( JWTAuthManager, "get_objects", @@ -2075,9 +1893,7 @@ async def test_auth_builder_uses_team_from_header_e2e(): return_value=(user_object, None, None, None, user_object.user_id), ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), ): mock_auth_jwt.return_value = { "sub": "user-1", @@ -2294,9 +2110,7 @@ async def test_auth_builder_rbac_team_loads_team_for_passthrough_allowlist(): return_value=(None, None, None, None, None), ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route", return_value=True, @@ -2325,9 +2139,7 @@ async def test_auth_builder_rbac_team_loads_team_for_passthrough_allowlist(): mock_get_team.assert_awaited_once() assert mock_get_team.await_args.kwargs["team_id"] == "team-rbac" user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"] - assert user_api_key_dict.team_metadata == { - "allowed_passthrough_routes": ["/my-pass-through"] - } + assert user_api_key_dict.team_metadata == {"allowed_passthrough_routes": ["/my-pass-through"]} @pytest.mark.asyncio @@ -2421,9 +2233,7 @@ async def test_auth_builder_admin_on_llm_route_honors_team_header(): patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "is_admin", return_value=True), - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_team, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_team, ): mock_auth_jwt.return_value = { "sub": "admin-user", @@ -2475,9 +2285,7 @@ async def test_auth_builder_admin_on_mgmt_route_ignores_team_header(): patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "is_admin", return_value=True), - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_team, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_team, ): mock_auth_jwt.return_value = { "sub": "admin-user", @@ -2531,9 +2339,7 @@ async def test_auth_builder_admin_on_llm_route_without_header_unchanged(): patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "is_admin", return_value=True), - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_team, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_team, ): mock_auth_jwt.return_value = { "sub": "admin-user", @@ -2577,9 +2383,7 @@ async def test_get_team_alias_with_nested_fields(): } # Test nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_alias_jwt_field="organization.team.name" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="organization.team.name") assert jwt_handler.get_team_alias(nested_token, None) == "engineering-team" # Test flat access (backward compatibility) @@ -2587,9 +2391,7 @@ async def test_get_team_alias_with_nested_fields(): assert jwt_handler.get_team_alias(nested_token, None) == "flat-team" # Test missing field returns default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_alias_jwt_field="nonexistent.field" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="nonexistent.field") assert jwt_handler.get_team_alias(nested_token, "default-team") == "default-team" # Test with team_alias_jwt_field not configured @@ -2620,9 +2422,7 @@ async def test_is_required_team_id_with_team_alias_field(): assert jwt_handler.is_required_team_id() is True # Both fields set - should return True - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_id_jwt_field="team_id", team_alias_jwt_field="team_name" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id", team_alias_jwt_field="team_name") assert jwt_handler.is_required_team_id() is True @@ -2654,9 +2454,7 @@ async def test_find_and_validate_specific_team_id_with_team_alias(): # Mock team object returned by get_team_object_by_alias team_object = LiteLLM_TeamTable(team_id="resolved-team-id", team_alias="my-team") - with patch( - "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock - ) as mock_get_by_alias: + with patch("litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock) as mock_get_by_alias: mock_get_by_alias.return_value = team_object team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id( @@ -2699,9 +2497,7 @@ async def test_find_and_validate_team_id_takes_precedence_over_name(): jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, - litellm_jwtauth=LiteLLM_JWTAuth( - team_id_jwt_field="team_id", team_alias_jwt_field="team_alias" - ), + litellm_jwtauth=LiteLLM_JWTAuth(team_id_jwt_field="team_id", team_alias_jwt_field="team_alias"), ) # Token with both team_id and team name @@ -2711,9 +2507,7 @@ async def test_find_and_validate_team_id_takes_precedence_over_name(): team_object = LiteLLM_TeamTable(team_id="direct-team-id") with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_by_id, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_by_id, patch( "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock, @@ -2793,9 +2587,7 @@ async def test_get_org_alias_with_nested_fields(): } # Test nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - org_alias_jwt_field="company.organization.name" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="company.organization.name") assert jwt_handler.get_org_alias(nested_token, None) == "acme-corp" # Test flat access @@ -2803,9 +2595,7 @@ async def test_get_org_alias_with_nested_fields(): assert jwt_handler.get_org_alias(nested_token, None) == "flat-org" # Test missing field returns default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - org_alias_jwt_field="nonexistent.field" - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="nonexistent.field") assert jwt_handler.get_org_alias(nested_token, "default-org") == "default-org" # Test with org_alias_jwt_field not configured @@ -2843,9 +2633,7 @@ async def test_get_objects_resolves_org_by_name(): models=[], ) - with patch( - "litellm.proxy.auth.handle_jwt.get_org_object_by_alias", new_callable=AsyncMock - ) as mock_get_by_alias: + with patch("litellm.proxy.auth.handle_jwt.get_org_object_by_alias", new_callable=AsyncMock) as mock_get_by_alias: mock_get_by_alias.return_value = org_object ( @@ -2923,9 +2711,7 @@ async def test_resolve_jwks_url_resolves_oidc_discovery_document(): litellm_jwtauth=LiteLLM_JWTAuth(), ) - discovery_url = ( - "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" - ) + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" mock_response = MagicMock() @@ -2956,9 +2742,7 @@ async def test_resolve_jwks_url_caches_resolved_jwks_uri(): litellm_jwtauth=LiteLLM_JWTAuth(), ) - discovery_url = ( - "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" - ) + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" mock_response = MagicMock() @@ -3102,9 +2886,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_notation(): error_msg = str(exc_info.value) # Should mention the bad field name and suggest the fix assert "roles.0" in error_msg, f"Expected field name in: {error_msg}" - assert "roles" in error_msg and "list" in error_msg, ( - f"Expected hint about using 'roles' instead: {error_msg}" - ) + assert "roles" in error_msg and "list" in error_msg, f"Expected hint about using 'roles' instead: {error_msg}" @pytest.mark.asyncio @@ -3132,9 +2914,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_index_notation() error_msg = str(exc_info.value) assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}" - assert "roles" in error_msg and "list" in error_msg, ( - f"Expected hint about using 'roles' instead: {error_msg}" - ) + assert "roles" in error_msg and "list" in error_msg, f"Expected hint about using 'roles' instead: {error_msg}" @pytest.mark.asyncio @@ -3244,9 +3024,7 @@ async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team( if len(user_teams) == 1 and get_team_object_return == "resolved_row": only = user_teams[0] team_table = LiteLLM_TeamTable(team_id=only) - membership = LiteLLM_TeamMembership( - user_id=user_id, team_id=only, litellm_budget_table=None - ) + membership = LiteLLM_TeamMembership(user_id=user_id, team_id=only, litellm_budget_table=None) get_team_return_value = team_table membership_return_value = membership else: @@ -3305,9 +3083,7 @@ async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team( ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -3324,9 +3100,7 @@ async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team( code = 404 if get_team_object_return == "http_404" else 500 mock_get_team.side_effect = HTTPException( status_code=code, - detail={ - "error": f"Team doesn't exist in db. Team={user_teams[0]}. Create team via `/team/new` call." - }, + detail={"error": f"Team doesn't exist in db. Team={user_teams[0]}. Create team via `/team/new` call."}, ) else: mock_get_team.return_value = get_team_return_value @@ -3435,9 +3209,7 @@ async def test_auth_builder_single_team_fallback_membership_outage_raises_instea ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -3478,9 +3250,7 @@ def _reset_unscoped_warning_flag(): JWTHandler._unscoped_jwt_warning_emitted = False -def test_build_decode_kwargs_no_env_disables_both_verifications( - monkeypatch, _reset_unscoped_warning_flag -): +def test_build_decode_kwargs_no_env_disables_both_verifications(monkeypatch, _reset_unscoped_warning_flag): monkeypatch.delenv("JWT_AUDIENCE", raising=False) monkeypatch.delenv("JWT_ISSUER", raising=False) @@ -3491,9 +3261,7 @@ def test_build_decode_kwargs_no_env_disables_both_verifications( assert kwargs["options"] == {"verify_aud": False, "verify_iss": False} -def test_build_decode_kwargs_audience_only_enables_aud_verification( - monkeypatch, _reset_unscoped_warning_flag -): +def test_build_decode_kwargs_audience_only_enables_aud_verification(monkeypatch, _reset_unscoped_warning_flag): monkeypatch.setenv("JWT_AUDIENCE", "my-proxy") monkeypatch.delenv("JWT_ISSUER", raising=False) @@ -3505,9 +3273,7 @@ def test_build_decode_kwargs_audience_only_enables_aud_verification( assert kwargs["options"] == {"verify_iss": False} -def test_build_decode_kwargs_issuer_only_enables_iss_verification( - monkeypatch, _reset_unscoped_warning_flag -): +def test_build_decode_kwargs_issuer_only_enables_iss_verification(monkeypatch, _reset_unscoped_warning_flag): monkeypatch.delenv("JWT_AUDIENCE", raising=False) monkeypatch.setenv("JWT_ISSUER", "https://idp.example.com/") @@ -3518,9 +3284,7 @@ def test_build_decode_kwargs_issuer_only_enables_iss_verification( assert kwargs["options"] == {"verify_aud": False} -def test_build_decode_kwargs_both_set_enables_full_verification( - monkeypatch, _reset_unscoped_warning_flag -): +def test_build_decode_kwargs_both_set_enables_full_verification(monkeypatch, _reset_unscoped_warning_flag): monkeypatch.setenv("JWT_AUDIENCE", "my-proxy") monkeypatch.setenv("JWT_ISSUER", "https://idp.example.com/") @@ -3532,9 +3296,7 @@ def test_build_decode_kwargs_both_set_enables_full_verification( assert kwargs["options"] is None -def test_build_decode_kwargs_warns_once_when_unscoped( - monkeypatch, _reset_unscoped_warning_flag, caplog -): +def test_build_decode_kwargs_warns_once_when_unscoped(monkeypatch, _reset_unscoped_warning_flag, caplog): """The warning about unscoped JWT auth should fire on the first call but not on every subsequent decode.""" import logging @@ -3550,17 +3312,12 @@ def test_build_decode_kwargs_warns_once_when_unscoped( matching = [ r for r in caplog.records - if "JWT auth is enabled" in r.getMessage() - and "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() + if "JWT auth is enabled" in r.getMessage() and "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() ] - assert len(matching) == 1, ( - f"Expected exactly one warning across 3 calls, got {len(matching)}" - ) + assert len(matching) == 1, f"Expected exactly one warning across 3 calls, got {len(matching)}" -def test_build_decode_kwargs_no_warning_when_scoped( - monkeypatch, _reset_unscoped_warning_flag, caplog -): +def test_build_decode_kwargs_no_warning_when_scoped(monkeypatch, _reset_unscoped_warning_flag, caplog): import logging monkeypatch.setenv("JWT_AUDIENCE", "my-proxy") @@ -3569,11 +3326,7 @@ def test_build_decode_kwargs_no_warning_when_scoped( JWTHandler._build_decode_kwargs() - matching = [ - r - for r in caplog.records - if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() - ] + matching = [r for r in caplog.records if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()] assert matching == [] @@ -3631,11 +3384,7 @@ async def test_find_team_with_model_access_unresolved_group_claim_returns_none( from litellm.router import Router - router = Router( - model_list=[ - {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} - ] - ) + router = Router(model_list=[{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}]) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) @@ -3707,9 +3456,7 @@ async def test_find_and_validate_specific_team_id_non_404_http_exception_propaga "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, ) as mock_get_team: - mock_get_team.side_effect = HTTPException( - status_code=status_code, detail="non-404 failure" - ) + mock_get_team.side_effect = HTTPException(status_code=status_code, detail="non-404 failure") with pytest.raises(HTTPException) as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( @@ -3782,9 +3529,7 @@ async def test_find_team_with_model_access_resolved_team_without_model_still_rai async def mock_get_team_object(*_args, **_kwargs): return team - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object - ) + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() @@ -3849,11 +3594,7 @@ async def test_find_team_with_model_access_unresolved_group_claim_default_raises from litellm.router import Router - router = Router( - model_list=[ - {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} - ] - ) + router = Router(model_list=[{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}]) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) @@ -3890,12 +3631,7 @@ def test_canonical_user_id_rebinds_to_legacy_uuid(): jwt_email = "matt@example.com" user_object = LiteLLM_UserTable(user_id=legacy_uuid, user_email=jwt_email) - assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id=jwt_email, user_object=user_object - ) - == legacy_uuid - ) + assert JWTAuthManager._canonical_user_id_from_db(user_id=jwt_email, user_object=user_object) == legacy_uuid def test_canonical_user_id_no_change_when_ids_match(): @@ -3903,28 +3639,20 @@ def test_canonical_user_id_no_change_when_ids_match(): same = "alice@example.com" user_object = LiteLLM_UserTable(user_id=same, user_email=same) - assert ( - JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object) - == same - ) + assert JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object) == same def test_canonical_user_id_returns_claim_when_no_user_object(): """No resolved row (e.g. upsert disabled / brand new) -> keep the claim.""" assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id="newcomer@example.com", user_object=None - ) + JWTAuthManager._canonical_user_id_from_db(user_id="newcomer@example.com", user_object=None) == "newcomer@example.com" ) def test_canonical_user_id_returns_none_when_claim_none_and_no_object(): """Defensive: no claim and no row -> stays None, never invents an id.""" - assert ( - JWTAuthManager._canonical_user_id_from_db(user_id=None, user_object=None) - is None - ) + assert JWTAuthManager._canonical_user_id_from_db(user_id=None, user_object=None) is None def test_canonical_user_id_no_change_when_db_user_id_falsy(): @@ -3934,10 +3662,7 @@ def test_canonical_user_id_no_change_when_db_user_id_falsy(): user_id = "" assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id="jwt@example.com", user_object=_Stub() - ) - == "jwt@example.com" + JWTAuthManager._canonical_user_id_from_db(user_id="jwt@example.com", user_object=_Stub()) == "jwt@example.com" ) @@ -3953,9 +3678,7 @@ async def test_auth_jwt_expired_token_raises_401_jwk_path(): jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() with ( - patch.object( - jwt_handler, "get_public_key", new_callable=AsyncMock - ) as mock_get_public_key, + patch.object(jwt_handler, "get_public_key", new_callable=AsyncMock) as mock_get_public_key, patch( "litellm.proxy.auth.handle_jwt.jwt.get_unverified_header", return_value={"kid": "test-kid"}, @@ -3991,9 +3714,7 @@ async def test_auth_jwt_expired_token_raises_401_pem_cert_path(): mock_cert.public_key.return_value.public_bytes.return_value = b"fake-key" with ( - patch.object( - jwt_handler, "get_public_key", new_callable=AsyncMock - ) as mock_get_public_key, + patch.object(jwt_handler, "get_public_key", new_callable=AsyncMock) as mock_get_public_key, patch( "litellm.proxy.auth.handle_jwt.jwt.get_unverified_header", return_value={"kid": "test-kid"}, @@ -4007,9 +3728,7 @@ async def test_auth_jwt_expired_token_raises_401_pem_cert_path(): side_effect=jwt_lib.ExpiredSignatureError("Signature has expired"), ), ): - mock_get_public_key.return_value = ( - "-----BEGIN CERTIFICATE-----\nfake\n-----END CERTIFICATE-----" - ) + mock_get_public_key.return_value = "-----BEGIN CERTIFICATE-----\nfake\n-----END CERTIFICATE-----" with pytest.raises(ProxyException) as exc_info: await jwt_handler.auth_jwt(token="expired.jwt.token") @@ -4123,9 +3842,7 @@ async def test_get_public_key_fetches_and_caches_jwks_response(): ) assert public_key == jwk - cached_keys = await cache.async_get_cache( - key="litellm_jwt_auth_keys_https://issuer.example.com/keys" - ) + cached_keys = await cache.async_get_cache(key="litellm_jwt_auth_keys_https://issuer.example.com/keys") assert cached_keys == [jwk] @@ -4370,9 +4087,7 @@ async def test_lowering_public_key_stale_ttl_stops_serving_a_copy_cached_under_t # The operator tightens the window and restarts; the cache, and its long-lived copy, survive. endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) - tightened = _get_jwt_handler_with_scripted_endpoint( - cache, endpoint, public_key_stale_ttl=lowered_stale_ttl - ) + tightened = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=lowered_stale_ttl) await cache.async_set_cache( key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}", value=time.time() - 7200, @@ -4720,9 +4435,7 @@ def test_get_jwks_url_for_issuer_falls_back_to_discovery_document(): jwks_url = jwt_handler._get_jwks_url_for_issuer(issuer_config=issuer_config) - assert ( - jwks_url == "https://issuer.example.com/tenant/.well-known/openid-configuration" - ) + assert jwks_url == "https://issuer.example.com/tenant/.well-known/openid-configuration" @pytest.mark.asyncio @@ -4746,9 +4459,7 @@ async def test_get_objects_team_membership_uses_rebound_user_id(): return None jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - user_id_jwt_field="email", user_id_upsert=True - ) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="email", user_id_upsert=True) with ( patch( @@ -4844,9 +4555,7 @@ async def test_multi_issuer_jwt_validates_selected_issuer_and_maps_claims( assert claims[JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == issuer_two assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-org" - assert jwt_handler.get_team_id(token=claims, default_value=None) == ( - "example-org/litellm-fork" - ) + assert jwt_handler.get_team_id(token=claims, default_value=None) == ("example-org/litellm-fork") @pytest.mark.asyncio @@ -4949,9 +4658,7 @@ async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch): claims = await jwt_handler.auth_jwt(token=token) - assert ( - jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace" - ) + assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace" @pytest.mark.asyncio @@ -4984,7 +4691,7 @@ async def test_multi_issuer_jwt_unknown_issuer_falls_back_to_global_jwks(monkeyp kid="issuer-key", ) - with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc: + with pytest.raises(Exception, match="Missing JWT Public Key URL from environment\\.") as exc: await jwt_handler.auth_jwt(token=token) assert "Missing JWT Public Key URL from environment." in str(exc.value) @@ -5058,7 +4765,7 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) kid=shared_kid, ) - with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc: + with pytest.raises(Exception, match="Validation fails: Signature verification failed") as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -5113,7 +4820,7 @@ def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception, match='must configure audience or set') as exc: + with pytest.raises(Exception, match="must configure audience or set") as exc: LiteLLM_JWTAuth( issuers=[ { @@ -5130,7 +4837,7 @@ def test_multi_issuer_jwt_rejects_audience_with_disable_audience_validation(): issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception, match='cannot set audience and disable_audience_validation=True') as exc: + with pytest.raises(Exception, match="cannot set audience and disable_audience_validation=True") as exc: LiteLLM_JWTAuth( issuers=[ { @@ -5142,9 +4849,7 @@ def test_multi_issuer_jwt_rejects_audience_with_disable_audience_validation(): ] ) - assert "cannot set audience and disable_audience_validation=True together" in str( - exc.value - ) + assert "cannot set audience and disable_audience_validation=True together" in str(exc.value) @pytest.mark.asyncio @@ -5197,21 +4902,15 @@ async def test_global_jwt_ignores_user_supplied_internal_claims(monkeypatch): claims = await jwt_handler.auth_jwt(token=token) - assert jwt_handler.get_user_id(token=claims, default_value=None) == ( - "real-user@example.com" - ) - assert jwt_handler.get_user_email(token=claims, default_value=None) == ( - "real-user@example.com" - ) + assert jwt_handler.get_user_id(token=claims, default_value=None) == ("real-user@example.com") + assert jwt_handler.get_user_email(token=claims, default_value=None) == ("real-user@example.com") assert jwt_handler.get_team_id(token=claims, default_value=None) == "real-team" assert jwt_handler.get_team_ids_from_jwt(token=claims) == [ "real-team", "secondary-team", ] assert jwt_handler.get_org_id(token=claims, default_value=None) == "real-org" - assert jwt_handler.get_end_user_id(token=claims, default_value=None) == ( - "real-end-user" - ) + assert jwt_handler.get_end_user_id(token=claims, default_value=None) == ("real-end-user") @pytest.mark.asyncio @@ -5251,15 +4950,11 @@ async def test_multi_issuer_jwt_strips_unmapped_internal_claims(monkeypatch): assert JWTHandler.LITELLM_TEAM_ID_CLAIM not in claims assert jwt_handler.get_user_id(token=claims, default_value=None) is None assert jwt_handler.get_team_id(token=claims, default_value=None) is None - assert jwt_handler.get_user_email(token=claims, default_value=None) == ( - "real-user@example.com" - ) + assert jwt_handler.get_user_email(token=claims, default_value=None) == ("real-user@example.com") @pytest.mark.asyncio -async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning( - monkeypatch, caplog -): +async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning(monkeypatch, caplog): import logging monkeypatch.delenv("JWT_AUDIENCE", raising=False) @@ -5309,11 +5004,7 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym JWTHandler._build_decode_kwargs() - matching = [ - r - for r in caplog.records - if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() - ] + matching = [r for r in caplog.records if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()] assert len(matching) == 1 @@ -5441,9 +5132,7 @@ async def test_resolve_db_team_fallback_skips_unresolvable_membership(): "expect_403", ), [ - pytest.param( - True, ["team_solo"], None, "team_solo", False, id="flag_on_single_db_team" - ), + pytest.param(True, ["team_solo"], None, "team_solo", False, id="flag_on_single_db_team"), pytest.param( True, ["team_a", "team_b"], @@ -5521,9 +5210,7 @@ async def test_auth_builder_db_team_fallback_when_jwt_has_no_team( async def call_auth_builder(): with ( - patch.object( - jwt_handler, "auth_jwt", new_callable=AsyncMock - ) as mock_auth_jwt, + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "get_rbac_role", return_value=None), patch.object(jwt_handler, "get_scopes", return_value=[]), @@ -5563,9 +5250,7 @@ async def test_auth_builder_db_team_fallback_when_jwt_has_no_team( ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -5669,9 +5354,7 @@ async def test_resolve_db_team_fallback_skips_team_without_model_access(): teams=["restricted_team", "allowed_team"], ) teams = { - "restricted_team": LiteLLM_TeamTable( - team_id="restricted_team", models=["claude-3"] - ), + "restricted_team": LiteLLM_TeamTable(team_id="restricted_team", models=["claude-3"]), "allowed_team": LiteLLM_TeamTable(team_id="allowed_team", models=["gpt-4"]), } @@ -5853,9 +5536,7 @@ async def _run_auth_builder_with_header_team( jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = jwt_auth_config with ( - patch.object( - jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value=token - ), + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value=token), patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "get_rbac_role", return_value=None), patch.object(jwt_handler, "get_scopes", return_value=[]), @@ -5874,9 +5555,7 @@ async def _run_auth_builder_with_header_team( new_callable=AsyncMock, return_value=None, ), - patch.object( - JWTAuthManager, "get_all_team_ids", return_value=allowed_team_ids - ), + patch.object(JWTAuthManager, "get_all_team_ids", return_value=allowed_team_ids), patch.object( JWTAuthManager, "get_objects", @@ -5885,9 +5564,7 @@ async def _run_auth_builder_with_header_team( ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -5914,9 +5591,7 @@ async def _run_auth_builder_with_header_team( @pytest.mark.asyncio -async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> ( - None -): +async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> None: """A provisional x-litellm-team-id naming a nonexistent team must produce the exact same 403 shape as one naming an existing team outside the caller's memberships. Letting get_team_object's 404 surface would give any @@ -5936,19 +5611,15 @@ async def test_auth_builder_header_team_not_found_matches_non_membership_denial( return LiteLLM_TeamTable(team_id=team_id) with pytest.raises(HTTPException) as missing_exc: - await _run_auth_builder_with_header_team( - config, token, "team_ghost", user_object, _team_lookup_404, set() - ) + await _run_auth_builder_with_header_team(config, token, "team_ghost", user_object, _team_lookup_404, set()) with pytest.raises(HTTPException) as outsider_exc: - await _run_auth_builder_with_header_team( - config, token, "team_other", user_object, team_exists, set() - ) + await _run_auth_builder_with_header_team(config, token, "team_other", user_object, team_exists, set()) assert missing_exc.value.status_code == 403 assert outsider_exc.value.status_code == 403 - assert missing_exc.value.detail.replace( - "team_ghost", "" - ) == outsider_exc.value.detail.replace("team_other", "") + assert missing_exc.value.detail.replace("team_ghost", "") == outsider_exc.value.detail.replace( + "team_other", "" + ) assert "exist" not in missing_exc.value.detail @@ -6145,9 +5816,7 @@ async def test_auth_builder_db_fallback_does_not_validate_rbac_team_against_db_m ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6299,9 +5968,7 @@ async def test_auth_builder_db_fallback_runs_when_only_team_id_default_set(): ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6383,9 +6050,7 @@ async def test_auth_builder_alias_only_token_resolves_alias_not_db_fallback(): ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6442,14 +6107,10 @@ async def test_find_and_validate_specific_team_id_alias_wins_over_team_id_defaul ) jwt_token = {"sub": "user-1", "team_alias": "my-team"} - alias_team = LiteLLM_TeamTable( - team_id="alias_resolved_team", team_alias="my-team" - ) + alias_team = LiteLLM_TeamTable(team_id="alias_resolved_team", team_alias="my-team") with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_by_id, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_by_id, patch( "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock, @@ -6496,9 +6157,7 @@ async def test_find_and_validate_specific_team_id_team_id_default_used_without_a default_team = LiteLLM_TeamTable(team_id="config_default_team") with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock - ) as mock_get_by_id, + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_by_id, patch( "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock, @@ -6567,9 +6226,7 @@ async def test_auth_builder_db_fallback_enforces_passthrough_route_access(): ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6723,18 +6380,14 @@ async def test_sync_user_role_and_teams_singular_claim_reconciles_memberships(): "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock, ) as mock_patch: - await JWTAuthManager.sync_user_role_and_teams( - jwt_handler, token, user, AsyncMock() - ) + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, AsyncMock()) mock_patch.assert_awaited_once() assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == { "team_stale_a", "team_stale_b", } - assert set(mock_patch.call_args.kwargs["teams_ids_to_add_user_to"]) == { - "team_primary" - } + assert set(mock_patch.call_args.kwargs["teams_ids_to_add_user_to"]) == {"team_primary"} assert user.teams == ["team_primary"] @@ -6793,9 +6446,7 @@ async def test_auth_builder_provisional_header_team_is_not_upserted(): ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6870,9 +6521,7 @@ async def test_auth_builder_header_cannot_override_rbac_team_under_db_fallback() ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -6922,9 +6571,7 @@ async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fa async def call(route: str): with ( - patch.object( - jwt_handler, "auth_jwt", new_callable=AsyncMock - ) as mock_auth_jwt, + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), patch.object(jwt_handler, "get_rbac_role", return_value=None), patch.object(jwt_handler, "get_scopes", return_value=[]), @@ -6952,9 +6599,7 @@ async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fa ), patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), patch.object(JWTAuthManager, "validate_object_id", return_value=True), - patch.object( - JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock - ), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), patch( "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, @@ -7312,14 +6957,10 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock, ) as mock_patch: - await JWTAuthManager.sync_user_role_and_teams( - jwt_handler, token, user, AsyncMock() - ) + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, AsyncMock()) mock_patch.assert_awaited_once() - assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == { - "team_existing" - } + assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == {"team_existing"} assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == [] assert user.teams == [] @@ -7544,8 +7185,12 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a if identity_only: identity = await JWTAuthManager.resolve_identity( - api_key=token, jwt_handler=jwt_handler, prisma_client=None, - user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, + api_key=token, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, ) assert identity.agent_id == "canonical-agent-id" return @@ -7580,8 +7225,12 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch if identity_only: with pytest.raises(HTTPException) as denial: await JWTAuthManager.resolve_identity( - api_key=token, jwt_handler=jwt_handler, prisma_client=None, - user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, + api_key=token, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, ) assert denial.value.status_code == 403 return @@ -7648,7 +7297,9 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc from litellm.proxy.management_endpoints import team_endpoints handler, token = _entra_signed_app_token( - monkeypatch, azp="canonical-agent-id", scope=LiteLLM_JWTAuth().admin_jwt_scope, + monkeypatch, + azp="canonical-agent-id", + scope=LiteLLM_JWTAuth().admin_jwt_scope, ) handler.bind_agent_lookup(_entra_agent_registry()) handler.litellm_jwtauth.team_id_upsert = True @@ -7660,10 +7311,16 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc resolve = JWTAuthManager.auth_builder if admission else JWTAuthManager.authorize_jwt result = await resolve( - api_key=token, jwt_handler=handler, request_data={}, general_settings={}, - route="/chat/completions", prisma_client=database, - user_api_key_cache=handler.user_api_key_cache, parent_otel_span=None, - proxy_logging_obj=MagicMock(), request_headers={"x-litellm-team-id": "new-team"}, + api_key=token, + jwt_handler=handler, + request_data={}, + general_settings={}, + route="/chat/completions", + prisma_client=database, + user_api_key_cache=handler.user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + request_headers={"x-litellm-team-id": "new-team"}, ) assert result["is_proxy_admin"] is True @@ -7688,7 +7345,12 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning cache: Final = UserApiKeyCache() cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk]) user_id: Final = f"admin-status-{existing_user}-{warm_cache}-{email}" - user: Final = LiteLLM_UserTable(user_id=user_id, user_email="admin@allowed.example", metadata={"scim_active": False}, organization_memberships=[]) + user: Final = LiteLLM_UserTable( + user_id=user_id, + user_email="admin@allowed.example", + metadata={"scim_active": False}, + organization_memberships=[], + ) if existing_user and warm_cache: cache.set_cache(user_id, user) database: Final = MagicMock() @@ -7701,7 +7363,9 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning prisma_client=database, user_api_key_cache=cache, litellm_jwtauth=LiteLLM_JWTAuth( - user_id_jwt_field="sub", user_id_upsert=True, user_email_jwt_field="email", + user_id_jwt_field="sub", + user_id_upsert=True, + user_email_jwt_field="email", user_allowed_email_domain="allowed.example", ), ) @@ -7709,12 +7373,22 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning monkeypatch.setenv("JWT_ISSUER", "https://admin.example") monkeypatch.setenv("JWT_AUDIENCE", "gateway") token: Final = _encode_rsa_jwt( - private_key, "https://admin.example", "gateway", "admin-status", + private_key, + "https://admin.example", + "gateway", + "admin-status", {"sub": user_id, "scope": "litellm_proxy_admin", **({"email": email} if email else {})}, ) result: Final = await JWTAuthManager.auth_builder( - api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache, - parent_otel_span=None, proxy_logging_obj=MagicMock(), request_data={}, general_settings={}, route="/user/info", + api_key=token, + jwt_handler=handler, + prisma_client=database, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + request_data={}, + general_settings={}, + route="/user/info", ) assert result["is_proxy_admin"] is True assert result["user_id"] == user_id @@ -7722,3 +7396,162 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning users.create.assert_not_awaited() if existing_user: assert users.find_unique.await_count == (0 if warm_cache else 1) + + +# --- B5: JWT algorithm allowlists (LIT-8429) --- + + +def _okp_keypair_and_jwk() -> "tuple[object, dict]": + import json + + from cryptography.hazmat.primitives.asymmetric import ed25519 + + private_key = ed25519.Ed25519PrivateKey.generate() + jwk = { + **json.loads(__import__("jwt").algorithms.OKPAlgorithm.to_jwk(private_key.public_key())), + "kid": "ed", + "alg": "EdDSA", + "use": "sig", + } + return private_key, jwk + + +def _eddsa_jwt(private_key: object, kid: str = "ed") -> str: + import jwt + + current_time = int(time.time()) + return jwt.encode( + {"sub": "test-subject", "iat": current_time, "exp": current_time + 300}, + private_key, # pyright: ignore[reportArgumentType] + algorithm="EdDSA", + headers={"kid": kid}, + ) + + +def test_approved_jwt_algorithm_literal_matches_tuple(): + from typing import get_args + + from litellm.proxy.auth.jwt_algorithms import APPROVED_JWT_ALGORITHMS, ApprovedJwtAlgorithm + + assert set(get_args(ApprovedJwtAlgorithm)) == set(APPROVED_JWT_ALGORITHMS) + + +def test_allowed_jwt_algorithms_drops_legacy_only_in_fips_mode(): + from litellm.proxy.auth.jwt_algorithms import allowed_jwt_algorithms + + assert "EdDSA" not in allowed_jwt_algorithms(True) + assert list(allowed_jwt_algorithms(False)) == JWTHandler.SUPPORTED_JWT_ALGORITHMS + + +@pytest.mark.parametrize( + "key,algorithms,expected", + [ + ({"kty": "RSA", "alg": "RS256", "kid": "a"}, ("RS256",), True), + ({"kty": "RSA", "alg": "HS256", "kid": "a"}, ("RS256", "ES256"), False), + ({"kty": "OKP", "alg": "EdDSA", "kid": "a"}, ("RS256", "ES256"), False), + ({"kty": "RSA", "kid": "a"}, ("RS256", "ES256"), True), + ({"kty": "EC", "kid": "a"}, ("RS256", "ES256"), True), + ({"kty": "EC", "kid": "a"}, ("RS256",), False), + ({"kty": "OKP", "kid": "a"}, ("RS256", "EdDSA"), True), + ({"kty": "OKP", "kid": "a"}, ("RS256", "ES256"), False), + ({"kty": "oct", "kid": "a"}, ("RS256", "HS256"), False), + ({"kid": "a"}, ("RS256",), False), + ], +) +def test_jwks_keys_for_filters_by_declared_or_inferred_algorithm(key, algorithms, expected): + from litellm.proxy.auth.jwt_algorithms import jwks_keys_for + + assert list(jwks_keys_for([key], algorithms)) == ([key] if expected else []) + + +def test_fips_mode_rejects_eddsa_token_but_accepts_rs256(): + import jwt + + ed_private, ed_jwk = _okp_keypair_and_jwk() + rsa_private, rsa_jwk = _rsa_keypair_and_jwk() + handler = JWTHandler(fips_mode=lambda: True) + + with pytest.raises(jwt.exceptions.InvalidAlgorithmError): + handler._decode_jwt_with_public_key( + token=_eddsa_jwt(ed_private), + public_key=ed_jwk, + audience=None, + disable_audience_validation=True, + ) + claims: Final = handler._decode_jwt_with_public_key( + token=_encode_rsa_jwt(rsa_private, "iss", "aud", "rsa"), + public_key=rsa_jwk, + audience="aud", + issuer="iss", + ) + assert claims["sub"] == "test-subject" + + +def test_eddsa_token_accepted_with_deprecation_log_outside_fips(caplog): + import logging + + from litellm.proxy.auth.handle_jwt import _log_eddsa_deprecation + + ed_private, ed_jwk = _okp_keypair_and_jwk() + handler = JWTHandler(fips_mode=lambda: False) + + _log_eddsa_deprecation.cache_clear() + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + claims: Final = handler._decode_jwt_with_public_key( + token=_eddsa_jwt(ed_private), + public_key=ed_jwk, + audience=None, + disable_audience_validation=True, + ) + assert claims["sub"] == "test-subject" + assert "EdDSA" in caplog.text and "deprecated" in caplog.text and "LITELLM_FIPS_MODE" in caplog.text + + _log_eddsa_deprecation.cache_clear() + caplog.clear() + rsa_private, rsa_jwk = _rsa_keypair_and_jwk() + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + handler._decode_jwt_with_public_key( + token=_encode_rsa_jwt(rsa_private, "iss", "aud", "rsa"), + public_key=rsa_jwk, + audience="aud", + issuer="iss", + ) + assert "deprecated" not in caplog.text + _log_eddsa_deprecation.cache_clear() + + +def _rsa_keypair_and_jwk() -> "tuple[object, dict]": + import json + + import jwt as jwt_lib + from cryptography.hazmat.primitives.asymmetric import rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk = { + **json.loads(jwt_lib.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())), + "kid": "rsa", + "alg": "RS256", + "use": "sig", + } + return private_key, jwk + + +@pytest.mark.asyncio +async def test_jwks_url_kid_matching_only_a_filtered_key_raises(): + _, ed_jwk = _okp_keypair_and_jwk() + jwks_url = "https://idp.example.com/jwks" + + fips_handler = JWTHandler(fips_mode=lambda: True) + fips_handler.user_api_key_cache = DualCache() + await fips_handler.user_api_key_cache.async_set_cache( + key=f"litellm_jwt_auth_keys_{jwks_url}", value=[ed_jwk], ttl=600 + ) + with pytest.raises(NoMatchingJWTPublicKeyError): + await fips_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="ed") + + normal_handler = JWTHandler(fips_mode=lambda: False) + normal_handler.user_api_key_cache = DualCache() + await normal_handler.user_api_key_cache.async_set_cache( + key=f"litellm_jwt_auth_keys_{jwks_url}", value=[ed_jwk], ttl=600 + ) + assert (await normal_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="ed"))["kid"] == "ed" diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index cb2276ab39d..1150fab69b4 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -174,9 +174,7 @@ def test_jwks_public_key_can_verify_signed_jwt(): def test_build_claims_standard_fields(): """_build_claims() populates iss, aud, iat, exp, nbf.""" - signer = _make_signer( - issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300 - ) + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) user_dict = _make_user_api_key_dict() data = {"mcp_tool_name": "get_weather"} @@ -347,17 +345,13 @@ async def test_hook_skips_non_mcp_call_types(): data=original_data, call_type=call_type, # type: ignore[arg-type] ) - assert "extra_headers" not in ( - result or {} - ), f"extra_headers should not be set for {call_type}" + assert "extra_headers" not in (result or {}), f"extra_headers should not be set for {call_type}" @pytest.mark.asyncio async def test_hook_signs_list_mcp_tools(): """async_pre_call_hook() signs JWT for list_mcp_tools with list scope.""" - signer = _make_signer( - issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300 - ) + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) user_dict = _make_user_api_key_dict(user_id="alice", team_id="backend") data = {"mcp_tool_name": "should_be_cleared"} @@ -382,9 +376,7 @@ async def test_hook_signs_list_mcp_tools(): @pytest.mark.asyncio async def test_signed_token_is_verifiable(): """The JWT injected by the hook can be verified against the JWKS public key.""" - signer = _make_signer( - issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300 - ) + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) user_dict = _make_user_api_key_dict(user_id="alice", team_id="backend") data = {"mcp_tool_name": "search"} @@ -600,9 +592,7 @@ async def test_channel_token_injected_when_configured(): assert isinstance(result, dict) assert "x-mcp-channel-token" in result["extra_headers"] - channel_token = result["extra_headers"]["x-mcp-channel-token"].removeprefix( - "Bearer " - ) + channel_token = result["extra_headers"]["x-mcp-channel-token"].removeprefix("Bearer ") channel_payload = _decode_unverified(channel_token) assert channel_payload["aud"] == "bedrock-gateway" @@ -818,9 +808,7 @@ def test_initialize_guardrail_passes_all_params(): litellm_params.issuer = "https://litellm.example.com" litellm_params.audience = "mcp-test" litellm_params.ttl_seconds = 120 - litellm_params.access_token_discovery_uri = ( - "https://idp.example.com/.well-known/openid-configuration" - ) + litellm_params.access_token_discovery_uri = "https://idp.example.com/.well-known/openid-configuration" litellm_params.token_introspection_endpoint = "https://idp.example.com/introspect" litellm_params.verify_issuer = "https://idp.example.com" litellm_params.verify_audience = "api://test" @@ -843,10 +831,7 @@ def test_initialize_guardrail_passes_all_params(): assert signer.issuer == "https://litellm.example.com" assert signer.audience == "mcp-test" assert signer.ttl_seconds == 120 - assert ( - signer.access_token_discovery_uri - == "https://idp.example.com/.well-known/openid-configuration" - ) + assert signer.access_token_discovery_uri == "https://idp.example.com/.well-known/openid-configuration" assert signer.token_introspection_endpoint == "https://idp.example.com/introspect" assert signer.verify_issuer == "https://idp.example.com" assert signer.verify_audience == "api://test" @@ -879,9 +864,7 @@ def _make_httpx_response(json_body: dict, status_code: int = 200): if status_code >= 400: from httpx import HTTPStatusError, Request, Response - mock_resp.raise_for_status.side_effect = HTTPStatusError( - "error", request=MagicMock(), response=MagicMock() - ) + mock_resp.raise_for_status.side_effect = HTTPStatusError("error", request=MagicMock(), response=MagicMock()) return mock_resp @@ -940,9 +923,7 @@ async def test_fetch_jwks_uses_cache_on_second_call(): @pytest.mark.asyncio async def test_get_oidc_discovery_caches_when_jwks_uri_present(): """_get_oidc_discovery caches the doc when jwks_uri is in the response.""" - signer = _make_signer( - access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration" - ) + signer = _make_signer(access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration") signer._oidc_discovery_doc = None # ensure fresh discovery_doc = { @@ -964,9 +945,7 @@ async def test_get_oidc_discovery_caches_when_jwks_uri_present(): @pytest.mark.asyncio async def test_get_oidc_discovery_does_not_cache_when_jwks_uri_absent(): """_get_oidc_discovery does NOT cache a doc that is missing jwks_uri.""" - signer = _make_signer( - access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration" - ) + signer = _make_signer(access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration") signer._oidc_discovery_doc = None bad_doc = {"issuer": "https://idp.example.com"} # no jwks_uri @@ -1069,9 +1048,7 @@ async def test_verify_incoming_jwt_raises_on_expired_token(): @pytest.mark.asyncio async def test_introspect_opaque_token_returns_claims_when_active(): """_introspect_opaque_token returns the introspection payload for active tokens.""" - signer = _make_signer( - token_introspection_endpoint="https://idp.example.com/introspect" - ) + signer = _make_signer(token_introspection_endpoint="https://idp.example.com/introspect") introspection_response = { "active": True, @@ -1095,9 +1072,7 @@ async def test_introspect_opaque_token_returns_claims_when_active(): @pytest.mark.asyncio async def test_introspect_opaque_token_raises_on_inactive_token(): """_introspect_opaque_token raises ExpiredSignatureError when active=false.""" - signer = _make_signer( - token_introspection_endpoint="https://idp.example.com/introspect" - ) + signer = _make_signer(token_introspection_endpoint="https://idp.example.com/introspect") fake_resp = _make_httpx_response({"active": False}) mock_client = MagicMock() @@ -1128,9 +1103,7 @@ async def test_hook_raises_401_when_jwt_verification_fails(): """async_pre_call_hook raises HTTP 401 when incoming JWT verification fails.""" from fastapi import HTTPException - signer = _make_signer( - access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration" - ) + signer = _make_signer(access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration") with patch.object( signer, @@ -1269,3 +1242,101 @@ async def test_inject_mcp_jwt_signs_for_tool_call_path(): scopes = set(decoded["scope"].split()) assert "mcp:tools/call" in scopes assert "mcp:tools/search_web:call" in scopes + + +# --------------------------------------------------------------------------- +# B5: incoming-JWT JWKS allowlist (LIT-8429) +# --------------------------------------------------------------------------- + + +def _okp_idp_key_and_token(now: int): + import json + + from cryptography.hazmat.primitives.asymmetric import ed25519 + + private_key = ed25519.Ed25519PrivateKey.generate() + jwk = { + **json.loads(jwt.algorithms.OKPAlgorithm.to_jwk(private_key.public_key())), + "kid": "ed", + "alg": "EdDSA", + } + token = jwt.encode( + {"sub": "idp-user", "iat": now, "exp": now + 300}, + private_key, + algorithm="EdDSA", + headers={"kid": "ed"}, + ) + return jwk, token + + +def _oct_idp_key_and_token(now: int): + secret = b"integration-hs256-client-secret-0123456789abcdef" + jwk = { + "kty": "oct", + "kid": "sym", + "alg": "HS256", + "k": base64.urlsafe_b64encode(secret).rstrip(b"=").decode(), + } + token = jwt.encode( + {"sub": "idp-user", "iat": now, "exp": now + 300}, + secret, + algorithm="HS256", + headers={"kid": "sym"}, + ) + return jwk, token + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_and_token", (_oct_idp_key_and_token, _okp_idp_key_and_token), ids=("oct-HS256", "OKP-EdDSA") +) +async def test_verify_incoming_jwt_rejects_jwks_without_approved_algorithms(key_and_token): + """A JWKS whose only keys use non-approved algorithms can never verify the incoming token.""" + jwks_key, incoming_token = key_and_token(int(time.time())) + signer = _make_signer( + access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration", + ) + with patch.object( + signer, + "_get_oidc_discovery", + new_callable=AsyncMock, + return_value={"jwks_uri": "https://idp.example.com/jwks"}, + ): + with patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer._fetch_jwks", + new_callable=AsyncMock, + return_value=[jwks_key], + ): + with pytest.raises(jwt.exceptions.PyJWKSetError, match="approved signing algorithm"): + await signer._verify_incoming_jwt(incoming_token) + + +@pytest.mark.asyncio +async def test_verify_incoming_jwt_ignores_filtered_keys_sharing_kid(): + """An approved RS256 key still verifies when the JWKS also carries a non-approved key with the same kid.""" + signer = _make_signer( + access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration", + ) + now = int(time.time()) + oct_jwk, _ = _oct_idp_key_and_token(now) + oct_jwk["kid"] = signer._kid + incoming_token = jwt.encode( + {"sub": "idp-user", "iat": now, "exp": now + 300}, + signer._private_key, + algorithm="RS256", + headers={"kid": signer._kid}, + ) + jwks = [oct_jwk, *signer.get_jwks()["keys"]] + with patch.object( + signer, + "_get_oidc_discovery", + new_callable=AsyncMock, + return_value={"jwks_uri": "https://idp.example.com/jwks"}, + ): + with patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer._fetch_jwks", + new_callable=AsyncMock, + return_value=jwks, + ): + payload = await signer._verify_incoming_jwt(incoming_token) + assert payload["sub"] == "idp-user" diff --git a/tests/test_litellm/types/test_mcp.py b/tests/test_litellm/types/test_mcp.py index 5450ec4aa48..de138bf662c 100644 --- a/tests/test_litellm/types/test_mcp.py +++ b/tests/test_litellm/types/test_mcp.py @@ -53,12 +53,12 @@ def test_has_header_matches_any_casing() -> None: @pytest.mark.parametrize( "target,expected", [ - ("https://upstream.example.com/other", False), # same origin + ("https://upstream.example.com/other", False), # same origin ("https://upstream.example.com:443/other", False), # explicit default port - ("https://attacker.example.com/collect", True), # different host - ("http://upstream.example.com/collect", True), # scheme downgrade, same host + ("https://attacker.example.com/collect", True), # different host + ("http://upstream.example.com/collect", True), # scheme downgrade, same host ("https://upstream.example.com:8443/other", True), # different port, same host - ("https://sub.upstream.example.com/x", True), # different host + ("https://sub.upstream.example.com/x", True), # different host ], ) def test_origin_is_scheme_host_and_port_not_host_alone(target: str, expected: bool) -> None: @@ -85,3 +85,29 @@ async def test_the_hook_drops_the_slot_only_once_the_origin_changes() -> None: foreign = httpx.Request("GET", "https://attacker.example.com/x", headers={"esb-oauth": "Bearer x"}) await hook(foreign) assert "esb-oauth" not in foreign.headers + + +def test_new_mcp_server_request_rejects_non_approved_client_assertion_signing_alg() -> None: + from pydantic import ValidationError + + from litellm.proxy._types import NewMCPServerRequest + + with pytest.raises(ValidationError) as exc: + NewMCPServerRequest( + transport="http", + url="https://mcp.example.com", + credentials={"client_assertion_signing_alg": "HS256"}, + ) + assert "client_assertion_signing_alg" in str(exc.value) + + +def test_new_mcp_server_request_accepts_approved_client_assertion_signing_alg() -> None: + from litellm.proxy._types import NewMCPServerRequest + + request = NewMCPServerRequest( + transport="http", + url="https://mcp.example.com", + credentials={"client_assertion_signing_alg": "ES256"}, + ) + assert request.credentials is not None + assert request.credentials["client_assertion_signing_alg"] == "ES256"