diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 3633f7b6127..970f046c3d6 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -12,6 +12,11 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +class OAuth2TokenValidationError(ValueError): + """The token-info or introspection endpoint did not vouch for the token: it rejected it, reported it + inactive, or could not be reached.""" + + class Oauth2Handler: """ Handles OAuth2 token validation. @@ -207,7 +212,7 @@ class Oauth2Handler: ) except httpx.HTTPStatusError as e: # This will catch any 4xx or 5xx errors - raise ValueError(f"Oauth 2.0 Token validation failed: {e}") + raise OAuth2TokenValidationError(f"Oauth 2.0 Token validation failed: {e}") except Exception as e: # This will catch any other errors (like network issues) - raise ValueError(f"An error occurred during token validation: {e}") + raise OAuth2TokenValidationError(f"An error occurred during token validation: {e}") diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index 625cde78495..c788a0f5e22 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -10,6 +10,7 @@ from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias +import jwt from fastapi import HTTPException from pydantic import SecretStr, TypeAdapter from starlette.datastructures import Headers @@ -45,6 +46,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -250,10 +252,25 @@ def _caused_by(exc: BaseException, error_type: type[BaseException]) -> bool: return any(isinstance(link, error_type) for link in _cause_chain(exc)) +def _signed_with(token: str, secret: str) -> bool: + """True when ``token`` is a JWT whose HS256 signature verifies with ``secret``, whatever its claims say.""" + try: + jwt.decode( + token, + secret, + algorithms=["HS256"], + options={"verify_exp": False, "verify_nbf": False, "verify_iat": False, "verify_aud": False}, + ) + except jwt.PyJWTError: + return False + return True + + async def _bearer_is_litellm_credential(token: str) -> bool: - """True when ``token`` is, or may be, a credential LiteLLM issued: the master key, a sealed gateway - bearer, or a virtual key the key table holds (a lookup fault counts as held). Only ``sk-`` tokens are - looked up, because no other token can be a virtual key.""" + """True when ``token`` is, or may be, a credential LiteLLM issued: the master key or a JWT it signed + (dashboard sessions, BYOK sessions), a sealed gateway bearer, or a virtual key the key table holds (a + lookup fault counts as held). Only ``sk-`` tokens are looked up, because no other token can be a + virtual key.""" from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module is_sealed_gateway_bearer, virtual_key_may_exist, @@ -263,7 +280,7 @@ async def _bearer_is_litellm_credential(token: str) -> bool: ) from litellm.proxy.proxy_server import master_key # noqa: PLC0415 # startup owns the master key - if master_key and secrets.compare_digest(token.encode(), master_key.encode()): + if master_key and (secrets.compare_digest(token.encode(), master_key.encode()) or _signed_with(token, master_key)): return True if is_sealed_gateway_bearer(token): return True @@ -298,16 +315,13 @@ def _identity_check(token: str, route: str, settings: Mapping[str, object]) -> _ return None -async def _verify_identity(check: _IdentityCheck, token: str, request: Request) -> None: - from litellm.proxy.auth.oauth2_check import Oauth2Handler # noqa: PLC0415 # proxy import cycle +async def _verify_identity(check: Literal["oauth2_proxy", "jwt"], token: str, request: Request) -> None: from litellm.proxy.auth.oauth2_proxy_hook import ( # noqa: PLC0415 # proxy import cycle handle_oauth2_proxy_request, ) from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration match check: - case "oauth2": - await Oauth2Handler.check_oauth2_token(token=token) case "oauth2_proxy": await handle_oauth2_proxy_request(request=request) case "jwt": @@ -320,8 +334,10 @@ async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: R """The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none. Custom auth decides alone, so any rejection is ``None``. Otherwise the identity check LiteLLM would run - (OAuth2 introspection, reverse-proxy headers, or a JWT) must pass first. Past that, only a user LiteLLM - does not know is ``None``; every other refusal (budget, blocked team, disallowed route) propagates.""" + must vouch for the caller. Reverse-proxy headers and a JWT are verified first, and OAuth2 introspection + runs once inside the full check, where a provider that does not vouch is ``None``. Past that, only a + user LiteLLM does not know is ``None``; every other refusal (budget, blocked team, disallowed route) + propagates.""" from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 # proxy import cycle from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module custom_auth_configured, @@ -340,14 +356,15 @@ async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: R ) if check is None: return None - try: - await _verify_identity(check, token, request) - except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer - return None + if check != "oauth2": + try: + await _verify_identity(check, token, request) + except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer + return None try: return await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: - if _caused_by(exc, UserNotFoundError): + if _caused_by(exc, UserNotFoundError) or _caused_by(exc, OAuth2TokenValidationError): return None raise diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index c772e2e9d09..9f00a3f2138 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -4176,6 +4176,12 @@ class MCPServerManager: referenced: Final = collect_env_var_references(strings=static_headers.values()) return bool(referenced & user_var_names) + async def resolved_static_headers( + self, server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None + ) -> dict[str, str] | None: + """``server.static_headers`` interpolated exactly as a tool call sends them.""" + return await self._resolve_static_headers_with_env_vars(server, user_api_key_auth) + async def _resolve_static_headers_with_env_vars( self, server: MCPServer, diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index d173fe60b82..6c79bc2d937 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -2104,6 +2104,7 @@ if MCP_AVAILABLE: url: str, auth_header: str, timeout: float = 5.0, + headers: Mapping[str, str] | None = None, ) -> tuple[int, str | None]: """JSON-RPC initialize-probe the upstream URL to check whether the token is accepted. @@ -2136,6 +2137,7 @@ if MCP_AVAILABLE: probe_headers: Final = { "Accept": "application/json, text/event-stream", **({"Authorization": auth_header} if auth_header else {}), + **(headers or {}), } try: resp: Final = await client.post( @@ -2154,6 +2156,30 @@ if MCP_AVAILABLE: verbose_logger.debug("_probe_upstream_auth: probe to %s failed (%s), allowing request through", url, exc) return 200, None + async def _probe_with_outbound_headers( + server: MCPServer, + auth_header: str, + caller_headers: Mapping[str, str], + user_api_key_auth: UserAPIKeyAuth | None, + ) -> tuple[int, str | None]: + """Probe with the headers a real tool call sends: the caller's prepared headers with the server's + resolved static headers on top, so a static ``Authorization`` wins here too. When those cannot be + resolved yet, such as a per-user variable the caller has not set, the probe fails open like a network + error and leaves the error to the call.""" + try: + static_headers: Final = await global_mcp_server_manager.resolved_static_headers(server, user_api_key_auth) + except Exception as exc: # noqa: BLE001 # the probe fails open on anything it cannot reproduce + verbose_logger.debug("Skipping upstream auth probe for %s (%s)", server.name, type(exc).__name__) + return 200, None + outbound: Final = merge_mcp_headers(extra_headers=caller_headers, static_headers=static_headers) or {} + authorization: Final = next( + (value for name, value in reversed(outbound.items()) if name.lower() == "authorization"), auth_header + ) + companion_headers: Final = {name: value for name, value in outbound.items() if name.lower() != "authorization"} + if not companion_headers: + return await _probe_upstream_auth(server.url or "", authorization) + return await _probe_upstream_auth(server.url or "", authorization, headers=companion_headers) + async def _check_passthrough_upstream_auth( scope: Scope, user_api_key_auth: UserAPIKeyAuth | None, @@ -2163,6 +2189,7 @@ if MCP_AVAILABLE: oauth2_headers: Mapping[str, str] | None, mcp_server_auth_headers: dict[str, dict[str, str]] | None, raw_headers: Mapping[str, str] | None, + requested_servers: Sequence[str] | None = None, ) -> None: """Probe pass-through and legacy delegate upstream servers in parallel before the MCP session starts. @@ -2189,7 +2216,10 @@ if MCP_AVAILABLE: ) # A legacy delegate is probed only on a single-server route, so one rejected token cannot 401 a # multi-server connect; its challenge echoes the requested name, matching the tokenless challenge. - single_requested_name: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None + requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers + single_requested_name: Final = ( + requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None + ) prepared_headers: Final = tuple( ( srv, @@ -2216,8 +2246,13 @@ if MCP_AVAILABLE: ) for srv, (auth, extra) in prepared_headers ) - probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple( - (srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name) + probe_targets: Final[tuple[tuple[MCPServer, str, str, Mapping[str, str]], ...]] = tuple( + ( + srv, + auth_header, + srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name, + headers or {}, + ) for srv, headers in probe_headers if ( auth_header := next( @@ -2229,9 +2264,12 @@ if MCP_AVAILABLE: return probe_results: Final = await asyncio.gather( - *[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets] + *[ + _probe_with_outbound_headers(srv, auth_header, headers, user_api_key_auth) + for srv, auth_header, _, headers in probe_targets + ] ) - for (_srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results): + for (_srv, _, challenge_server_name, _), (probe_status, _) in zip(probe_targets, probe_results): if probe_status == 401: # Token is missing or expired: keep pass-through clients on the # protected-resource discovery flow so they re-authorize against @@ -2315,6 +2353,7 @@ if MCP_AVAILABLE: if context.mcp_server_auth_headers is not None else None, raw_headers=context.raw_headers, + requested_servers=context.mcp_servers, ) return None diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index 7225b854181..68f309b4926 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1599,9 +1599,7 @@ class TestMCPRequestHandler: team_id="test-team-id", ) - with patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" - ) as mock_user_api_key_auth: + with patch("litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth") as mock_user_api_key_auth: mock_user_api_key_auth.return_value = mock_auth_result # Call the method @@ -2775,6 +2773,40 @@ class TestMCPDelegateAuthToUpstream: assert auth_result.mcp_forwards_caller_authorization is False assert oauth2_headers is None + @pytest.mark.parametrize( + "signing_secret, expires_in, forwarded", + [ + ("sk-master-secret", timedelta(hours=1), False), + ("sk-master-secret", timedelta(hours=-1), False), + ("upstream-idp-signing-secret", timedelta(hours=1), True), + ], + ids=["master-key-signed", "master-key-signed-expired", "signed-elsewhere"], + ) + async def test_delegate_master_key_signed_jwt_is_never_forwarded( + self, signing_secret: str, expires_in: timedelta, forwarded: bool + ) -> None: + """LiteLLM signs dashboard and OAuth tokens with the master key, so such a JWT is LiteLLM's own even + once expired: it takes normal admission and never reaches the upstream. A JWT signed by anyone + else is the upstream's token.""" + import jwt + + from litellm.proxy import proxy_server + + token: Final = jwt.encode( + {"sub": "dashboard-user", "exp": datetime.now(timezone.utc) + expires_in}, signing_secret, algorithm="HS256" + ) + admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="dashboard-user")) + with ( + patch.object(proxy_server, "master_key", "sk-master-secret"), + TestMCPDelegateAuthToUpstream._delegate_admission(admission), + ): + auth_result, *_ = await MCPRequestHandler.process_mcp_request( + TestMCPDelegateAuthToUpstream._bearer_scope(token) + ) + + assert auth_result.mcp_forwards_caller_authorization is forwarded + assert admission.await_count == (0 if forwarded else 1) + @pytest.mark.parametrize( "token, ui_token_lookup", [ @@ -2912,28 +2944,26 @@ class TestMCPDelegateAuthToUpstream: ): yield - @pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"]) @pytest.mark.parametrize( "verify_effect, allowed_server_ids, expected_user_id", [ (None, ["test-server-id"], "idp-user"), (None, [], None), - (ValueError("Token is not active"), ["test-server-id"], None), + (ValueError("Missing trusted proxy header"), ["test-server-id"], None), ], ids=["granted-user", "user-not-granted", "identity-rejected"], ) - async def test_delegate_idp_identity_follows_the_jwt_rules( + async def test_delegate_proxy_header_identity_follows_the_jwt_rules( self, - mechanism: Literal["oauth2", "oauth2_proxy"], verify_effect: Exception | None, allowed_server_ids: Sequence[str], expected_user_id: str | None, ) -> None: - """OAuth2 introspection and trusted reverse-proxy headers admit a granted user by name, and a user - who is not granted or not vouched for anonymously. The bearer is forwarded either way.""" + """Trusted reverse-proxy headers admit a granted user by name, and a user who is not granted or not + vouched for anonymously. The bearer is forwarded either way.""" admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user")) with ( - TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock(side_effect=verify_effect)), + TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2_proxy", AsyncMock(side_effect=verify_effect)), TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids), ): auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request( @@ -2944,6 +2974,55 @@ class TestMCPDelegateAuthToUpstream: assert auth_result.mcp_forwards_caller_authorization is True assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"} + @pytest.mark.parametrize( + "active, allowed_server_ids, expected_user_id", + [(True, ["test-server-id"], "idp-user"), (True, [], None), (False, ["test-server-id"], None)], + ids=["granted-user", "user-not-granted", "token-inactive"], + ) + async def test_delegate_oauth2_introspects_once_and_follows_the_jwt_rules( + self, active: bool, allowed_server_ids: Sequence[str], expected_user_id: str | None + ) -> None: + """The provider is asked about the token once per request. A token it vouches for admits a granted + user by name, and one it reports inactive is forwarded anonymously instead of failing the request.""" + import httpx + + from litellm.proxy import proxy_server + from litellm.types.mcp import MCPAuth + + async def introspect(outbound: httpx.Request, **kwargs: object) -> httpx.Response: + return httpx.Response(200, request=outbound, json={"active": active, "sub": "idp-user"}) + + with ( + patch.dict(proxy_server.general_settings, {"enable_oauth2_auth": True}), + patch.object(proxy_server, "premium_user", True), + patch.object(proxy_server, "user_custom_auth", None), + patch.object(proxy_server, "master_key", "sk-fixture-master"), + patch.object(proxy_server, "prisma_client", None), + patch.object(proxy_server, "user_api_key_cache", DualCache()), + patch.dict( + os.environ, + { + "OAUTH_TOKEN_INFO_ENDPOINT": "https://idp.fixture.test/introspect", + "OAUTH_CLIENT_ID": "fixture-client", + "OAUTH_CLIENT_SECRET": "fixture-client-secret", + }, + ), + patch("httpx.AsyncClient.send", side_effect=introspect) as network, + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = TestMCPDelegateAuthToUpstream._make_server( + auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True + ) + mock_mgr.get_allowed_mcp_servers = AsyncMock(return_value=list(allowed_server_ids)) + auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request( + TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token") + ) + + assert network.call_count == 1 + assert auth_result.user_id == expected_user_id + assert auth_result.mcp_forwards_caller_authorization is True + assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"} + @pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"]) async def test_delegate_idp_user_refusal_propagates(self, mechanism: Literal["oauth2", "oauth2_proxy"]) -> None: """Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply.""" @@ -2966,7 +3045,9 @@ class TestMCPDelegateAuthToUpstream: introspect: Final = AsyncMock() with ( TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False), - TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))), + TestMCPDelegateAuthToUpstream._delegate_admission( + AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user")) + ), ): auth_result, *_ = await MCPRequestHandler.process_mcp_request( TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token") @@ -8005,9 +8086,7 @@ class TestMCPDcrBridgeDualCredential: new_callable=AsyncMock, return_value=UserAPIKeyAuth(user_id=presented_user, api_key="any-hash"), ), - patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr, + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), ): mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server() @@ -8058,7 +8137,9 @@ class TestMCPDcrBridgeDualCredential: @pytest.mark.parametrize("root_path", ["", "/litellm"]) @pytest.mark.parametrize("invalid_token", [False, True]) def test_gateway_challenge_metadata_contains_the_public_resource( - public_path: str, root_path: str, invalid_token: bool, + public_path: str, + root_path: str, + invalid_token: bool, ) -> None: from mcp.shared.auth_utils import check_resource_allowed from starlette.requests import Request @@ -8071,8 +8152,12 @@ def test_gateway_challenge_metadata_contains_the_public_resource( ) scope: Final[Scope] = { - "type": "http", "method": "POST", "scheme": "https", "path": "/mcp", - "root_path": root_path, "_original_path": f"{root_path}{public_path}", + "type": "http", + "method": "POST", + "scheme": "https", + "path": "/mcp", + "root_path": root_path, + "_original_path": f"{root_path}{public_path}", "headers": [(b"host", b"gateway.example.com")], } request: Final = Request(scope) @@ -8081,8 +8166,7 @@ def test_gateway_challenge_metadata_contains_the_public_resource( metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0] toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp") assert metadata_url == ( - f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" - + ("" if toolset else "/mcp") + f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" + ("" if toolset else "/mcp") ) metadata: Final = ( oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request) @@ -8285,9 +8369,7 @@ class TestAggregateGatewayDcrChallenge: ): with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr, + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): mock_mgr.get_mcp_server_by_name.return_value = server with pytest.raises(HTTPException) as exc_info: @@ -8333,9 +8415,7 @@ class TestAggregateGatewayDcrChallenge: with ( patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}), patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr, + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): mock_mgr.get_mcp_server_by_name.return_value = server with pytest.raises(HTTPException) as exc_info: @@ -8366,9 +8446,7 @@ class TestAggregateGatewayDcrChallenge: ): with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr, + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): mock_mgr.get_mcp_server_by_name.return_value = resolved mock_mgr.get_mcp_server_by_identifier.return_value = resolved @@ -8413,9 +8491,7 @@ class TestAggregateGatewayDcrChallenge: (None, None), ] for resolved, expected in cases: - with patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr: + with patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr: mock_mgr.get_mcp_server_by_identifier.return_value = resolved assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None @@ -11311,14 +11387,15 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session from litellm.proxy import proxy_server from litellm.types.mcp_server.mcp_toolset import MCPToolset - permission = LiteLLM_ObjectPermissionTable( - object_permission_id="current-policy", mcp_toolsets=["session-toolset"] - ) + permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_toolsets=["session-toolset"]) toolset: Final = MCPToolset( - toolset_id="session-toolset", toolset_name="session-toolset", + toolset_id="session-toolset", + toolset_name="session-toolset", tools=[{"server_id": "current-server", "tool_name": "read"}], ) - tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))) + tables: Final = SimpleNamespace( + litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset])) + ) monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) current = UserAPIKeyAuth( @@ -11345,10 +11422,13 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session caller.mcp_toolset_id = "session-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) - assert refreshed.object_permission == permission.model_copy(update={ - "mcp_servers": ["current-server"], "mcp_toolsets": [], - "mcp_tool_permissions": {"current-server": ["read"]}, - }) + assert refreshed.object_permission == permission.model_copy( + update={ + "mcp_servers": ["current-server"], + "mcp_toolsets": [], + "mcp_tool_permissions": {"current-server": ["read"]}, + } + ) assert refreshed.object_permission_id == "current-policy" assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ( "new-team", @@ -12452,9 +12532,7 @@ async def test_personal_key_guardrail_auth_preserves_native_tool_name_matching_s from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: - return _personal_key_team_granting_tools( - team_id, ("math-multiply",) if team_id == "team-b" else ("other",) - ) + return _personal_key_team_granting_tools(team_id, ("math-multiply",) if team_id == "team-b" else ("other",)) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) @@ -12493,9 +12571,7 @@ async def test_personal_key_mcp_billing_inherits_adopted_organization_limits( ) async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: - return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy( - update={"organization_id": "org-1"} - ) + return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"}) org_loader: Final = AsyncMock(return_value=organization) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) @@ -12541,16 +12617,12 @@ async def test_personal_key_mcp_billing_keeps_key_organization_limits( ) async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: - return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy( - update={"organization_id": "org-1"} - ) + return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"}) org_loader: Final = AsyncMock(return_value=organization) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) - auth: Final = personal_key_owner_teams.model_copy( - update={"org_id": "org-k", "organization_rpm_limit": 9} - ) + auth: Final = personal_key_owner_teams.model_copy(update={"org_id": "org-k", "organization_rpm_limit": 9}) server: Final = global_mcp_server_manager.registry["srv-b"] with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server): @@ -12682,9 +12754,7 @@ async def test_personal_key_mcp_billing_uses_fresh_team_and_membership_for_stric } ) - async def _get_team_membership( - user_id: str, team_id: str, **kwargs: object - ) -> LiteLLM_TeamMembership: + async def _get_team_membership(user_id: str, team_id: str, **kwargs: object) -> LiteLLM_TeamMembership: database_value: Final = bool(kwargs.get("check_db_only")) return LiteLLM_TeamMembership( user_id=user_id, diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index 1d78573b1a1..77526095f21 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -7335,16 +7335,19 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual @pytest.mark.asyncio -async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None: # test-quality-ok: the credential sent to the upstream probe is the security contract and has no return value +@pytest.mark.usefixtures("httpx_transport") +async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None: """With a separate x-litellm-api-key, the preflight probes the upstream with its own token and never with the LiteLLM key.""" + import respx + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server.server import ( _check_passthrough_upstream_auth, ) - server = _delegate_auth_mcp_server() - scope = _delegate_scope( + server: Final = _delegate_auth_mcp_server() + scope: Final = _delegate_scope( [ (b"x-litellm-api-key", b"sk-litellm-proxy-key"), (b"authorization", b"Bearer upstream-token"), @@ -7356,11 +7359,9 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), - patch( # test-quality-ok: the probe's outbound credential is the security contract under test - "litellm.proxy.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(200, None)), - ) as probe, + respx.mock as upstream, ): + probe: Final = upstream.post(server.url).respond(200) await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), @@ -7371,7 +7372,10 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"}, ) - assert probe.await_args.args == ("http://upstream:9401/mcp", "Bearer upstream-token") + assert probe.call_count == 1 + sent_headers: Final = probe.calls[0].request.headers + assert sent_headers["Authorization"] == "Bearer upstream-token" + assert all("sk-litellm-proxy-key" not in value for value in sent_headers.values()) def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]: @@ -7416,7 +7420,9 @@ async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_chall assert exc_info.value.status_code == 401 challenge: Final = exc_info.value.headers["www-authenticate"] assert 'error="invalid_token"' in challenge - assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + assert ( + 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + ) @pytest.mark.asyncio @@ -7426,6 +7432,131 @@ async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_ assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + "static_headers, expected_authorization", + [ + ({"X-Tenant": "tenant-1"}, "Bearer expired-upstream-token"), + ({"X-Tenant": "tenant-1", "Authorization": "Bearer server-wide-token"}, "Bearer server-wide-token"), + ], + ids=["companion-header", "static-authorization-wins"], +) +async def test_legacy_delegate_probe_sends_the_headers_a_tool_call_sends( + static_headers: dict[str, str], expected_authorization: str +) -> None: + """An upstream that needs a static header next to the token sees it in the probe too, so a token the + upstream accepts on real calls is not rejected before the session opens.""" + import respx + + server: Final = _delegate_auth_mcp_server().model_copy(update={"static_headers": static_headers}) + with ( + patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + respx.mock as upstream, + ): + probe: Final = upstream.post(server.url).mock( + side_effect=lambda request: httpx.Response(200 if request.headers.get("X-Tenant") == "tenant-1" else 401) + ) + assert await _expired_delegate_token_preflight(["delegate_test"]) is None + + assert probe.call_count == 1 + assert probe.calls[0].request.headers["Authorization"] == expected_authorization + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_user_values() -> None: + """A static header built from a per-user value the caller has not set cannot be reproduced, so the + probe lets the request through and the tool call reports the missing value.""" + import respx + + server: Final = _delegate_auth_mcp_server().model_copy( + update={ + "static_headers": {"X-Tenant-Token": "${TENANT_TOKEN}"}, + "env_vars": [{"name": "TENANT_TOKEN", "value": "", "scope": "user"}], + } + ) + with ( + patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + respx.mock as upstream, + ): + probe: Final = upstream.post(server.url).respond( + 401, headers={"www-authenticate": 'Bearer error="invalid_token"'} + ) + assert await _expired_delegate_token_preflight(["delegate_test"]) is None + + assert probe.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + "requested, expected_status", + [(("delegate_test",), 401), (("delegate_test", "other_server"), None)], + ids=["single-server", "two-requested-one-reachable"], +) +async def test_modern_interaction_probes_legacy_delegate_only_for_a_single_requested_server( + requested: tuple[str, ...], expected_status: int | None +) -> None: + """The single-server rule follows what the caller asked for, not what it can reach, so a request + naming two servers is not failed by one expired delegate token.""" + import respx + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.mcp_server import server + from litellm.proxy.mcp_server.contracts import OperationContext + + target: Final = _delegate_auth_mcp_server() + scope: Final = _delegate_scope( + [ + (b"authorization", b"Bearer expired-upstream-token"), + (b"mcp-protocol-version", b"2026-07-28"), + (b"mcp-method", b"tools/call"), + (b"mcp-name", b"delegate-add"), + ] + ) + context: Final = OperationContext( + _caller=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}), + mcp_servers=requested, + oauth2_headers={"Authorization": "Bearer expired-upstream-token"}, + raw_headers={"authorization": "Bearer expired-upstream-token"}, + ) + body: Final = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "delegate-add", + "arguments": {"a": 2, "b": 3}, + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, + }, + }, + } + ).encode() + with ( + patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new=AsyncMock(return_value=[target]), + ), + respx.mock as upstream, + ): + upstream.post(target.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'}) + outcome: Final = await asyncio.gather( + server._preflight_modern_interaction(scope, body, context), return_exceptions=True + ) + + assert getattr(outcome[0], "status_code", None) == expected_status + + @pytest.mark.parametrize( "probe_status, expected_status", [(200, None), (401, 401), (403, 403)], @@ -8969,8 +9100,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_unknown_name_fails_ ] with patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp." - "MCPRequestHandler.get_mcp_servers_from_access_groups", + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", new_callable=AsyncMock, return_value=[], ): @@ -9028,8 +9158,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_known_alias_returns ] with patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp." - "MCPRequestHandler.get_mcp_servers_from_access_groups", + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", new_callable=AsyncMock, return_value=[], ): @@ -9061,8 +9190,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_mixed_known_and_unk ] with patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp." - "MCPRequestHandler.get_mcp_servers_from_access_groups", + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", new_callable=AsyncMock, return_value=[], ): @@ -9094,8 +9222,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_access_group_resolv ] with patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp." - "MCPRequestHandler.get_mcp_servers_from_access_groups", + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", new_callable=AsyncMock, return_value=["id-b"], ): @@ -9363,15 +9490,11 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion( is_otel_v2_enabled.cache_clear() notified: Final = asyncio.Event() - async def call_tool( - ctx: ServerRequestContext[object, object], params: CallToolRequestParams - ) -> CallToolResult: + async def call_tool(ctx: ServerRequestContext[object, object], params: CallToolRequestParams) -> CallToolResult: with tracer.start_as_current_span("tool-handler"): return CallToolResult(content=[TextContent(type="text", text="done")]) - async def notification( - ctx: ServerRequestContext[object, object], params: NotificationParams - ) -> HandlerResult: + async def notification(ctx: ServerRequestContext[object, object], params: NotificationParams) -> HandlerResult: with tracer.start_as_current_span("notification-handler"): notified.set() return None @@ -9389,51 +9512,85 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion( try: async with manager.run(): async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client: - initialized: Final = await client.post(initial_url, headers=headers, json={ - "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": { - "protocolVersion": "2025-11-25", "capabilities": {}, - "clientInfo": {"name": "trace-exclusion", "version": "1"}, - }, - }) - assert initialized.status_code == 200, initialized.text - assert any( - span.name == "initialize" for span in exporter.get_finished_spans() - ) is not initialized_excluded - exporter.clear() - session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]} - ready: Final = await client.post(initial_url, headers=session_headers, - json={"jsonrpc": "2.0", "method": "notifications/initialized"}) - assert ready.status_code == 202, ready.text - response: Final = await client.post(current_url, headers=session_headers, json={ - "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { - "name": "probe", "arguments": {}, "_meta": { - "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + initialized: Final = await client.post( + initial_url, + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "clientInfo": {"name": "trace-exclusion", "version": "1"}, }, }, - }) + ) + assert initialized.status_code == 200, initialized.text + assert ( + any(span.name == "initialize" for span in exporter.get_finished_spans()) is not initialized_excluded + ) + exporter.clear() + session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]} + ready: Final = await client.post( + initial_url, headers=session_headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"} + ) + assert ready.status_code == 202, ready.text + response: Final = await client.post( + current_url, + headers=session_headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "probe", + "arguments": {}, + "_meta": { + "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + }, + }, + }, + ) assert response.status_code == 200 and '"text":"done"' in response.text, response.text names: Final = {span.name for span in exporter.get_finished_spans()} assert ("tools/call probe" in names) is initialized_excluded assert ("tool-handler" in names) is initialized_excluded exporter.clear() - notification_response: Final = await client.post(current_url, headers=session_headers, - json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"}) + notification_response: Final = await client.post( + current_url, + headers=session_headers, + json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"}, + ) assert notification_response.status_code == 202, notification_response.text await notified.wait() - assert any( - span.name == "notification-handler" for span in exporter.get_finished_spans() - ) is initialized_excluded + assert ( + any(span.name == "notification-handler" for span in exporter.get_finished_spans()) + is initialized_excluded + ) exporter.clear() - modern_response: Final = await client.post(current_url, headers={ - "Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28", - "MCP-Method": "tools/call", "MCP-Name": "probe", - }, json={ - "jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": { - "name": "probe", "arguments": {}, "_meta": { - PROTOCOL_VERSION_META_KEY: "2026-07-28", CLIENT_CAPABILITIES_META_KEY: {}, + modern_response: Final = await client.post( + current_url, + headers={ + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2026-07-28", + "MCP-Method": "tools/call", + "MCP-Name": "probe", + }, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "probe", + "arguments": {}, + "_meta": { + PROTOCOL_VERSION_META_KEY: "2026-07-28", + CLIENT_CAPABILITIES_META_KEY: {}, + }, }, }, - }) + ) assert modern_response.status_code == 200, modern_response.text assert '"text":"done"' in modern_response.text, modern_response.text modern_names: Final = {span.name for span in exporter.get_finished_spans()} @@ -9536,8 +9693,7 @@ async def test_get_allowed_mcp_servers_includes_active_servers_submitted_by_user return_value={"submitted-1": submitted_server}, ), patch( - "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp." - "MCPRequestHandler.get_allowed_mcp_servers", + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_allowed_mcp_servers", AsyncMock(return_value=[]), ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), @@ -10497,9 +10653,7 @@ async def test_handle_list_tools_all_faults_on_continuation_stays_degraded( outcomes={"a": ServerListFault(tag="upstream_error", status_code=429)}, ) - result: Final = await _list_tools_with_listing( - _mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1") - ) + result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1")) assert isinstance(result, ListToolsResult) assert result.tools == []