feat(mcp): union team-inherited MCP grants across all of a user's teams for keyless admission

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

@@ -1,3 +1,4 @@
+import asyncio
 import re
 from datetime import datetime, timezone
 from typing import Dict, List, Optional, Set, Tuple, cast
@@ -1593,10 +1594,86 @@ class MCPRequestHandler:
 
     @staticmethod
     async def _get_allowed_mcp_servers_for_team(
-        user_api_key_auth: Optional[UserAPIKeyAuth] = None,
-    ) -> List[str]:
+        user_api_key_auth: UserAPIKeyAuth | None = None,
+    ) -> list[str]:
         """
-        Get allowed MCP servers for a team.
+        Get allowed MCP servers a caller inherits from team membership.
+
+        For a key-based caller the ``team_id`` on the auth is the one team, and the result
+        is that team's grants (byte-identical to before this method learned about multiple
+        teams). For a user-subject caller admitted WITHOUT a key (the gateway DCR session
+        bearer and the bridge user-envelope, which carry a ``user_id`` and no ``api_key`` or
+        ``team_id``), a ``UserAPIKeyAuth`` can only pin one team while the user may belong to
+        many, so the inherited grant is the UNION across every team the user belongs to.
+        Without this a signed-in user would see only servers granted to them directly and
+        none granted through their teams, which is how servers are meant to be shared
+        (assign teams, not individuals). Key-based auth never enters the union branch, so its
+        access is unchanged.
+        """
+        team_ids = await MCPRequestHandler._team_ids_for_mcp_grant(user_api_key_auth)
+        if not team_ids:
+            return []
+        per_team = await asyncio.gather(
+            *(
+                MCPRequestHandler._allowed_mcp_servers_for_single_team(team_id, user_api_key_auth)
+                for team_id in team_ids
+            )
+        )
+        return list({server for servers in per_team for server in servers})
+
+    @staticmethod
+    async def _team_ids_for_mcp_grant(user_api_key_auth: UserAPIKeyAuth | None) -> list[str]:
+        """The team ids whose MCP grants a caller inherits.
+
+        A key-based caller (``api_key`` set) or any caller with an explicit ``team_id`` uses
+        that single team, so key auth is unchanged. Only a keyless user-subject caller (no
+        ``api_key``, no ``team_id``, a ``user_id``) fans out to the user's full team list,
+        resolved once from the live user record. The ``UI_TEAM_ID`` sentinel resolves to no
+        teams exactly as before."""
+        if user_api_key_auth is None:
+            return []
+        if user_api_key_auth.team_id:
+            return [] if user_api_key_auth.team_id == UI_TEAM_ID else [user_api_key_auth.team_id]
+        if user_api_key_auth.api_key is not None or not user_api_key_auth.user_id:
+            return []
+        return await MCPRequestHandler._resolve_user_team_ids(user_api_key_auth.user_id, user_api_key_auth)
+
+    @staticmethod
+    async def _resolve_user_team_ids(user_id: str, user_api_key_auth: UserAPIKeyAuth) -> list[str]:
+        """The distinct team ids a user belongs to, from the live user record. Returns [] on
+        no DB, a missing user, or any resolution failure so a lookup blip narrows access
+        rather than raising; the caller's direct grants still apply."""
+        from litellm.proxy.auth.auth_checks import get_user_object
+        from litellm.proxy.proxy_server import (
+            prisma_client,
+            proxy_logging_obj,
+            user_api_key_cache,
+        )
+
+        if prisma_client is None:
+            return []
+        try:
+            user_object = await get_user_object(
+                user_id=user_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                user_id_upsert=False,
+                parent_otel_span=user_api_key_auth.parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+        except Exception as e:  # noqa: BLE001  # a team-resolution blip narrows access, never raises
+            verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {str(e)}")
+            return []
+        if user_object is None or not user_object.teams:
+            return []
+        return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID))
+
+    @staticmethod
+    async def _allowed_mcp_servers_for_single_team(
+        team_id: str,
+        user_api_key_auth: UserAPIKeyAuth | None,
+    ) -> list[str]:
+        """Allowed MCP servers granted by ONE team.
 
         Unions two sources:
         - Legacy team.object_permission (mcp_servers, mcp_access_groups,
@@ -1620,17 +1697,15 @@ class MCPRequestHandler:
                 user_api_key_cache,
             )
 
-            if user_api_key_auth is None or not user_api_key_auth.team_id or prisma_client is None:
+            if not team_id or team_id == UI_TEAM_ID or prisma_client is None:
                 return []
 
-            if user_api_key_auth.team_id == UI_TEAM_ID:
-                return []
-
-            team_obj: Optional[LiteLLM_TeamTable] = await get_team_object(
-                team_id=user_api_key_auth.team_id,
+            parent_otel_span = user_api_key_auth.parent_otel_span if user_api_key_auth is not None else None
+            team_obj: LiteLLM_TeamTable | None = await get_team_object(
+                team_id=team_id,
                 prisma_client=prisma_client,
                 user_api_key_cache=user_api_key_cache,
-                parent_otel_span=user_api_key_auth.parent_otel_span,
+                parent_otel_span=parent_otel_span,
                 proxy_logging_obj=proxy_logging_obj,
             )
             if team_obj is None:

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

@@ -6235,3 +6235,97 @@ class TestGatewaySessionAdmission:
             with pytest.raises((HTTPException, ProxyException)):
                 await MCPRequestHandler.process_mcp_request(self._scope(token, path="/mcp/github"))
         mock_auth.assert_called_once()
+
+
+@pytest.mark.asyncio
+class TestUserSubjectTeamUnion:
+    """_get_allowed_mcp_servers_for_team unions across ALL a user's teams for a keyless
+    user-subject caller (the gateway DCR session bearer and bridge user-envelope), while a
+    key-based caller keeps its single-team behavior byte-identically."""
+
+    def _team(self, team_id, mcp_servers):
+        from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable
+
+        return LiteLLM_TeamTable(
+            team_id=team_id,
+            access_group_ids=[],
+            object_permission=LiteLLM_ObjectPermissionTable(
+                object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers
+            ),
+        )
+
+    @contextlib.contextmanager
+    def _patch(self, *, teams_by_id, user_teams=None):
+        async def _get_team_object(team_id, **kw):
+            return teams_by_id.get(team_id)
+
+        async def _get_user_object(user_id, **kw):
+            return MagicMock(user_id=user_id, teams=user_teams or [])
+
+        with (
+            patch("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object),
+            patch("litellm.proxy.auth.auth_checks.get_user_object", _get_user_object),
+            patch("litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", AsyncMock(return_value=[])),
+            patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
+            patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
+        ):
+            yield
+
+    async def test_keyless_user_unions_servers_across_all_their_teams(self):
+        teams = {"team-a": self._team("team-a", ["srv1", "srv2"]), "team-b": self._team("team-b", ["srv2", "srv3"])}
+        auth = UserAPIKeyAuth(user_id="sso-user", api_key=None)
+        with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
+        assert set(result) == {"srv1", "srv2", "srv3"}
+
+    async def test_key_based_caller_uses_single_team_only(self):
+        """A key-based caller (api_key set) with a team_id sees ONLY that team, even though the
+        same user belongs to other teams: key auth must be byte-identical to before."""
+        teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2", "srv3"])}
+        auth = UserAPIKeyAuth(user_id="sso-user", api_key="sk-hash", team_id="team-a")
+        with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
+        assert set(result) == {"srv1"}
+
+    async def test_keyless_user_with_explicit_team_id_uses_that_team_only(self):
+        """A keyless caller that already pins a team_id (not the user-subject fan-out shape)
+        resolves only that team; the union is strictly for the no-team-id user-subject case."""
+        teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])}
+        auth = UserAPIKeyAuth(user_id="sso-user", api_key=None, team_id="team-a")
+        with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
+        assert set(result) == {"srv1"}
+
+    async def test_keyless_user_with_no_teams_gets_nothing_from_teams(self):
+        auth = UserAPIKeyAuth(user_id="lonely-user", api_key=None)
+        with self._patch(teams_by_id={}, user_teams=[]):
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
+        assert result == []
+
+    async def test_ui_session_team_id_still_resolves_to_nothing(self):
+        from litellm.proxy._types import UI_TEAM_ID
+
+        auth = UserAPIKeyAuth(user_id="dash-user", api_key="sk-hash", team_id=UI_TEAM_ID)
+        with self._patch(teams_by_id={}, user_teams=["team-a"]):
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
+        assert result == []
+
+    async def test_team_ids_helper_gates_on_shape(self):
+        from litellm.proxy._types import UI_TEAM_ID
+
+        # key-based with team -> that team
+        assert await MCPRequestHandler._team_ids_for_mcp_grant(
+            UserAPIKeyAuth(api_key="sk", team_id="t1", user_id="u")
+        ) == ["t1"]
+        # keyless user-subject, no team -> resolved from user record
+        with self._patch(teams_by_id={}, user_teams=["t2", "t3"]):
+            assert await MCPRequestHandler._team_ids_for_mcp_grant(
+                UserAPIKeyAuth(api_key=None, user_id="u")
+            ) == ["t2", "t3"]
+        # keyless, no user_id -> nothing
+        assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == []
+        # UI sentinel -> nothing
+        assert await MCPRequestHandler._team_ids_for_mcp_grant(
+            UserAPIKeyAuth(api_key="sk", team_id=UI_TEAM_ID, user_id="u")
+        ) == []