From 327515a3ba526f93ff88d716377045d593fd34b9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:09:52 -0700 Subject: [PATCH] fix(mcp): preserve credential authority in DCR bridge authentication (#42563) * fix(mcp): admit dcr_bridge envelope alongside an explicit litellm credential and mint under jwt principals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(mcp): suppress LIT002 on concrete dict header payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): mint bridge envelope for jwt mapped to a key without a user_id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): mint and admit bridge envelopes under the master key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): bind mapped JWT envelopes to stored key tokens * fix(mcp): preserve master envelope scope enforcement * fix(mcp): reject bridge minting that loses JWT restrictions --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 152 +++++++++- .../mcp_server/bridge_token_flow.py | 65 ++++- .../auth/test_user_api_key_auth_mcp.py | 257 ++++++++++++++++- .../mcp_server/test_discoverable_endpoints.py | 264 ++++++++++++++++++ 4 files changed, 716 insertions(+), 22 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 d6478421834..43d62ba8b22 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 @@ -42,6 +42,7 @@ from litellm.proxy._types import ( SpecialMCPServerName, SpecialMCPServerNames, UserAPIKeyAuth, + hash_token, user_api_key_has_admin_view, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( @@ -182,6 +183,20 @@ def _is_litellm_auth_admission_error(exc: Exception) -> bool: return False +def _explicit_credential_matches_envelope( + explicit_auth: UserAPIKeyAuth, + presented_token: str, + identity: EnvelopeIdentity, +) -> bool: + """Match the stored key hash or user ID, including token-only mapped JWT keys.""" + match identity.subject_type: + case "key_hash": + return identity.subject in (hash_token(presented_token), explicit_auth.token) + case "user_id": + return explicit_auth.user_id is not None and explicit_auth.user_id == identity.subject + return assert_never(identity.subject_type) + + def _has_client_supplied_mcp_auth( mcp_auth_header: str | None, mcp_server_auth_headers: dict[str, dict[str, str]] | None, @@ -475,6 +490,31 @@ class MCPRequestHandler: # Only OAuth metadata routes registered under /.well-known/ are public. if request_route.startswith("/.well-known/"): validated_user_api_key_auth = UserAPIKeyAuth() + elif ( + has_explicit_litellm_key + and oauth2_headers + and is_bridge_envelope_shaped(oauth2_headers["Authorization"]) + and ( + dual_bridge_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 + ): + ( + validated_user_api_key_auth, + mcp_server_auth_headers, + ) = await MCPRequestHandler._admit_dcr_bridge_dual_credential( + server=dual_bridge_target.server, + requested_name=dual_bridge_target.requested_name, + authorization_value=oauth2_headers["Authorization"], + litellm_api_key=litellm_api_key, + mcp_server_auth_headers=mcp_server_auth_headers, + request=request, + route=request_route, + ) elif has_explicit_litellm_key: # An explicit x-litellm-api-key is always a LiteLLM credential, even # for a delegated server, so validate it: identity / spend / rate @@ -782,6 +822,45 @@ class MCPRequestHandler: higher-priority alias slot, pairing the admitted identity with an attacker's upstream credential; the alias-keyed injection overwrites any such caller value. """ + result: Final = await MCPRequestHandler._open_dcr_bridge_envelope( + server=server, + requested_name=requested_name, + authorization_value=authorization_value, + request=request, + route=route, + ) + header_key: Final = 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: Final = await MCPRequestHandler._reload_admitted_principal(result.identity) + await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) + injected: Final = { # mutable-ok: mcp_server_auth_headers contract requires concrete dicts + header_key: { # mutable-ok: concrete dict header payload + "Authorization": result.upstream_authorization.get_secret_value() + } + } + new_headers: Final = { # mutable-ok: merged header map must stay a concrete dict + **(mcp_server_auth_headers or {}), # mutable-ok: empty-dict fallback for the merge + **injected, + } + return admitted, new_headers + + @staticmethod + async def _open_dcr_bridge_envelope( + server: MCPServer, + requested_name: str, + authorization_value: str, + request: Request, + route: str, + ) -> BridgeEnvelopeAdmitted: + """Open a bridge envelope after the pre-DB gates, or fail closed with the scope's challenge. + + Shared by the envelope-only arm (:meth:`_admit_dcr_bridge_delegate`) and the dual-credential + arm (:meth:`_admit_dcr_bridge_dual_credential`): both require master_key, run the same + proxy-wide pre-DB checks the standard pipeline applies before any key lookup, and resolve + the envelope's crypto. Returns only the ``BridgeEnvelopeAdmitted`` result; an invalid, + expired, tampered, or non-envelope value raises the requested scope's ``invalid_token`` + challenge instead.""" from litellm.proxy.proxy_server import master_key if not master_key: @@ -793,20 +872,67 @@ class MCPRequestHandler: result: Final = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id) match result: case BridgeEnvelopeAdmitted(): - header_key: Final = 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: Final = await MCPRequestHandler._reload_admitted_principal(result.identity) - await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) - injected: Final = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} - new_headers: Final = {**(mcp_server_auth_headers or {}), **injected} - return admitted, new_headers + return result case BridgeEnvelopeInvalid() | NotBridgeEnvelope(): raise MCPRequestHandler._dcr_bridge_invalid_token_challenge( requested_name=requested_name, request=request ) - case _: - assert_never(result) + return assert_never(result) + + @staticmethod + async def _admit_dcr_bridge_dual_credential( + server: MCPServer, + requested_name: str, + authorization_value: str, + litellm_api_key: str, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + request: Request, + route: str, + ) -> tuple[UserAPIKeyAuth, dict[str, dict[str, str]] | None]: + """Admit a request carrying BOTH an explicit litellm credential and a bridge envelope. + + MCP clients send ``x-litellm-api-key`` on every request, including the ``tools/list`` that + follows the ``/{server}/token`` mint, so the envelope arrives alongside the key rather than + alone. The explicit credential is validated first (its own pipeline, so a bad key keeps the + normal 401/403), then the envelope is opened and its sealed identity must match the explicit + credential's principal — a mismatch is a 403, never a fallback onto either credential alone. + On a match the explicit credential's ``UserAPIKeyAuth`` is the admission context (key + permissions, budgets, rate limits) and the sealed upstream token is injected under the + server's per-server auth-header key, while the leak-defense chokepoint strips the envelope + ``Authorization`` itself from egress.""" + presented_token: Final = _get_bearer_token_or_received_api_key(litellm_api_key) + explicit_auth: Final = await user_api_key_auth(api_key=f"Bearer {presented_token}", request=request) + result: Final = await MCPRequestHandler._open_dcr_bridge_envelope( + server=server, + requested_name=requested_name, + authorization_value=authorization_value, + request=request, + route=route, + ) + if not _explicit_credential_matches_envelope( + explicit_auth=explicit_auth, + presented_token=presented_token, + identity=result.identity, + ): + raise HTTPException( + status_code=403, + detail={ # mutable-ok: HTTPException detail payload requires a concrete dict + "error": "oauth_principal_mismatch" + }, + ) + header_key: Final = 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") + injected: Final = { # mutable-ok: mcp_server_auth_headers contract requires concrete dicts + header_key: { # mutable-ok: concrete dict header payload + "Authorization": result.upstream_authorization.get_secret_value() + } + } + new_headers: Final = { # mutable-ok: merged header map must stay a concrete dict + **(mcp_server_auth_headers or {}), # mutable-ok: empty-dict fallback for the merge + **injected, + } + return explicit_auth, new_headers @staticmethod async def _admit_dcr_bridge_authorization( @@ -1091,9 +1217,15 @@ class MCPRequestHandler: 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._experimental.mcp_server.bridge_token_flow import ( + master_key_admin_auth, # noqa: PLC0415 # inline import avoids a module-load circular import + ) from litellm.proxy.auth.auth_checks import get_key_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + admin: Final = master_key_admin_auth(key_hash) + if admin is not None: + return admin if prisma_client is None: raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") try: diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 6ae33cc1629..e5d11271c67 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -209,6 +209,31 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR return await _reload_active_key_by_hash(hash_token(token)) +def master_key_admin_auth(key_hash: str) -> "UserAPIKeyAuth | None": + from litellm.constants import ( # noqa: PLC0415 # inline import avoids a module-load circular import + LITELLM_PROXY_MASTER_KEY_ALIAS, + ) + from litellm.proxy._types import ( # noqa: PLC0415 # inline import avoids a module-load circular import + LitellmUserRoles, + UserAPIKeyAuth, + hash_token, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + litellm_proxy_admin_name, + master_key, + ) + + if not master_key or not secrets.compare_digest(key_hash, hash_token(master_key)): + return None + auth: Final = UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id=litellm_proxy_admin_name, + ) + auth.via_virtual_key = True + return auth + + async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure": """Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state, returning the resolved key or a precise failure. Shared by the token request's presented-key @@ -233,6 +258,8 @@ async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResol user_api_key_cache, ) + if (admin := master_key_admin_auth(key_hash)) is not None: + return _ResolvedKey(key_hash=key_hash, key=admin) if prisma_client is None: return "unresolvable" try: @@ -475,7 +502,7 @@ async def _resolve_jwt_auth( proxy_logging_obj=proxy_logging_obj, ) if isinstance(mapped, UserAPIKeyAuth): - return None if await _key_owner_scim_deactivated(mapped) or not _active_key_user_id(mapped) else mapped + return None if await _key_owner_scim_deactivated(mapped) or not _key_is_active(mapped) else mapped if mapped is not None: return None if write_route is None: @@ -593,6 +620,7 @@ def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenG _BridgeMintError = Literal[ "no_identity", + "jwt_client_policy_unsupported", "invalid_refresh", "identity_unavailable", "identity_faulted", @@ -633,6 +661,13 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: "this server issues a gateway-bound credential; complete the interactive sign-in, or " "send a litellm credential (x-litellm-api-key or Authorization) on the token request", ) + case "jwt_client_policy_unsupported": + status, code, desc = ( + 400, + "invalid_request", + "JWT bridge minting is not supported with a claim-based MCP client allowlist; " + "the bridge credential cannot preserve the signed client identity", + ) case "invalid_refresh": status, code, desc = ( 400, @@ -736,11 +771,18 @@ async def _prepare_bridge_mint( Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway - authorization code) and mints a user subject. The scripted two-header client presents a litellm key - on the token request instead, so its identity is the active key's hash and mints a key_hash subject. - A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; - neither source present is ``no_identity``. The refresh_token grant has its own phase-1 + authorization code) and mints a user subject. The scripted two-header client presents a litellm + credential (a virtual key or a JWT) on the token request instead: a key mints a key_hash subject, + while a JWT resolves through the same auth path as admission and mints a key_hash subject when it + maps to a virtual key. An unmapped JWT is rejected because a user subject cannot preserve its + JWT-specific authorization restrictions. A JWT client-claim allowlist also prevents JWT minting: + the envelope cannot retain the signed client identity for subsequent allowlist checks. A missing or invalid + presented key keeps its resolution origin so the mapper statuses it truthfully; neither source + present is ``no_identity``. The refresh_token grant has its own phase-1 (:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope.""" + from litellm.proxy._experimental.mcp_server.client_allowlist import ( # noqa: PLC0415 # keep mint policy dependencies local + load_mcp_client_allowlist, + ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import envelope_keys_from_master_key, ) @@ -748,7 +790,10 @@ async def _prepare_bridge_mint( key_hash_identity, user_identity, ) + from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle + from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + general_settings, master_key, ) @@ -758,6 +803,16 @@ async def _prepare_bridge_mint( if bridge_identity is not None: identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) return _BridgeMintReady(identity=identity, keys=keys) + presented_token: Final = _litellm_key_from_request(request) + if presented_token is not None and JWTHandler.is_jwt(presented_token): + client_allowlist: Final = load_mcp_client_allowlist(general_settings) + if client_allowlist is not None and client_allowlist.jwt_field is not None: + return "jwt_client_policy_unsupported" + resolved_jwt: Final = await _resolve_jwt_auth(request, presented_token, None) + if isinstance(resolved_jwt, UserAPIKeyAuth) and resolved_jwt.token: + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved_jwt.token) + return _BridgeMintReady(identity=identity, keys=keys) + return "no_identity" resolved: Final = await _resolve_active_litellm_key(request) if not isinstance(resolved, _ResolvedKey): return _key_resolution_failure_to_mint_error(resolved) 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 0012be23428..05ed53df8e8 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 @@ -6529,6 +6529,53 @@ class TestMCPDcrBridgeDelegateAdmission: assert exc_info.value.status_code == 503 + async def test_reload_admitted_key_returns_admin_for_master_key_hash(self): + """An envelope sealed under the master key has no DB row to reload; the reload resolves it + to the PROXY_ADMIN auth context (api_key is the alias, never the hash) rather than failing. + A hash that is NOT the master key's still reaches the prisma gate and fails the same as + before (500 with no database connection).""" + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy._types import LitellmUserRoles, hash_token + + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + admitted = await MCPRequestHandler._reload_admitted_key(hash_token(self._MASTER_KEY)) + assert admitted.user_role == LitellmUserRoles.PROXY_ADMIN + assert admitted.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler._reload_admitted_key("not-the-master-hash") + assert exc_info.value.status_code == 500 + + @pytest.mark.parametrize( + "flag_enabled, scope, expected", + [(True, "scoped", []), (False, "scoped", ["public"]), (True, "unscoped", ["public"])], + ) + async def test_master_envelope_respects_allow_all_scope(self, flag_enabled, scope, expected): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy._types import hash_token + + manager = MCPServerManager() + with patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY): + admitted = await MCPRequestHandler._reload_admitted_key(hash_token(self._MASTER_KEY)) + with ( + patch.object(manager, "get_allow_all_keys_server_ids", return_value=["public"]), + patch.object(manager, "_get_active_submitted_mcp_server_ids_for_user", new=AsyncMock(return_value=[])), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_user", + new=AsyncMock(return_value=["granted"] if scope == "scoped" else []), + ), + ): + servers = await manager.get_allowed_mcp_servers( + admitted, + access=MCPServerAccess(server_ids=(), scope=scope), + general_settings={"mcp_allow_all_keys_respects_mcp_scope": flag_enabled}, + ) + assert servers == expected + 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 @@ -6959,10 +7006,12 @@ class TestMCPDcrBridgeDelegateAdmission: assert exc_info.value.status_code == 403 assert not exc_info.value.headers - 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.""" + async def test_explicit_litellm_key_matching_envelope_admits_under_explicit_key(self): + """The dual-credential arm: an explicit x-litellm-api-key paired with an envelope sealing the + SAME key hash admits under the explicit key's auth context AND injects the sealed upstream + token for egress. When the envelope seals a different principal the request is a 403 instead + (covered by the mismatch tests), and the explicit key never silently drops the envelope the + way the pre-fix ordering did.""" envelope = self._mint_bridge_envelope() scope = { "type": "http", @@ -6975,7 +7024,7 @@ class TestMCPDcrBridgeDelegateAdmission: } async def mock_user_api_key_auth(api_key, request): - return UserAPIKeyAuth(api_key=api_key, user_id="litellm-key-user") + return UserAPIKeyAuth(api_key=self._KEY_HASH, user_id="litellm-key-user") with ( patch( @@ -6997,9 +7046,10 @@ class TestMCPDcrBridgeDelegateAdmission: mock_auth.assert_called_once() assert mock_auth.call_args.kwargs["api_key"] == "Bearer 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 == {} + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } 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 @@ -7139,6 +7189,199 @@ class TestMCPDcrBridgeDelegateAdmission: assert exc_info.value.status_code == 500 +@pytest.mark.asyncio +class TestMCPDcrBridgeDualCredential: + """Dual-credential arm: ``x-litellm-api-key`` alongside an ``llm_env_`` bearer on a + DCR-bridge ``oauth_delegate`` route (issue #38208). + + Real MCP clients send their litellm key on every request, so the envelope minted at + ``/{server}/token`` arrives paired with the key rather than alone. The explicit credential + is the admission context and the envelope supplies the upstream token, but only when both + name the same principal; a mismatch is a 403, an invalid envelope is the scope's + ``invalid_token`` challenge, and the envelope itself never reaches egress. + """ + + _DELEGATE = TestMCPDcrBridgeDelegateAdmission + _MASTER_KEY = TestMCPDcrBridgeDelegateAdmission._MASTER_KEY + _KEY_HASH = TestMCPDcrBridgeDelegateAdmission._KEY_HASH + + @staticmethod + def _dual_scope(envelope: str, explicit_key: str): + return { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [ + (b"authorization", f"Bearer {envelope}".encode("latin-1")), + (b"x-litellm-api-key", explicit_key.encode("latin-1")), + ], + } + + @pytest.mark.parametrize("dual_credential", [False, True]) + async def test_admission_rejects_server_without_routable_name(self, dual_credential): + envelope = self._DELEGATE._mint_bridge_envelope() + server = self._DELEGATE._bridge_delegate_server(server_name=None) + admission = ( + MCPRequestHandler._admit_dcr_bridge_dual_credential( + server=server, + requested_name="bridge_delegate_server", + authorization_value=f"Bearer {envelope}", + litellm_api_key="sk-explicit-key", + mcp_server_auth_headers=None, + request=self._DELEGATE._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + if dual_credential + else MCPRequestHandler._admit_dcr_bridge_delegate( + server=server, + requested_name="bridge_delegate_server", + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=None, + request=self._DELEGATE._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new=AsyncMock(return_value=self._DELEGATE._reloaded_key()), + ), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._DELEGATE._patch_key_reload() as reload_key, + pytest.raises(HTTPException) as exc_info, + ): + await admission + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Server misconfigured: MCP server has no routable name" + reload_key.assert_not_awaited() + + @pytest.mark.parametrize("mapped_jwt", [False, True]) + async def test_dual_credential_matching_key_admits_under_explicit_key_and_forwards_upstream_token(self, mapped_jwt): + """The reported bug: before the fix this request validated the key and dropped the + envelope, so egress forwarded no upstream credential and the upstream 401 yielded + ``tools: []``. Now the explicit key's auth context wins admission AND the sealed + upstream token is injected per-server, while the envelope bearer is scrubbed from + every egress header context.""" + envelope = self._DELEGATE._mint_bridge_envelope(key_hash=self._KEY_HASH) + explicit_auth = self._DELEGATE._reloaded_key( + api_key=None if mapped_jwt else self._KEY_HASH, + token=self._KEY_HASH, + user_id=None if mapped_jwt else "explicit-key-user", + ) + presented_token = "aaa.bbb.ccc" if mapped_jwt else "sk-explicit-key" + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + return_value=explicit_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), + self._DELEGATE._patch_key_reload() as get_key_object, + ): + mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server() + ( + auth_result, + _mcp_auth_header, + _mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = await MCPRequestHandler.process_mcp_request(self._dual_scope(envelope, presented_token)) + + mock_auth.assert_awaited_once() + assert mock_auth.await_args.kwargs["api_key"] == f"Bearer {presented_token}" + assert auth_result is explicit_auth + get_key_object.assert_not_awaited() + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + assert oauth2_headers is None + assert all("llm_env_" not in str(v) for v in raw_headers.values()) + + async def test_dual_credential_principal_mismatch_is_403(self): + """An envelope minted under one key presented alongside a different key must not admit: + the request names two different principals, so it fails closed with + ``oauth_principal_mismatch`` rather than falling back onto either credential.""" + envelope = self._DELEGATE._mint_bridge_envelope(key_hash=self._KEY_HASH) + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + return_value=self._DELEGATE._reloaded_key(api_key="a-different-key-hash"), + ), + 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._DELEGATE._patch_key_reload(), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._dual_scope(envelope, "sk-other-key")) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == {"error": "oauth_principal_mismatch"} + + async def test_dual_credential_user_subject_envelope_matches_on_user_id(self): + """An interactive (user_id) envelope pairs with an explicit credential whose resolved + user_id is the same user; a different user is a 403, never a silent admit.""" + for presented_user, expected_status in (("sso-user-7", None), ("sso-user-9", 403)): + envelope = self._DELEGATE._mint_bridge_envelope(user_id="sso-user-7") + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + return_value=UserAPIKeyAuth(user_id=presented_user, api_key="any-hash"), + ), + 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._DELEGATE._bridge_delegate_server() + if expected_status is not None: + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._dual_scope(envelope, "sk-key")) + assert exc_info.value.status_code == expected_status + assert exc_info.value.detail == {"error": "oauth_principal_mismatch"} + else: + ( + auth_result, + _h, + _s, + mcp_server_auth_headers, + _o, + _r, + ) = await MCPRequestHandler.process_mcp_request(self._dual_scope(envelope, "sk-key")) + assert auth_result.user_id == "sso-user-7" + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + + async def test_dual_credential_invalid_envelope_is_401_challenge_not_silent_admit(self): + """A tampered envelope next to a perfectly valid key must still fail closed with the + scope's ``invalid_token`` challenge; the explicit key alone never unlocks a bridge + server's upstream token.""" + envelope = self._DELEGATE._mint_bridge_envelope(key_hash=self._KEY_HASH) + tampered = envelope[:-4] + "AAAA" + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + return_value=self._DELEGATE._reloaded_key(), + ), + 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._DELEGATE._patch_key_reload(), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._dual_scope(tampered, "sk-explicit-key")) + + assert exc_info.value.status_code == 401 + assert "invalid_token" in str(exc_info.value.headers) + + @pytest.mark.asyncio class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index c1d5cedeba0..f9a0075e530 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -6725,6 +6725,270 @@ async def test_bridge_mint_unresolvable_identity_is_500_before_upstream(): post.assert_not_called() +async def _exchange_for_bridge_server_with_jwt(jwt_auth_result, upstream_body=None): + """Drive exchange_token_with_server for a bridge oauth_delegate authorization_code request whose + presented credential is JWT-shaped, with _resolve_jwt_auth stubbed to a given result. Returns + (response, post_mock) so a test can assert the minted envelope's sealed identity or the mapped + error status.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + request = _bridge_mock_request() + request.headers = {"x-litellm-api-key": "aaa.bbb.ccc"} + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body or { + "access_token": "UP", + "token_type": "Bearer", + "expires_in": 3600, + } + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_jwt_auth", + new=AsyncMock(return_value=jwt_auth_result), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + return response, fake_http_client.post + + +@pytest.mark.asyncio +async def test_bridge_mint_unmapped_jwt_is_rejected_before_upstream(): + from litellm.proxy.auth.handle_jwt import JWTIdentity + + response, post = await _exchange_for_bridge_server_with_jwt( + JWTIdentity(user_id="jwt-user-5", user_object=None, agent_id=None) + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_jwt_mapped_to_virtual_key_seals_key_hash_subject(): + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.proxy._types import UserAPIKeyAuth + + response, _post = await _exchange_for_bridge_server_with_jwt( + UserAPIKeyAuth(token="mapped-key-hash-99", user_id="mapped-user") + ) + assert response.status_code == 200 + token = json.loads(response.body)["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), "bridge_srv") + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.subject_type == "key_hash" + assert opened.identity.subject == "mapped-key-hash-99" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("jwt_claims", [{"client_id": "allowed"}, {"client_id": "denied"}, {}]) +async def test_bridge_mint_jwt_cannot_drop_signed_client_policy(jwt_claims): + from litellm.proxy._types import UserAPIKeyAuth + + settings = { + "mcp_allowed_clients": [{"alias": "Allowed", "value": "allowed"}], + "mcp_client_id_header": "x-client-id", + "litellm_jwtauth": {"mcp_client_id_jwt_field": "client_id"}, + } + with patch("litellm.proxy.proxy_server.general_settings", settings): + response, post = await _exchange_for_bridge_server_with_jwt( + UserAPIKeyAuth(token="mapped-key-hash-99", jwt_claims=jwt_claims) + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + assert "signed client identity" in json.loads(response.body)["error_description"] + post.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "settings", + [ + {}, + {"litellm_jwtauth": {"mcp_client_id_jwt_field": "client_id"}}, + { + "mcp_allowed_clients": [{"alias": "Allowed", "value": "allowed"}], + "mcp_client_id_header": "x-client-id", + }, + ], +) +async def test_bridge_mint_mapped_jwt_without_signed_client_policy(settings): + from litellm.proxy._types import UserAPIKeyAuth + + with patch("litellm.proxy.proxy_server.general_settings", settings): + response, post = await _exchange_for_bridge_server_with_jwt( + UserAPIKeyAuth(token="mapped-key-hash-99", jwt_claims={"client_id": "allowed"}) + ) + assert response.status_code == 200 + assert json.loads(response.body)["access_token"].startswith("llm_env_") + post.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_bridge_mint_jwt_with_no_resolved_identity_is_400_before_upstream(): + """A JWT that resolves to nothing (or to an identity with no user_id) cannot back an envelope: + the mint returns 400 invalid_request WITHOUT consuming the single-use code upstream, matching + the no-credential path rather than hashing the raw JWT string.""" + response, post = await _exchange_for_bridge_server_with_jwt(None) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + post.assert_not_called() + + from litellm.proxy.auth.handle_jwt import JWTIdentity + + response, post = await _exchange_for_bridge_server_with_jwt( + JWTIdentity(user_id=None, user_object=None, agent_id=None) + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + post.assert_not_called() + + +def _jwt_auth_patches(mapped_key): + from contextlib import ExitStack + + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.set_cache(key=jwt_key_mapping_cache_key("sub", "mapped-client"), value=mapped_key.token) + cache.set_cache(key=mapped_key.token, value=mapped_key) + handler = MagicMock() + handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="sub") + handler.auth_jwt = AsyncMock(return_value={"sub": "mapped-client"}) + stack = ExitStack() + stack.enter_context(patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True})) + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) + stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", object())) + stack.enter_context(patch("litellm.proxy.proxy_server.jwt_handler", handler)) + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache", cache)) + stack.enter_context( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._key_owner_scim_deactivated", + new=AsyncMock(return_value=False), + ) + ) + return stack + + +@pytest.mark.asyncio +async def test_jwt_mapped_to_service_account_key_without_user_id_resolves(): + """A JWT mapped to a team or service-account virtual key (no user_id) is still an active + credential: _resolve_jwt_auth returns the mapped key, and the mint seals a key_hash-subject + envelope rather than 400ing with no_identity.""" + from litellm.proxy._experimental.mcp_server import bridge_token_flow + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth + + mapped_key = UserAPIKeyAuth(user_id=None, token="svc-key-hash-1") + request = _bridge_mock_request() + request.headers = {"x-litellm-api-key": "aaa.bbb.ccc"} + with ( + _jwt_auth_patches(mapped_key), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + resolved = await bridge_token_flow._resolve_jwt_auth(request, "aaa.bbb.ccc", None) + assert isinstance(resolved, UserAPIKeyAuth) + assert resolved.token == mapped_key.token + assert resolved.api_key is None + assert resolved.user_id is None + + mint = await bridge_token_flow._prepare_bridge_mint( + request=request, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + ) + assert isinstance(mint, bridge_token_flow._BridgeMintReady) + assert mint.identity.subject_type == "key_hash" + assert mint.identity.subject == "svc-key-hash-1" + + +@pytest.mark.asyncio +async def test_jwt_mapped_to_blocked_key_is_rejected(): + """The relaxed gate is still active-state gated: a JWT mapped to a blocked virtual key resolves + to None, so the mint cannot seal an envelope under it.""" + from litellm.proxy._experimental.mcp_server import bridge_token_flow + from litellm.proxy._types import UserAPIKeyAuth + + mapped_key = UserAPIKeyAuth(user_id=None, token="blocked-key-hash-1", blocked=True) + request = _bridge_mock_request() + request.headers = {"x-litellm-api-key": "aaa.bbb.ccc"} + with ( + _jwt_auth_patches(mapped_key), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + resolved = await bridge_token_flow._resolve_jwt_auth(request, "aaa.bbb.ccc", None) + assert resolved is None + + mint = await bridge_token_flow._prepare_bridge_mint( + request=request, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + ) + assert mint == "no_identity" + + +@pytest.mark.asyncio +async def test_master_key_at_token_endpoint_mints_key_hash_envelope(): + """The master key has no row in LiteLLM_VerificationTokenTable, but it is the proxy's root + credential: presented at the bridge /token endpoint it must mint a key_hash-subject envelope + (sealed under hash_token(master_key)) even with no database connection at all. A presented key + that is NOT the master key still hits the unresolvable gate when prisma is down, unchanged.""" + from litellm.proxy._experimental.mcp_server import bridge_token_flow + from litellm.proxy._types import hash_token + from litellm.types.mcp import MCPAuth + + master = "sk-test-master-key-mint-0000" + request = _bridge_mock_request() + request.headers = {"x-litellm-api-key": master} + with ( + patch("litellm.proxy.proxy_server.master_key", master), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + mint = await bridge_token_flow._prepare_bridge_mint( + request=request, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + ) + assert isinstance(mint, bridge_token_flow._BridgeMintReady) + assert mint.identity.subject_type == "key_hash" + assert mint.identity.subject == hash_token(master) + + other = _bridge_mock_request() + other.headers = {"x-litellm-api-key": "sk-not-the-master-key"} + with ( + patch("litellm.proxy.proxy_server.master_key", master), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + mint = await bridge_token_flow._prepare_bridge_mint( + request=other, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + ) + assert mint == "identity_unresolvable" + + @pytest.mark.asyncio async def test_bridge_mint_upstream_expired_lifetime_is_502(): """An upstream token response reporting an already-elapsed lifetime (a parseable non-positive