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:
parent
5eafe1c1fc
commit
909a5f597a
@ -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(
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user