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 19f508e6af9..92ebd248af0 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 @@ -8,6 +8,10 @@ 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 ( + NoServers, + parse_mcp_server_grant, +) from litellm.proxy._types import ( LiteLLM_TeamTable, ProxyException, @@ -646,9 +650,8 @@ class MCPRequestHandler: # The key explicitly opted out of every MCP server. This overrides # team inheritance and additive grants (mirrors no-default-models). - if ( - SpecialMCPServerNames.no_mcp_servers.value - in allowed_mcp_servers_for_key + if isinstance( + parse_mcp_server_grant(allowed_mcp_servers_for_key), NoServers ): return [] @@ -1070,8 +1073,9 @@ class MCPRequestHandler: # Sentinel opt-out: surface it unexpanded so the caller can short-circuit # to zero servers instead of inheriting the team. - if SpecialMCPServerNames.no_mcp_servers.value in ( - key_object_permission.mcp_servers or [] + if isinstance( + parse_mcp_server_grant(key_object_permission.mcp_servers or []), + NoServers, ): return [SpecialMCPServerNames.no_mcp_servers.value] diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1b86551972a..e28d0ef578a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -16,6 +16,8 @@ import time from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse +from typing_extensions import assert_never + import anyio from fastapi import HTTPException from httpx import HTTPStatusError @@ -49,6 +51,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.permission_grant import ( + AllServers, + ExplicitServers, + NoServers, + parse_mcp_server_grant, +) from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, ) @@ -1367,9 +1375,9 @@ class MCPServerManager: key_object_permission = ( user_api_key_auth.object_permission if user_api_key_auth else None ) - if key_object_permission is not None and ( - SpecialMCPServerNames.no_mcp_servers.value - in (key_object_permission.mcp_servers or []) + if key_object_permission is not None and isinstance( + parse_mcp_server_grant(key_object_permission.mcp_servers or []), + NoServers, ): return [] @@ -4132,6 +4140,30 @@ class MCPServerManager: if server.available_on_public_internet or server.server_id in public_ids ] + def _resolve_permission_identifier( + self, identifier: str, registry: Dict[str, MCPServer] + ) -> Tuple[str, ...]: + if identifier in registry: + return (identifier,) + matches = tuple( + server_id + for server_id, server in registry.items() + if server.alias == identifier + or server.server_name == identifier + or server.name == identifier + ) + if matches: + return matches + # %r quotes and escapes control chars so an admin-controlled identifier + # with newlines cannot forge log lines. + verbose_logger.debug( + "MCP permission entry %r does not resolve to any known " + "server (config + DB union). Passing through — the " + "downstream access check will deny it if it's stale.", + identifier, + ) + return (identifier,) + def expand_permission_list(self, identifiers: List[str]) -> List[str]: """ Expand a permission list of server_ids/names/aliases into concrete @@ -4147,38 +4179,31 @@ class MCPServerManager: The ``all-proxy-mcps`` sentinel expands to every server in the live registry, so an entity granted it stays in sync as servers are added - or removed without re-editing its permission list. + or removed without re-editing its permission list. The + ``no-mcp-servers`` sentinel surfaces unexpanded so opt-out-aware callers + can short-circuit before inheriting any other scope. """ if not identifiers: return [] registry = self.get_registry() - if SpecialMCPServerNames.all_proxy_mcp_servers.value in identifiers: - return list(registry.keys()) - expanded: Set[str] = set() - for identifier in identifiers: - if identifier in registry: - expanded.add(identifier) - continue - matches: List[str] = [ - server_id - for server_id, server in registry.items() - if server.alias == identifier - or server.server_name == identifier - or server.name == identifier - ] - if matches: - expanded.update(matches) - else: - # %r quotes and escapes control chars so an admin-controlled - # identifier with newlines cannot forge log lines. - verbose_logger.debug( - "MCP permission entry %r does not resolve to any known " - "server (config + DB union). Passing through — the " - "downstream access check will deny it if it's stale.", - identifier, + grant = parse_mcp_server_grant(identifiers) + match grant: + case AllServers(): + return list(registry.keys()) + case NoServers(): + return [SpecialMCPServerNames.no_mcp_servers.value] + case ExplicitServers(identifiers=explicit): + return list( + frozenset( + server_id + for identifier in explicit + for server_id in self._resolve_permission_identifier( + identifier, registry + ) + ) ) - expanded.add(identifier) - return list(expanded) + case _: + assert_never(grant) def expand_tool_permissions( self, diff --git a/litellm/proxy/_experimental/mcp_server/permission_grant.py b/litellm/proxy/_experimental/mcp_server/permission_grant.py new file mode 100644 index 00000000000..0b7778cf2a6 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/permission_grant.py @@ -0,0 +1,54 @@ +"""Parse raw ``mcp_servers`` permission lists into a tagged grant union. + +This module is pure: it only knows about the sentinel strings via the +``SpecialMCPServerNames`` enum and never touches the server manager, the DB, or +the proxy server. Downstream resolvers turn an ``ExplicitServers`` set of +identifiers into concrete deployments. +""" + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Union + +from litellm.proxy._types import SpecialMCPServerNames + +MCP_GRANT_SENTINELS: frozenset[str] = frozenset( + { + SpecialMCPServerNames.no_mcp_servers.value, + SpecialMCPServerNames.all_proxy_mcp_servers.value, + } +) + + +@dataclass(frozen=True, slots=True) +class AllServers: + """Grant every MCP server on the proxy, including ones added later.""" + + +@dataclass(frozen=True, slots=True) +class NoServers: + """Block all MCP servers.""" + + +@dataclass(frozen=True, slots=True) +class ExplicitServers: + """Grant a specific, unresolved set of server_ids / aliases / names.""" + + identifiers: frozenset[str] + + +MCPServerGrant = Union[AllServers, 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. + """ + values = frozenset(raw) + if SpecialMCPServerNames.no_mcp_servers.value in values: + return NoServers() + if SpecialMCPServerNames.all_proxy_mcp_servers.value in values: + return AllServers() + return ExplicitServers(values) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 0a61fc5d3dd..e7bbf97bb1a 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -7,11 +7,18 @@ import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union from fastapi import HTTPException, status +from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import SpecialMCPServerNames +from litellm.proxy._experimental.mcp_server.permission_grant import ( + MCP_GRANT_SENTINELS, + AllServers, + ExplicitServers, + NoServers, + parse_mcp_server_grant, +) from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import MCPServerRepository @@ -286,12 +293,15 @@ def _rewrite_object_permission_mcp_servers( if not isinstance(mcp_servers, list): return - normalized_servers: List[str] = [] - for identifier in mcp_servers: - if identifier == SpecialMCPServerNames.no_mcp_servers.value: - normalized_servers.append(SpecialMCPServerNames.no_mcp_servers.value) - continue - normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, []))) + normalized_servers = [ + resolved + for identifier in mcp_servers + for resolved in ( + [identifier] + if identifier in MCP_GRANT_SENTINELS + else sorted(identifier_to_server_ids.get(identifier, [])) + ) + ] object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers) @@ -355,30 +365,55 @@ async def _resolve_team_allowed_mcp_servers( - Direct mcp_servers list - Servers from mcp_access_groups - Server IDs referenced in mcp_tool_permissions keys - """ - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[ - str - ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( - team_object_permission.mcp_access_groups or [] - ) - raw_tool_perms = team_object_permission.mcp_tool_permissions or {} - if isinstance(raw_tool_perms, str): - raw_tool_perms = json.loads(raw_tool_perms) - tool_perm_servers: List[str] = list(raw_tool_perms.keys()) - raw_servers = set(direct_servers + access_group_servers + tool_perm_servers) - resolved_servers = await _resolve_mcp_server_identifiers_to_ids( - identifiers=raw_servers, - prisma_client=prisma_client, - ) - unresolved_servers = { - server_id for server_id in raw_servers if not resolved_servers.get(server_id) - } - return _flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers + The grant on the direct list decides everything: ``all-proxy-mcps`` resolves + to every registered server, ``no-mcp-servers`` blocks all of them. + """ + grant = parse_mcp_server_grant(team_object_permission.mcp_servers or []) + match grant: + case NoServers(): + return set() + case AllServers(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return set(global_mcp_server_manager.get_registry().keys()) + case ExplicitServers(identifiers): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + access_group_servers = ( + await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] + ) + ) + raw_tool_perms_value = team_object_permission.mcp_tool_permissions or {} + raw_tool_perms = ( + json.loads(raw_tool_perms_value) + if isinstance(raw_tool_perms_value, str) + else raw_tool_perms_value + ) + raw_servers = ( + set(identifiers) + | set(access_group_servers) + | set(raw_tool_perms.keys()) + ) + resolved_servers = await _resolve_mcp_server_identifiers_to_ids( + identifiers=raw_servers, + prisma_client=prisma_client, + ) + unresolved_servers = { + server_id + for server_id in raw_servers + if not resolved_servers.get(server_id) + } + return ( + _flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers + ) + case _: + assert_never(grant) def _get_allow_all_keys_server_ids() -> Set[str]: @@ -426,17 +461,21 @@ def _extract_requested_mcp_server_ids( if not object_permission or not isinstance(object_permission, dict): return set() - server_ids: Set[str] = set() mcp_servers = object_permission.get("mcp_servers") - if isinstance(mcp_servers, list): - server_ids.update(mcp_servers) - server_ids.discard(SpecialMCPServerNames.no_mcp_servers.value) + explicit_servers = ( + frozenset(s for s in mcp_servers if s not in MCP_GRANT_SENTINELS) + if isinstance(mcp_servers, list) + else frozenset() + ) mcp_tool_permissions = object_permission.get("mcp_tool_permissions") - if isinstance(mcp_tool_permissions, dict): - server_ids.update(mcp_tool_permissions.keys()) + tool_permission_servers = ( + frozenset(mcp_tool_permissions.keys()) + if isinstance(mcp_tool_permissions, dict) + else frozenset() + ) - return server_ids + return set(explicit_servers | tool_permission_servers) def _extract_requested_mcp_access_groups( @@ -493,6 +532,34 @@ async def validate_key_mcp_servers_against_team( requested_toolsets = _extract_requested_mcp_toolsets(object_permission) + key_mcp_servers = ( + object_permission.get("mcp_servers") + if isinstance(object_permission, dict) + else None + ) + key_grant = ( + parse_mcp_server_grant(key_mcp_servers) + if isinstance(key_mcp_servers, list) + else ExplicitServers(frozenset()) + ) + team_grant = ( + parse_mcp_server_grant(team_obj.object_permission.mcp_servers or []) + if team_obj is not None and team_obj.object_permission is not None + else ExplicitServers(frozenset()) + ) + + if isinstance(key_grant, AllServers) and not teamless_admin_assignment: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + "all-proxy-mcps cannot be granted directly on a key. A key in a " + "team inherits the team's MCP grant at request time; only a proxy " + "admin may assign all-proxy-mcps to a teamless key." + ) + }, + ) + # Nothing to validate if not requested_servers and not requested_access_groups and not requested_toolsets: return object_permission @@ -532,30 +599,32 @@ async def validate_key_mcp_servers_against_team( identifier_to_server_ids ) - allowed_servers = all_allowed_servers - if teamless_admin_assignment: - allowed_servers = all_allowed_servers | active_requested_servers - - disallowed_servers = active_requested_servers - allowed_servers - if disallowed_servers: - if team_obj is not None: - team_id = team_obj.team_id - detail = ( - f"Key requests MCP servers not allowed by team '{team_id}': " - f"{sorted(disallowed_servers)}. " - f"Team allows: {sorted(team_allowed_servers)}. " - f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}." - ) - else: - detail = ( - f"Key is not in a team. Only globally available (allow_all_keys) MCP servers " - f"can be assigned: {sorted(allow_all_keys_servers)}. " - f"Disallowed servers: {sorted(disallowed_servers)}." - ) - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": detail}, + if not isinstance(team_grant, AllServers): + allowed_servers = ( + all_allowed_servers | active_requested_servers + if teamless_admin_assignment + else all_allowed_servers ) + disallowed_servers = active_requested_servers - allowed_servers + if disallowed_servers: + if team_obj is not None: + team_id = team_obj.team_id + detail = ( + f"Key requests MCP servers not allowed by team '{team_id}': " + f"{sorted(disallowed_servers)}. " + f"Team allows: {sorted(team_allowed_servers)}. " + f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}." + ) + else: + detail = ( + f"Key is not in a team. Only globally available (allow_all_keys) MCP servers " + f"can be assigned: {sorted(allow_all_keys_servers)}. " + f"Disallowed servers: {sorted(disallowed_servers)}." + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": detail}, + ) # Validate requested access groups (must be subset of team's access groups) if requested_access_groups: 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 dc5a2783a91..1d87af54c8c 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 @@ -160,6 +160,55 @@ class TestMCPServerManager: # blanket "return everything". assert manager.expand_permission_list(["srv-a"]) == ["srv-a"] + async def test_expand_permission_list_resolution_branches(self): + """Pin every non-all-proxy branch of expand_permission_list: empty, + concrete id passthrough, alias/server_name expansion, and unknown + passthrough. Guards the ExplicitServers resolution loop.""" + manager = MCPServerManager() + await manager.add_server( + LiteLLM_MCPServerTable( + server_id="srv-x", + alias="my-alias", + server_name="my-server-name", + description="", + url=None, + transport=MCPTransport.stdio, + command="python", + args=["-m", "server"], + env={}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ) + + assert manager.expand_permission_list([]) == [] + assert manager.expand_permission_list(["srv-x"]) == ["srv-x"] + assert manager.expand_permission_list(["my-alias"]) == ["srv-x"] + assert manager.expand_permission_list(["my-server-name"]) == ["srv-x"] + assert manager.expand_permission_list(["ghost"]) == ["ghost"] + + async def test_expand_permission_list_no_mcp_servers_passthrough(self): + """The no-mcp-servers sentinel is not a registered server, so today it + passes through unchanged. Callers (key opt-out, team expansion) rely on + this literal surviving expand_permission_list, so pin it.""" + 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.no_mcp_servers.value + assert manager.expand_permission_list([sentinel]) == [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 new file mode 100644 index 00000000000..430d9bcfb5b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py @@ -0,0 +1,63 @@ +import dataclasses + +import pytest + +from litellm.proxy._experimental.mcp_server.permission_grant import ( + AllServers, + ExplicitServers, + MCPServerGrant, + NoServers, + parse_mcp_server_grant, +) +from litellm.proxy._types import SpecialMCPServerNames + +NO_MCP = SpecialMCPServerNames.no_mcp_servers.value +ALL_PROXY = SpecialMCPServerNames.all_proxy_mcp_servers.value + + +def test_concrete_only_returns_explicit_servers() -> None: + grant = parse_mcp_server_grant(["srv-a", "srv-b", "alias-c"]) + assert grant == ExplicitServers(frozenset({"srv-a", "srv-b", "alias-c"})) + + +def test_all_proxy_sentinel_returns_all_servers() -> None: + assert parse_mcp_server_grant([ALL_PROXY]) == AllServers() + + +def test_all_proxy_sentinel_dominates_concrete_ids() -> None: + assert parse_mcp_server_grant([ALL_PROXY, "srv-a"]) == AllServers() + + +def test_no_mcp_sentinel_returns_no_servers() -> None: + assert parse_mcp_server_grant([NO_MCP]) == NoServers() + + +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_empty_list_returns_empty_explicit_servers() -> None: + assert parse_mcp_server_grant([]) == ExplicitServers(frozenset()) + + +def test_grants_are_frozen() -> None: + grant = ExplicitServers(frozenset({"srv-a"})) + with pytest.raises(dataclasses.FrozenInstanceError): + grant.identifiers = frozenset({"other"}) # pyright: ignore[reportAttributeAccessIssue] + + +def test_grants_equal_by_value() -> None: + assert AllServers() == AllServers() + assert NoServers() == NoServers() + assert ExplicitServers(frozenset({"a"})) == ExplicitServers(frozenset({"a"})) + assert AllServers() != NoServers() + + +def test_union_alias_admits_each_variant() -> None: + grants: tuple[MCPServerGrant, ...] = ( + AllServers(), + NoServers(), + ExplicitServers(frozenset({"a"})), + ) + assert len(grants) == 3 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 2b38d732e9d..f5fa670676e 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 @@ -9,7 +9,7 @@ sys.path.insert(0, os.path.abspath("../../../..")) from unittest.mock import AsyncMock, MagicMock, patch -from litellm.proxy._types import LiteLLM_ObjectPermissionTable +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerNames from litellm.proxy.management_helpers.object_permission_utils import ( _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, @@ -890,3 +890,211 @@ async def test_validate_search_tools_raises_when_not_subset(): team_obj=_make_team_obj_search(search_tools=["t1"]), ) assert exc.value.status_code == 403 + + +# ---- Tests for the all-proxy-mcps grant ---- + +ALL_PROXY = SpecialMCPServerNames.all_proxy_mcp_servers.value +NO_MCP = SpecialMCPServerNames.no_mcp_servers.value + + +def test_extract_requested_mcp_server_ids_excludes_all_proxy_sentinel(): + obj_perm = {"mcp_servers": [ALL_PROXY, "server-1"]} + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"} + + +def test_rewrite_object_permission_mcp_servers_preserves_all_proxy_sentinel(): + obj_perm = {"mcp_servers": [ALL_PROXY, "alias-1"]} + _rewrite_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}}) + assert obj_perm["mcp_servers"] == [ALL_PROXY, "server-1"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1", "server-2"), +) +async def test_resolve_team_allowed_mcp_servers_all_proxy_expands_to_registry(): + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = [ALL_PROXY] + mock_perm.mcp_access_groups = [] + mock_perm.mcp_tool_permissions = {} + result = await _resolve_team_allowed_mcp_servers(mock_perm) + assert result == {"server-1", "server-2"} + + +@pytest.mark.asyncio +async def test_resolve_team_allowed_mcp_servers_no_mcp_servers_blocks_all(): + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = [NO_MCP] + 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() + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1", "server-2"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_team_all_proxy_grants_any_real_server( + mock_access_groups, mock_allow_all +): + """A team scoped to all-proxy-mcps lets a key request any real server without a 403.""" + team_obj = _make_team_obj(mcp_servers=[ALL_PROXY]) + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1"]}, + team_obj=team_obj, + ) # Must not raise + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("registry-server"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._resolve_mcp_server_identifiers_to_ids", + new_callable=AsyncMock, + return_value={"db-only": {"db-only"}}, +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_team_all_proxy_allows_server_outside_registry( + mock_access_groups, mock_allow_all, mock_resolve +): + """An all-proxy team must allow a key's server that resolves via the DB even when + it is absent from the in-memory registry keys; the validator short-circuits on the + team's all-proxy grant instead of enumerating the registry.""" + team_obj = _make_team_obj(mcp_servers=[ALL_PROXY]) + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["db-only"]}, + team_obj=team_obj, + ) # Must not raise + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_all_proxy_under_explicit_team_raises( + mock_access_groups, mock_allow_all +): + """A key may not self-grant all-proxy-mcps when its team only allows specific servers.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [ALL_PROXY]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_all_proxy_in_team_raises_even_when_team_all_proxy( + mock_access_groups, mock_allow_all +): + """A key in a team may not set all-proxy-mcps directly; it inherits the team's + grant at request time. Persisting the literal on the key would let it keep + all-server access if the team's grant were later revoked, so it is rejected.""" + team_obj = _make_team_obj(mcp_servers=[ALL_PROXY]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [ALL_PROXY]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_teamless_admin_can_assign_all_proxy( + mock_access_groups, mock_allow_all +): + """A proxy admin assigning all-proxy-mcps to a teamless key is allowed and preserved.""" + obj_perm = {"mcp_servers": [ALL_PROXY]} + await validate_key_mcp_servers_against_team( + object_permission=obj_perm, + team_obj=None, + is_proxy_admin=True, + ) + assert obj_perm["mcp_servers"] == [ALL_PROXY] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_teamless_non_admin_all_proxy_raises( + mock_access_groups, mock_allow_all +): + """A non-admin teamless key cannot self-grant all-proxy-mcps; it is rejected, not dropped.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [ALL_PROXY]}, + team_obj=None, + is_proxy_admin=False, + ) + assert exc_info.value.status_code == 403