From 70005a2a395de56e7f5c6abc03f19b6d4ece7ca1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Jul 2026 03:24:27 +0000 Subject: [PATCH] fix: validate all-mcp-servers sentinel on key-level permission check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../object_permission_utils.py | 42 ++++++++- .../test_object_permission_utils.py | 94 +++++++++++++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index a21e585a2b9..ed8219d8931 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -419,6 +419,17 @@ def _extract_requested_mcp_server_ids( return server_ids +def _has_all_mcp_servers_sentinel( + object_permission: Optional[ObjectPermissionDict], +) -> bool: + if not object_permission or not isinstance(object_permission, dict): + return False + mcp_servers = object_permission.get("mcp_servers") + if isinstance(mcp_servers, list): + return SpecialMCPServerNames.all_mcp_servers.value in mcp_servers + return False + + def _extract_requested_mcp_access_groups( object_permission: Optional[ObjectPermissionDict], ) -> Set[str]: @@ -468,13 +479,42 @@ async def validate_key_mcp_servers_against_team( Raises HTTPException(403) if validation fails. """ teamless_admin_assignment = team_obj is None and is_proxy_admin + has_all_sentinel = _has_all_mcp_servers_sentinel(object_permission) requested_servers = _extract_requested_mcp_server_ids(object_permission) requested_access_groups = _extract_requested_mcp_access_groups(object_permission) requested_toolsets = _extract_requested_mcp_toolsets(object_permission) + if has_all_sentinel: + if teamless_admin_assignment: + pass + elif team_obj is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + "Key is not in a team. The 'all-mcp-servers' sentinel " + "cannot be assigned to personal keys by non-admin callers." + ) + }, + ) + else: + team_op = team_obj.object_permission + team_mcp_servers = (team_op.mcp_servers or []) if team_op else [] + if SpecialMCPServerNames.all_mcp_servers.value not in team_mcp_servers: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"Key requests 'all-mcp-servers' but team '{team_obj.team_id}' " + f"does not grant the all-mcp-servers sentinel. " + f"Team allows: {sorted(team_mcp_servers)}." + ) + }, + ) + # Nothing to validate - if not requested_servers and not requested_access_groups and not requested_toolsets: + if not requested_servers and not requested_access_groups and not requested_toolsets and not has_all_sentinel: return object_permission allow_all_keys_servers = _get_allow_all_keys_server_ids() 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 8d543d19bbb..a6d65b3ebd6 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 @@ -13,6 +13,7 @@ from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_ObjectPer from litellm.proxy.management_helpers.object_permission_utils import ( _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, + _has_all_mcp_servers_sentinel, _resolve_team_allowed_mcp_servers, _rewrite_object_permission_mcp_servers, _set_object_permission, @@ -668,6 +669,99 @@ async def test_validate_stale_mcp_server_ids_are_removed_from_object_permission( assert object_permission["mcp_servers"] == [] +# ---- Tests for all-mcp-servers sentinel validation in validate_key_mcp_servers_against_team ---- + + +def test_has_all_mcp_servers_sentinel_true(): + assert _has_all_mcp_servers_sentinel({"mcp_servers": ["all-mcp-servers"]}) is True + + +def test_has_all_mcp_servers_sentinel_false(): + assert _has_all_mcp_servers_sentinel({"mcp_servers": ["server-1"]}) is False + assert _has_all_mcp_servers_sentinel(None) is False + assert _has_all_mcp_servers_sentinel({}) is False + + +@pytest.mark.asyncio +async def test_validate_all_sentinel_teamless_non_admin_raises(): + """Non-admin, teamless key with all-mcp-servers sentinel must be rejected""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["all-mcp-servers"]}, + team_obj=None, + is_proxy_admin=False, + ) + assert exc_info.value.status_code == 403 + assert "all-mcp-servers" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("s1", "s2"), +) +@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_all_sentinel_teamless_admin_passes( + mock_access_groups, mock_allow_all +): + """Proxy admin assigning all-mcp-servers to a teamless key should pass""" + result = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["all-mcp-servers"]}, + team_obj=None, + is_proxy_admin=True, + ) + assert result is not None + assert "all-mcp-servers" in result["mcp_servers"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("s1", "s2"), +) +@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_all_sentinel_key_in_team_with_sentinel_passes( + mock_access_groups, mock_allow_all +): + """Key with all-mcp-servers in a team that also has the sentinel should pass""" + team_obj = _make_team_obj(mcp_servers=["all-mcp-servers"]) + result = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["all-mcp-servers"]}, + team_obj=team_obj, + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_validate_all_sentinel_key_in_team_without_sentinel_raises(): + """Key with all-mcp-servers in a team that does NOT have the sentinel must be rejected""" + team_obj = _make_team_obj(mcp_servers=["server-1", "server-2"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["all-mcp-servers"]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + assert "all-mcp-servers" in str(exc_info.value.detail) + assert "team-1" in str(exc_info.value.detail) + + @pytest.mark.asyncio @patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",