feat(mcp): admit gateway DCR session bearers at the aggregate /mcp scope

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -25,6 +25,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti
 from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
     EnvelopeIdentity,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
+    is_session_bearer_shaped,
+)
 from litellm.proxy._types import (
     UI_TEAM_ID,
     LiteLLM_TeamTable,
@@ -124,6 +127,17 @@ def _has_client_supplied_mcp_auth(
     return bool(mcp_auth_header) or bool(mcp_server_auth_headers)
 
 
+def _is_aggregate_mcp_scope(route: str, mcp_servers: list[str] | None) -> bool:
+    """True when a request targets the aggregate ``/mcp`` endpoint rather than any named
+    server. Named targets arrive either through ``x-mcp-servers`` (``mcp_servers``) or a
+    path segment (``/mcp/{server}`` / ``/{server}/mcp``); the aggregate scope has neither.
+    The gateway-DCR session arm and challenge fire only here, so a per-server flow is never
+    affected."""
+    if mcp_servers:
+        return False
+    return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0
+
+
 def _is_aggregate_gateway_dcr_challenge_scope(
     route: str,
     mcp_servers: list[str] | None,
@@ -141,11 +155,9 @@ def _is_aggregate_gateway_dcr_challenge_scope(
     client. Fails closed to the original admission error otherwise."""
     if not _is_litellm_auth_admission_error(exc):
         return False
-    if mcp_servers:
-        return False
     if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers):
         return False
-    return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0
+    return _is_aggregate_mcp_scope(route, mcp_servers)
 
 
 def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException:
@@ -362,6 +374,22 @@ class MCPRequestHandler:
                 request=request,
                 route=request_route,
             )
+        elif (
+            is_mcp_gateway_dcr_enabled()
+            and _is_aggregate_mcp_scope(request_route, mcp_servers)
+            and oauth2_headers
+            and is_session_bearer_shaped(oauth2_headers["Authorization"])
+        ):
+            # A gateway DCR session bearer at the aggregate /mcp scope: open the
+            # identity-only session token and admit under the live litellm user it
+            # references. A session-shaped bearer that does not open fails closed with
+            # the aggregate invalid_token challenge; a non-session bearer never reaches
+            # here (is_session_bearer_shaped is false) and falls through to the oauth2 arm.
+            validated_user_api_key_auth = await MCPRequestHandler._admit_gateway_session(
+                authorization_value=oauth2_headers["Authorization"],
+                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
@@ -626,6 +654,59 @@ class MCPRequestHandler:
             case _:
                 assert_never(result)
 
+    @staticmethod
+    async def _admit_gateway_session(
+        authorization_value: str,
+        request: Request,
+        route: str,
+    ) -> UserAPIKeyAuth:
+        """Open a gateway DCR session bearer and admit the live litellm user it references.
+
+        The custody sibling of :meth:`_admit_dcr_bridge_delegate`: the session token seals
+        no upstream credential (those are vaulted per user and resolved at egress), so this
+        admits identity only and injects no per-server header. The token's signature proves
+        the user signed in when it was minted, but authorization is resolved fresh here, the
+        sealed ``user_id`` reloads the current user record through the SAME
+        :meth:`_reload_admitted_user` the bridge user-subject path uses, and the admitted
+        identity runs through the centralized policy gate, so the user's present team, org,
+        budget, and SCIM state gate the request rather than a snapshot frozen at mint time.
+
+        Fails closed with the aggregate ``invalid_token`` challenge on an expired, tampered,
+        or foreign token, on a refresh token presented at the tool edge, and when the
+        referenced user is missing, deactivated, or rejected by the policy gate. The
+        pre-DB gates (size, IP, route allowlist) run first, mirroring the bridge arm and the
+        standard pipeline, so a caller blocked by IP or route is turned away before any
+        crypto or DB read."""
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
+            NotSessionBearer,
+            SessionBearerAdmitted,
+            SessionBearerInvalid,
+            resolve_session_bearer,
+            session_keys_from_master_key,
+        )
+        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 = session_keys_from_master_key(master_key)
+        result = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc))
+        match result:
+            case SessionBearerAdmitted():
+                admitted = await MCPRequestHandler._reload_admitted_user(result.principal.user_id)
+                await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route)
+                return admitted
+            case SessionBearerInvalid():
+                raise _aggregate_gateway_dcr_challenge(request, invalid_token=True)
+            case NotSessionBearer():
+                # is_session_bearer_shaped gated entry, so a non-session bearer here means a
+                # session-shaped-but-empty value; fail closed with the same challenge.
+                raise _aggregate_gateway_dcr_challenge(request, invalid_token=True)
+            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

tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py

@@ -6080,3 +6080,165 @@ class TestAggregateGatewayDcrChallenge:
             with pytest.raises(ProxyException) as exc_info:
                 await MCPRequestHandler.process_mcp_request(self._scope())
         assert str(exc_info.value.code) == "500"
+
+
+@pytest.mark.asyncio
+class TestGatewaySessionAdmission:
+    """The aggregate /mcp session-bearer admission arm (mcp_gateway_dcr). A valid session
+    token admits under the LIVE litellm user it references; an invalid/expired/refresh/foreign
+    token fails closed with the aggregate invalid_token challenge; the arm fires ONLY at the
+    aggregate scope with the flag on, never for named servers or per-server flows."""
+
+    _MASTER_KEY = "sk-gateway-session-admission-master-key"
+    _FLAG = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled"
+
+    def _session_bearer(self, user_id="sso-user-42", client_id="llm_dcrc_abc"):
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
+            session_keys_from_master_key,
+        )
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
+            SessionPrincipal,
+            mint_session_token,
+            mint_session_refresh_token,
+        )
+
+        keys = session_keys_from_master_key(self._MASTER_KEY)
+        principal = SessionPrincipal(user_id=user_id, client_id=client_id)
+        return mint_session_token, mint_session_refresh_token, principal, keys
+
+    def _access_token(self, **kw):
+        from datetime import datetime, timezone
+
+        mint, _refresh, principal, keys = self._session_bearer(**kw)
+        return mint(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
+
+    def _scope(self, bearer, path="/mcp", extra_headers=()):
+        return {
+            "type": "http",
+            "method": "POST",
+            "path": path,
+            "headers": [(b"host", b"testserver"), (b"authorization", f"Bearer {bearer}".encode()), *extra_headers],
+        }
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _patch_user_reload(*, user_id, active=True):
+        get_user_object = AsyncMock(
+            return_value=MagicMock(
+                user_id=user_id,
+                metadata={"scim_active": active} if not active else {"scim_active": True},
+                user_role=None,
+                object_permission=None,
+                object_permission_id=None,
+            )
+        )
+        with (
+            patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object),
+            patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
+        ):
+            yield get_user_object
+
+    async def test_valid_session_admits_under_live_user_at_aggregate_scope(self):
+        token = self._access_token(user_id="sso-user-42")
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+            patch(
+                "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+                new_callable=AsyncMock,
+            ) as mock_auth,
+            self._patch_user_reload(user_id="sso-user-42") as get_user_object,
+        ):
+            auth_result, _h, _servers, mcp_server_auth_headers, _o, _r = await MCPRequestHandler.process_mcp_request(
+                self._scope(token)
+            )
+        assert get_user_object.await_args.kwargs["user_id"] == "sso-user-42"
+        assert auth_result.user_id == "sso-user-42"
+        mock_auth.assert_not_called()
+        # Identity-only admission injects no per-server upstream credential (unlike the
+        # bridge envelope arm); the headers dict is whatever the request carried, here empty.
+        assert not mcp_server_auth_headers
+
+    async def test_expired_session_fails_closed_with_invalid_token_challenge(self):
+        from datetime import datetime, timezone
+
+        mint, _refresh, principal, keys = self._session_bearer()
+        token = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await MCPRequestHandler.process_mcp_request(self._scope(token))
+        assert exc_info.value.status_code == 401
+        assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"]
+
+    async def test_tampered_session_fails_closed(self):
+        token = self._access_token()
+        tampered = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb")
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await MCPRequestHandler.process_mcp_request(self._scope(tampered))
+        assert exc_info.value.status_code == 401
+
+    async def test_refresh_token_is_not_admitted_at_the_tool_edge(self):
+        from datetime import datetime, timezone
+
+        _mint, refresh, principal, keys = self._session_bearer()
+        refresh_token = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await MCPRequestHandler.process_mcp_request(self._scope(refresh_token))
+        assert exc_info.value.status_code == 401
+
+    async def test_foreign_key_session_fails_closed(self):
+        token = self._access_token()
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", "sk-a-totally-different-master-key"),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await MCPRequestHandler.process_mcp_request(self._scope(token))
+        assert exc_info.value.status_code == 401
+
+    async def test_arm_does_not_fire_when_flag_off(self):
+        """Flag off: a session-shaped bearer is treated as an ordinary bearer and hits the
+        oauth2 arm, which validates it as a litellm credential and fails it there (not the
+        session arm). Proven by user_api_key_auth being called, unlike the flag-on path."""
+        token = self._access_token()
+        with (
+            patch(self._FLAG, return_value=False),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+            patch(
+                "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+                new_callable=AsyncMock,
+                side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401),
+            ) as mock_auth,
+        ):
+            with pytest.raises((HTTPException, ProxyException)):
+                await MCPRequestHandler.process_mcp_request(self._scope(token))
+        mock_auth.assert_called_once()
+
+    async def test_arm_does_not_fire_for_named_server(self):
+        """A session-shaped bearer aimed at a named server (path scope) does not enter the
+        aggregate arm; it is treated as an ordinary bearer on that server."""
+        token = self._access_token()
+        with (
+            patch(self._FLAG, return_value=True),
+            patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
+            patch(
+                "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+                new_callable=AsyncMock,
+                side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401),
+            ) as mock_auth,
+        ):
+            with pytest.raises((HTTPException, ProxyException)):
+                await MCPRequestHandler.process_mcp_request(self._scope(token, path="/mcp/github"))
+        mock_auth.assert_called_once()