diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index e7ffef0e4e3..a5a8275a6ad 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,5 +1,6 @@ import re -from typing import Dict, List, Optional, Set, Tuple, cast +from datetime import datetime, timezone +from typing import Dict, List, Optional, Set, Tuple, assert_never, cast from fastapi import HTTPException from starlette.datastructures import Headers @@ -7,6 +8,14 @@ from starlette.requests import Request from starlette.types import Scope from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + BridgeEnvelopeInvalid, + NotBridgeEnvelope, + envelope_keys_from_master_key, + is_bridge_envelope_shaped, + resolve_bridge_envelope, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -23,6 +32,7 @@ from litellm.repositories.table_repositories import ( AgentsRepository, MCPServerRepository, ) +from litellm.types.mcp_server.mcp_server_manager import MCPServer def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]: @@ -226,6 +236,27 @@ class MCPRequestHandler: client_ip=IPAddressUtils.get_mcp_client_ip(request), ): validated_user_api_key_auth = UserAPIKeyAuth() + elif ( + ( + bridge_delegate_target := MCPRequestHandler._single_dcr_bridge_delegate_target( + path=request_route, + mcp_servers=mcp_servers, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + ) + is not None + and oauth2_headers + and is_bridge_envelope_shaped(oauth2_headers["Authorization"]) + ): + # A single DCR-bridge oauth_delegate target carrying an envelope-shaped + # Authorization: open the envelope, admit under its recovered identity, and + # inject the inner upstream token for egress. A non-envelope bearer on the same + # server is NOT admitted here — it falls through to the oauth2 arm, which 401s. + validated_user_api_key_auth, mcp_server_auth_headers = MCPRequestHandler._admit_dcr_bridge_delegate( + server=bridge_delegate_target, + authorization_value=oauth2_headers["Authorization"], + mcp_server_auth_headers=mcp_server_auth_headers, + ) elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and @@ -432,6 +463,70 @@ class MCPRequestHandler: return False return True + @staticmethod + def _single_dcr_bridge_delegate_target( + path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str] + ) -> Optional[MCPServer]: + """The one DCR-bridge ``oauth_delegate`` server this request targets, or ``None``. + + Returns the server only when EXACTLY ONE target resolves and it is both + ``is_oauth_delegate`` and ``is_dcr_bridge``. Fails closed (``None``) on a + multi-target request, an unresolved target, or a non-matching server, so the + envelope admission arm never fires for an aggregate scope or a server that did not + opt into the bridge. Mirrors :meth:`_target_servers_are_true_passthrough`. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + target_names = MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers) + if len(target_names) != 1: + return None + server = global_mcp_server_manager.get_mcp_server_by_name(target_names[0], client_ip=client_ip) + if server is None or not server.is_oauth_delegate or not server.is_dcr_bridge: + return None + # Egress resolves the injected per-server token only by alias / server_name; a server with + # neither cannot receive the forwarded token, so fail closed rather than admit-and-drop. + if not (server.server_name or server.alias): + return None + return server + + @staticmethod + def _admit_dcr_bridge_delegate( + server: MCPServer, + authorization_value: str, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + ) -> Tuple[UserAPIKeyAuth, Optional[Dict[str, Dict[str, str]]]]: + """Open the bridge envelope and admit the caller under its recovered identity. + + The envelope's signature is itself the proof the user authenticated when it was + minted, so the recovered ``user_id`` is admitted without any re-validation. The + inner upstream token is injected under the server's per-server auth-header key so + egress forwards it via the ``PassthroughConfig`` override; the envelope + ``Authorization`` the leak-defense strips never reaches the upstream. A new headers + dict is returned rather than mutating the input. Fails closed with a 401 on an + invalid or expired envelope. + """ + from litellm.proxy.proxy_server import master_key + + if not master_key: + raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + + keys = envelope_keys_from_master_key(master_key) + result = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id) + match result: + case BridgeEnvelopeAdmitted(): + header_key = server.server_name or server.alias + if header_key is None: + raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") + injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} + new_headers = {**(mcp_server_auth_headers or {}), **injected} + return UserAPIKeyAuth(user_id=result.identity.user_id), new_headers + case BridgeEnvelopeInvalid() | NotBridgeEnvelope(): + raise HTTPException(status_code=401, detail="Invalid or expired credential") + case _: + assert_never(result) + @staticmethod def _resolve_target_server_names(path: str, mcp_servers_header: Optional[List[str]]) -> List[str]: """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index ebaa6bc7cc0..1002c33783a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,6 +1,7 @@ import json import os import sys +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -4866,3 +4867,428 @@ async def test_get_allowed_mcp_servers_team_all_proxy_key_scoped_to_one_end_to_e finally: for sid in ("srv-x", "srv-y"): global_mcp_server_manager.registry.pop(sid, None) + + +@pytest.mark.asyncio +class TestMCPDcrBridgeDelegateAdmission: + """Admission-side arm for a DCR-bridge ``oauth_delegate`` client that authenticates with + a single envelope bearer (LIT-4338). + + The arm fires only for a single ``is_dcr_bridge`` ``is_oauth_delegate`` target carrying an + envelope-shaped Authorization. It opens the litellm-signed envelope, admits under the + recovered identity WITHOUT re-validating (the signature is the proof), and injects the inner + upstream token under the server's per-server auth-header key so egress forwards it. Everything + else must stay on its existing admission path. + """ + + _MASTER_KEY = "sk-bridge-master-key-for-envelope-derivation" + + @staticmethod + def _bridge_delegate_server(server_name="bridge_delegate_server", dcr_bridge=True, alias=None): + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id="bridge-server-id", + name=server_name or "bridge-fallback-name", + server_name=server_name, + alias=alias, + transport="http", + auth_type=MCPAuth.oauth_delegate, + dcr_bridge=dcr_bridge, + ) + + @classmethod + def _mint_bridge_envelope( + cls, + *, + user_id="envelope-user-42", + server_id="bridge-server-id", + access_token="inner-upstream-access-token", + token_type="Bearer", + expires_in=1800, + minted_at=None, + master_key=None, + ): + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, + SealedEnvelope, + UpstreamTokenGrant, + mint_envelope, + ) + from pydantic import SecretStr + + keys = envelope_keys_from_master_key(master_key or cls._MASTER_KEY) + now = minted_at or datetime.now(timezone.utc) + sealed = mint_envelope( + identity=EnvelopeIdentity(user_id=user_id, server_id=server_id), + grant=UpstreamTokenGrant( + access_token=SecretStr(access_token), + token_type=token_type, + expires_in=expires_in, + ), + keys=keys, + now=now, + ) + assert isinstance(sealed, SealedEnvelope), sealed + return sealed.token.get_secret_value() + + async def test_valid_envelope_admits_identity_and_injects_inner_token(self): + """A valid envelope on a single dcr_bridge oauth_delegate server admits under the envelope's + identity WITHOUT re-validating (user_api_key_auth is never called) and injects the inner + upstream token, keyed by the server name, for egress forwarding.""" + envelope = self._mint_bridge_envelope(user_id="envelope-user-42") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + ( + auth_result, + _mcp_auth_header, + _mcp_servers, + mcp_server_auth_headers, + _oauth2_headers, + _raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # Signature is the proof of prior authentication: identity admitted, no re-validation. + assert auth_result.user_id == "envelope-user-42" + mock_auth.assert_not_called() + # Inner upstream token injected under the per-server key so egress forwards it. + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + + async def test_alias_only_server_injects_under_alias_egress_can_resolve(self): + """When server_name is None, the inner token must be keyed under the alias (which egress + resolves), never under server.name (which egress never looks up), so the forwarded token is + not silently dropped.""" + envelope = self._mint_bridge_envelope() + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + patch("litellm.proxy._experimental.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._bridge_delegate_server( + server_name=None, alias="bridge_alias" + ) + (_auth, _h, _s, mcp_server_auth_headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert mcp_server_auth_headers == {"bridge_alias": {"Authorization": "Bearer inner-upstream-access-token"}} + + async def test_server_with_no_alias_or_server_name_is_not_admitted_via_bridge_arm(self): + """A bridge server egress cannot route to (no alias and no server_name) must not take the + envelope arm; it fails closed to normal oauth2 admission rather than admitting and dropping + the inner token under an unresolvable key.""" + envelope = self._mint_bridge_envelope() + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=401, detail="Invalid key"), + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server(server_name=None, alias=None) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_called_once() + + async def test_expired_envelope_fails_closed_401(self): + """An envelope whose exp is in the past must fail closed with a 401, never fall through to + anonymous admission.""" + expired = self._mint_bridge_envelope( + expires_in=60, + minted_at=datetime.now(timezone.utc) - timedelta(hours=2), + ) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {expired}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_not_called() + + async def test_envelope_minted_for_a_different_server_fails_closed_401(self): + """An envelope sealed for another server_id must be rejected when presented to this server, + so a captured or misrouted envelope cannot forward one server's upstream credential to + another. The signature verifies, but the server binding does not.""" + wrong_server = self._mint_bridge_envelope(server_id="some-other-server-id") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {wrong_server}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_not_called() + + async def test_envelope_under_wrong_master_key_fails_closed_401(self): + """An envelope-shaped bearer whose signature does not verify under the proxy's derived keys + (e.g. minted against a different master_key, or tampered) must fail closed with a 401.""" + foreign = self._mint_bridge_envelope(master_key="a-different-master-key-entirely") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {foreign}".encode("latin-1"))], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_not_called() + + async def test_non_envelope_bearer_on_bridge_server_falls_through_to_oauth2_arm(self): + """A plain (non-envelope) bearer on the same bridge server must NOT be admitted by the + envelope arm: it falls through to the oauth2 arm, which validates it as a LiteLLM key and + 401s here. Proves the arm is gated on envelope shape, not merely on the target being a + bridge server.""" + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", b"Bearer plain-upstream-bearer-not-an-envelope")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + # The envelope arm was skipped, so the oauth2 arm ran and validated the bearer. + mock_auth.assert_called_once() + + async def test_explicit_litellm_key_wins_over_envelope_arm(self): + """An explicit x-litellm-api-key is always a LiteLLM credential and its arm precedes the + envelope arm: user_api_key_auth validates the key and NO inner token is injected, even + though the Authorization header carries a valid envelope.""" + envelope = self._mint_bridge_envelope() + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [ + (b"x-litellm-api-key", b"sk-explicit-litellm-key"), + (b"authorization", f"Bearer {envelope}".encode("latin-1")), + ], + } + + async def mock_user_api_key_auth(api_key, request): + return UserAPIKeyAuth(api_key=api_key, user_id="litellm-key-user") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server() + ( + auth_result, + _mcp_auth_header, + _mcp_servers, + mcp_server_auth_headers, + _oauth2_headers, + _raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + mock_auth.assert_called_once() + assert mock_auth.call_args.kwargs["api_key"] == "sk-explicit-litellm-key" + # The explicit-key arm admitted; the envelope arm never ran, so no inner token is injected. + assert auth_result.user_id == "litellm-key-user" + assert mcp_server_auth_headers == {} + + async def test_non_bridge_oauth_delegate_server_does_not_take_envelope_arm(self): + """An oauth_delegate server that is NOT a DCR bridge (``dcr_bridge`` unset) must not take the + envelope arm even for an envelope-shaped bearer: is_dcr_bridge is False, so the gate returns + None and admission falls through to the oauth2 arm (which 401s here).""" + envelope = self._mint_bridge_envelope() + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/plain_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ) as mock_auth, + patch("litellm.proxy._experimental.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._bridge_delegate_server( + server_name="plain_delegate_server", dcr_bridge=False + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + # Not admitted by the envelope arm — the oauth2 arm ran instead. + mock_auth.assert_called_once() + + async def test_multi_target_including_bridge_server_does_not_take_envelope_arm(self): + """A multi-target request that includes the bridge server must not take the envelope arm: + the gate requires exactly one target, so it returns None and admission falls through.""" + from litellm.types.mcp import MCPAuth + + envelope = self._mint_bridge_envelope() + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"authorization", f"Bearer {envelope}".encode("latin-1")), + (b"x-mcp-servers", b"bridge_delegate_server,other_server"), + ], + } + + def mock_lookup(name, client_ip=None): + if name == "bridge_delegate_server": + return self._bridge_delegate_server() + return TestMCPDelegateAuthToUpstream._make_server(auth_type=MCPAuth.api_key) + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ) as mock_auth, + patch("litellm.proxy._experimental.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.side_effect = mock_lookup + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_called_once() + + async def test_admit_helper_returns_new_headers_without_mutating_input(self): + """Unit: ``_admit_dcr_bridge_delegate`` must return a NEW headers dict that preserves the + caller's existing per-server entries and adds the injected inner token, never mutating the + input dict.""" + envelope = self._mint_bridge_envelope(user_id="unit-user") + existing = {"other_server": {"Authorization": "Bearer someone-elses-token"}} + + with patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY): + auth_result, new_headers = MCPRequestHandler._admit_dcr_bridge_delegate( + server=self._bridge_delegate_server(), + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=existing, + ) + + assert auth_result.user_id == "unit-user" + # Input untouched. + assert existing == {"other_server": {"Authorization": "Bearer someone-elses-token"}} + # New dict carries both the pre-existing entry and the injected inner token. + assert new_headers is not existing + assert new_headers == { + "other_server": {"Authorization": "Bearer someone-elses-token"}, + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"}, + } + + async def test_admit_helper_raises_500_when_master_key_missing(self): + """Unit: without a configured master_key the gateway cannot derive envelope keys, so + admission raises a 500 rather than silently admitting.""" + envelope = self._mint_bridge_envelope() + with patch("litellm.proxy.proxy_server.master_key", None): + with pytest.raises(HTTPException) as exc_info: + MCPRequestHandler._admit_dcr_bridge_delegate( + server=self._bridge_delegate_server(), + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=None, + ) + assert exc_info.value.status_code == 500