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..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 @@ -1,12 +1,23 @@ import re +from datetime import datetime, timezone 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 +import litellm 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, @@ -17,12 +28,17 @@ 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, 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 +242,29 @@ 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 = await MCPRequestHandler._admit_dcr_bridge_delegate( + 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 # LiteLLM credential, so a failed validation is a genuine 401/403 and @@ -432,6 +471,247 @@ 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 + 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, 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, 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 + 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 + + 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: + case BridgeEnvelopeAdmitted(): + 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) + 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 + case BridgeEnvelopeInvalid() | NotBridgeEnvelope(): + raise HTTPException(status_code=401, detail="Invalid or expired credential") + 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. + + 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). 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 + + 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 + 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. + + 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 (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 + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + try: + 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 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 authorization checks over the admitted identity. + + 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``: 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, + request_data=await _read_request_body(request=request), + route=route, + ) + 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 + MCPRequestHandler._raise_503_if_db_unavailable(e) + 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 + 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 ebaa6bc7cc04..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 @@ -1,6 +1,8 @@ +import contextlib import json import os import sys +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,6 +17,8 @@ MCPRequestHandler, ) from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + ProxyException, SpecialHeaders, SpecialMCPServerNames, UserAPIKeyAuth, @@ -4866,3 +4870,895 @@ 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, 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" + + @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, + ) + + _KEY_HASH = "hashed-litellm-key-abc123" + + @classmethod + def _mint_bridge_envelope( + cls, + *, + key_hash=None, + 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(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), + 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() + + @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, 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) + patchers = [ + 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()), + ] + 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 + 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", + 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), + 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() + ( + auth_result, + _mcp_auth_header, + _mcp_servers, + mcp_server_auth_headers, + _oauth2_headers, + _raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # 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_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_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 + + _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_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_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_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.""" + 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 + 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), + 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" + ) + (_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_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, + 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 + 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() + existing = {"other_server": {"Authorization": "Bearer someone-elses-token"}} + + 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, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + + 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: + await MCPRequestHandler._admit_dcr_bridge_delegate( + 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 + + 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, + request=self._mcp_request(), + route="/mcp/bridge_delegate_server", + ) + assert exc_info.value.status_code == 500 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"