scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(key_management): count team unified access group MCP servers when validating key MCP grants (#41231)

Browse source 
* 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>
This commit is contained in:
devin-ai-integration[bot] 2026-09-24 23:34:15 -07:00 • committed by GitHub
parent e319bf270c
commit 8b3939da25
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 159 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/proxy/management_helpers/object_permission_utils.py
View file

@ -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.
"""

132
tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py
View file

@ -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",