From 27d1974e2fccd5eca826a01b5e38a23ae5299d65 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 24 Sep 2026 18:47:44 -0500 Subject: [PATCH] fix(mcp): cap an agent key's tools at what the invoking user and team may call (#42478) * fix(mcp): cap an agent key's tools at what the invoking user and team may call The invoking user's and team's x-litellm-user-id / x-litellm-team-id, echoed back by the agent, already narrowed which MCP servers the agent key could reach, but not which tools on those servers. An agent granted every tool on a server kept them all when acting for a user who may only call a subset. The caller's team and user tool grants now intersect the agent's tool list on each server, mirroring the servers axis, so the headers only ever narrow. Co-Authored-By: Claude Fable 5.1 * test(mcp): pick the caller principal explicitly instead of getattr in the tool grant stub Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): return immutable tool sequences from the agent caller tool ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Claude Fable 5.1 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 48 ++++++- .../auth/test_user_api_key_auth_mcp.py | 123 +++++++++++++++++- 2 files changed, 163 insertions(+), 8 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 43d62ba8b22..cdf52e6dc8d 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 @@ -2057,7 +2057,7 @@ class MCPRequestHandler: @staticmethod async def _get_team_object_permission( user_api_key_auth: UserAPIKeyAuth | None = None, - ): + ) -> LiteLLM_ObjectPermissionTable | None: """ Get team object_permission - automatically loaded by get_team_object() in main auth flow. @@ -2289,6 +2289,10 @@ class MCPRequestHandler: ) ) + allowed_tools = _as_list( + await MCPRequestHandler._apply_agent_caller_tool_ceiling(allowed_tools, server_id, user_api_key_auth) + ) + return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source ) @@ -3170,6 +3174,48 @@ class MCPRequestHandler: return list(user_tools) return list(set(allowed_tools) & set(user_tools)) + @staticmethod + async def _apply_agent_caller_tool_ceiling( + allowed_tools: Sequence[str] | None, + server_id: str, + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> Sequence[str] | None: + """Narrow an agent key's tools on ``server_id`` to those the invoking user and team (echoed back + by the agent as ``x-litellm-user-id`` / ``x-litellm-team-id``) may call: the echoed team's tool + grants when it names any on this server, then the echoed user's own tool entitlement. The tools + axis twin of ``_apply_agent_caller_ceiling``, so the headers only ever narrow. Denies every tool + on the server when the caller's team cannot be loaded, since a caller we cannot resolve must not + read as unrestricted.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None + if caller_auth is None: + return allowed_tools + try: + team_obj_perm: Final = await MCPRequestHandler._get_team_object_permission(caller_auth) + team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(team_obj_perm, server_id) + except Exception as e: # noqa: BLE001 # an unresolved caller team must deny, not widen + verbose_logger.warning( + "MCP agent caller team tool ceiling unresolvable, denying tools on %r: %s", server_id, e + ) + return () + 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 + ) + team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools) + team_capped: Final = ( + allowed_tools + if team_tools is None + else tuple(team_tools) + if allowed_tools is None + else tuple(frozenset(allowed_tools) & frozenset(team_tools)) + ) + return await MCPRequestHandler._apply_user_tool_ceiling(team_capped, server_id, caller_auth) + @staticmethod async def _apply_end_user_tool_ceiling( allowed_tools: Sequence[str] | 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 05ed53df8e8..24f09e79dbb 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 @@ -2,6 +2,7 @@ import contextlib import json import os from datetime import datetime, timedelta, timezone +from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -4284,6 +4285,105 @@ class TestAgentMCPPermissions: ): assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == [] + @staticmethod + def _tool_grants(grants: dict[str, dict[str, list[str]]], keyed_by: Literal["team_id", "user_id"]) -> AsyncMock: + """Object permissions keyed by the ``team_id`` or ``user_id`` being asked about; anyone else has none.""" + + async def by_principal(user_api_key_auth: UserAPIKeyAuth | None = None) -> LiteLLM_ObjectPermissionTable | None: + assert user_api_key_auth is not None + principal = (user_api_key_auth.team_id if keyed_by == "team_id" else user_api_key_auth.user_id) or "" + tools = grants.get(principal) + if tools is None: + return None + return LiteLLM_ObjectPermissionTable(object_permission_id=f"perm-{principal}", mcp_tool_permissions=tools) + + return AsyncMock(side_effect=by_principal) + + @contextlib.contextmanager + def _caller_tool_levels( + self, team_grants: dict[str, dict[str, list[str]]], user_grants: dict[str, dict[str, list[str]]] + ): + with ( + patch.object( # test-quality-ok: the level loaders read proxy_server globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=None + ), + patch.object( # test-quality-ok: same seam, keyed by which team is being asked about + MCPRequestHandler, "_get_team_object_permission", self._tool_grants(team_grants, keyed_by="team_id") + ), + patch.object( # test-quality-ok: same seam, keyed by which user is being asked about + MCPRequestHandler, "_get_user_object_permission", self._tool_grants(user_grants, keyed_by="user_id") + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=None) + ), + ): + yield + + async def test_agent_key_acting_for_a_user_only_sees_the_tools_that_user_may_call(self): + """The agent's key may call every tool on server-a, the invoking team grants two of them and the + invoking user only one, so on that user's behalf the agent sees exactly that one tool.""" + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + with self._caller_tool_levels( + team_grants={"callers": {"server-a": ["read_wiki_structure", "ask_wiki_question"]}}, + user_grants={"alice": {"server-a": ["read_wiki_structure", "read_wiki_contents"]}}, + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("server-a", agent_key) == [ + "read_wiki_structure" + ] + + async def test_agent_key_acting_for_a_user_without_tool_grants_keeps_its_own_tools(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + with self._caller_tool_levels(team_grants={}, user_grants={}): + assert await MCPRequestHandler.get_allowed_tools_for_server("server-a", agent_key) is None + + async def test_agent_key_acting_for_a_team_is_capped_at_that_teams_tools_on_the_server(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + with self._caller_tool_levels( + team_grants={"callers": {"server-a": ["ask_wiki_question"], "server-b": ["other"]}}, user_grants={} + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("server-a", agent_key) == ["ask_wiki_question"] + assert await MCPRequestHandler.get_allowed_tools_for_server("server-c", agent_key) is None + + async def test_agent_key_not_acting_for_anyone_ignores_the_caller_tool_ceiling(self): + agent_key = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + + with self._caller_tool_levels( + team_grants={"callers": {"server-a": ["ask_wiki_question"]}}, + user_grants={"alice": {"server-a": ["read_wiki_structure"]}}, + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("server-a", agent_key) is None + + async def test_agent_key_acting_for_a_caller_whose_team_is_unreadable_gets_no_tools(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + async def only_the_callers_team_is_unreadable( + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> LiteLLM_ObjectPermissionTable | None: + if user_api_key_auth is not None and user_api_key_auth.team_id == "callers": + raise RuntimeError("db down") + return None + + with ( + patch.object( # test-quality-ok: the level loaders read proxy_server globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=None + ), + patch.object( # test-quality-ok: same seam; the agent's own team resolves, the caller's team does not + MCPRequestHandler, + "_get_team_object_permission", + AsyncMock(side_effect=only_the_callers_team_is_unreadable), + ), + patch.object( # test-quality-ok: same seam + MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=None) + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=None) + ), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("server-a", agent_key) == [] + async def test_get_allowed_mcp_servers_agent_intersection(self): """Key/team allow [server_1, server_2]; agent allows [server_1]. Result = [server_1].""" user_api_key_auth = UserAPIKeyAuth( @@ -4354,8 +4454,7 @@ class TestAgentMCPPermissions: assert result == frozenset({"ag-server-id"}) assert asked == ["agent-ag"] assert ( - await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k"), resolve) - is None + await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k"), resolve) is None ) assert asked == ["agent-ag"] @@ -4492,7 +4591,9 @@ class TestAgentMCPPermissions: stack.enter_context(patcher) stack.enter_context( patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here - MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + AsyncMock(return_value=["server-a", "server-b"]), ) ) stack.enter_context( @@ -4518,7 +4619,9 @@ class TestAgentMCPPermissions: await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) stack.enter_context( patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here - MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + AsyncMock(return_value=["server-a", "server-b"]), ) ) stack.enter_context( @@ -4542,9 +4645,15 @@ class TestAgentMCPPermissions: with contextlib.ExitStack() as stack: for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): stack.enter_context(patcher) - server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth) - server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth) - server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth) + server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-a", user_api_key_auth + ) + server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-b", user_api_key_auth + ) + server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-c", user_api_key_auth + ) assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"] assert server_b_tools == ["tool_b"]