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 1d1a8fc6fcac..d4dbe5551115 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 @@ -126,6 +127,21 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +MCP_ADMITTED_USER_SUBJECT_METADATA = "mcp_admitted_user_subject" +"""Marker key stamped into ``UserAPIKeyAuth.metadata`` by ``_reload_admitted_user`` for a +subject admitted keyless through the gateway session / bridge user path. It is what lets +``_team_ids_for_mcp_grant`` union the user's teams for exactly those admissions without also +broadening JWT auth, which produces a structurally identical keyless auth.""" + + +def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth) -> bool: + """True when this auth is a keyless subject admitted by the gateway session / bridge user + path (stamped at admission), as opposed to a JWT or other keyless auth that merely lacks a + ``team_id``.""" + metadata = user_api_key_auth.metadata + return isinstance(metadata, dict) and metadata.get(MCP_ADMITTED_USER_SUBJECT_METADATA) is True + + 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 @@ -807,6 +823,7 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: org_id=user_object.organization_id, object_permission=object_permission, object_permission_id=user_object.object_permission_id, + metadata={MCP_ADMITTED_USER_SUBJECT_METADATA: True}, ) @staticmethod @@ -1590,10 +1607,91 @@ async def _get_allowed_mcp_servers_for_key( @staticmethod async def _get_allowed_mcp_servers_for_team( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: - """ - Get allowed MCP servers for a team. + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> list[str]: + """ + Get allowed MCP servers a caller inherits from team membership. + + For a caller with a ``team_id`` (every key-based caller) the result is that one team's + grants, byte-identical to before this method learned about multiple teams. For a + subject admitted keyless through the gateway DCR session or bridge user path, 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). The union is gated on the admission marker + ``_team_ids_for_mcp_grant`` checks, NOT on ``api_key is None``, so JWT auth (also + keyless, also possibly team-less) keeps its prior behavior and is not silently + broadened. + """ + 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 caller with an explicit ``team_id`` (every key-based caller, and any auth that pins + a team) uses that single team, so key auth is byte-identical. The fan-out to the + user's full team list happens ONLY for a subject admitted keyless through the gateway + session or bridge user path, which ``_reload_admitted_user`` stamps with + ``MCP_ADMITTED_USER_SUBJECT_METADATA``. Gating on that positive marker rather than on + ``api_key is None`` is deliberate: JWT auth also produces a keyless ``user_id`` auth + with no ``team_id``, and it must keep its prior behavior (no team-inherited grants) + rather than silently gaining the union across every team the user belongs to. 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 not user_api_key_auth.user_id or not _is_mcp_admitted_user_subject(user_api_key_auth): + 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, @@ -1617,17 +1715,15 @@ async def _get_allowed_mcp_servers_for_team( 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 94b52153a4dc..4d8ae8cb14ac 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 @@ -6218,3 +6218,119 @@ async def test_arm_does_not_fire_for_named_server(self): 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 + + @staticmethod + def _admitted_subject(user_id): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCP_ADMITTED_USER_SUBJECT_METADATA, + ) + + return UserAPIKeyAuth(user_id=user_id, api_key=None, metadata={MCP_ADMITTED_USER_SUBJECT_METADATA: True}) + + 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 = self._admitted_subject("sso-user") + 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 = self._admitted_subject("lonely-user") + 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 subject admitted by the gateway/bridge path (marked), no team -> resolved from record + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant(self._admitted_subject("u")) == ["t2", "t3"] + # keyless, no user_id -> nothing + assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == [] + # keyless with a user_id but NOT admission-marked (JWT auth) -> nothing (unchanged behavior) + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key=None, user_id="jwt-user") + ) == [] + # UI sentinel -> nothing + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key="sk", team_id=UI_TEAM_ID, user_id="u") + ) == [] + + async def test_jwt_keyless_user_without_team_claim_does_not_union(self): + """Regression for the review finding: a JWT-authenticated caller is also keyless with a + user_id and (with no team claim) no team_id, but it is NOT admission-marked, so it must + keep its prior behavior of inheriting no team grants rather than silently gaining the + union across every team the user belongs to.""" + teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])} + jwt_auth = UserAPIKeyAuth(user_id="jwt-user", api_key=None) # no admission marker + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(jwt_auth) + assert result == []