fix(mcp): resolve key.access_group_ids → MCP servers (ungated) (#29195)

* 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.
This commit is contained in:
ryan-crabbe-berri 2026-05-29 15:07:12 -07:00 committed by GitHub
parent 5eafe1c1fc
commit 909a5f597a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 201 additions and 21 deletions

View File

@ -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(

View File

@ -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)