scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(mcp): enforce end user mcp_tool_permissions on tools/list and tools/call (#40865)

Browse source 
* fix(mcp): apply end user mcp_tool_permissions as a tool ceiling on tools/list and tools/call

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(mcp): restore scoped session admission coverage dropped by mistake

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-12 10:33:09 -07:00 • committed by GitHub
parent e4f59a953c
commit 0c98afa780
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 111 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2162,6 +2162,10 @@ class MCPRequestHandler:
# No team restrictions → use key restrictions
allowed_tools = cast(list[str], key_tools)
allowed_tools = _as_list(
await MCPRequestHandler._apply_end_user_tool_ceiling(allowed_tools, server_id, user_api_key_auth)
)
allowed_tools = _as_list(
await MCPRequestHandler._apply_user_tool_ceiling(
allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source
@ -3027,6 +3031,38 @@ class MCPRequestHandler:
return list(user_tools)
return list(set(allowed_tools) & set(user_tools))
@staticmethod
async def _apply_end_user_tool_ceiling(
allowed_tools: Sequence[str] | None,
server_id: str,
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> Sequence[str] | None:
"""Narrow a key/team tool allowlist by the end user's (customer's) tool entitlement."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.proxy_server import prisma_client
if user_api_key_auth is None or not user_api_key_auth.end_user_id or prisma_client is None:
return allowed_tools
object_permissions: Final = await MCPRequestHandler._get_end_user_object_permission(
user_api_key_auth, prisma_client
)
if object_permissions is None:
return allowed_tools
end_user_direct_tools: Final = global_mcp_server_manager.expand_tool_permissions(
object_permissions.mcp_tool_permissions
).get(server_id)
end_user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(object_permissions, server_id)
end_user_tools: Final = MCPRequestHandler._union_tool_grants(end_user_direct_tools, end_user_toolset_tools)
if end_user_tools is None:
return allowed_tools
if allowed_tools is None:
return list(end_user_tools)
return list(set(allowed_tools) & set(end_user_tools))
# Sentinel stored in cache when an agent has no object_permission, so we
# don't re-query the DB on every MCP request for that agent.
_AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__"

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

@ -9266,17 +9266,24 @@ def _agent_prisma(object_permission_id=None, side_effect=None):
@contextlib.contextmanager
def _entitlement_fault_globals(prisma_client=None):
def _entitlement_fault_globals(prisma_client=None, user_api_key_cache=None):
from litellm.caching.dual_cache import DualCache
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client or MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache or DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", _proxy_logging_with_awaitable_hooks()),
):
yield
def _proxy_logging_with_awaitable_hooks():
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
return proxy_logging_obj
@pytest.mark.asyncio
class TestEntitlementFaultSemantics:
"""Each entitlement level distinguishes two fault classes for a KEY-authenticated caller.
@ -9412,6 +9419,71 @@ class TestEntitlementFaultSemantics:
assert set(allowed) == {"srv1"}
async def _cache_with_end_user(end_user_id, *, mcp_tool_permissions=None, object_permission_id=None):
"""A real DualCache already holding the end user row, so ``get_end_user_object`` answers from
cache and no ``litellm.`` internal has to be patched. ``object_permission_id`` without a
permission body models a row that NAMES an entitlement the DB then fails to serve."""
from litellm.caching.dual_cache import DualCache
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key
cache = DualCache()
await cache.async_set_cache(
key=end_user_cache_key(end_user_id),
value=LiteLLM_EndUserTable(
user_id=end_user_id,
blocked=False,
object_permission_id=object_permission_id or ("op-eu" if mcp_tool_permissions else None),
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions
)
if mcp_tool_permissions
else None,
),
)
return cache
@pytest.mark.asyncio
class TestEndUserToolCeiling:
"""The end user (customer) level narrows the TOOLS axis exactly as it narrows the servers axis,
so `object_permission.mcp_tool_permissions` on `/customer/new` is enforced, not just stored."""
async def test_end_user_tool_permissions_intersect_key_tools(self):
auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1")
cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]})
with _entitlement_fault_globals(user_api_key_cache=cache):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert tools == ["tool_a"]
async def test_end_user_tool_permissions_become_allowlist_when_key_is_unrestricted(self):
auth = _key_auth_reaching("srv1", end_user_id="eu-1")
cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]})
with _entitlement_fault_globals(user_api_key_cache=cache):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert tools == ["tool_a"]
async def test_end_user_tool_permissions_on_another_server_place_no_ceiling(self):
auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1")
cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv2": ["tool_z"]})
with _entitlement_fault_globals(user_api_key_cache=cache):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert sorted(tools) == ["tool_a", "tool_b"]
async def test_end_user_named_but_unloadable_permission_denies_tools(self):
auth = _key_auth_reaching("srv1", tools=["tool_a"], end_user_id="eu-1")
cache = await _cache_with_end_user("eu-1", object_permission_id="op-eu")
with _entitlement_fault_globals(user_api_key_cache=cache):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert tools == [], "an end-user entitlement we know exists but cannot read must deny its tools"
async def test_no_end_user_row_places_no_tool_ceiling(self):
auth = _key_auth_reaching("srv1", tools=["tool_a"], end_user_id="eu-1")
with _entitlement_fault_globals(user_api_key_cache=await _cache_with_end_user("someone-else")):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert tools == ["tool_a"]
@pytest.mark.asyncio
class TestScopedSessionAdmission:
"""LIT-4917: a session bearer sealed to one server (RFC 8707 resource at authorize)