scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

fix(mcp): cap an agent key's tools at what the invoking user and team may call (#42478)

Browse source 
* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-09-24 18:47:44 -05:00 • committed by GitHub
parent dc83a9c979
commit 27d1974e2f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 163 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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