From b82fb75292432d2b37e45c589d486331d212ebbe Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:51:06 -0700 Subject: [PATCH] feat(mcp): union team-inherited MCP grants across all of a user's teams for keyless admission --- .../mcp_server/auth/user_api_key_auth_mcp.py | 95 +++++++++++++++++-- .../auth/test_user_api_key_auth_mcp.py | 94 ++++++++++++++++++ 2 files changed, 179 insertions(+), 10 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 0ef7416b677..4965b2da1ed 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,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: 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 33e78027872..25f7feb099b 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 @@ -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") + ) == []