scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix: enforce MCP toolsets attached to a team, org, or internal user (#38488)

Browse source 
* fix: enforce MCP toolsets attached to a team, org, or internal user

object_permission.mcp_toolsets was resolved into servers and tools only at
the key level; every other principal read mcp_tool_permissions and silently
ignored its toolsets. A team/org/user toolset alongside a server grant was
inert (all tools callable), a toolset alone granted nothing, and an inert
team toolset let the org server list substitute for the empty team result,
handing the caller every org server.

Resolve toolsets at each level that resolves mcp_tool_permissions, union
their servers into that level's granted server set, and count a declared
key/team toolset toward has_lower_level_mcp_restrictions so the org list
can only cap, never substitute, even when the toolset resolves empty.

Resolves LIT-5749

* fix: deny when a team's declared MCP toolset cannot be resolved

The team server resolver swallowed UnloadableEntitlementError into an empty
list, so a dangling team toolset dropped the team ceiling instead of denying,
unlike the org and user paths. Re-raise it so the top-level resolver denies.

Also anchor the test-quality suppression comments on the patch opener lines
the gate reads, with per-seam reasons.
This commit is contained in:
yucheng-berri 2026-08-28 12:06:57 -07:00 • committed by GitHub
parent 4d9025c2bf
commit 2ef77f30e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 596 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

152
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1,5 +1,5 @@
import re
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
@ -67,6 +67,9 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({})
def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: resolver returns a list
"""Widen a read-only allowlist back to the mutable list the resolver's own contract returns,
preserving the ``None`` that means "no restriction"."""
@ -1497,7 +1500,11 @@ class MCPRequestHandler:
team_set: Final = set(allowed_mcp_servers_for_team)
grants_set: Final = set(key_access_group_grants)
has_lower_level_mcp_restrictions = bool(key_set or team_set or grants_set)
# A DECLARED toolset restricts even when it resolves to no servers: the org
# ceiling below may only cap it, never substitute the org's full server list.
has_lower_level_mcp_restrictions = bool(key_set or team_set or grants_set) or (
await MCPRequestHandler._key_or_team_declares_toolsets(user_api_key_auth)
)
# 1. Key/team ceiling. An empty set means "this level does not restrict".
if not team_set:
@ -1941,6 +1948,105 @@ class MCPRequestHandler:
return team_obj.object_permission
@staticmethod
async def _toolset_tool_permissions(
object_permission: LiteLLM_ObjectPermissionTable | None,
) -> Mapping[str, Sequence[str]]:
"""The ``server_id -> tool names`` grants of this permission row's toolsets, empty when it
declares none. The shared resolver for the team, org, and internal-user levels, so a toolset
behaves identically wherever it is attached.
RAISES ``UnloadableEntitlementError`` when the row DECLARES toolsets but resolution yields
nothing (deleted or unknown ids, a swallowed DB fault, or a toolset with no tools): that is a
KNOWN restriction with unknown contents, and every caller already turns this error into deny
rather than letting the level read as unrestricted."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
if object_permission is None or not object_permission.mcp_toolsets:
return _EMPTY_TOOLSET_GRANTS
resolved: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=object_permission.mcp_toolsets
)
if not resolved:
raise UnloadableEntitlementError(
f"declared mcp_toolsets {object_permission.mcp_toolsets!r} resolved to no grants"
)
return resolved
@staticmethod
async def _toolset_tools_for_server(
object_permission: LiteLLM_ObjectPermissionTable | None,
server_id: str,
) -> Sequence[str] | None:
"""Tool names this row's toolsets grant on ``server_id``, ``None`` when its toolsets place
no restriction on that server (it declares no toolsets, or none of them name it)."""
return (await MCPRequestHandler._toolset_tool_permissions(object_permission)).get(server_id)
@staticmethod
def _union_tool_grants(
direct: Sequence[str] | None,
via_toolsets: Sequence[str] | None,
) -> Sequence[str] | None:
"""Union of one level's direct tool grants and its toolset-granted tools on one server,
``None`` when neither source restricts (allow-all from this level)."""
if direct is None and via_toolsets is None:
return None
return tuple({*(direct or ()), *(via_toolsets or ())})
@staticmethod
async def _key_object_permission_hydrated(
user_api_key_auth: UserAPIKeyAuth,
) -> LiteLLM_ObjectPermissionTable | None:
"""The key's object_permission, loading it by ``object_permission_id`` when the main auth
flow cached the key with the relation unhydrated (its loader swallows a failed read and
caches the partial object)."""
loaded: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth)
if loaded is not None or not user_api_key_auth.object_permission_id:
return loaded
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,
)
if prisma_client is None:
return None
return 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,
)
@staticmethod
async def _key_or_team_declares_toolsets(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
"""Whether the key or its team GRANTS any toolset, resolvable or not. A declared toolset is
a lower-level restriction even when it resolves to no servers (deleted or unknown ids), so the
org ceiling may only cap it; reading an empty resolution as "no restriction" would substitute
the org's entire server list for the narrowest grant an operator can write.
Falls back to the DB when the auth object carries ``object_permission_id`` unhydrated (the
main auth flow swallows a failed load and caches the partial object). An INDETERMINATE fault
answers False — no gate, org substitution as before the fault — mirroring how the org ceiling
keeps key auth open on a fault it cannot classify."""
if user_api_key_auth is None:
return False
try:
key_obj_perm: Final = await MCPRequestHandler._key_object_permission_hydrated(user_api_key_auth)
if key_obj_perm is not None and key_obj_perm.mcp_toolsets:
return True
if not user_api_key_auth.team_id:
return False
team_obj_perm: Final = await MCPRequestHandler._get_team_object_permission(user_api_key_auth)
return bool(team_obj_perm is not None and team_obj_perm.mcp_toolsets)
except Exception as e: # noqa: BLE001 # indeterminate fault: no gate, as before this level existed
verbose_logger.warning("Failed to check declared MCP toolsets, org ceiling unchanged: %s", e)
return False
@staticmethod
async def get_allowed_tools_for_server(
server_id: str,
@ -2004,12 +2110,17 @@ class MCPRequestHandler:
if key_direct_tools is not None or key_toolset_tools is not None
else None
)
team_tools: Final = (
team_direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id)
if team_obj_perm
else None
)
# Tools granted through the team's toolsets restrict this server exactly
# as the team's direct tool permissions do, mirroring the key path above
team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(team_obj_perm, server_id)
team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools)
# Apply same inheritance logic as get_allowed_mcp_servers
if team_tools:
if key_tools:
@ -2094,11 +2205,13 @@ class MCPRequestHandler:
e,
)
return allowed_tools
org_tools: Final = (
org_direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(org_obj_perm.mcp_tool_permissions).get(server_id)
if org_obj_perm and org_obj_perm.mcp_tool_permissions
else None
)
org_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(org_obj_perm, server_id)
org_tools: Final = MCPRequestHandler._union_tool_grants(org_direct_tools, org_toolset_tools)
if org_tools is not None:
allowed_tools = (
list(set(allowed_tools) & set(org_tools)) if allowed_tools is not None else list(org_tools)
@ -2340,7 +2453,8 @@ class MCPRequestHandler:
async def _team_granted_servers(team_obj: LiteLLM_TeamTable, team_access_group_servers: list[str]) -> set[str]:
"""The raw MCP-server set a team grants (before any org ceiling): its object_permission (direct
``mcp_servers``, the ``all_proxy_servers`` sentinel → the full registry, legacy access groups,
tool-perm-referenced servers) unioned with its unified ``access_group_ids`` servers."""
tool-perm-referenced servers, toolset-referenced servers) unioned with its unified
``access_group_ids`` servers."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
@ -2357,6 +2471,7 @@ class MCPRequestHandler:
set(global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []))
| set(legacy_access_group_servers)
| set(global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys())
| (await MCPRequestHandler._toolset_tool_permissions(object_permissions)).keys()
| set(team_access_group_servers)
)
@ -2415,6 +2530,8 @@ class MCPRequestHandler:
servers: Final = await MCPRequestHandler._team_granted_servers(team_obj, team_access_group_servers)
return list(servers)
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
raise
verbose_logger.warning("Failed to get allowed MCP servers for team: %s", e)
return []
@ -2546,7 +2663,13 @@ class MCPRequestHandler:
global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys()
)
all_servers: Final = direct_mcp_servers + access_group_servers + tool_perm_servers
# servers referenced by the org's toolset grants are part of the org ceiling,
# exactly as servers referenced by its inline tool permissions are
toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(object_permissions)
all_servers: Final = tuple(
{*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants}
)
return list(set(all_servers))
except Exception as e:
# None = ceiling UNRESOLVED, distinct from [] = org places no restriction. Collapsing them
@ -2740,8 +2863,8 @@ class MCPRequestHandler:
``[]`` means this human places no restriction (allow-all from this level); ``None`` means the
ceiling is UNRESOLVED, which the caller denies on. Servers named only under
``mcp_tool_permissions`` count as entitled, exactly as they do for a key or a team, so
granting one tool never requires naming its server twice.
``mcp_tool_permissions`` or reached through ``mcp_toolsets`` count as entitled, exactly as
they do for a key or a team, so granting one tool never requires naming its server twice.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
@ -2759,7 +2882,8 @@ class MCPRequestHandler:
tool_perm_servers: Final = list(
global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys()
)
return list(set(direct_mcp_servers + access_group_servers + tool_perm_servers))
toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(object_permissions)
return tuple({*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants})
except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling"
verbose_logger.warning("Failed to get allowed MCP servers for user: %s", e)
return None
@ -2860,12 +2984,14 @@ class MCPRequestHandler:
verbose_logger.warning("MCP user tool ceiling unresolvable, denying tools on %r: %s", server_id, e)
return []
if object_permissions is None or not object_permissions.mcp_tool_permissions:
if object_permissions is None:
return allowed_tools
user_tools = global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).get(
server_id
)
user_direct_tools: Final = global_mcp_server_manager.expand_tool_permissions(
object_permissions.mcp_tool_permissions
).get(server_id)
user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(object_permissions, server_id)
user_tools: Final = MCPRequestHandler._union_tool_grants(user_direct_tools, user_toolset_tools)
if user_tools is None:
return allowed_tools
if allowed_tools is None:

463
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -504,6 +504,448 @@ class TestMCPRequestHandler:
assert result is None
# ------------------------------------------------------------------
# LIT-5749: toolsets attached to a TEAM, ORG, or internal USER must be
# enforced exactly like inline tool allowlists, on both axes
# ------------------------------------------------------------------
async def test_team_toolset_restricts_tools_on_granted_server(self):
"""A team's toolset must narrow the server's tools on list and on call,
unioned with the team's direct tool grants, mirroring the key path"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-1")
team_object_permission = self._toolset_only_object_permission(["toolset-1"])
team_object_permission.mcp_tool_permissions = {"server-a": ["direct_tool"]}
mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels", "read_thread"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission)
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
allowed = await MCPRequestHandler.get_allowed_tools_for_server(
server_id="server-a", user_api_key_auth=user_api_key_auth
)
send_message_allowed = await MCPRequestHandler.is_tool_allowed_for_server(
tool_name="send_message", server_id="server-a", user_api_key_auth=user_api_key_auth
)
toolset_tool_allowed = await MCPRequestHandler.is_tool_allowed_for_server(
tool_name="read_thread", server_id="server-a", user_api_key_auth=user_api_key_auth
)
assert allowed is not None
assert set(allowed) == {"direct_tool", "search_channels", "read_thread"}
assert send_message_allowed is False
assert toolset_tool_allowed is True
async def test_team_toolset_only_restricts_tools_without_direct_grants(self):
"""A team whose ONLY tool grant is a toolset must not fall through to
allow-all; every tool the toolset does not name is refused"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-1")
team_object_permission = self._toolset_only_object_permission(["toolset-1"])
mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission)
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
allowed = await MCPRequestHandler.get_allowed_tools_for_server(
server_id="server-a", user_api_key_auth=user_api_key_auth
)
assert allowed == ["search_channels"]
async def test_team_granted_servers_include_toolset_servers(self):
"""The team's raw server grant must include servers reached only through
its toolsets, so a toolset-only team still lists its server"""
team_object_permission = self._toolset_only_object_permission(["toolset-1"])
team_obj = MagicMock()
team_obj.object_permission = team_object_permission
mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"], "server-b": ["get_doc"]})
with (
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
):
servers = await MCPRequestHandler._team_granted_servers(team_obj, [])
assert servers == {"server-a", "server-b"}
async def test_team_toolset_only_does_not_inherit_org_full_server_list(self):
"""The reported amplifier: a team whose only MCP grant is a toolset must
CAP the org list to the toolset's server, never inherit the org's full list"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-1", org_id="org-1")
team_object_permission = self._toolset_only_object_permission(["toolset-1"])
team_obj = MagicMock()
team_obj.blocked = False
team_obj.object_permission = team_object_permission
team_obj.access_group_ids = []
team_obj.organization_id = "org-1"
mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch( # test-quality-ok: team-server resolution requires the proxy's module-global prisma client
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
"litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team_obj)
),
patch( # test-quality-ok: access-group lookup hits the DB, not under test here
"litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
AsyncMock(return_value=[]),
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])
),
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler,
"_get_allowed_mcp_servers_for_org",
AsyncMock(return_value=["server-a", "server-x"]),
),
):
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == ["server-a"]
async def test_declared_toolset_resolving_empty_still_blocks_org_substitution(self):
"""A DECLARED toolset that resolves to nothing (deleted/unknown ids) is
still a lower-level restriction: the org list may cap it, never replace it"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", org_id="org-1")
key_object_permission = self._toolset_only_object_permission(["toolset-gone"])
mock_manager = self._mock_manager_with_toolsets({})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission
),
patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here
MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler,
"_get_allowed_mcp_servers_for_org",
AsyncMock(return_value=["server-x", "server-y"]),
),
):
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == []
async def test_team_dangling_toolset_denies_key_own_grants(self):
"""A team toolset that cannot be resolved must deny on the SERVER axis too,
not silently drop the team ceiling and pass the key's own grants through"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-1")
key_object_permission = self._toolset_only_object_permission([])
key_object_permission.mcp_toolsets = None
key_object_permission.mcp_servers = ["server-key-own"]
team_obj = MagicMock()
team_obj.blocked = False
team_obj.object_permission = self._toolset_only_object_permission(["toolset-gone"])
team_obj.access_group_ids = []
team_obj.organization_id = None
mock_manager = self._mock_manager_with_toolsets({})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission
),
patch( # test-quality-ok: team-server resolution requires the proxy's module-global prisma client
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
"litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team_obj)
),
patch( # test-quality-ok: access-group lookup hits the DB, not under test here
"litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
AsyncMock(return_value=[]),
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])
),
):
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == []
async def test_org_toolset_restricts_tools_on_granted_server(self):
"""An org's toolset must act as the org tool ceiling, unioned with the
org's direct tool permissions"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", org_id="org-1")
org_object_permission = self._toolset_only_object_permission(["toolset-1"])
mock_manager = self._mock_manager_with_toolsets({"server-a": ["read_tool_1", "read_tool_2"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)
),
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission)
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
allowed = await MCPRequestHandler.get_allowed_tools_for_server(
server_id="server-a", user_api_key_auth=user_api_key_auth
)
write_tool_allowed = await MCPRequestHandler.is_tool_allowed_for_server(
tool_name="write_tool", server_id="server-a", user_api_key_auth=user_api_key_auth
)
assert allowed is not None
assert set(allowed) == {"read_tool_1", "read_tool_2"}
assert write_tool_allowed is False
async def test_org_toolset_servers_join_org_ceiling(self):
"""Servers reached only through the org's toolsets are part of the org
ceiling, exactly as servers named by its inline tool permissions"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", org_id="org-1")
org_object_permission = self._toolset_only_object_permission(["toolset-1"])
mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission)
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth)
assert result == ["server-a"]
async def test_user_toolset_restricts_tools(self):
"""An internal user's toolset must narrow tools like their inline
mcp_tool_permissions: intersecting a lower-level list, or becoming the
allowlist when no lower level restricts"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="user-1")
user_object_permission = self._toolset_only_object_permission(["toolset-1"])
mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_1", "tool_2"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission)
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
becomes_allowlist = await MCPRequestHandler._apply_user_tool_ceiling(None, "server-a", user_api_key_auth)
intersected = await MCPRequestHandler._apply_user_tool_ceiling(
["tool_1", "other_tool"], "server-a", user_api_key_auth
)
untouched_server = await MCPRequestHandler._apply_user_tool_ceiling(
["any_tool"], "server-without-toolset", user_api_key_auth
)
assert becomes_allowlist is not None and set(becomes_allowlist) == {"tool_1", "tool_2"}
assert intersected == ["tool_1"]
assert untouched_server == ["any_tool"]
async def test_user_toolset_servers_count_as_entitled(self):
"""Servers reached only through the user's toolsets count toward the
user's entitlement, so a toolset-only user ceiling caps to that server"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="user-1")
user_object_permission = self._toolset_only_object_permission(["toolset-1"])
mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_1"]})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission)
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
entitled = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth)
capped, restricts = await MCPRequestHandler._apply_user_server_ceiling(
["server-a", "server-b"], user_api_key_auth
)
assert list(entitled) == ["server-a"]
assert capped == ("server-a",)
assert restricts is True
async def test_team_declared_toolset_resolving_empty_denies_tools(self):
"""A team toolset whose ids resolve to nothing (deleted/unknown) is a KNOWN restriction
with unknown contents: tools on the granted server deny instead of falling open"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-1")
team_object_permission = self._toolset_only_object_permission(["toolset-deleted"])
mock_manager = self._mock_manager_with_toolsets({})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission)
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
allowed = await MCPRequestHandler.get_allowed_tools_for_server(
server_id="server-a", user_api_key_auth=user_api_key_auth
)
assert allowed == []
async def test_org_declared_toolset_resolving_empty_denies_servers(self):
"""An org whose only MCP grant is an unresolvable toolset must deny, never read as
'org places no restriction' and leave the caller uncapped"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", org_id="org-1")
org_object_permission = self._toolset_only_object_permission(["toolset-deleted"])
mock_manager = self._mock_manager_with_toolsets({})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission)
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
with pytest.raises(Exception, match="resolved to no grants"):
await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth)
async def test_user_declared_toolset_resolving_empty_still_places_ceiling(self):
"""An admin (or any user) whose row declares an unresolvable toolset keeps a ceiling:
the entitlement reads UNRESOLVED (deny), never 'no restriction'"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="user-1")
user_object_permission = self._toolset_only_object_permission(["toolset-deleted"])
mock_manager = self._mock_manager_with_toolsets({})
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission)
),
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
),
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
mock_manager,
),
):
entitled = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth)
places_ceiling = await MCPRequestHandler._user_places_mcp_ceiling(user_api_key_auth)
assert entitled is None
assert places_ceiling is True
async def test_declares_toolsets_gate_falls_back_to_db_for_unhydrated_key(self):
"""The main auth flow can cache a key with object_permission_id set but object_permission
unloaded; the declared-toolsets gate must fetch the row rather than answer False"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", object_permission_id="op-1")
key_object_permission = self._toolset_only_object_permission(["toolset-1"])
with (
patch( # test-quality-ok: team-server resolution requires the proxy's module-global prisma client
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
"litellm.proxy.auth.auth_checks.get_object_permission",
AsyncMock(return_value=key_object_permission),
),
):
declares = await MCPRequestHandler._key_or_team_declares_toolsets(user_api_key_auth)
assert declares is True
async def test_declares_toolsets_gate_swallows_team_lookup_fault(self):
"""An indeterminate fault while checking the team must answer False (org substitution
unchanged, matching base fault behavior), never escape as deny-all"""
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-gone")
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the DB team loader to drive the real team-server resolution path
MCPRequestHandler,
"_get_team_object_permission",
AsyncMock(side_effect=Exception("team lookup blew up")),
),
):
declares = await MCPRequestHandler._key_or_team_declares_toolsets(user_api_key_auth)
assert declares is False
async def test_declares_toolsets_gate_skips_team_lookup_for_teamless_key(self):
user_api_key_auth = UserAPIKeyAuth(api_key="test-key")
with (
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_key_object_permission", return_value=None
),
patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam
MCPRequestHandler, "_get_team_object_permission", AsyncMock()
) as team_lookup,
):
declares = await MCPRequestHandler._key_or_team_declares_toolsets(user_api_key_auth)
assert declares is False
team_lookup.assert_not_awaited()
async def test_permission_inheritance_edge_cases(self):
"""Test edge cases in permission inheritance"""
@ -1104,10 +1546,12 @@ class TestMCPOAuth2AuthFlow:
async def mock_user_api_key_auth(api_key, request):
return UserAPIKeyAuth(api_key=api_key, user_id="test-user")
with patch( # test-quality-ok: capturing the exact api_key handed to key validation is the regression under test
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
side_effect=mock_user_api_key_auth,
) as mock_auth:
with (
patch( # test-quality-ok: capturing the exact api_key handed to key validation is the regression under test
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
side_effect=mock_user_api_key_auth,
) as mock_auth
):
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope)
mock_auth.assert_called_once()
@ -4161,6 +4605,7 @@ class TestOrgMCPPermissions:
auth = self._make_auth(org_id="org-123")
mock_perm = MagicMock()
mock_perm.mcp_toolsets = None # a bare MagicMock attr reads as a DECLARED toolset and now denies
mock_perm.mcp_servers = ["org_server_1", "org_server_2"]
mock_perm.mcp_access_groups = []
mock_perm.mcp_tool_permissions = {}
@ -4186,6 +4631,7 @@ class TestOrgMCPPermissions:
auth = self._make_auth(org_id="org-123")
mock_perm = MagicMock()
mock_perm.mcp_toolsets = None # a bare MagicMock attr reads as a DECLARED toolset and now denies
mock_perm.mcp_servers = []
mock_perm.mcp_access_groups = ["group-a"]
mock_perm.mcp_tool_permissions = {}
@ -4211,6 +4657,7 @@ class TestOrgMCPPermissions:
auth = self._make_auth(org_id="org-123")
mock_perm = MagicMock()
mock_perm.mcp_toolsets = None # a bare MagicMock attr reads as a DECLARED toolset and now denies
mock_perm.mcp_servers = []
mock_perm.mcp_access_groups = []
mock_perm.mcp_tool_permissions = {"tool_only_server": ["tool_x"]}
@ -4251,6 +4698,7 @@ class TestOrgMCPPermissions:
key_perm.mcp_tool_permissions = {"server_1": ["tool_a", "tool_b", "tool_c"]}
org_perm = MagicMock()
org_perm.mcp_toolsets = None # a bare MagicMock attr reads as a DECLARED toolset and now denies
org_perm.mcp_tool_permissions = {"server_1": ["tool_a", "tool_b"]}
with (
@ -4281,6 +4729,7 @@ class TestOrgMCPPermissions:
key_perm.mcp_tool_permissions = {"server_1": ["tool_a", "tool_b"]}
org_perm = MagicMock()
org_perm.mcp_toolsets = None # a bare MagicMock attr reads as a DECLARED toolset and now denies
org_perm.mcp_tool_permissions = {}
with (
@ -6129,7 +6578,9 @@ class TestMCPDcrBridgeDelegateAdmission:
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling challenge tests
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), # test-quality-ok: envelope keys derive from the proxy master_key module global
patch( # test-quality-ok: envelope keys derive from the proxy master_key module global
"litellm.proxy.proxy_server.master_key", self._MASTER_KEY
),
):
mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server(
server_name="bridge_name", alias="bridge_alias"
@ -8438,7 +8889,7 @@ class TestUserMCPEntitlement:
result = await MCPRequestHandler._get_allowed_mcp_servers_for_user(self._auth())
finally:
global_mcp_server_manager.registry.pop("srv-a", None)
assert result == ["srv-a"]
assert list(result) == ["srv-a"]
async def test_places_ceiling_is_true_when_unresolvable(self):
"""``_user_places_mcp_ceiling`` gates the admin shortcut that hands over the whole registry, so