From 909a5f597afc3076fea6ed0a996ad1792e125911 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 29 May 2026 15:07:12 -0700 Subject: [PATCH] =?UTF-8?q?fix(mcp):=20resolve=20key.access=5Fgroup=5Fids?= =?UTF-8?q?=20=E2=86=92=20MCP=20servers=20(ungated)=20(#29195)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): resolve key.access_group_ids → MCP servers (ungated) A teamless virtual key whose unified access_group_ids grant an MCP-granting access group now sees and can call that server instead of getting an empty list / 403. The key path previously read only the legacy object_permission; this folds key.access_group_ids into the key's base scope (ungated), mirroring can_key_call_model's fallback. The gated assigned_*-checked override is unchanged. * fix(mcp): expand name/alias entries in key access-group servers access_mcp_server_ids may hold server names/aliases, not just ids. The new ungated key path now runs them through expand_permission_list at the source, so the early-return and union branches both surface resolved ids — matching the legacy object_permission path and the gated extras path. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 75 ++++++--- .../auth/test_user_api_key_auth_mcp.py | 147 ++++++++++++++++++ 2 files changed, 201 insertions(+), 21 deletions(-) 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 97d3a8cf5c..6fc0a97838 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 @@ -925,39 +925,67 @@ class MCPRequestHandler: async def _get_allowed_mcp_servers_for_key( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: + """ + Get allowed MCP servers for a key (the key's own scope). + + Unions two sources: + - Legacy key.object_permission (mcp_servers, mcp_access_groups, + mcp_tool_permissions). + - Unified key.access_group_ids → access_group.access_mcp_server_ids. + Mirrors the ungated fallback in can_key_call_model — the group is + attached to the key itself, so it grants the key's own scope (no + assigned_key_ids re-check). The gated, team-ceiling-busting override + lives in _get_key_access_group_mcp_server_extras. + """ + if user_api_key_auth is None: + return [] try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.auth.auth_checks import ( + _get_mcp_server_ids_from_access_groups, + get_object_permission, + ) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + # Unified key.access_group_ids → MCP servers (ungated: the group is + # attached to the key, so it grants the key's own scope). Entries in + # access_mcp_server_ids may be server_ids OR names/aliases, so expand + # to ids here — matching the legacy object_permission path below. + key_access_group_servers = global_mcp_server_manager.expand_permission_list( + await _get_mcp_server_ids_from_access_groups( + access_group_ids=user_api_key_auth.access_group_ids or [], + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + # Get key object permission (already loaded in main auth flow, or fetch from DB) key_object_permission = MCPRequestHandler._get_key_object_permission( user_api_key_auth ) if ( key_object_permission is None - and user_api_key_auth and user_api_key_auth.object_permission_id + and prisma_client is not None ): - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, + key_object_permission = await get_object_permission( + object_permission_id=user_api_key_auth.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - - if prisma_client is not None: - key_object_permission = await get_object_permission( - object_permission_id=user_api_key_auth.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) if key_object_permission is None: - return [] + return list(set(key_access_group_servers)) # Permission entries may be server_ids OR names/aliases — expand to ids. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list( key_object_permission.mcp_servers or [] ) @@ -977,7 +1005,12 @@ class MCPRequestHandler: ) # Combine all lists - all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + all_servers = ( + direct_mcp_servers + + access_group_servers + + tool_perm_servers + + key_access_group_servers + ) return list(set(all_servers)) except Exception as e: verbose_logger.warning( 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 1499f7e474..ceefc41052 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 @@ -3641,3 +3641,150 @@ async def test_get_allowed_mcp_servers_includes_team_access_group_extras_end_to_ ): result = await MCPRequestHandler.get_allowed_mcp_servers(auth) assert result == ["srv-stripe"] + + +@pytest.mark.asyncio +async def test_key_access_group_ids_resolves_mcp_servers_ungated(): + """A teamless key whose unified access_group_ids grant an MCP server sees it + even though the group lists the key in NEITHER assigned_key_ids NOR + assigned_team_ids — the group is attached to the key, so it grants the key's + own scope (ungated, mirroring can_key_call_model's fallback).""" + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + access_group_ids=["mcp-premium"], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=["srv-stripe"], + ) as mock_resolver, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) + + assert result == ["srv-stripe"] + mock_resolver.assert_called_once() + assert mock_resolver.call_args.kwargs["access_group_ids"] == ["mcp-premium"] + + +@pytest.mark.asyncio +async def test_key_access_group_ids_union_with_object_permission(): + """When both legacy key.object_permission and unified key.access_group_ids + grant MCP servers, the final list is their union.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + global_mcp_server_manager.registry["srv-direct"] = MCPServer( + server_id="srv-direct", + name="srv-direct", + server_name="srv-direct", + url="https://srv-direct.example.com", + transport=MCPTransport.http, + ) + try: + perms = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-1", + mcp_servers=["srv-direct"], + mcp_access_groups=[], + vector_stores=[], + ) + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + access_group_ids=["mcp-premium"], + object_permission=perms, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=["srv-stripe"], + ), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) + + assert set(result) == {"srv-direct", "srv-stripe"} + finally: + global_mcp_server_manager.registry.pop("srv-direct", None) + + +@pytest.mark.asyncio +async def test_key_access_group_ids_empty_returns_no_extras(): + """Empty key.access_group_ids and no object_permission → resolver called with + [], short-circuits without DB access, returns [].""" + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + access_group_ids=[], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ) as mock_resolver, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) + + assert result == [] + mock_resolver.assert_called_once() + assert mock_resolver.call_args.kwargs["access_group_ids"] == [] + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_key_access_group_base_end_to_end(): + """End-to-end bug repro: a teamless key has an MCP-granting access group on + its access_group_ids, but the group lists the key in NEITHER assigned_key_ids + NOR assigned_team_ids. The gated extras path returns [] (no override), yet the + ungated base key path grants the server → the key sees it through + get_allowed_mcp_servers.""" + auth = UserAPIKeyAuth( + token="test-token", + api_key="sk-test", + access_group_ids=["mcp-group"], + ) + # Group grants the server but admits neither this key nor its (absent) team. + fake_ag = _fake_mcp_access_group( + access_group_id="mcp-group", + access_mcp_server_ids=["srv-deepwiki"], + assigned_team_ids=[], + assigned_key_ids=[], + ) + + patches = _patch_proxy_server_globals_for_mcp() + [ + # Ungated base resolver used by _get_allowed_mcp_servers_for_key. + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=["srv-deepwiki"], + ), + # Gated path (_get_key_access_group_mcp_server_extras) resolves the group + # via get_access_object; empty assigned_* → it contributes nothing. + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + _start_patches(patches) + try: + # Sanity: the gated extras path alone denies (the old behavior). + extras = await MCPRequestHandler._get_key_access_group_mcp_server_extras(auth) + assert extras == [] + + # But the key now sees the server via the ungated base path. + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert result == ["srv-deepwiki"] + finally: + _stop_patches(patches)