From 8b3939da2524c9ac8ab3b0bedcab6f4379e68a5a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:34:15 -0700 Subject: [PATCH] fix(key_management): count team unified access group MCP servers when validating key MCP grants (#41231) * fix(key_management): count team unified access group MCP servers when validating key MCP grants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: drop explanatory suffix from pyright suppression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(key_management): restore reason on pyright suppression for LIT004 gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(key_management): assert union result and suppress TQ008 on litellm-internal patches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: verify unified MCP team grants through real resolvers --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../object_permission_utils.py | 33 ++++- .../test_object_permission_utils.py | 132 ++++++++++++++++++ 2 files changed, 159 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 4aaa77f8d45..b12a689429d 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -565,17 +565,35 @@ async def _get_team_allowed_mcp_servers( """ Get the full set of MCP server IDs a team allows. - If team has no object_permission or no MCP config, returns empty set - (meaning only allow_all_keys servers are permitted). + Combines servers granted via the team's object_permission with servers + granted via the team's unified access groups (access_group_ids). If the + team grants neither, returns empty set (meaning only allow_all_keys + servers are permitted). """ if team_obj is None: return set() + from litellm.proxy.auth.auth_checks import ( + _get_mcp_server_ids_from_access_groups, # pyright: ignore[reportPrivateUsage] # same resolver runtime MCP auth calls + ) + + access_group_servers: Final = await _get_mcp_server_ids_from_access_groups( + access_group_ids=team_obj.access_group_ids or [], + prisma_client=prisma_client, + ) + resolved_access_group_servers: Final = await _resolve_mcp_server_identifiers_to_ids( + identifiers=set(access_group_servers), + prisma_client=prisma_client, + ) + unified_servers: Final = _flatten_resolved_mcp_server_ids(resolved_access_group_servers) | { + server for server in access_group_servers if not resolved_access_group_servers.get(server) + } + team_object_permission: Final = team_obj.object_permission if team_object_permission is None: - return set() + return unified_servers - return await _resolve_team_allowed_mcp_servers( + return unified_servers | await _resolve_team_allowed_mcp_servers( team_object_permission=team_object_permission, prisma_client=prisma_client, ) @@ -650,14 +668,17 @@ async def validate_key_mcp_servers_against_team( Rules: - If key is in a team: key's mcp_servers must be a subset of - (team's allowed servers + allow_all_keys servers) + (team's allowed servers + allow_all_keys servers), where the team's + allowed servers include servers granted via the team's unified + access groups - If key is NOT in a team and the caller is a proxy admin: any server or access group may be assigned. A proxy admin can already reach every MCP server, and runtime access is granted directly from the key's own object_permission, so the key is scoped to exactly what the admin selected - If key is NOT in a team and the caller is not a proxy admin: key's mcp_servers must only contain allow_all_keys servers - - If team has no MCP config: key can only use allow_all_keys servers + - If team has no MCP config (no object_permission and no unified + access groups): key can only use allow_all_keys servers Raises HTTPException(403) if validation fails. """ 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 078315c2bf8..2fba39b30f6 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 @@ -1,12 +1,16 @@ import json +from collections.abc import Iterator +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from litellm.proxy._types import ( + LiteLLM_AccessGroupTable, LiteLLM_ObjectPermissionBase, LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTableCachedObj, ObjectPermissionDict, SpecialMCPServerName, ) @@ -217,10 +221,12 @@ def _make_team_obj( mcp_servers=None, mcp_access_groups=None, mcp_tool_permissions=None, + access_group_ids=None, ): """Create a mock team object with the given MCP permissions.""" mock_team = MagicMock() mock_team.team_id = team_id + mock_team.access_group_ids = access_group_ids or [] if ( mcp_servers is not None @@ -541,6 +547,132 @@ async def test_validate_team_no_mcp_config_blocks_all( assert exc_info.value.status_code == 403 +@pytest.fixture +def unified_mcp_prisma() -> Iterator[MagicMock]: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + prisma: Final = MagicMock() + prisma.db.litellm_accessgrouptable.find_unique = AsyncMock( + return_value=LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group one", + access_mcp_server_ids=["server-1"], + ) + ) + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + manager: Final = _make_mock_mcp_manager( + "server-1", + "server-2", + servers=[_make_mock_mcp_server("server-1", alias="server-alias")], + ) + manager.config_mcp_servers = {} + manager.get_allow_all_keys_server_ids.return_value = [] + with ( + patch( # test-quality-ok: management helpers read this module singleton without a registry injection seam + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=manager, + ), + patch( # test-quality-ok: unified group resolver obtains its cache from the proxy singleton + "litellm.proxy.proxy_server.user_api_key_cache", + new=UserApiKeyCache(), + ), + ): + yield prisma + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("group_identifier", "requested_identifier"), + [("server-1", "server-1"), ("server-alias", "server-1"), ("server-1", "server-alias")], +) +async def test_validate_key_servers_granted_via_team_unified_access_group_pass( + unified_mcp_prisma: MagicMock, + group_identifier: str, + requested_identifier: str, +) -> None: + unified_mcp_prisma.db.litellm_accessgrouptable.find_unique.return_value = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group one", + access_mcp_server_ids=[group_identifier], + ) + team: Final = LiteLLM_TeamTableCachedObj(team_id="team-1", access_group_ids=["ag-1"]) + result: Final = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": [requested_identifier]}, + team_obj=team, + prisma_client=unified_mcp_prisma, + ) + assert result == {"mcp_servers": [requested_identifier]} + + +@pytest.mark.asyncio +async def test_validate_key_servers_outside_team_unified_access_group_rejected( + unified_mcp_prisma: MagicMock, +) -> None: + team: Final = LiteLLM_TeamTableCachedObj(team_id="team-1", access_group_ids=["ag-1"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-2"]}, + team_obj=team, + prisma_client=unified_mcp_prisma, + ) + assert exc_info.value.status_code == 403 + assert "server-2" in str(exc_info.value.detail) + assert "Team allows: ['server-1']" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_team_allowed_servers_union_object_permission_and_unified_access_group( + unified_mcp_prisma: MagicMock, +) -> None: + team: Final = LiteLLM_TeamTableCachedObj( + team_id="team-1", + access_group_ids=["ag-1"], + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["server-2"]), + ) + result: Final = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1", "server-2"]}, + team_obj=team, + prisma_client=unified_mcp_prisma, + ) + assert result == {"mcp_servers": ["server-1", "server-2"]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group_state", ["empty", "missing", "unresolved"]) +async def test_team_unified_access_group_without_servers_preserves_direct_grants( + unified_mcp_prisma: MagicMock, + group_state: Literal["empty", "missing", "unresolved"], +) -> None: + unified_mcp_prisma.db.litellm_accessgrouptable.find_unique.return_value = ( + LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="empty or stale group", + access_mcp_server_ids=["deleted-server"] if group_state == "unresolved" else [], + ) + if group_state != "missing" + else None + ) + team: Final = LiteLLM_TeamTableCachedObj( + team_id="team-1", + access_group_ids=["ag-1"], + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["server-2"]), + ) + allowed: Final = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-2"]}, + team_obj=team, + prisma_client=unified_mcp_prisma, + ) + assert allowed == {"mcp_servers": ["server-2"]} + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1"]}, + team_obj=team, + prisma_client=unified_mcp_prisma, + ) + assert exc_info.value.status_code == 403 + assert "['server-1']. Team allows:" in str(exc_info.value.detail) + + @pytest.mark.asyncio @patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",