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 bd8dfea3621..d5461f01ed8 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 @@ -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: 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 aa6ddbfb49d..0144fbb17dd 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 @@ -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