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 92ebd248af0..238b174dd4f 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 @@ -9,6 +9,7 @@ from starlette.types import Scope from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy._experimental.mcp_server.permission_grant import ( + AllTeamServers, NoServers, parse_mcp_server_grant, ) @@ -647,12 +648,11 @@ class MCPRequestHandler: user_api_key_auth ) ) + key_grant = parse_mcp_server_grant(allowed_mcp_servers_for_key) # The key explicitly opted out of every MCP server. This overrides # team inheritance and additive grants (mirrors no-default-models). - if isinstance( - parse_mcp_server_grant(allowed_mcp_servers_for_key), NoServers - ): + if isinstance(key_grant, NoServers): return [] allowed_mcp_servers_for_team = ( @@ -677,7 +677,11 @@ class MCPRequestHandler: has_lower_level_mcp_restrictions = bool(key_set or team_set or grants_set) # 1. Key/team ceiling. An empty set means "this level does not restrict". - if not team_set: + if isinstance(key_grant, AllTeamServers): + # The key explicitly requested its team's full grant: resolve to + # exactly the team set, failing closed when the team grants nothing. + base = team_set + elif not team_set: base = key_set # no team restriction elif not key_set: base = team_set # key has no own perms → inherits team @@ -1071,13 +1075,13 @@ class MCPRequestHandler: if key_object_permission is None: return [] - # Sentinel opt-out: surface it unexpanded so the caller can short-circuit - # to zero servers instead of inheriting the team. - if isinstance( - parse_mcp_server_grant(key_object_permission.mcp_servers or []), - NoServers, - ): + # Sentinels are surfaced unexpanded so the caller can short-circuit: + # no-mcp-servers -> zero servers, all-team-mcps -> the team's own set. + key_grant = parse_mcp_server_grant(key_object_permission.mcp_servers or []) + if isinstance(key_grant, NoServers): return [SpecialMCPServerNames.no_mcp_servers.value] + if isinstance(key_grant, AllTeamServers): + return [SpecialMCPServerNames.all_team_mcp_servers.value] # Permission entries may be server_ids OR names/aliases — expand to ids. direct_mcp_servers = global_mcp_server_manager.expand_permission_list( diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index bc1068d77d0..2571982328f 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -17,10 +17,13 @@ from litellm.proxy._types import ( MCPEnvVarScope, MCPSubmissionsSummary, NewMCPServerRequest, - SpecialMCPServerName, UpdateMCPServerRequest, UserAPIKeyAuth, ) +from litellm.proxy._experimental.mcp_server.permission_grant import ( + AllTeamServers, + parse_mcp_server_grant, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _get_salt_key, decrypt_value_helper, @@ -494,9 +497,10 @@ async def get_all_mcp_servers_for_user( ) mcp_server_ids.update(token_mcp_servers) - # check for special team membership + # all-team-mcps expands to the key's team servers, so discovery matches + # what the access resolver grants at request time. if ( - SpecialMCPServerName.all_team_servers in mcp_server_ids + isinstance(parse_mcp_server_grant(token_mcp_servers), AllTeamServers) and user.team_id is not None ): team_mcp_servers = await get_mcp_servers_by_team( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e28d0ef578a..7221e11b838 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -53,6 +53,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.permission_grant import ( AllServers, + AllTeamServers, ExplicitServers, NoServers, parse_mcp_server_grant, @@ -4190,6 +4191,8 @@ class MCPServerManager: match grant: case AllServers(): return list(registry.keys()) + case AllTeamServers(): + return [SpecialMCPServerNames.all_team_mcp_servers.value] case NoServers(): return [SpecialMCPServerNames.no_mcp_servers.value] case ExplicitServers(identifiers=explicit): diff --git a/litellm/proxy/_experimental/mcp_server/permission_grant.py b/litellm/proxy/_experimental/mcp_server/permission_grant.py index 0b7778cf2a6..95c02e29af6 100644 --- a/litellm/proxy/_experimental/mcp_server/permission_grant.py +++ b/litellm/proxy/_experimental/mcp_server/permission_grant.py @@ -16,6 +16,7 @@ MCP_GRANT_SENTINELS: frozenset[str] = frozenset( { SpecialMCPServerNames.no_mcp_servers.value, SpecialMCPServerNames.all_proxy_mcp_servers.value, + SpecialMCPServerNames.all_team_mcp_servers.value, } ) @@ -25,6 +26,16 @@ class AllServers: """Grant every MCP server on the proxy, including ones added later.""" +@dataclass(frozen=True, slots=True) +class AllTeamServers: + """Grant every MCP server the holder's team can reach. + + Key-only: it is resolved against the team at request time, so it tracks the + team grant and fails closed (zero servers) when the key has no team or the + team grants nothing. + """ + + @dataclass(frozen=True, slots=True) class NoServers: """Block all MCP servers.""" @@ -37,18 +48,21 @@ class ExplicitServers: identifiers: frozenset[str] -MCPServerGrant = Union[AllServers, NoServers, ExplicitServers] +MCPServerGrant = Union[AllServers, AllTeamServers, NoServers, ExplicitServers] def parse_mcp_server_grant(raw: Iterable[str]) -> MCPServerGrant: """Resolve precedence once, most-restrictive first. - ``no-mcp-servers`` blocks everything and beats ``all-proxy-mcps``; absent any - sentinel the entries are explicit identifiers. + ``no-mcp-servers`` blocks everything; ``all-team-mcps`` caps to the team and + so beats the broader ``all-proxy-mcps``; absent any sentinel the entries are + explicit identifiers. """ values = frozenset(raw) if SpecialMCPServerNames.no_mcp_servers.value in values: return NoServers() + if SpecialMCPServerNames.all_team_mcp_servers.value in values: + return AllTeamServers() if SpecialMCPServerNames.all_proxy_mcp_servers.value in values: return AllServers() return ExplicitServers(values) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d3279b29a06..a9d08f6d8d0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1229,11 +1229,6 @@ from litellm.models.team import LiteLLM_ModelTable as LiteLLM_ModelTable # noqa # MCP Types -class SpecialMCPServerName(str, enum.Enum): - all_team_servers = "all-team-mcpservers" - all_proxy_servers = "all-proxy-mcpservers" - - class MCPApprovalStatus(str, enum.Enum): pending_review = "pending_review" active = "active" @@ -2986,6 +2981,7 @@ class SpecialModelNames(enum.Enum): class SpecialMCPServerNames(enum.Enum): no_mcp_servers = "no-mcp-servers" all_proxy_mcp_servers = "all-proxy-mcps" + all_team_mcp_servers = "all-team-mcps" class SpecialProxyStrings(enum.Enum): diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4ab308990b5..cdc9abf26b4 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -163,11 +163,13 @@ if MCP_AVAILABLE: MCPUserEnvVarsStatus, NewMCPServerRequest, RejectMCPServerRequest, - SpecialMCPServerName, UpdateMCPServerRequest, UserAPIKeyAuth, UserMCPManagementMode, ) + from litellm.proxy._experimental.mcp_server.permission_grant import ( + MCP_GRANT_SENTINELS, + ) from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, user_api_key_auth, @@ -1451,11 +1453,9 @@ if MCP_AVAILABLE: }, ) - # Block reserved special server IDs - if ( - SpecialMCPServerName.all_team_servers == payload.server_id - or SpecialMCPServerName.all_proxy_servers == payload.server_id - ): + # Block reserved special server IDs — a grant sentinel must never be + # usable as a real server_id. + if payload.server_id in MCP_GRANT_SENTINELS: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={ diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 7f5b0176c3e..a2807dea96e 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._experimental.mcp_server.permission_grant import ( MCP_GRANT_SENTINELS, AllServers, + AllTeamServers, ExplicitServers, MCPServerGrant, NoServers, @@ -372,7 +373,10 @@ async def _resolve_team_allowed_mcp_servers( """ grant = parse_mcp_server_grant(team_object_permission.mcp_servers or []) match grant: - case NoServers(): + case NoServers() | AllTeamServers(): + # no-mcp-servers blocks everything; all-team-mcps is key-only and + # meaningless on a team (a team cannot grant "its own team's + # servers"), so a misconfigured team fails closed to no servers. return set() case AllServers(): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -697,6 +701,17 @@ async def validate_key_mcp_servers_against_team( ) }, ) + if isinstance(key_grant, AllTeamServers) and team_obj is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + "all-team-mcps can only be granted to a key that belongs to a " + "team. It resolves to the team's MCP servers at request time, so " + "a teamless key would reach zero servers." + ) + }, + ) if not any((requested_servers, requested_access_groups, requested_toolsets)): return object_permission 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 5b7c4948c56..83ea385051b 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 @@ -280,6 +280,65 @@ class TestMCPRequestHandler: assert result == [SpecialMCPServerNames.no_mcp_servers.value] + @pytest.mark.parametrize( + "team_servers,expected", + [ + (["team_server1", "team_server2"], ["team_server1", "team_server2"]), + # Fail closed: a team that grants nothing leaves an all-team-mcps key + # with zero servers rather than inheriting the whole proxy. + ([], []), + ], + ) + async def test_all_team_mcps_sentinel_resolves_to_team_set( + self, team_servers, expected + ): + """A key scoped to all-team-mcps resolves to exactly its team's servers, + tracking the team grant and never leaking the sentinel marker.""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [ + SpecialMCPServerNames.all_team_mcp_servers.value + ] + + with patch.object( + MCPRequestHandler, + "_get_key_object_permission", + return_value=key_object_permission, + ), patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=team_servers, + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert sorted(result) == sorted(expected) + assert SpecialMCPServerNames.all_team_mcp_servers.value not in result + + async def test_get_allowed_mcp_servers_for_key_returns_all_team_marker(self): + """_get_allowed_mcp_servers_for_key surfaces the all-team-mcps sentinel + unexpanded so the caller can map it onto the team set, ignoring any other + entries on the key.""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [ + SpecialMCPServerNames.all_team_mcp_servers.value, + "some-other-server", + ] + + with patch.object( + MCPRequestHandler, + "_get_key_object_permission", + return_value=key_object_permission, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key( + user_api_key_auth + ) + + assert result == [SpecialMCPServerNames.all_team_mcp_servers.value] + async def test_permission_inheritance_edge_cases(self): """Test edge cases in permission inheritance""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 1d87af54c8c..0e8338f2260 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -209,6 +209,32 @@ class TestMCPServerManager: sentinel = SpecialMCPServerNames.no_mcp_servers.value assert manager.expand_permission_list([sentinel]) == [sentinel] + async def test_expand_permission_list_all_team_mcps_passthrough(self): + """all-team-mcps cannot be expanded here (it needs team context, which + the registry-only manager lacks), so it is surfaced unexpanded for the + key resolver to map onto the team set. Pin that the literal survives and + is never confused with a registered server.""" + manager = MCPServerManager() + await manager.add_server( + LiteLLM_MCPServerTable( + server_id="srv-x", + alias="srv-x", + description="", + url=None, + transport=MCPTransport.stdio, + command="python", + args=["-m", "server"], + env={}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ) + sentinel = SpecialMCPServerNames.all_team_mcp_servers.value + assert manager.expand_permission_list([sentinel]) == [sentinel] + # Mixed with a real id the sentinel still dominates (most-restrictive + # grant), so the named server is NOT additionally resolved. + assert manager.expand_permission_list([sentinel, "srv-x"]) == [sentinel] + async def test_create_mcp_client_stdio(self): """Test creating MCP client for stdio transport""" manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py index 430d9bcfb5b..f21a4165f0e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py @@ -4,6 +4,7 @@ import pytest from litellm.proxy._experimental.mcp_server.permission_grant import ( AllServers, + AllTeamServers, ExplicitServers, MCPServerGrant, NoServers, @@ -13,6 +14,7 @@ from litellm.proxy._types import SpecialMCPServerNames NO_MCP = SpecialMCPServerNames.no_mcp_servers.value ALL_PROXY = SpecialMCPServerNames.all_proxy_mcp_servers.value +ALL_TEAM = SpecialMCPServerNames.all_team_mcp_servers.value def test_concrete_only_returns_explicit_servers() -> None: @@ -32,11 +34,31 @@ def test_no_mcp_sentinel_returns_no_servers() -> None: assert parse_mcp_server_grant([NO_MCP]) == NoServers() +def test_all_team_sentinel_returns_all_team_servers() -> None: + assert parse_mcp_server_grant([ALL_TEAM]) == AllTeamServers() + + +def test_all_team_sentinel_dominates_concrete_ids() -> None: + assert parse_mcp_server_grant([ALL_TEAM, "srv-a"]) == AllTeamServers() + + def test_block_all_wins_over_all_proxy_regardless_of_order() -> None: assert parse_mcp_server_grant([NO_MCP, ALL_PROXY]) == NoServers() assert parse_mcp_server_grant([ALL_PROXY, NO_MCP]) == NoServers() +def test_block_all_wins_over_all_team_regardless_of_order() -> None: + assert parse_mcp_server_grant([NO_MCP, ALL_TEAM]) == NoServers() + assert parse_mcp_server_grant([ALL_TEAM, NO_MCP]) == NoServers() + + +def test_all_team_beats_all_proxy_regardless_of_order() -> None: + # all-team caps to the team and so is the more restrictive grant; it wins + # over the broader all-proxy when both somehow appear. + assert parse_mcp_server_grant([ALL_PROXY, ALL_TEAM]) == AllTeamServers() + assert parse_mcp_server_grant([ALL_TEAM, ALL_PROXY]) == AllTeamServers() + + def test_empty_list_returns_empty_explicit_servers() -> None: assert parse_mcp_server_grant([]) == ExplicitServers(frozenset()) @@ -49,15 +71,18 @@ def test_grants_are_frozen() -> None: def test_grants_equal_by_value() -> None: assert AllServers() == AllServers() + assert AllTeamServers() == AllTeamServers() assert NoServers() == NoServers() assert ExplicitServers(frozenset({"a"})) == ExplicitServers(frozenset({"a"})) assert AllServers() != NoServers() + assert AllServers() != AllTeamServers() def test_union_alias_admits_each_variant() -> None: grants: tuple[MCPServerGrant, ...] = ( AllServers(), + AllTeamServers(), NoServers(), ExplicitServers(frozenset({"a"})), ) - assert len(grants) == 3 + assert len(grants) == 4 diff --git a/tests/test_litellm/proxy/db/mcp_server/test_db.py b/tests/test_litellm/proxy/db/mcp_server/test_db.py index ff6400ac5b7..b03042b9bd1 100644 --- a/tests/test_litellm/proxy/db/mcp_server/test_db.py +++ b/tests/test_litellm/proxy/db/mcp_server/test_db.py @@ -1,6 +1,7 @@ import json import os import sys +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,8 +9,57 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.proxy._experimental.mcp_server.db import get_mcp_servers_by_team +from litellm.proxy._experimental.mcp_server.db import ( + get_all_mcp_servers_for_user, + get_mcp_servers_by_team, +) +from litellm.proxy._types import SpecialMCPServerNames, UserAPIKeyAuth + +DB_MODULE = "litellm.proxy._experimental.mcp_server.db" def test_fetch_mcp_servers_by_team(): assert True == True + + +@pytest.mark.asyncio +@patch(f"{DB_MODULE}.get_mcp_servers", new_callable=AsyncMock) +@patch(f"{DB_MODULE}.get_mcp_servers_by_team", new_callable=AsyncMock) +@patch(f"{DB_MODULE}.get_mcp_servers_by_verificationtoken", new_callable=AsyncMock) +async def test_get_all_mcp_servers_for_user_all_team_mcps_expands_to_team( + mock_token_servers, mock_team_servers, mock_get_servers +): + """A key scoped to all-team-mcps surfaces its team's servers in discovery, + matching the access resolver. Pins that the live sentinel value drives the + expansion (the removed legacy all-team-mcpservers value would not).""" + mock_token_servers.return_value = [ + SpecialMCPServerNames.all_team_mcp_servers.value + ] + mock_team_servers.return_value = ["team-srv-1", "team-srv-2"] + mock_get_servers.return_value = [] + + user = UserAPIKeyAuth(api_key="k", user_id="u", team_id="team-1") + await get_all_mcp_servers_for_user(MagicMock(), user) + + mock_team_servers.assert_awaited_once() + queried_ids = set(mock_get_servers.call_args.args[1]) + assert {"team-srv-1", "team-srv-2"} <= queried_ids + + +@pytest.mark.asyncio +@patch(f"{DB_MODULE}.get_mcp_servers", new_callable=AsyncMock) +@patch(f"{DB_MODULE}.get_mcp_servers_by_team", new_callable=AsyncMock) +@patch(f"{DB_MODULE}.get_mcp_servers_by_verificationtoken", new_callable=AsyncMock) +async def test_get_all_mcp_servers_for_user_explicit_servers_skip_team_expansion( + mock_token_servers, mock_team_servers, mock_get_servers +): + """Without the all-team-mcps sentinel, discovery never pulls in the team's + servers — proving the expansion is gated on the sentinel, not unconditional.""" + mock_token_servers.return_value = ["explicit-srv"] + mock_get_servers.return_value = [] + + user = UserAPIKeyAuth(api_key="k", user_id="u", team_id="team-1") + await get_all_mcp_servers_for_user(MagicMock(), user) + + mock_team_servers.assert_not_awaited() + assert set(mock_get_servers.call_args.args[1]) == {"explicit-srv"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index f40904e234d..bdeea7023ee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( LitellmUserRoles, MCPTransport, NewMCPServerRequest, + SpecialMCPServerNames, UpdateMCPServerRequest, UserAPIKeyAuth, ) @@ -2590,6 +2591,49 @@ class TestUpdateMCPServer: assert result.alias == "Updated Test Server" +class TestAddMCPServerReservedIds: + """A grant sentinel must never be usable as a real MCP server_id — otherwise + a real server could shadow the permission sentinel and corrupt grant parsing. + Pins that the guard tracks the live SpecialMCPServerNames values.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "reserved_id", + [ + SpecialMCPServerNames.no_mcp_servers.value, + SpecialMCPServerNames.all_proxy_mcp_servers.value, + SpecialMCPServerNames.all_team_mcp_servers.value, + ], + ) + async def test_create_rejects_grant_sentinel_server_id(self, reserved_id): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + payload = NewMCPServerRequest( + server_id=reserved_id, + alias="echo", + url="https://echo.example.com/mcp", + transport=MCPTransport.http, + ) + admin = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ) + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + MagicMock(), + ), + pytest.raises(HTTPException) as exc_info, + ): + await add_mcp_server(payload=payload, user_api_key_dict=admin) + assert exc_info.value.status_code == 400 + + class TestAddMCPServerAtomicity: """A committed MCP server must survive a post-write registry refresh failure. diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index f5fa670676e..66bf4d36c40 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -1098,3 +1098,67 @@ async def test_validate_teamless_non_admin_all_proxy_raises( is_proxy_admin=False, ) assert exc_info.value.status_code == 403 + + +# ---- Tests for the all-team-mcps grant ---- + +ALL_TEAM = SpecialMCPServerNames.all_team_mcp_servers.value + + +def test_extract_requested_mcp_server_ids_excludes_all_team_sentinel(): + obj_perm = {"mcp_servers": [ALL_TEAM, "server-1"]} + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"} + + +@pytest.mark.asyncio +async def test_validate_key_all_team_in_team_allowed_and_preserved(): + """A key in a team may carry all-team-mcps: it resolves to the team grant at + request time, so the validator accepts it and leaves the sentinel intact + (so it keeps tracking the team) without enumerating servers.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + obj_perm = {"mcp_servers": [ALL_TEAM]} + result = await validate_key_mcp_servers_against_team( + object_permission=obj_perm, + team_obj=team_obj, + ) + assert result == obj_perm + assert obj_perm["mcp_servers"] == [ALL_TEAM] + + +@pytest.mark.asyncio +async def test_validate_key_all_team_teamless_raises(): + """all-team-mcps on a teamless key is rejected, not silently resolved to + zero: there is no team to track, so it is a meaningless grant.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [ALL_TEAM]}, + team_obj=None, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_validate_key_all_team_teamless_admin_still_raises(): + """Unlike all-proxy-mcps (which a proxy admin may assign to a teamless key), + all-team-mcps stays team-only even for an admin — it has no team to resolve + against.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [ALL_TEAM]}, + team_obj=None, + is_proxy_admin=True, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_resolve_team_allowed_mcp_servers_all_team_returns_empty(): + """A team carrying all-team-mcps is nonsensical (a team cannot grant 'its own + team's servers' recursively), so the team resolver fails closed to empty + rather than the registry.""" + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = [ALL_TEAM] + mock_perm.mcp_access_groups = ["group-a"] + mock_perm.mcp_tool_permissions = {"server-x": ["tool1"]} + result = await _resolve_team_allowed_mcp_servers(mock_perm) + assert result == set() diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json index 4490673f839..875564acba8 100644 --- a/ui/litellm-dashboard/eslint-metrics.json +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -1,5 +1,5 @@ { - "@typescript-eslint/no-explicit-any": 2030, + "@typescript-eslint/no-explicit-any": 2024, "complexity": 128, "max-depth": 61 } diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx index f1b33a3534d..0c601679dd7 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx @@ -3,7 +3,11 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import MCPServerSelector from "./MCPServerSelector"; -import { ALL_PROXY_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { + ALL_PROXY_MCPS_SENTINEL, + ALL_TEAM_MCPS_SENTINEL, + NO_MCP_SERVERS_SENTINEL, +} from "../mcp_tools/constants"; vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn(), @@ -45,15 +49,22 @@ const mockUseMCPServers = vi.mocked(useMCPServers); const mockUseMCPAccessGroups = vi.mocked(useMCPAccessGroups); const mockUseMCPToolsets = vi.mocked(useMCPToolsets); +const setupMcpHooks = (servers: { server_id: string; server_name: string }[]) => { + vi.clearAllMocks(); + mockUseMCPServers.mockReturnValue({ data: servers, isLoading: false } as unknown as ReturnType); + mockUseMCPAccessGroups.mockReturnValue({ + data: [], + isLoading: false, + } as unknown as ReturnType); + mockUseMCPToolsets.mockReturnValue({ + data: [], + isLoading: false, + } as unknown as ReturnType); +}; + describe("MCPServerSelector no-mcp-servers option", () => { beforeEach(() => { - vi.clearAllMocks(); - mockUseMCPServers.mockReturnValue({ - data: [{ server_id: "srv-1", server_name: "Server One" }], - isLoading: false, - } as any); - mockUseMCPAccessGroups.mockReturnValue({ data: [], isLoading: false } as any); - mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); + setupMcpHooks([{ server_id: "srv-1", server_name: "Server One" }]); }); const optionByValue = (value: string) => @@ -101,13 +112,7 @@ describe("MCPServerSelector no-mcp-servers option", () => { describe("MCPServerSelector all-proxy-mcps option", () => { beforeEach(() => { - vi.clearAllMocks(); - mockUseMCPServers.mockReturnValue({ - data: [{ server_id: "srv-1", server_name: "Server One" }], - isLoading: false, - } as any); - mockUseMCPAccessGroups.mockReturnValue({ data: [], isLoading: false } as any); - mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); + setupMcpHooks([{ server_id: "srv-1", server_name: "Server One" }]); }); const optionByValue = (value: string) => @@ -152,3 +157,71 @@ describe("MCPServerSelector all-proxy-mcps option", () => { expect(optionByValue(ALL_PROXY_MCPS_SENTINEL)?.disabled).toBe(false); }); }); + +describe("MCPServerSelector all-team-mcps option", () => { + beforeEach(() => { + setupMcpHooks([{ server_id: "srv-1", server_name: "Server One" }]); + }); + + const optionByValue = (value: string) => + Array.from(screen.getByTestId("mcp-select").querySelectorAll("option")).find( + (o) => (o as HTMLOptionElement).value === value, + ) as HTMLOptionElement | undefined; + + it("hides All Team MCPs when no team is selected, even with allowAllTeamMcps", () => { + renderWithProviders( + , + ); + expect(optionByValue(ALL_TEAM_MCPS_SENTINEL)).toBeUndefined(); + }); + + it("shows All Team MCPs once a team is selected", () => { + renderWithProviders( + , + ); + expect(optionByValue(ALL_TEAM_MCPS_SENTINEL)).toBeDefined(); + }); + + it("emits an exclusive sentinel when All Team MCPs is selected", async () => { + const onChange = vi.fn(); + renderWithProviders( + , + ); + await userEvent.selectOptions(screen.getByTestId("mcp-select"), [ALL_TEAM_MCPS_SENTINEL]); + + expect(onChange).toHaveBeenCalledWith({ servers: [ALL_TEAM_MCPS_SENTINEL], accessGroups: [], toolsets: [] }); + }); + + it("makes the three exclusive sentinels mutually exclusive", () => { + renderWithProviders( + , + ); + expect(optionByValue(ALL_TEAM_MCPS_SENTINEL)?.disabled).toBe(false); + expect(optionByValue(NO_MCP_SERVERS_SENTINEL)?.disabled).toBe(true); + expect(optionByValue("srv-1")?.disabled).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index c966fd61e28..fee616387d6 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -3,7 +3,11 @@ import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers" import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { Select } from "antd"; import React from "react"; -import { ALL_PROXY_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { + ALL_PROXY_MCPS_SENTINEL, + ALL_TEAM_MCPS_SENTINEL, + NO_MCP_SERVERS_SENTINEL, +} from "@/components/mcp_tools/constants"; interface MCPServerSelectorProps { onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; @@ -19,28 +23,34 @@ interface MCPServerSelectorProps { teamId?: string | null; allowNoMcpServers?: boolean; allowAllProxyMcps?: boolean; + allowAllTeamMcps?: boolean; } const TOOLSET_PREFIX = "toolset:"; -type MCPSelection = { servers: string[]; accessGroups: string[]; toolsets: string[] }; +type ExclusiveSentinel = { value: string; label: string; caption: string; color: string }; -// "No MCP Servers" and "All Proxy MCPs" are mutually exclusive with everything -// else. The most restrictive ("block all") wins, mirroring the backend -// precedence where the no-mcp-servers sentinel overrides every grant. -const resolveExclusiveSelection = ( - selected: string[], +// Each entry grants a whole class of servers, so it is mutually exclusive with +// every other pick. Ordered most-restrictive first to mirror the backend grant +// precedence (block-all > team > proxy); the first match wins. "All Team MCPs" +// only appears once a team is chosen, since it resolves against that team at +// request time and is meaningless on a teamless key. +const buildExclusiveSentinels = ( allowNoMcpServers: boolean, + allowAllTeamMcps: boolean, allowAllProxyMcps: boolean, -): MCPSelection | null => { - if (allowNoMcpServers && selected.includes(NO_MCP_SERVERS_SENTINEL)) { - return { servers: [NO_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }; - } - if (allowAllProxyMcps && selected.includes(ALL_PROXY_MCPS_SENTINEL)) { - return { servers: [ALL_PROXY_MCPS_SENTINEL], accessGroups: [], toolsets: [] }; - } - return null; -}; + teamId: string | null | undefined, +): ExclusiveSentinel[] => [ + ...(allowNoMcpServers + ? [{ value: NO_MCP_SERVERS_SENTINEL, label: "No MCP Servers", caption: "Block all", color: "#8c8c8c" }] + : []), + ...(allowAllTeamMcps && !!teamId + ? [{ value: ALL_TEAM_MCPS_SENTINEL, label: "All Team MCPs", caption: "Team's servers", color: "#13c2c2" }] + : []), + ...(allowAllProxyMcps + ? [{ value: ALL_PROXY_MCPS_SENTINEL, label: "All Proxy MCPs", caption: "Every server", color: "#1890ff" }] + : []), +]; const MCPServerSelector: React.FC = ({ onChange, @@ -52,6 +62,7 @@ const MCPServerSelector: React.FC = ({ teamId, allowNoMcpServers = false, allowAllProxyMcps = false, + allowAllTeamMcps = false, }) => { const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId); const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(); @@ -101,15 +112,21 @@ const MCPServerSelector: React.FC = ({ ...(value?.toolsets || []).map((id) => `${TOOLSET_PREFIX}${id}`), ]; - const hasNoMcpServersSelected = allowNoMcpServers && selectedValues.includes(NO_MCP_SERVERS_SENTINEL); - const hasAllProxyMcpsSelected = allowAllProxyMcps && selectedValues.includes(ALL_PROXY_MCPS_SENTINEL); - const hasExclusiveSelected = hasNoMcpServersSelected || hasAllProxyMcpsSelected; + const exclusiveSentinels = buildExclusiveSentinels( + allowNoMcpServers, + allowAllTeamMcps, + allowAllProxyMcps, + teamId, + ); + + const selectedExclusive = exclusiveSentinels.find((s) => selectedValues.includes(s.value)) ?? null; + const hasExclusiveSelected = selectedExclusive !== null; // Handle selection const handleChange = (selected: string[]) => { - const exclusive = resolveExclusiveSelection(selected, allowNoMcpServers, allowAllProxyMcps); + const exclusive = exclusiveSentinels.find((s) => selected.includes(s.value)); if (exclusive) { - onChange(exclusive); + onChange({ servers: [exclusive.value], accessGroups: [], toolsets: [] }); return; } const toolsetsSelected = selected @@ -121,38 +138,22 @@ const MCPServerSelector: React.FC = ({ onChange({ servers, accessGroups: accessGroupsSelected, toolsets: toolsetsSelected }); }; - const renderExclusiveOptions = () => [ - ...(allowAllProxyMcps - ? [ - -
- All Proxy MCPs - Every server -
- , - ] - : []), - ...(allowNoMcpServers - ? [ - -
- No MCP Servers - Block all -
- , - ] - : []), - ]; + const renderExclusiveOptions = () => + exclusiveSentinels.map((sentinel) => ( + +
+ {sentinel.label} + + {sentinel.caption} + +
+ + )); return (
@@ -168,7 +169,7 @@ const MCPServerSelector: React.FC = ({ style={{ width: "100%" }} disabled={disabled} filterOption={(input, option) => { - if (option?.value === NO_MCP_SERVERS_SENTINEL || option?.value === ALL_PROXY_MCPS_SENTINEL) return true; + if (exclusiveSentinels.some((s) => s.value === option?.value)) return true; const searchText = options.find((opt) => opt.value === option?.value)?.searchText || ""; return searchText.toLowerCase().includes(input.toLowerCase()); }} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/constants.ts b/ui/litellm-dashboard/src/components/mcp_tools/constants.ts index 4eb9b2f6558..6c8b87a5975 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/constants.ts +++ b/ui/litellm-dashboard/src/components/mcp_tools/constants.ts @@ -3,3 +3,6 @@ export const NO_MCP_SERVERS_SENTINEL = "no-mcp-servers"; // Must match the backend SpecialMCPServerNames.all_proxy_mcp_servers enum value. export const ALL_PROXY_MCPS_SENTINEL = "all-proxy-mcps"; + +// Must match the backend SpecialMCPServerNames.all_team_mcp_servers enum value. +export const ALL_TEAM_MCPS_SENTINEL = "all-team-mcps"; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 2a448d63300..057522b6479 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -32,7 +32,7 @@ import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/Budg import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { Team } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; -import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_TEAM_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; import NotificationsManager from "../molecules/notifications_manager"; import { @@ -1402,6 +1402,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp teamId={selectedCreateKeyTeam?.team_id ?? null} placeholder="Select MCP servers or access groups (optional)" allowNoMcpServers + allowAllTeamMcps /> @@ -1423,7 +1424,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp accessToken={accessToken} selectedServers={( form.getFieldValue("allowed_mcp_servers_and_groups")?.servers || [] - ).filter((s: string) => s !== NO_MCP_SERVERS_SENTINEL)} + ).filter((s: string) => s !== NO_MCP_SERVERS_SENTINEL && s !== ALL_TEAM_MCPS_SENTINEL)} toolPermissions={form.getFieldValue("mcp_tool_permissions") || {}} onChange={(toolPerms) => form.setFieldsValue({ mcp_tool_permissions: toolPerms })} /> diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx index 8a3149a9719..a224f8261f1 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx @@ -3,7 +3,7 @@ import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import MCPServerPermissions from "./MCPServerPermissions"; import * as networking from "../networking"; -import { ALL_PROXY_MCPS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_PROXY_MCPS_SENTINEL, ALL_TEAM_MCPS_SENTINEL } from "../mcp_tools/constants"; vi.mock("../networking"); @@ -271,6 +271,26 @@ describe("MCPServerPermissions", () => { expect(screen.queryByText(ALL_PROXY_MCPS_SENTINEL)).not.toBeInTheDocument(); }); + it("should display the all-team-mcps grant instead of listing the sentinel as a server", async () => { + /** + * The all-team-mcps sentinel renders the "team's servers" grant state rather + * than a literal server row, so a viewer sees the key tracks its team. + */ + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + + render( + , + ); + + expect(await screen.findByText(/All team MCP servers/)).toBeInTheDocument(); + expect(screen.queryByText(ALL_TEAM_MCPS_SENTINEL)).not.toBeInTheDocument(); + }); + it("should handle multiple servers with different tool permissions", async () => { /** * Tests that multiple servers can each have their own tool permissions diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index 079f6609320..be45e702d50 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -4,7 +4,7 @@ import { ServerIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/ import { Tooltip } from "antd"; import { fetchMCPServers, fetchMCPToolsets } from "../networking"; import { MCPServer, MCPToolset } from "../mcp_tools/types"; -import { ALL_PROXY_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_PROXY_MCPS_SENTINEL, ALL_TEAM_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; interface MCPServerPermissionsProps { mcpServers: string[]; @@ -97,11 +97,17 @@ export function MCPServerPermissions({ const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); const grantsAllProxyMcps = mcpServers.includes(ALL_PROXY_MCPS_SENTINEL); + const grantsAllTeamMcps = mcpServers.includes(ALL_TEAM_MCPS_SENTINEL); // Merge servers and access groups into one list const mergedItems = [ ...mcpServers - .filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCPS_SENTINEL) + .filter( + (server) => + server !== NO_MCP_SERVERS_SENTINEL && + server !== ALL_PROXY_MCPS_SENTINEL && + server !== ALL_TEAM_MCPS_SENTINEL, + ) .map((server) => ({ type: "server", value: server })), ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), ]; @@ -112,8 +118,11 @@ export function MCPServerPermissions({
MCP Servers - - {blocksAllMcpServers ? "Blocked" : grantsAllProxyMcps ? "All" : totalCount} + + {blocksAllMcpServers ? "Blocked" : grantsAllProxyMcps ? "All" : grantsAllTeamMcps ? "Team" : totalCount}
@@ -131,6 +140,14 @@ export function MCPServerPermissions({ All proxy MCP servers — access to every MCP server registered on the proxy, including ones added later
+ ) : grantsAllTeamMcps ? ( +
+ + + All team MCP servers — access to every MCP server this key's team can reach, tracking the team as its + grant changes + +
) : totalCount > 0 ? (
{mergedItems.map((item, index) => { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index e991db8069e..988293b49bb 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -19,7 +19,7 @@ import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; import { KeyResponse } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; -import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_TEAM_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; import NotificationsManager from "../molecules/notifications_manager"; import { getPromptsList, modelAvailableCall, tagListCall } from "../networking"; @@ -626,8 +626,10 @@ export function KeyEditView({ onChange={(val) => form.setFieldValue("mcp_servers_and_groups", val)} value={form.getFieldValue("mcp_servers_and_groups")} accessToken={accessToken || ""} + teamId={keyData.team_id ?? null} placeholder="Select MCP servers or access groups (optional)" allowNoMcpServers + allowAllTeamMcps /> @@ -648,7 +650,7 @@ export function KeyEditView({ s !== NO_MCP_SERVERS_SENTINEL, + (s: string) => s !== NO_MCP_SERVERS_SENTINEL && s !== ALL_TEAM_MCPS_SENTINEL, )} toolPermissions={form.getFieldValue("mcp_tool_permissions") || {}} onChange={(toolPerms) => form.setFieldsValue({ mcp_tool_permissions: toolPerms })}