From 55eea3ac36fdb263279eea4e3fe776c94b7e6dab Mon Sep 17 00:00:00 2001 From: Yucheng Low Date: Wed, 26 Aug 2026 23:50:34 -0700 Subject: [PATCH 1/8] 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 --- .../mcp_server/auth/user_api_key_auth_mcp.py | 150 ++++++- .../auth/test_user_api_key_auth_mcp.py | 379 +++++++++++++++++- 2 files changed, 510 insertions(+), 19 deletions(-) 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..799f8dc7d19 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) ) @@ -2546,7 +2661,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 +2861,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 +2880,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 +2982,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..c63576e150f 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,364 @@ 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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object( + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission) + ), + patch( + "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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object( + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission) + ), + patch( + "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( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team_obj)), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), + patch.object( + 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(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])), + patch.object(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object( + 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_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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch.object( + MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission) + ), + patch( + "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( + MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission) + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch( + "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( + MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission) + ), + patch( + "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( + MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission) + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch( + "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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object( + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=team_object_permission) + ), + patch( + "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( + MCPRequestHandler, "_get_org_object_permission", AsyncMock(return_value=org_object_permission) + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch( + "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( + MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=user_object_permission) + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch( + "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("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object( + 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(MCPRequestHandler, "_get_key_object_permission", return_value=None), + patch.object(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 +1462,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 +4521,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 +4547,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 +4573,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 +4614,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 +4645,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 +6494,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( + "litellm.proxy.proxy_server.master_key", self._MASTER_KEY + ), # test-quality-ok: envelope keys derive from the proxy master_key module global ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server( server_name="bridge_name", alias="bridge_alias" @@ -8438,7 +8805,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 From 987e7623c9ef74cf9febf0e708420194fc86569d Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 27 Aug 2026 10:16:23 -0700 Subject: [PATCH 2/8] 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. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 + .../auth/test_user_api_key_auth_mcp.py | 184 +++++++++++++----- 2 files changed, 136 insertions(+), 50 deletions(-) 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 799f8dc7d19..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 @@ -2530,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 [] 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 c63576e150f..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 @@ -518,11 +518,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels", "read_thread"]}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object( + 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( + 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, ), @@ -550,11 +552,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object( + 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( + 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, ), @@ -574,11 +578,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"], "server-b": ["get_doc"]}) with ( - patch( + 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(MCPRequestHandler, "_get_mcp_servers_from_access_groups", 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=[]) + ), ): servers = await MCPRequestHandler._team_granted_servers(team_obj, []) @@ -597,19 +603,27 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team_obj)), - patch( + 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( + 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(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), - patch.object( + 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"]), @@ -627,15 +641,23 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])), - patch.object(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), - patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), - patch( + 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( + 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"]), @@ -645,6 +667,46 @@ class TestMCPRequestHandler: 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""" @@ -653,12 +715,16 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["read_tool_1", "read_tool_2"]}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), - patch.object( + 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( + 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, ), @@ -682,11 +748,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["search_channels"]}) with ( - patch.object( + 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(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), - patch( + 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, ), @@ -704,10 +772,10 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_1", "tool_2"]}) with ( - patch.object( + 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( + 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, ), @@ -732,11 +800,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_1"]}) with ( - patch.object( + 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(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), - patch( + 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, ), @@ -758,11 +828,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({}) with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object( + 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( + 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, ), @@ -781,11 +853,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({}) with ( - patch.object( + 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(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), - patch( + 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, ), @@ -801,11 +875,13 @@ class TestMCPRequestHandler: mock_manager = self._mock_manager_with_toolsets({}) with ( - patch.object( + 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(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), - patch( + 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, ), @@ -823,8 +899,10 @@ class TestMCPRequestHandler: key_object_permission = self._toolset_only_object_permission(["toolset-1"]) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch( + 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), ), @@ -839,8 +917,10 @@ class TestMCPRequestHandler: user_api_key_auth = UserAPIKeyAuth(api_key="test-key", team_id="team-gone") with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object( + 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")), @@ -854,8 +934,12 @@ class TestMCPRequestHandler: user_api_key_auth = UserAPIKeyAuth(api_key="test-key") with ( - patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None), - patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock()) as team_lookup, + 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) @@ -6494,9 +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( + patch( # test-quality-ok: envelope keys derive from the proxy master_key module global "litellm.proxy.proxy_server.master_key", self._MASTER_KEY - ), # test-quality-ok: envelope keys derive from the proxy master_key module global + ), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server( server_name="bridge_name", alias="bridge_alias" From ceea12e434107b310445171b04dd76967a3636e0 Mon Sep 17 00:00:00 2001 From: Yucheng Low Date: Fri, 28 Aug 2026 10:26:32 -0700 Subject: [PATCH 3/8] test: e2e coverage for MCP toolset enforcement at team, org, and user levels Co-Authored-By: Claude Fable 5 --- tests/e2e/coverage_registry/mcp.yaml | 54 ++++++ tests/e2e/mcp/datadog_mcp.py | 7 +- tests/e2e/mcp/mcp_client.py | 152 +++++++++++++-- .../mcp/test_mcp_toolset_enforcement_e2e.py | 181 ++++++++++++++++++ tests/e2e/models.py | 4 + 5 files changed, 382 insertions(+), 16 deletions(-) create mode 100644 tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..d0f060dc07d 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,57 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.team_toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:2453" + rationale: "A toolset attached to a team's object_permission narrows the team key's tools/list to exactly the tools it names (LIT-5749: stored but inert before the fix)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.team_toolset_denied_outside + module: mcp + tier: P0 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "A team key calling a tool on the granted server but outside the team's toolset is refused; the named tool stays callable (LIT-5749)" + fail_before_fix: proven +- id: mcp.list_tools.api_key.org_toolset_scoped + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:1449" + rationale: "A toolset on an org's object_permission caps a member team's key that declares no MCP grants of its own (LIT-5749: the org substitution amplifier)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.org_toolset_denied_outside + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "An org-inherited key calling outside the org's toolset is refused while the named tool stays callable (LIT-5749)" + fail_before_fix: proven +- id: mcp.list_tools.api_key.user_toolset_scoped + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:2859" + rationale: "A toolset on an internal user's row ceils the servers the user's own key grants; user-level toolsets narrow, never grant (LIT-5749)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.user_toolset_denied_outside + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "A key whose user row holds a toolset is refused calling outside it even though the key itself grants the whole server (LIT-5749)" + fail_before_fix: proven diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..54c5899812b 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the Datadog remote MCP server. `allowed_tools` defaults to the + search-logs slice; pass None to expose the full core toolset (needed when a + test must prove narrowing, so a second tool has to exist to be denied).""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=list(allowed_tools) if allowed_tools is not None else None, mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..8bd61e04335 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -20,7 +20,19 @@ from pydantic import BaseModel, ConfigDict, Field, RootModel from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import ( + KeyGenerateBody, + ObjectPermission, + OrgDeleteBody, + OrgNewBody, + OrgNewResponse, + TeamDeleteBody, + TeamNewBody, + TeamNewResponse, + UserDeleteBody, + UserNewBody, + UserNewResponse, +) from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,6 +58,20 @@ class McpServerNewResponse(BaseModel): server_id: str +class McpToolsetToolSpec(BaseModel): + server_id: str + tool_name: str + + +class McpToolsetNewBody(BaseModel): + toolset_name: str + tools: list[McpToolsetToolSpec] + + +class McpToolsetNewResponse(BaseModel): + toolset_id: str + + class McpServerRow(BaseModel): server_id: str alias: str | None = None @@ -74,9 +100,7 @@ class McpToolsListResponse(BaseModel): def tool_names_for_server(self, server_id: str) -> frozenset[str]: return frozenset( - tool.name - for tool in self.tools - if tool.mcp_info is not None and tool.mcp_info.server_id == server_id + tool.name for tool in self.tools if tool.mcp_info is not None and tool.mcp_info.server_id == server_id ) def tool_name_containing(self, server_id: str, needle: str) -> str | None: @@ -225,6 +249,7 @@ class McpClient: mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, models: list[str] | None = None, + team_id: str | None = None, ) -> str: object_permission = ( ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) @@ -235,10 +260,114 @@ class McpClient: KeyGenerateBody( models=models if models is not None else [], user_id=user_id, + team_id=team_id, object_permission=object_permission, ) ) + def create_toolset(self, *, name: str, server_id: str, tool_names: list[str]) -> str: + """Create a named toolset over `server_id` (admin), returning its id.""" + return unwrap( + self.proxy.transport.post( + "/v1/mcp/toolset", + headers=self.proxy.transport.master, + json=McpToolsetNewBody( + toolset_name=name, + tools=[McpToolsetToolSpec(server_id=server_id, tool_name=tool) for tool in tool_names], + ), + response_type=McpToolsetNewResponse, + ) + ).toolset_id + + def delete_toolset(self, toolset_id: str) -> None: + _ = self.proxy.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + def create_team( + self, + *, + alias: str, + object_permission: ObjectPermission | None = None, + organization_id: str | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/team/new", + headers=self.proxy.transport.master, + json=TeamNewBody( + team_alias=alias, + organization_id=organization_id, + object_permission=object_permission, + ), + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + _ = self.proxy.transport.post( + "/team/delete", + headers=self.proxy.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def create_org( + self, + *, + alias: str, + object_permission: ObjectPermission | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/organization/new", + headers=self.proxy.transport.master, + json=OrgNewBody( + organization_alias=alias, + object_permission=object_permission, + ), + response_type=OrgNewResponse, + ) + ).organization_id + + def delete_org(self, organization_id: str) -> None: + _ = self.proxy.transport.delete( + "/organization/delete", + headers=self.proxy.transport.master, + json=OrgDeleteBody(organization_ids=[organization_id]), + response_type=NoBody, + ) + + def create_user( + self, + *, + user_email: str, + object_permission: ObjectPermission | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/user/new", + headers=self.proxy.transport.master, + json=UserNewBody( + user_email=user_email, + user_role="internal_user", + object_permission=object_permission, + ), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + def list_tools(self, key: str) -> Result[McpToolsListResponse]: return self.proxy.transport.get( "/mcp-rest/tools/list", @@ -316,13 +445,10 @@ class McpClient: if isinstance(last, UnknownApiError) and last.status_code == 403: return last if not _is_mcp_not_synced(last, tool_name=name): - raise AssertionError( - f"ungranted key's tools/call was not 403 access_denied: {last}" - ) + raise AssertionError(f"ungranted key's tools/call was not 403 access_denied: {last}") if time.monotonic() >= deadline: raise AssertionError( - f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; " - f"last result: {last}" + f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; last result: {last}" ) time.sleep(self.proxy.poll_interval) @@ -368,9 +494,7 @@ class McpClient: return self.proxy.transport.post( "/mcp-rest/tools/call", headers=ApiKeyHeaders(x_litellm_api_key=key), - json=McpCallToolBody( - name=name, arguments=dict(arguments), server_id=server_id - ), + json=McpCallToolBody(name=name, arguments=dict(arguments), server_id=server_id), response_type=McpCallToolResponse, ) @@ -403,9 +527,7 @@ def _is_mcp_not_synced( # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message) if tool_name is not None: - return ( - re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None - ) + return re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..b6c1f66328d --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,181 @@ +"""Live e2e: an MCP toolset attached to a team, org, or internal user narrows a +real MCP server to exactly the tools the toolset names (LIT-5749). + +An admin registers the Datadog remote MCP server with its full core toolset (no +`allowed_tools` cap, so more than one tool exists to be denied), creates a +toolset naming only `search_datadog_logs` on it, and attaches that toolset at +one principal level per test. A control key granted the whole server first +proves the upstream is alive and serves a second tool, so a later denial is an +authorization decision rather than a dead server. The scoped key must then see +exactly the toolset's one tool on `tools/list`, be refused (403) calling any +other tool on the same server, and still successfully call the granted tool. + +Levels covered, one per test: a team holding server + toolset, a team that +inherits its grant from an org holding server + toolset, and an internal user +whose toolset ceils the servers their own key grants. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import DD_SEARCH_FROM, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ObjectPermission + +pytestmark = pytest.mark.e2e + + +def _open_server_and_toolset(client: McpClient, resources: ResourceManager) -> tuple[str, str]: + """Register the DD server uncapped and a toolset naming only the search-logs + tool on it, returning (server_id, toolset_id).""" + server_id = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + toolset_id = client.create_toolset( + name=f"e2e-search-only-{unique_marker()}", + server_id=server_id, + tool_names=[SEARCH_LOGS_TOOL], + ) + resources.defer(lambda: client.delete_toolset(toolset_id)) + return server_id, toolset_id + + +def _control_tools(client: McpClient, resources: ResourceManager, server_id: str) -> tuple[str, str]: + """A fully-granted control key's view of the server: the fully-qualified + search-logs tool name and one tool outside the toolset. Proves the upstream + is alive and actually serves something a toolset can deny.""" + control_key = client.generate_key(user_id=f"e2e-ts-control-{unique_marker()}", mcp_servers=[server_id]) + resources.defer(lambda: client.proxy.delete_key(control_key)) + granted_tool = client.await_tool(control_key, server_id, SEARCH_LOGS_TOOL) + all_tools = unwrap(client.list_tools(control_key)).tool_names_for_server(server_id) + outside = sorted(all_tools - {granted_tool}) + assert outside, ( + f"the uncapped Datadog core toolset served only {all_tools}; a toolset " + f"cannot be proven to narrow a one-tool server" + ) + return granted_tool, outside[0] + + +def _assert_toolset_ceiling( + client: McpClient, + scoped_key: str, + *, + server_id: str, + granted_tool: str, + outside_tool: str, +) -> None: + """The scoped key sees exactly the toolset's tool, is 403-refused on a tool + outside it, and can still execute the granted one.""" + _ = client.await_tool(scoped_key, server_id, SEARCH_LOGS_TOOL) + listed = unwrap(client.list_tools(scoped_key)).tool_names_for_server(server_id) + assert listed == frozenset({granted_tool}), ( + f"toolset-scoped key must list exactly {granted_tool!r}; the toolset ceiling leaked: {sorted(listed)}" + ) + + denied = client.await_call_tool_denied(scoped_key, server_id=server_id, name=outside_tool, arguments={}) + assert denied.status_code == 403 + + result = client.await_call_tool( + scoped_key, + server_id=server_id, + name=granted_tool, + arguments={"query": f"service:e2e-toolset-{unique_marker()}", "from": DD_SEARCH_FROM}, + ) + assert result.is_error is not True, f"the toolset-granted tool must stay callable; upstream said: {result.all_text}" + + +class TestMcpToolsetEnforcementPerLevel: + @pytest.mark.covers("mcp.list_tools.api_key.team_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.team_toolset_denied_outside") + def test_team_toolset_narrows_team_key( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + team_id = client.create_team( + alias=f"e2e-ts-team-{unique_marker()}", + object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_team(team_id)) + team_key = client.generate_key( + user_id=f"e2e-ts-team-user-{unique_marker()}", + mcp_servers=None, + team_id=team_id, + ) + resources.defer(lambda: client.proxy.delete_key(team_key)) + + _assert_toolset_ceiling( + client, + team_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) + + @pytest.mark.covers("mcp.list_tools.api_key.org_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.org_toolset_denied_outside") + def test_org_toolset_caps_inherited_team_key( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + org_id = client.create_org( + alias=f"e2e-ts-org-{unique_marker()}", + object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_org(org_id)) + # The team declares no MCP grants of its own; whatever its key can reach + # comes from the org, so the org's toolset must cap it. + team_id = client.create_team(alias=f"e2e-ts-org-team-{unique_marker()}", organization_id=org_id) + resources.defer(lambda: client.delete_team(team_id)) + team_key = client.generate_key( + user_id=f"e2e-ts-org-user-{unique_marker()}", + mcp_servers=None, + team_id=team_id, + ) + resources.defer(lambda: client.proxy.delete_key(team_key)) + + _assert_toolset_ceiling( + client, + team_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) + + @pytest.mark.covers("mcp.list_tools.api_key.user_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.user_toolset_denied_outside") + def test_user_toolset_ceils_own_key_grant( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + # The user's row holds only the toolset; their key grants the whole + # server. The user-level toolset is a ceiling over the key's grant. + user_id = client.create_user( + user_email=f"e2e-ts-user-{unique_marker()}@example.com", + object_permission=ObjectPermission(mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_user(user_id)) + user_key = client.generate_key(user_id=user_id, mcp_servers=[server_id]) + resources.defer(lambda: client.proxy.delete_key(user_key)) + + _assert_toolset_ceiling( + client, + user_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 95a02b58824..331ea3fbb55 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -52,6 +52,7 @@ class KeyMetadata(BaseModel): class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -922,6 +923,7 @@ class TeamNewBody(BaseModel): team_id: str | None = None organization_id: str | None = None metadata: TeamMetadata | None = None + object_permission: ObjectPermission | None = None class TeamNewResponse(BaseModel): @@ -979,6 +981,7 @@ class UserNewBody(BaseModel): user_email: str user_role: UserRole user_id: str | None = None + object_permission: ObjectPermission | None = None class UserNewResponse(BaseModel): @@ -1029,6 +1032,7 @@ class UserListResponse(BaseModel): class OrgNewBody(BaseModel): organization_alias: str models: list[str] = [] + object_permission: ObjectPermission | None = None class OrgNewResponse(BaseModel): From 9966b8855cabd4884725ff86147169bb758feb77 Mon Sep 17 00:00:00 2001 From: Yucheng Low Date: Fri, 28 Aug 2026 10:34:46 -0700 Subject: [PATCH 4/8] test: poll the exact tool set to the shared deadline instead of trusting one replica's sync Co-Authored-By: Claude Fable 5 --- .../mcp/test_mcp_toolset_enforcement_e2e.py | 38 +++++++++++++------ 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py index b6c1f66328d..fe10bf7c0e9 100644 --- a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -17,8 +17,9 @@ whose toolset ceils the servers their own key grants. from __future__ import annotations -import pytest +import time +import pytest from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker from e2e_http import unwrap @@ -59,6 +60,25 @@ def _control_tools(client: McpClient, resources: ResourceManager, server_id: str return granted_tool, outside[0] +def _await_exact_tools(client: McpClient, key: str, server_id: str, expected: frozenset[str]) -> None: + """Poll tools/list until the key's view of `server_id` is exactly `expected`. + + A single successful listing proves one replica has synced the grant rows; the + next request may land on another. Polling the exact set to the shared deadline + is the propagation barrier, and a genuine enforcement leak stays red because + the leaked set never converges to `expected`.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + listed = unwrap(client.list_tools(key)).tool_names_for_server(server_id) + if listed == expected: + return + if time.monotonic() >= deadline: + raise AssertionError( + f"toolset-scoped key must list exactly {sorted(expected)}; the toolset ceiling leaked: {sorted(listed)}" + ) + time.sleep(client.proxy.poll_interval) + + def _assert_toolset_ceiling( client: McpClient, scoped_key: str, @@ -68,12 +88,10 @@ def _assert_toolset_ceiling( outside_tool: str, ) -> None: """The scoped key sees exactly the toolset's tool, is 403-refused on a tool - outside it, and can still execute the granted one.""" - _ = client.await_tool(scoped_key, server_id, SEARCH_LOGS_TOOL) - listed = unwrap(client.list_tools(scoped_key)).tool_names_for_server(server_id) - assert listed == frozenset({granted_tool}), ( - f"toolset-scoped key must list exactly {granted_tool!r}; the toolset ceiling leaked: {sorted(listed)}" - ) + outside it, and can still execute the granted one. Every leg polls to the + shared deadline so it cannot flake on a replica that has not synced the + just-created grant rows yet.""" + _await_exact_tools(client, scoped_key, server_id, frozenset({granted_tool})) denied = client.await_call_tool_denied(scoped_key, server_id=server_id, name=outside_tool, arguments={}) assert denied.status_code == 403 @@ -133,8 +151,7 @@ class TestMcpToolsetEnforcementPerLevel: object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), ) resources.defer(lambda: client.delete_org(org_id)) - # The team declares no MCP grants of its own; whatever its key can reach - # comes from the org, so the org's toolset must cap it. + # the team declares no MCP grants of its own, so only the org's cap applies team_id = client.create_team(alias=f"e2e-ts-org-team-{unique_marker()}", organization_id=org_id) resources.defer(lambda: client.delete_team(team_id)) team_key = client.generate_key( @@ -162,8 +179,7 @@ class TestMcpToolsetEnforcementPerLevel: server_id, toolset_id = _open_server_and_toolset(client, resources) granted_tool, outside_tool = _control_tools(client, resources, server_id) - # The user's row holds only the toolset; their key grants the whole - # server. The user-level toolset is a ceiling over the key's grant. + # the key grants the whole server; the user row's toolset must ceil it user_id = client.create_user( user_email=f"e2e-ts-user-{unique_marker()}@example.com", object_permission=ObjectPermission(mcp_toolsets=[toolset_id]), From 854266d56282af174b66298b127d29fdffe2186f Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:41:03 -0700 Subject: [PATCH 5/8] test(mcp): reuse the successful catalog snapshot for toolset controls --- tests/e2e/mcp/mcp_client.py | 9 +++++++-- tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py | 12 ++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 2f6f905fb51..71e2ad43c67 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -278,8 +278,13 @@ class McpClient: return self.await_tool_entry(key, server_id, needle).name def await_tool_entry(self, key: str, server_id: str, needle: str) -> McpToolEntry: + tool = self.await_tool_catalog(key, server_id, needle).tool_containing(server_id, needle) + assert tool is not None + return tool + + def await_tool_catalog(self, key: str, server_id: str, needle: str) -> McpToolsListResponse: """Poll tools/list until `server_id` serves a tool matching `needle`, and - return the successful tool snapshot. Fails at poll_timeout. + return that catalog snapshot. Fails at poll_timeout. /v1/mcp/server returns as soon as the DB row is written, but the gateway runs the initialize + tools/list handshake against the upstream lazily on @@ -293,7 +298,7 @@ class McpClient: if isinstance(result, Success): tool = result.data.tool_containing(server_id, needle) if tool is not None: - return tool + return result.data if time.monotonic() >= deadline: raise AssertionError( f"server {server_id} never served a tool matching {needle!r} within " diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py index 1c1f725fff4..1deafdf0ecc 100644 --- a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -75,8 +75,10 @@ class TestMcpToolsetEnforcement: client.await_registered(server_id) catalog_key: Final = _key(client, resources, "catalog", server_id=server_id) - known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL) - catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id) + snapshot: Final = client.await_tool_catalog(catalog_key, server_id, SEARCH_LOGS_TOOL) + known_wire: Final = snapshot.tool_name_containing(server_id, SEARCH_LOGS_TOOL) + assert known_wire is not None + catalog: Final = snapshot.tool_names_for_server(server_id) assert len(catalog) > 2, ( f"the Datadog core toolset must serve more tools than the toolset names, or the " f"restriction has nothing to hide; got {sorted(catalog)}" @@ -149,8 +151,10 @@ def _assert_principal_toolset(client: McpClient, resources: ResourceManager, pri for transport in client.proxy.replicas_for("/mcp-rest/tools/list").values(): replica = McpClient(proxy=replace(client.proxy, transport=transport)) - granted = replica.await_tool_entry(control_key, server_id, SEARCH_LOGS_TOOL) - catalog = unwrap(replica.list_tools(control_key)).tool_names_for_server(server_id) + snapshot = replica.await_tool_catalog(control_key, server_id, SEARCH_LOGS_TOOL) + granted = snapshot.tool_containing(server_id, SEARCH_LOGS_TOOL) + assert granted is not None + catalog = snapshot.tool_names_for_server(server_id) outside = sorted(catalog - {granted.name}) assert outside, f"uncapped upstream must expose a tool outside the grant: {catalog}" expected = frozenset({granted.name}) From 68abd0d867355cc71f054e6c69e0311d6714a2b9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:44:29 -0700 Subject: [PATCH 6/8] test(mcp): assert the toolset-specific denial reason --- tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py index 1deafdf0ecc..135830c52ed 100644 --- a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -161,7 +161,7 @@ def _assert_principal_toolset(client: McpClient, resources: ResourceManager, pri assert replica.await_tools(scoped_key, server_id, expected=expected) == expected denied = replica.call_tool(scoped_key, server_id=server_id, name=outside[0], arguments={}) assert isinstance(denied, UnknownApiError) and denied.status_code == 403, denied - assert "access_denied" in denied.body, denied.body + assert "is not allowed for your key/team on server" in denied.body, denied.body arguments = {"query": f"service:e2e-toolset-{unique_marker()}", "from": DD_SEARCH_FROM} granted.assert_arguments_are_documented(arguments) result = unwrap(replica.call_tool(scoped_key, server_id=server_id, name=granted.name, arguments=arguments)) From a6759ff3c32b9ec3a3545d3d828b2356c5204487 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 00:50:20 +0000 Subject: [PATCH 7/8] docs(e2e): cite the Datadog schema source and scope the health row Co-Authored-By: bot_apk --- tests/e2e/coverage_registry/README.md | 2 +- tests/e2e/mcp/datadog_mcp.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index 8bf71f7f430..f52856a4125 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -120,7 +120,7 @@ pending. Integration contract IDs belong in `tests/integration/contracts.json`, | Requirement | Existing regression or owning suite | Remaining acceptance and owner | Required before | |---|---|---|---| -| Principal discovery | `mcp/test_mcp_key_access_e2e.py`, `mcp/test_mcp_access_group_e2e.py`, `mcp/test_mcp_toolset_enforcement_e2e.py` | Run key/team/org/user controls on every configured replica; native/REST parity remains LIT-4506 | Phase 0 and affected authorization changes | +| Principal discovery | `mcp/test_mcp_key_access_e2e.py`, `mcp/test_mcp_access_group_e2e.py`, `mcp/test_mcp_toolset_enforcement_e2e.py` | Run key/team/org/user controls on every configured replica; the E2E health check intersects with test-owned servers and proves grant visibility only, non-disclosure of unrelated servers is proven by integration contract other.mcp.health.restricted_keys_intersect_grants_in_both_modes (#41731); native/REST parity remains LIT-4506 | Phase 0 and affected authorization changes | | No self-attached unauthorized grants | Management key authorization tests | Read back unchanged server/toolset/access-group grants after rejected writes, LIT-4502 | Affected grant capability activation | | UI/API parity | Admin MCP UI suite | Same non-admin actor and permissions across both surfaces, LIT-3644 / LIT-4506 | Affected UI capability activation | | Server identity routing | Saved-server lifecycle integration and resolver tests | Cold routing, duplicate/unprefixed names, LIT-4500 | Routing changes | diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index 352b4446cfd..e16f62d7be6 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -1,4 +1,9 @@ -"""Shared helpers for e2e tests that register the real Datadog remote MCP server.""" +"""Shared helpers for e2e tests that register the real Datadog remote MCP server. + +The search_datadog_logs arguments these suites send (query, from, to, max_tokens) come from +https://docs.datadoghq.com/mcp_server/tools/ (checked 2026-09-19). McpToolEntry.assert_arguments_are_documented +compares them with the schema the gateway advertises live, so a failure there after a Datadog rename is stale, +not broken litellm code""" from __future__ import annotations From 8a1c76549cfbd13e85a42e50aa36ddaaeea2f338 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 19 Sep 2026 00:52:59 +0000 Subject: [PATCH 8/8] test(mcp): credit the absorbed inventory and schema guard work Records co-authorship for the #34055 inventory and #35405 input-schema guard consolidated in this branch Co-authored-by: mubashir1osmani Co-Authored-By: bot_apk