From 46977d6e4c0fcdf2f2d381821963853f759997f1 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 13:38:22 -0700 Subject: [PATCH 1/9] feat(mcp): admit dcr_bridge oauth_delegate clients via a single envelope bearer --- .../mcp_server/auth/user_api_key_auth_mcp.py | 97 +++- .../auth/test_user_api_key_auth_mcp.py | 426 ++++++++++++++++++ 2 files changed, 522 insertions(+), 1 deletion(-) 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 e7ffef0e4e36..a5a8275a6ad6 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.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 @@ 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 @@ async def mock_body(): 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 @@ def _target_servers_are_true_passthrough( 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 ebaa6bc7cc04..1002c33783a8 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 From 50cc2c01cfbcea16ba3cdb8a233aac10ce571ac2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 17:06:41 -0700 Subject: [PATCH 2/9] fix(mcp): admit dcr_bridge envelopes under the live key, not a frozen identity The bridge envelope sealed only user_id/server_id, and admission fabricated a UserAPIKeyAuth(user_id=...) with no object_permission, team_id, org_id, or key identity. Downstream MCP permission checks read the missing restrictions as unrestricted, so a caller holding a valid envelope for a restricted key could reach tools and servers that key was never granted, and a revoked key kept working until the envelope expired. Bind the hashed authorizing key into the envelope identity and reload the live UserAPIKeyAuth by it at admission via get_key_object, failing closed with a 401 when the key is missing, blocked, or expired. Authorization is resolved fresh per request instead of frozen at mint time, so current key/team/org and tool restrictions plus revocation are enforced. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 73 +++++-- .../outbound_credentials/envelope.py | 22 ++- .../auth/test_user_api_key_auth_mcp.py | 183 ++++++++++++++++-- .../test_bridge_credentials.py | 6 +- .../outbound_credentials/test_envelope.py | 16 +- 5 files changed, 252 insertions(+), 48 deletions(-) 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 a5a8275a6ad6..adb1812059b5 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 @@ -252,7 +252,7 @@ async def mock_body(): # 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( + validated_user_api_key_auth, mcp_server_auth_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( server=bridge_delegate_target, authorization_value=oauth2_headers["Authorization"], mcp_server_auth_headers=mcp_server_auth_headers, @@ -492,20 +492,23 @@ def _single_dcr_bridge_delegate_target( return server @staticmethod - def _admit_dcr_bridge_delegate( + async 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. + """Open the bridge envelope and admit the caller under the live key it references. + + The envelope's signature proves the user authenticated when it was minted, but + authorization is resolved fresh here rather than trusted from the envelope: the + sealed ``key_hash`` reloads the current ``UserAPIKeyAuth`` record, so the key's + present team/org/object-permission restrictions and its revocation state gate the + request instead of a snapshot frozen at mint time. 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, or + when the referenced key is missing, blocked, or expired. """ from litellm.proxy.proxy_server import master_key @@ -519,14 +522,60 @@ def _admit_dcr_bridge_delegate( 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") + admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) 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 + return admitted, new_headers case BridgeEnvelopeInvalid() | NotBridgeEnvelope(): raise HTTPException(status_code=401, detail="Invalid or expired credential") case _: assert_never(result) + @staticmethod + async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: + """Reload the live key record an admitted envelope references. + + Resolving the current ``UserAPIKeyAuth`` (cache first, then DB) is what stops the + envelope from carrying frozen authority: the key's present team/org/object-permission + restrictions ride on the returned object, and a key that has since been deleted, + blocked, or expired fails closed with a 401 here rather than being admitted as an + unrestricted identity. ``get_key_object`` raises for a hash with no key row; a + blocked or expired row is rejected explicitly because ``get_key_object`` resolves a + row without applying those checks (the main ``user_api_key_auth`` pipeline enforces + them downstream, which this admission path bypasses). + """ + from litellm.proxy.auth.auth_checks import get_key_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") + try: + key_object = await get_key_object( + hashed_token=key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if not MCPRequestHandler._admitted_key_is_active(key_object): + raise HTTPException(status_code=401, detail="Invalid or expired credential") + return key_object + + @staticmethod + def _admitted_key_is_active(key_object: UserAPIKeyAuth) -> bool: + """False when the referenced key is blocked or past its expiry, so a revoked key + cannot be admitted through its still-unexpired envelope. Mirrors the active-key gate + the bridge token endpoint applies at mint time.""" + if key_object.blocked is True: + return False + expires = key_object.expires + if expires is None: + return True + expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + return expiry >= datetime.now(timezone.utc) + @staticmethod def _resolve_target_server_names(path: str, mcp_servers_header: Optional[List[str]]) -> List[str]: """ diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 298bc8d98ccf..517c2ef5c8f2 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -10,7 +10,7 @@ Wire shape: ``llm_env_`` + an HS256 JWT (same signing approach as the BYOK session bearer in ``byok_oauth_endpoints.py``). Registered claims are ``iss``/``iat``/``exp``; -custom claims are ``user_id``, ``server_id``, and ``grant``, where ``grant`` is the +custom claims are ``server_id``, ``key_hash``, and ``grant``, where ``grant`` is the upstream token grant serialized to JSON, encrypted with the repo's symmetric encryption helpers (``encrypt_value``/``decrypt_value`` from ``encrypt_decrypt_utils`` — the same family ``encrypt_value_helper`` applies to @@ -68,11 +68,19 @@ class EnvelopeIdentity(BaseModel): - """The litellm identity the envelope binds the inner grant to.""" + """The litellm identity the envelope binds the inner grant to. + + ``key_hash`` is the hashed litellm key that authorized the mint, never a raw + credential (and the edge rejects a bare hash presented as a bearer). Admission + reloads the live key record by it, so the key's current team/org/object-permission + restrictions and its revocation state are enforced at use time rather than frozen at + mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed + across a server boundary. + """ model_config = ConfigDict(frozen=True) - user_id: str = Field(min_length=1) server_id: str = Field(min_length=1) + key_hash: str = Field(min_length=1) class UpstreamTokenGrant(BaseModel): @@ -174,7 +182,7 @@ class DecryptFailed(BaseModel): class _EnvelopeClaims(BaseModel): """Decoded-claims boundary that pins the exact shape :func:`mint_envelope` emits. - ``user_id``/``server_id`` mirror the ``min_length`` constraints of + ``server_id``/``key_hash`` mirror the ``min_length`` constraints of :class:`EnvelopeIdentity` so any claim set that validates here also constructs an identity, keeping :func:`open_envelope` raise-free: a correctly signed JWT with an empty identity claim fails here and maps to ``MalformedPayload``. @@ -191,8 +199,8 @@ class _EnvelopeClaims(BaseModel): iss: str iat: int exp: int - user_id: str = Field(min_length=1) server_id: str = Field(min_length=1) + key_hash: str = Field(min_length=1) grant: str = Field(min_length=1) @@ -227,8 +235,8 @@ def mint_envelope( iss=ENVELOPE_ISSUER, iat=int(now.timestamp()), exp=int(expires_at.timestamp()), - user_id=identity.user_id, server_id=identity.server_id, + key_hash=identity.key_hash, grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), ) token = ENVELOPE_PREFIX + jwt.encode( @@ -273,7 +281,7 @@ def open_envelope( if not isinstance(grant, UpstreamTokenGrant): return grant return OpenedEnvelope( - identity=EnvelopeIdentity(user_id=claims.user_id, server_id=claims.server_id), + identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash), grant=grant, ) 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 1002c33783a8..9873f3ac7c59 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,3 +1,4 @@ +import contextlib import json import os import sys @@ -16,6 +17,8 @@ MCPRequestHandler, ) from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + ProxyException, SpecialHeaders, SpecialMCPServerNames, UserAPIKeyAuth, @@ -4875,10 +4878,12 @@ class TestMCPDcrBridgeDelegateAdmission: 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. + envelope-shaped Authorization. It opens the litellm-signed envelope, reloads the live key + record the sealed ``key_hash`` references so the caller is admitted under the key's current + authorization context (team/org/object-permission) and revocation state, and injects the inner + upstream token under the server's per-server auth-header key so egress forwards it. A key that + is missing, blocked, or expired fails closed with a 401. Everything else must stay on its + existing admission path. """ _MASTER_KEY = "sk-bridge-master-key-for-envelope-derivation" @@ -4898,11 +4903,13 @@ def _bridge_delegate_server(server_name="bridge_delegate_server", dcr_bridge=Tru dcr_bridge=dcr_bridge, ) + _KEY_HASH = "hashed-litellm-key-abc123" + @classmethod def _mint_bridge_envelope( cls, *, - user_id="envelope-user-42", + key_hash=None, server_id="bridge-server-id", access_token="inner-upstream-access-token", token_type="Bearer", @@ -4924,7 +4931,7 @@ def _mint_bridge_envelope( 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), + identity=EnvelopeIdentity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), grant=UpstreamTokenGrant( access_token=SecretStr(access_token), token_type=token_type, @@ -4936,18 +4943,51 @@ def _mint_bridge_envelope( 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") + @staticmethod + def _reloaded_key(**overrides): + """A live key record as ``get_key_object`` would return it: carries real authorization + context (key identity, team, org, and an object-permission restricting MCP servers) so a + test can prove admission admits under THAT context rather than a blank identity.""" + defaults = dict( + user_id="envelope-user-42", + api_key=TestMCPDcrBridgeDelegateAdmission._KEY_HASH, + team_id="team-restricted", + org_id="org-restricted", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-1", mcp_servers=["only-this-server"] + ), + ) + defaults.update(overrides) + return UserAPIKeyAuth(**defaults) + + @staticmethod + @contextlib.contextmanager + def _patch_key_reload(*, return_value=None, side_effect=None): + """Patch the live-key reload dependencies used by ``_reload_admitted_key``: the + ``get_key_object`` lookup plus the ``prisma_client`` / ``user_api_key_cache`` globals it + reads. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was + the reload key.""" + get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) + with ( + patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_key_object + + async def test_valid_envelope_reloads_live_key_and_admits_its_authorization_context(self): + """A valid envelope admits under the LIVE key record the sealed key_hash references, not a + blank identity: the reload is keyed by that exact hash, and the admitted auth carries the + key's current team/org/object-permission (the MCP tool/server restrictions the finding was + about). The heavyweight ``user_api_key_auth`` pipeline is still never invoked. The inner + upstream token is injected under the per-server key for egress.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) 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", @@ -4955,6 +4995,7 @@ async def test_valid_envelope_admits_identity_and_injects_inner_token(self): ) 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), + self._patch_key_reload(return_value=self._reloaded_key()) as get_key_object, ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() ( @@ -4966,14 +5007,101 @@ async def test_valid_envelope_admits_identity_and_injects_inner_token(self): _raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) - # Signature is the proof of prior authentication: identity admitted, no re-validation. + # The live key was reloaded by the exact hash the envelope sealed. + assert get_key_object.await_args.kwargs["hashed_token"] == self._KEY_HASH + # Admission carries the reloaded key's authorization context, not a blank UserAPIKeyAuth. assert auth_result.user_id == "envelope-user-42" + assert auth_result.team_id == "team-restricted" + assert auth_result.org_id == "org-restricted" + assert auth_result.object_permission is not None + assert auth_result.object_permission.mcp_servers == ["only-this-server"] + # The full raw-key auth pipeline is still bypassed for the envelope arm. 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_revoked_key_envelope_fails_closed_401(self): + """An envelope whose key has since been deleted must fail closed: ``get_key_object`` raises + for the missing row, so admission 401s instead of admitting the caller as an unrestricted + identity. This is the core regression for the dropped-authorization-context finding.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + revoked = self._patch_key_reload( + side_effect=ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type="token_not_found_in_db", + param="key", + code=401, + ) + ) + + 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), + revoked, + ): + 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_blocked_key_envelope_fails_closed_401(self): + """A reloaded key that is blocked must fail closed with a 401, so revoking a key by blocking + it takes effect immediately for any envelope still holding its hash.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(blocked=True)), + ): + 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 + + async def test_expired_key_record_fails_closed_401(self): + """A reloaded key past its expiry must fail closed with a 401, distinct from an expired + envelope: even a still-valid envelope cannot outlive the key it was minted under.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + expired_at = datetime.now(timezone.utc) - timedelta(hours=1) + + with ( + 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), + self._patch_key_reload(return_value=self._reloaded_key(expires=expired_at)), + ): + 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 + 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 @@ -4985,7 +5113,6 @@ async def test_alias_only_server_injects_under_alias_egress_can_resolve(self): "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", @@ -4993,6 +5120,7 @@ async def test_alias_only_server_injects_under_alias_egress_can_resolve(self): ), 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), + self._patch_key_reload(return_value=self._reloaded_key()), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server( server_name=None, alias="bridge_alias" @@ -5260,11 +5388,14 @@ 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") + envelope = self._mint_bridge_envelope() 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( + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(user_id="unit-user")), + ): + auth_result, new_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=existing, @@ -5286,7 +5417,23 @@ async def test_admit_helper_raises_500_when_master_key_missing(self): 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( + await 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 + + async def test_admit_helper_raises_500_when_no_db_connection(self): + """Unit: with a valid envelope but no database to reload the key from, admission raises a 500 + rather than admitting on unresolved authorization.""" + envelope = self._mint_bridge_envelope() + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler._admit_dcr_bridge_delegate( server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index dca9dc193142..82e8e2aae89e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -34,7 +34,7 @@ _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" -_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") _SERVER_ID = _IDENTITY.server_id @@ -138,7 +138,7 @@ def test_resolve_envelope_minted_for_another_server_is_invalid(): captured or misrouted envelope cannot forward one server's upstream credential to another. The valid access token stays sealed; the mismatch alone fails the resolve.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - other_server_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-OTHER") + other_server_identity = EnvelopeIdentity(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash) token = _sealed_token(keys, identity=other_server_identity) result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) @@ -155,7 +155,7 @@ def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - unicode_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-café") + unicode_identity = EnvelopeIdentity(server_id="srv-café", key_hash=_IDENTITY.key_hash) token = _sealed_token(keys, identity=unicode_identity) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index 71de206aa1e0..b44f3f84cc95 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -51,7 +51,7 @@ _WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" -_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") def _full_grant() -> UpstreamTokenGrant: @@ -137,12 +137,12 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims(): def test_claim_layout_and_no_plaintext_token_in_envelope(): token = _sealed_token(_full_grant()) claims = _unverified_claims(token) - assert set(claims) == {"iss", "iat", "exp", "user_id", "server_id", "grant"} + assert set(claims) == {"iss", "iat", "exp", "server_id", "key_hash", "grant"} assert claims["iss"] == ENVELOPE_ISSUER assert claims["iat"] == int(_NOW.timestamp()) assert claims["exp"] == int(_NOW.timestamp()) + 600 - assert claims["user_id"] == "user-123" assert claims["server_id"] == "srv-456" + assert claims["key_hash"] == "hashed-key-123" assert _ACCESS_TOKEN not in token assert _ACCESS_TOKEN not in json.dumps(claims) assert _REFRESH_TOKEN not in json.dumps(claims) @@ -226,11 +226,11 @@ def test_wrong_issuer_is_malformed_payload(): def test_missing_identity_claim_is_malformed_payload(): claims = _unverified_claims(_sealed_token(_full_grant())) - forged = _forge({key: value for key, value in claims.items() if key != "user_id"}) + forged = _forge({key: value for key, value in claims.items() if key != "key_hash"}) assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) -@pytest.mark.parametrize("identity_claim", ["user_id", "server_id"]) +@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"]) def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): claims = _unverified_claims(_sealed_token(_full_grant())) forged = _forge({**claims, identity_claim: ""}) @@ -463,9 +463,9 @@ def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): def test_empty_identity_and_key_fields_are_rejected_at_construction(): with pytest.raises(ValidationError): - EnvelopeIdentity(user_id="", server_id="srv-456") + EnvelopeIdentity(server_id="", key_hash="hashed-key-123") with pytest.raises(ValidationError): - EnvelopeIdentity(user_id="user-123", server_id="") + EnvelopeIdentity(server_id="srv-456", key_hash="") with pytest.raises(ValidationError): EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) with pytest.raises(ValidationError): @@ -484,4 +484,4 @@ def test_public_models_are_frozen(): with pytest.raises(ValidationError): opened.grant = _minimal_grant() with pytest.raises(ValidationError): - _IDENTITY.user_id = "someone-else" + _IDENTITY.key_hash = "someone-elses-hash" From c59f16f42ea54e799c8d73c5875165d21c84fe02 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 17:58:43 -0700 Subject: [PATCH 3/9] fix(mcp): enforce team block and alias-priority token injection on bridge admission Two follow-ups on the envelope admission arm flagged in review. Team revocation bypass: _reload_admitted_key checked only the key's own blocked/expires, so blocking a key's team left every envelope minted under it live until expiry. Reload the team and reject a blocked team, mirroring common_checks, so a team block revokes its envelopes immediately. Caller-overridable upstream token: egress resolves the per-server auth header alias-first, but injection keyed under server_name, so for a server with a distinct alias a caller-forwarded x-mcp-{alias}-authorization sat at the higher-priority slot and paired the admitted identity with an attacker's upstream credential. Inject under alias-first so the sealed token owns the slot egress resolves. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 41 ++++++++++++-- .../auth/test_user_api_key_auth_mcp.py | 56 +++++++++++++++++-- 2 files changed, 89 insertions(+), 8 deletions(-) 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 adb1812059b5..bee621b3721f 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 @@ -508,7 +508,13 @@ async def _admit_dcr_bridge_delegate( ``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, or - when the referenced key is missing, blocked, or expired. + when the referenced key is missing, blocked, or expired, or its team is blocked. + + The sealed token is keyed alias-first, matching the order egress resolves + (``lookup_mcp_server_auth_in_headers`` tries ``alias`` before ``server_name``). Keying + under ``server_name`` would leave a caller-supplied ``x-mcp-{alias}-authorization`` at the + higher-priority alias slot, pairing the admitted identity with an attacker's upstream + credential; the alias-keyed injection overwrites any such caller value. """ from litellm.proxy.proxy_server import master_key @@ -519,7 +525,7 @@ async def _admit_dcr_bridge_delegate( 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 + header_key = server.alias or server.server_name if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) @@ -533,7 +539,7 @@ async def _admit_dcr_bridge_delegate( @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: - """Reload the live key record an admitted envelope references. + """Reload the live key record an admitted envelope references and re-check live policy. Resolving the current ``UserAPIKeyAuth`` (cache first, then DB) is what stops the envelope from carrying frozen authority: the key's present team/org/object-permission @@ -542,7 +548,9 @@ async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: unrestricted identity. ``get_key_object`` raises for a hash with no key row; a blocked or expired row is rejected explicitly because ``get_key_object`` resolves a row without applying those checks (the main ``user_api_key_auth`` pipeline enforces - them downstream, which this admission path bypasses). + them downstream, which this admission path bypasses). The key's team is reloaded and + rejected when blocked, mirroring ``common_checks``, so blocking a team revokes every + envelope minted under its keys rather than leaving them live until expiry. """ from litellm.proxy.auth.auth_checks import get_key_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -559,8 +567,33 @@ async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: raise HTTPException(status_code=401, detail="Invalid or expired credential") from None if not MCPRequestHandler._admitted_key_is_active(key_object): raise HTTPException(status_code=401, detail="Invalid or expired credential") + await MCPRequestHandler._reject_if_admitted_team_blocked(key_object) return key_object + @staticmethod + async def _reject_if_admitted_team_blocked(key_object: UserAPIKeyAuth) -> None: + """Reload the key's team and fail closed with a 401 when it is blocked or no longer + resolves. ``get_key_object`` returns the key row without any team validation, so this + admission path applies the same live team-block gate ``common_checks`` runs on the + standard auth pipeline; without it, blocking a team would not revoke envelopes already + minted under its keys until they expired.""" + team_id = key_object.team_id + if not team_id: + return + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + try: + team_object = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if team_object.blocked is True: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + @staticmethod def _admitted_key_is_active(key_object: UserAPIKeyAuth) -> bool: """False when the referenced key is blocked or past its expiry, so a revoked key 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 9873f3ac7c59..d0d75231b593 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 @@ -4962,14 +4962,17 @@ def _reloaded_key(**overrides): @staticmethod @contextlib.contextmanager - def _patch_key_reload(*, return_value=None, side_effect=None): + def _patch_key_reload(*, return_value=None, side_effect=None, team_blocked=False): """Patch the live-key reload dependencies used by ``_reload_admitted_key``: the - ``get_key_object`` lookup plus the ``prisma_client`` / ``user_api_key_cache`` globals it - reads. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was - the reload key.""" + ``get_key_object`` lookup, the ``get_team_object`` lookup its team-block gate runs, and the + ``prisma_client`` / ``user_api_key_cache`` globals they read. The team resolves unblocked by + default; ``team_blocked=True`` simulates an admin blocking the key's team. Yields the + ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key.""" get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) + get_team_object = AsyncMock(return_value=MagicMock(blocked=team_blocked)) with ( patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), + patch("litellm.proxy.auth.auth_checks.get_team_object", get_team_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ): @@ -5102,6 +5105,29 @@ async def test_expired_key_record_fails_closed_401(self): assert exc_info.value.status_code == 401 + async def test_blocked_team_envelope_fails_closed_401(self): + """Blocking the key's TEAM must revoke its envelopes immediately: the reloaded key is active + but its team is blocked, so admission 401s. Without the live team re-check, a caller could + keep executing tools after an admin blocked the team, until the envelope expired.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(), team_blocked=True), + ): + 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 + 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 @@ -5129,6 +5155,28 @@ async def test_alias_only_server_injects_under_alias_egress_can_resolve(self): assert mcp_server_auth_headers == {"bridge_alias": {"Authorization": "Bearer inner-upstream-access-token"}} + async def test_sealed_token_wins_over_caller_forwarded_alias_header(self): + """When a bridge server has both a server_name and a distinct alias, the sealed inner token + must occupy the alias slot, the identifier egress resolves first. Otherwise a caller who + forwards x-mcp-{alias}-authorization keeps that entry at the higher-priority slot and pairs + the admitted identity with an attacker-chosen upstream credential.""" + envelope = self._mint_bridge_envelope() + attacker_forwarded = {"bridge_alias": {"Authorization": "Bearer ATTACKER-UPSTREAM-TOKEN"}} + + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key()), + ): + _auth, new_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( + server=self._bridge_delegate_server(server_name="bridge_name", alias="bridge_alias"), + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=attacker_forwarded, + ) + + # The sealed token owns the alias slot, overwriting the caller's value; the attacker token + # survives nowhere egress would resolve. + assert new_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 From 9e6d6e509c79484d2772102dcc392232718fbdc9 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 10:56:05 -0700 Subject: [PATCH 4/9] fix(mcp): route bridge admission through the centralized policy gate and mirror the SCIM owner check --- .../mcp_server/auth/user_api_key_auth_mcp.py | 86 ++++++--- .../auth/test_user_api_key_auth_mcp.py | 182 +++++++++++++++++- 2 files changed, 234 insertions(+), 34 deletions(-) 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 bee621b3721f..998aa4846c6e 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 @@ -26,7 +26,11 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.auth.user_api_key_auth import ( + _run_centralized_common_checks, + user_api_key_auth, +) +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl from litellm.repositories.table_repositories import ( AgentsRepository, @@ -256,6 +260,8 @@ async def mock_body(): server=bridge_delegate_target, authorization_value=oauth2_headers["Authorization"], mcp_server_auth_headers=mcp_server_auth_headers, + request=request, + route=request_route, ) elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real @@ -496,19 +502,24 @@ async def _admit_dcr_bridge_delegate( server: MCPServer, authorization_value: str, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + request: Request, + route: str, ) -> Tuple[UserAPIKeyAuth, Optional[Dict[str, Dict[str, str]]]]: """Open the bridge envelope and admit the caller under the live key it references. The envelope's signature proves the user authenticated when it was minted, but authorization is resolved fresh here rather than trusted from the envelope: the - sealed ``key_hash`` reloads the current ``UserAPIKeyAuth`` record, so the key's - present team/org/object-permission restrictions and its revocation state gate the - request instead of a snapshot frozen at mint time. The inner upstream token is - injected under the server's per-server auth-header key so egress forwards it via the + sealed ``key_hash`` reloads the current ``UserAPIKeyAuth`` record, and the admitted + identity then runs through the standard pipeline's centralized policy gate, so the + key's present restrictions and revocation state gate the request instead of a + snapshot frozen at mint time. 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, or - when the referenced key is missing, blocked, or expired, or its team is blocked. + when the referenced key is missing, blocked, or expired, its owner is + SCIM-deactivated, or the centralized policy gate rejects it (blocked team or + project, org or budget limits). The sealed token is keyed alias-first, matching the order egress resolves (``lookup_mcp_server_auth_in_headers`` tries ``alias`` before ``server_name``). Keying @@ -529,6 +540,7 @@ async def _admit_dcr_bridge_delegate( if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) + await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} new_headers = {**(mcp_server_auth_headers or {}), **injected} return admitted, new_headers @@ -548,9 +560,11 @@ async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: unrestricted identity. ``get_key_object`` raises for a hash with no key row; a blocked or expired row is rejected explicitly because ``get_key_object`` resolves a row without applying those checks (the main ``user_api_key_auth`` pipeline enforces - them downstream, which this admission path bypasses). The key's team is reloaded and - rejected when blocked, mirroring ``common_checks``, so blocking a team revokes every - envelope minted under its keys rather than leaving them live until expiry. + them downstream, which this admission path bypasses). The owner's SCIM state is the + other builder-inline check mirrored here, so IdP offboarding revokes every envelope + minted under the user's keys rather than leaving them live until expiry. Team, + project, org, and budget state are NOT re-checked here; the caller runs the admitted + identity through ``_enforce_admitted_live_policy`` for those. """ from litellm.proxy.auth.auth_checks import get_key_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -567,33 +581,57 @@ async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: raise HTTPException(status_code=401, detail="Invalid or expired credential") from None if not MCPRequestHandler._admitted_key_is_active(key_object): raise HTTPException(status_code=401, detail="Invalid or expired credential") - await MCPRequestHandler._reject_if_admitted_team_blocked(key_object) + await MCPRequestHandler._reject_if_admitted_owner_scim_deactivated(key_object) return key_object @staticmethod - async def _reject_if_admitted_team_blocked(key_object: UserAPIKeyAuth) -> None: - """Reload the key's team and fail closed with a 401 when it is blocked or no longer - resolves. ``get_key_object`` returns the key row without any team validation, so this - admission path applies the same live team-block gate ``common_checks`` runs on the - standard auth pipeline; without it, blocking a team would not revoke envelopes already - minted under its keys until they expired.""" - team_id = key_object.team_id - if not team_id: + async def _reject_if_admitted_owner_scim_deactivated(key_object: UserAPIKeyAuth) -> None: + """Fail closed with a 401 when the key's owning user was deactivated via SCIM. + + The standard pipeline enforces this inline in ``_user_api_key_auth_builder`` rather + than in ``common_checks``, so the centralized policy gate does not cover it; without + this mirror, IdP offboarding would leave the user's already-minted envelopes live + until expiry. A failed user lookup skips the gate, matching the builder.""" + if key_object.user_id is None: return - from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache try: - team_object = await get_team_object( - team_id=team_id, + user_object = await get_user_object( + user_id=key_object.user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + user_id_upsert=False, ) - except (ProxyException, HTTPException): - raise HTTPException(status_code=401, detail="Invalid or expired credential") from None - if team_object.blocked is True: + except Exception as e: # noqa: BLE001 # mirror the builder's fail-open user lookup; DB errors are of any type + verbose_logger.debug(f"bridge admission: user lookup failed, skipping SCIM gate: {e}") + user_object = None + if user_object is None or not isinstance(user_object.metadata, dict): + return + if user_object.metadata.get("scim_active") is False: raise HTTPException(status_code=401, detail="Invalid or expired credential") + @staticmethod + async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Request, route: str) -> None: + """Run the standard pipeline's single authorization point over the admitted identity. + + ``_run_centralized_common_checks`` is the same gate ``user_api_key_auth`` applies + after every builder path, so the envelope identity gets team-block, project-block, + org, and budget enforcement identical to the same key presented directly, and any + policy dimension added to the standard pipeline applies here without this arm + mirroring it. Every failure maps to the arm's uniform 401 so a caller probing with + a stolen envelope learns nothing about why it stopped working.""" + try: + await _run_centralized_common_checks( + user_api_key_auth_obj=admitted, + request=request, + request_data=await _read_request_body(request=request), + route=route, + ) + except Exception: # noqa: BLE001 # common_checks raises bare Exception for blocked states; narrowing would fail open + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + @staticmethod def _admitted_key_is_active(key_object: UserAPIKeyAuth) -> bool: """False when the referenced key is blocked or past its expiry, so a revoked key 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 d0d75231b593..58ef95f53a6e 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 @@ -4962,22 +4962,57 @@ def _reloaded_key(**overrides): @staticmethod @contextlib.contextmanager - def _patch_key_reload(*, return_value=None, side_effect=None, team_blocked=False): - """Patch the live-key reload dependencies used by ``_reload_admitted_key``: the - ``get_key_object`` lookup, the ``get_team_object`` lookup its team-block gate runs, and the - ``prisma_client`` / ``user_api_key_cache`` globals they read. The team resolves unblocked by - default; ``team_blocked=True`` simulates an admin blocking the key's team. Yields the - ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key.""" + def _patch_key_reload(*, return_value=None, side_effect=None, team_blocked=False, owner=None, project_object=None): + """Patch the live-policy dependencies of the admission arm: the ``get_key_object`` reload, + the ``prisma_client`` / ``user_api_key_cache`` globals, and optionally the live objects the + policy gates re-check. ``team_blocked=True`` patches the centralized gate's + ``get_team_object`` at the ``user_api_key_auth`` namespace it actually calls; ``owner`` + patches the SCIM gate's ``get_user_object`` (``auth_checks`` namespace); ``project_object`` + patches the centralized gate's ``get_project_object``. Unpatched lookups hit the MagicMock + prisma and are swallowed (``_safe_fetch`` / the SCIM gate's fail-open), so their checks + skip. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was + the reload key.""" get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) - get_team_object = AsyncMock(return_value=MagicMock(blocked=team_blocked)) - with ( + patchers = [ patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), - patch("litellm.proxy.auth.auth_checks.get_team_object", get_team_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - ): + ] + if team_blocked: + patchers.append( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + AsyncMock(return_value=MagicMock(blocked=True)), + ) + ) + if owner is not None: + patchers.append(patch("litellm.proxy.auth.auth_checks.get_user_object", AsyncMock(return_value=owner))) + if project_object is not None: + patchers.append( + patch( + "litellm.proxy.auth.user_api_key_auth.get_project_object", + AsyncMock(return_value=project_object), + ) + ) + with contextlib.ExitStack() as stack: + for patcher in patchers: + stack.enter_context(patcher) yield get_key_object + @staticmethod + def _mcp_request(path="/mcp/bridge_delegate_server"): + """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how + ``process_mcp_request`` builds one from the ASGI scope with a stubbed empty JSON body.""" + from starlette.requests import Request + + request = Request(scope={"type": "http", "method": "POST", "path": path, "headers": [], "query_string": b""}) + + async def mock_body(): + return b"{}" + + request.body = mock_body + return request + async def test_valid_envelope_reloads_live_key_and_admits_its_authorization_context(self): """A valid envelope admits under the LIVE key record the sealed key_hash references, not a blank identity: the reload is keyed by that exact hash, and the admitted auth carries the @@ -5128,6 +5163,86 @@ async def test_blocked_team_envelope_fails_closed_401(self): assert exc_info.value.status_code == 401 + async def test_scim_deactivated_owner_envelope_fails_closed_401(self): + """SCIM-deactivating the key's OWNER must revoke the user's envelopes immediately: the + standard pipeline rejects every key of a deactivated user inline in the builder, so the + admission arm mirrors that gate. Without it, IdP offboarding would leave the offboarded + user's already-minted envelopes executing tools until they expired.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload( + return_value=self._reloaded_key(), + owner=MagicMock(metadata={"scim_active": False}), + ), + ): + 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 + + async def test_scim_active_owner_envelope_admits(self): + """A SCIM-ACTIVE owner must still be admitted: the gate rejects only an explicit + ``scim_active: False``, so SCIM-managed users whose accounts are in good standing keep + working (and non-SCIM deployments, which never set the flag, are untouched).""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload( + return_value=self._reloaded_key(), + owner=MagicMock(metadata={"scim_active": True}), + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _, _, _, _, _) = await MCPRequestHandler.process_mcp_request(scope) + + assert auth_result.user_id == "envelope-user-42" + + async def test_blocked_project_envelope_fails_closed_401(self): + """Blocking the key's PROJECT must revoke its envelopes immediately: the admitted identity + runs through the standard pipeline's centralized policy gate, which rejects a blocked + project exactly as it would for the same key presented directly. This is the regression for + the project half of the revocation finding; the gate also covers future policy dimensions + without the admission arm mirroring them one by one.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload( + return_value=self._reloaded_key(project_id="project-restricted"), + project_object=MagicMock(blocked=True), + ), + ): + 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 + 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 @@ -5171,12 +5286,53 @@ async def test_sealed_token_wins_over_caller_forwarded_alias_header(self): server=self._bridge_delegate_server(server_name="bridge_name", alias="bridge_alias"), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=attacker_forwarded, + request=self._mcp_request(), + route="/mcp/bridge_name", ) # The sealed token owns the alias slot, overwriting the caller's value; the attacker token # survives nowhere egress would resolve. assert new_headers == {"bridge_alias": {"Authorization": "Bearer inner-upstream-access-token"}} + @pytest.mark.parametrize( + "server_name,alias", + [ + (None, "bridge_alias"), + ("bridge_delegate_server", None), + ("bridge_name", "bridge_alias"), + ], + ids=["alias_only", "server_name_only", "both"], + ) + async def test_injection_key_agrees_with_egress_lookup(self, server_name, alias): + """Round-trip the injected headers through the REAL egress resolver for every admissible + server shape: whatever identifier the admission arm keys the sealed token under, + ``lookup_mcp_server_auth_in_headers`` called the way egress calls it (alias first, then + server_name) must recover exactly that token. This pins the agreement between the two key + hierarchies so neither side can drift and silently drop the forwarded token.""" + from litellm.proxy._experimental.mcp_server.utils import lookup_mcp_server_auth_in_headers + + envelope = self._mint_bridge_envelope() + server = self._bridge_delegate_server(server_name=server_name, alias=alias) + + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key()), + ): + _auth, new_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( + server=server, + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=None, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + + resolved = lookup_mcp_server_auth_in_headers( + new_headers, + alias=server.alias, + server_name=server.server_name, + ) + assert resolved == {"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 @@ -5447,6 +5603,8 @@ async def test_admit_helper_returns_new_headers_without_mutating_input(self): server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=existing, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", ) assert auth_result.user_id == "unit-user" @@ -5469,6 +5627,8 @@ async def test_admit_helper_raises_500_when_master_key_missing(self): server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=None, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", ) assert exc_info.value.status_code == 500 @@ -5485,5 +5645,7 @@ async def test_admit_helper_raises_500_when_no_db_connection(self): server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=None, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", ) assert exc_info.value.status_code == 500 From 0e6ed18bfca9e1426c5158c3732fb70ed3f384a6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 11:46:41 -0700 Subject: [PATCH 5/9] fix(mcp): import assert_never from typing_extensions for Python 3.10 --- .../_experimental/mcp_server/auth/user_api_key_auth_mcp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 998aa4846c6e..d1d48ca7362a 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,11 +1,12 @@ import re from datetime import datetime, timezone -from typing import Dict, List, Optional, Set, Tuple, assert_never, cast +from typing import Dict, List, Optional, Set, Tuple, cast from fastapi import HTTPException from starlette.datastructures import Headers from starlette.requests import Request from starlette.types import Scope +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( From ea64ef7a2afa05be5ba2bb5059305bf221131613 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 12:19:52 -0700 Subject: [PATCH 6/9] fix(mcp): surface real status from bridge admission policy gate instead of flattening to 401 Over-budget rendered 401 (should be 429), model-access and other typed failures collapsed to 401, and a transient DB outage was masked as an auth error. Mirror UserAPIKeyAuthExceptionHandler: budget maps to 429, a sub-check's own HTTPException/ProxyException keeps its status, a DB outage is a retryable 503, and only a genuinely unresolvable failure stays the fail-closed 401. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 32 ++++++++++++-- .../auth/test_user_api_key_auth_mcp.py | 44 +++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) 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 d1d48ca7362a..ddeab918fcc5 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 @@ -8,6 +8,7 @@ from starlette.types import Scope from typing_extensions import assert_never +import litellm from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -592,7 +593,10 @@ async def _reject_if_admitted_owner_scim_deactivated(key_object: UserAPIKeyAuth) The standard pipeline enforces this inline in ``_user_api_key_auth_builder`` rather than in ``common_checks``, so the centralized policy gate does not cover it; without this mirror, IdP offboarding would leave the user's already-minted envelopes live - until expiry. A failed user lookup skips the gate, matching the builder.""" + until expiry. A failed user lookup skips the gate (fail-open), matching the builder: + this is the one deliberately fail-open check in an otherwise fail-closed arm, so a + transient DB outage during this lookup admits the request rather than rejecting it, + keeping parity with how the standard pipeline treats the same lookup failure.""" if key_object.user_id is None: return from litellm.proxy.auth.auth_checks import get_user_object @@ -621,8 +625,19 @@ async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Reque after every builder path, so the envelope identity gets team-block, project-block, org, and budget enforcement identical to the same key presented directly, and any policy dimension added to the standard pipeline applies here without this arm - mirroring it. Every failure maps to the arm's uniform 401 so a caller probing with - a stolen envelope learns nothing about why it stopped working.""" + mirroring it. + + Failures surface with the status the standard pipeline would give them, mirroring + ``UserAPIKeyAuthExceptionHandler``: an over-budget identity is a 429, a sub-check that + raised its own ``HTTPException``/``ProxyException`` keeps that status, a transient + database outage is a retryable 503, and only a genuinely unresolvable failure (a + blocked team/project raises a bare ``Exception``, same as the standard pipeline's + fallback) becomes the fail-closed 401. Collapsing every failure to 401 was misleading: + it told an over-budget but validly-authenticated caller their credential was invalid, + which on a DCR client reads as broken auth and can trigger a pointless re-authorize + loop that cannot fix a budget problem, and it masked a DB outage as an auth error.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + try: await _run_centralized_common_checks( user_api_key_auth_obj=admitted, @@ -630,7 +645,16 @@ async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Reque request_data=await _read_request_body(request=request), route=route, ) - except Exception: # noqa: BLE001 # common_checks raises bare Exception for blocked states; narrowing would fail open + except (HTTPException, ProxyException): + raise + except litellm.BudgetExceededError as e: + raise HTTPException(status_code=getattr(e, "status_code", 429), detail=str(e)) from None + except Exception as e: # noqa: BLE001 # untyped gate failure: retryable 503 for a DB outage, else fail closed 401 + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise HTTPException( + status_code=503, + detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + ) from None raise HTTPException(status_code=401, detail="Invalid or expired credential") from None @staticmethod 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 58ef95f53a6e..f9234cecde4f 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 @@ -5243,6 +5243,50 @@ async def test_blocked_project_envelope_fails_closed_401(self): assert exc_info.value.status_code == 401 + _POLICY_GATE = ( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp._run_centralized_common_checks" + ) + + async def _enforce_with_gate_error(self, error): + """Drive _enforce_admitted_live_policy with the centralized gate raising ``error`` and return + the HTTPException the arm maps it to.""" + with patch(self._POLICY_GATE, new=AsyncMock(side_effect=error)): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler._enforce_admitted_live_policy( + admitted=UserAPIKeyAuth(user_id="envelope-user-42"), + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + return exc_info.value + + async def test_over_budget_admission_surfaces_429_not_401(self): + """A validly-authenticated but over-budget identity surfaces the standard pipeline's 429, not + a misleading 401. Flattening budget to 401 told the caller their credential was invalid, which + on a DCR client reads as broken auth and triggers a re-authorize that cannot fix a budget + problem. Regression for the status-flattening finding on the live-policy gate.""" + import litellm + + mapped = await self._enforce_with_gate_error(litellm.BudgetExceededError(current_cost=10.0, max_budget=1.0)) + assert mapped.status_code == 429 + + async def test_db_outage_during_policy_surfaces_503_not_401(self): + """A transient database outage during the live-policy gate surfaces a retryable 503, not a 401 + that masks the outage as an auth failure and tells a valid caller to re-authenticate.""" + mapped = await self._enforce_with_gate_error(ConnectionError("could not reach database server")) + assert mapped.status_code == 503 + + async def test_blocked_state_bare_exception_stays_401(self): + """A blocked team/project raises a bare Exception (no status) in common_checks, which the + standard pipeline renders as 401; the arm keeps failing those closed as 401, never a 500.""" + mapped = await self._enforce_with_gate_error(Exception("Team=team-x is blocked.")) + assert mapped.status_code == 401 + + async def test_subcheck_httpexception_status_preserved(self): + """A sub-check that raises its own HTTPException (e.g. a 403 model-access denial) keeps that + status through the arm rather than being flattened to 401.""" + mapped = await self._enforce_with_gate_error(HTTPException(status_code=403, detail="model not allowed")) + assert mapped.status_code == 403 + 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 From f5f03cbd635456f48eb72a65567b0e90596bbb97 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 12:46:09 -0700 Subject: [PATCH 7/9] fix(mcp): map a DB outage during bridge key reload to a retryable 503 get_key_object's raw transport error propagated uncaught out of _reload_admitted_key as an opaque 500; classify it via the shared _raise_503_if_db_unavailable helper (also used by the live-policy gate) so a database outage is a retryable 503, while a key-not-found ProxyException stays the fail-closed 401. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 25 +++++++++++++------ .../auth/test_user_api_key_auth_mcp.py | 22 ++++++++++++++++ 2 files changed, 40 insertions(+), 7 deletions(-) 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 ddeab918fcc5..e785b6b8da4b 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 @@ -581,11 +581,28 @@ async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: ) except (ProxyException, HTTPException): raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + except Exception as e: # noqa: BLE001 # a DB outage during reload is a retryable 503, not an opaque 500 + MCPRequestHandler._raise_503_if_db_unavailable(e) + raise if not MCPRequestHandler._admitted_key_is_active(key_object): raise HTTPException(status_code=401, detail="Invalid or expired credential") await MCPRequestHandler._reject_if_admitted_owner_scim_deactivated(key_object) return key_object + @staticmethod + def _raise_503_if_db_unavailable(e: Exception) -> None: + """Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the + caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure + (401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``, + which renders a service-unavailable database error as 503 on the standard pipeline.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise HTTPException( + status_code=503, + detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + ) from None + @staticmethod async def _reject_if_admitted_owner_scim_deactivated(key_object: UserAPIKeyAuth) -> None: """Fail closed with a 401 when the key's owning user was deactivated via SCIM. @@ -636,8 +653,6 @@ async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Reque it told an over-budget but validly-authenticated caller their credential was invalid, which on a DCR client reads as broken auth and can trigger a pointless re-authorize loop that cannot fix a budget problem, and it masked a DB outage as an auth error.""" - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler - try: await _run_centralized_common_checks( user_api_key_auth_obj=admitted, @@ -650,11 +665,7 @@ async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Reque except litellm.BudgetExceededError as e: raise HTTPException(status_code=getattr(e, "status_code", 429), detail=str(e)) from None except Exception as e: # noqa: BLE001 # untyped gate failure: retryable 503 for a DB outage, else fail closed 401 - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): - raise HTTPException( - status_code=503, - detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", - ) from None + MCPRequestHandler._raise_503_if_db_unavailable(e) raise HTTPException(status_code=401, detail="Invalid or expired credential") from None @staticmethod 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 f9234cecde4f..da3e3250c59e 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 @@ -5275,6 +5275,28 @@ async def test_db_outage_during_policy_surfaces_503_not_401(self): mapped = await self._enforce_with_gate_error(ConnectionError("could not reach database server")) assert mapped.status_code == 503 + async def test_db_outage_during_key_reload_surfaces_503_not_500(self): + """A DB outage while reloading the admitted key surfaces a retryable 503, not the opaque 500 a + raw get_key_object transport error would otherwise propagate as, and not a 401 that masks the + outage as an auth failure. Regression for the reload-path exception gap.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(side_effect=ConnectionError("could not reach database server")), + ): + 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 == 503 + async def test_blocked_state_bare_exception_stays_401(self): """A blocked team/project raises a bare Exception (no status) in common_checks, which the standard pipeline renders as 401; the arm keeps failing those closed as 401, never a 500.""" From 688f535bbf1abc9b4b69704ccd76ba1e3bd06a65 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 13:34:23 -0700 Subject: [PATCH 8/9] fix(mcp): run the route gate on bridge admission so allowed_routes are enforced The envelope arm reloaded the identity and ran _run_centralized_common_checks but skipped RouteChecks.should_call_route, which the standard pipeline runs between the builder and common_checks. Because the centralized checks treat MCP as an inference route and never re-check allowed_routes, a key barred from MCP routes could mint an envelope at the token endpoint (not itself an MCP route) and replay it against MCP. Run the route gate before admitting, and clear the request-scoped budget_reservation, matching the wrapper's sequence; a disallowed route now surfaces the gate's own 403. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 36 +++++++++++-------- .../auth/test_user_api_key_auth_mcp.py | 25 +++++++++++++ 2 files changed, 47 insertions(+), 14 deletions(-) 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 e785b6b8da4b..f63819b1822a 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 @@ -636,24 +636,32 @@ async def _reject_if_admitted_owner_scim_deactivated(key_object: UserAPIKeyAuth) @staticmethod async def _enforce_admitted_live_policy(admitted: UserAPIKeyAuth, request: Request, route: str) -> None: - """Run the standard pipeline's single authorization point over the admitted identity. + """Run the standard pipeline's authorization checks over the admitted identity. - ``_run_centralized_common_checks`` is the same gate ``user_api_key_auth`` applies - after every builder path, so the envelope identity gets team-block, project-block, - org, and budget enforcement identical to the same key presented directly, and any - policy dimension added to the standard pipeline applies here without this arm - mirroring it. + Mirrors the ``user_api_key_auth`` wrapper between the builder and its return: clear the + request-scoped ``budget_reservation`` on the reloaded identity, run the route gate + (``RouteChecks.should_call_route``) to enforce the identity's ``allowed_routes`` and any + disabled/admin-only route, then run ``_run_centralized_common_checks`` (the same gate every + builder path funnels through) for team-block, project-block, org, and budget. The route gate + closes a bypass: a key barred from MCP routes could otherwise mint an envelope at the token + endpoint (not itself an MCP route) and replay it against MCP, because the centralized checks + treat MCP as an inference route and never re-check ``allowed_routes``. Failures surface with the status the standard pipeline would give them, mirroring - ``UserAPIKeyAuthExceptionHandler``: an over-budget identity is a 429, a sub-check that - raised its own ``HTTPException``/``ProxyException`` keeps that status, a transient - database outage is a retryable 503, and only a genuinely unresolvable failure (a - blocked team/project raises a bare ``Exception``, same as the standard pipeline's - fallback) becomes the fail-closed 401. Collapsing every failure to 401 was misleading: - it told an over-budget but validly-authenticated caller their credential was invalid, - which on a DCR client reads as broken auth and can trigger a pointless re-authorize - loop that cannot fix a budget problem, and it masked a DB outage as an auth error.""" + ``UserAPIKeyAuthExceptionHandler``: a disallowed route is the route gate's own 403, an + over-budget identity is a 429, a sub-check that raised its own ``HTTPException``/ + ``ProxyException`` keeps that status, a transient database outage is a retryable 503, and + only a genuinely unresolvable failure (a blocked team/project raises a bare ``Exception``, + same as the standard pipeline's fallback) becomes the fail-closed 401. Collapsing every + failure to 401 was misleading: it told an over-budget but validly-authenticated caller their + credential was invalid, which on a DCR client reads as broken auth and can trigger a + pointless re-authorize loop that cannot fix a budget problem, and it masked a DB outage as an + auth error.""" + from litellm.proxy.auth.route_checks import RouteChecks + + admitted.budget_reservation = None try: + RouteChecks.should_call_route(route=route, valid_token=admitted, request=request) await _run_centralized_common_checks( user_api_key_auth_obj=admitted, request=request, 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 da3e3250c59e..eefaaf1dc994 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 @@ -5297,6 +5297,31 @@ async def test_db_outage_during_key_reload_surfaces_503_not_500(self): assert exc_info.value.status_code == 503 + async def test_envelope_for_key_barred_from_mcp_routes_is_rejected_403(self): + """A key whose allowed_routes exclude MCP must not reach tools via an envelope: the arm runs + RouteChecks.should_call_route before admitting, exactly as the standard pipeline does between + the builder and common_checks. A route-restricted key can mint an envelope at the token + endpoint (not itself an MCP route) and would otherwise replay it against MCP, because the + centralized checks treat MCP as an inference route and never re-check allowed_routes; the + route gate rejects it with its own 403.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(allowed_routes=["/chat/completions"])), + ): + 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 == 403 + async def test_blocked_state_bare_exception_stays_401(self): """A blocked team/project raises a bare Exception (no status) in common_checks, which the standard pipeline renders as 401; the arm keeps failing those closed as 401, never a 500.""" From 1c3c1af529b6edfc63161c58a7172f72d25ff55b Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:39:57 -0700 Subject: [PATCH 9/9] fix(mcp): run proxy-wide pre-DB gates on bridge envelope admission The envelope arm bypasses user_api_key_auth, so it never ran pre_db_read_auth_checks (request-size and body-safety limits, the IP allowlist, and the general_settings route allowlist) that the normal MCP admission path runs before any key lookup. A caller blocked by IP or a disallowed proxy route could be admitted through an envelope where the same principal on the normal path is rejected. Run those gates before the envelope crypto, mirroring the pipeline's pre-DB ordering; a blocked IP or route surfaces its own 403. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 21 ++++++++++++++ .../auth/test_user_api_key_auth_mcp.py | 28 +++++++++++++++++-- 2 files changed, 46 insertions(+), 3 deletions(-) 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 f63819b1822a..e300a22e5db3 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 @@ -534,6 +534,8 @@ async def _admit_dcr_bridge_delegate( if not master_key: raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) + keys = envelope_keys_from_master_key(master_key) result = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id) match result: @@ -551,6 +553,25 @@ async def _admit_dcr_bridge_delegate( case _: assert_never(result) + @staticmethod + async def _run_pre_db_read_auth_checks(request: Request, route: str) -> None: + """Run the proxy-wide gates ``user_api_key_auth`` applies before any key lookup: the + request-size and body-safety limits, the IP allowlist, and the ``general_settings`` + route allowlist. The envelope arm bypasses ``user_api_key_auth`` (it opens the envelope + and reloads the identity itself), so without this a caller blocked by IP or hitting a + proxy route the allowlist forbids would be admitted through an envelope where the same + principal presented on the normal MCP admission path would be rejected. Runs before the + envelope crypto so a disallowed caller is turned away before any work, mirroring the + standard pipeline's pre-DB ordering. Violations raise the gate's own status (an IP or + route block is a 403, an oversized body its own limit error).""" + from litellm.proxy.auth.auth_utils import pre_db_read_auth_checks + + await pre_db_read_auth_checks( + request=request, + request_data=await _read_request_body(request=request), + route=route, + ) + @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: """Reload the live key record an admitted envelope references and re-check live policy. 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 eefaaf1dc994..c785ac577f7d 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 @@ -5243,9 +5243,7 @@ async def test_blocked_project_envelope_fails_closed_401(self): assert exc_info.value.status_code == 401 - _POLICY_GATE = ( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp._run_centralized_common_checks" - ) + _POLICY_GATE = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp._run_centralized_common_checks" async def _enforce_with_gate_error(self, error): """Drive _enforce_admitted_live_policy with the centralized gate raising ``error`` and return @@ -5322,6 +5320,30 @@ async def test_envelope_for_key_barred_from_mcp_routes_is_rejected_403(self): assert exc_info.value.status_code == 403 + async def test_envelope_rejected_by_proxy_wide_pre_db_gates_403(self): + """The envelope arm runs the same proxy-wide pre-DB gates user_api_key_auth applies before any + key lookup (request size, body safety, IP allowlist, general_settings route allowlist). Here + the proxy route allowlist forbids MCP, so the envelope is turned away with a 403 before the + identity is even reloaded, closing the gap where an envelope bypassed the IP/route allowlists + the normal MCP admission path enforces.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + 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.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch("litellm.proxy.proxy_server.general_settings", {"allowed_routes": ["/chat/completions"]}), + ): + 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 == 403 + async def test_blocked_state_bare_exception_stays_401(self): """A blocked team/project raises a bare Exception (no status) in common_checks, which the standard pipeline renders as 401; the arm keeps failing those closed as 401, never a 500."""