From 0f52da9232aa0604bb0114782d676b319f021ff2 Mon Sep 17 00:00:00 2001 From: lewis Date: Sat, 17 Oct 2026 21:11:21 -0400 Subject: [PATCH] fix(proxy): enforce customer MCP denial before open grants --- .../mcp_server/auth/user_api_key_auth_mcp.py | 43 +++++-- .../mcp_server/mcp_server_manager.py | 10 +- .../mcp_server/test_mcp_server_manager.py | 111 ++++++++++++++++++ 3 files changed, 151 insertions(+), 13 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 235580c18ad..baac6b0cbd3 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 @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timezone -from types import MappingProxyType +from types import EllipsisType, MappingProxyType from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import HTTPException @@ -464,6 +464,7 @@ def _admission_failure_fallback( class MCPServerAccess: server_ids: tuple[str, ...] scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped" + end_user_denies_mcp: bool = False @dataclass(frozen=True, slots=True) @@ -1830,7 +1831,21 @@ class MCPRequestHandler: Returns: List[str]: List of allowed MCP servers by server id """ - from litellm.proxy.proxy_server import general_settings + from litellm.proxy.proxy_server import general_settings, prisma_client + + try: + end_user_permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if user_api_key_auth and user_api_key_auth.end_user_id and prisma_client is not None + else None + ) + except UnloadableEntitlementError as error: + verbose_logger.warning("Denying MCP access, customer entitlement unreadable: %s", error) + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) + if end_user_permission is not None and SpecialMCPServerNames.no_mcp_servers.value in ( + end_user_permission.mcp_servers or [] + ): + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) if managed_agent_policy(user_api_key_auth) is not None: from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers @@ -1911,7 +1926,7 @@ class MCPRequestHandler: ######################################################### if user_api_key_auth and user_api_key_auth.end_user_id: allowed_mcp_servers_for_end_user: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( - user_api_key_auth + user_api_key_auth, object_permission=end_user_permission ) # If end_user has explicit MCP server permissions, apply intersection @@ -3428,6 +3443,8 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_end_user( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + object_permission: LiteLLM_ObjectPermissionTable | None | EllipsisType = ..., ) -> list[str]: """ Get allowed MCP servers for an end user. @@ -3445,16 +3462,18 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - object_permission: Final = await MCPRequestHandler._get_end_user_object_permission( - user_api_key_auth, prisma_client + permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if object_permission is ... + else object_permission ) - if object_permission is None: + if permission is None: return [] - if SpecialMCPServerNames.no_mcp_servers.value in (object_permission.mcp_servers or []): + if SpecialMCPServerNames.no_mcp_servers.value in (permission.mcp_servers or []): return [SpecialMCPServerNames.no_mcp_servers.value] toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( - object_permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy ) try: @@ -3464,24 +3483,24 @@ class MCPRequestHandler: ) direct_mcp_servers: Final = global_mcp_server_manager.expand_permission_list( - object_permission.mcp_servers or [] + permission.mcp_servers or [] ) # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( - object_permission.mcp_access_groups or [], + permission.mcp_access_groups or [], requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) # servers referenced in tool permissions should also be accessible tool_perm_servers: Final = list( - global_mcp_server_manager.expand_tool_permissions(object_permission.mcp_tool_permissions).keys() + global_mcp_server_manager.expand_tool_permissions(permission.mcp_tool_permissions).keys() ) all_servers: Final = frozenset( (*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants) ) - return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, object_permission)) + return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, permission)) except Exception as e: raise UnloadableEntitlementError( f"end user {user_api_key_auth.end_user_id!r} has an object permission " diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b724a61cd94..b83e5a1b496 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3691,6 +3691,14 @@ class MCPServerManager: ) try: + customer_access: Final = ( + access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + if user_api_key_auth and user_api_key_auth.end_user_id + else None + ) + if customer_access is not None and customer_access.end_user_denies_mcp: + return [] + # Admin view with no explicit object permission and no entitlement ceiling resolves the # whole registry, for keys AND admitted session subjects alike (one predicate owns the # question). Seeded into the union rather than returned early so the session resource @@ -3699,7 +3707,7 @@ class MCPServerManager: resolved_access: Final = ( MCPServerAccess(server_ids=()) if admin_unscoped - else access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + else customer_access or access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) ) resolved_server_ids: Final = ( set(self.get_registry().keys()) if admin_unscoped else set(resolved_access.server_ids) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 40901d1dd96..9a0129a06b4 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20292,3 +20292,114 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["manager", "mcp", "rest"]) +@pytest.mark.parametrize("admin", [False, True]) +@pytest.mark.parametrize("permission_state", ["denied", "stale-grants", "unreadable"]) +async def test_customer_denial_blocks_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + surface: Literal["manager", "mcp", "rest"], + admin: bool, + permission_state: Literal["denied", "stale-grants", "unreadable"], +) -> None: + from starlette.requests import Request + + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations, rest_endpoints + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="customer-denial", + mcp_servers=["no-mcp-servers"], + mcp_tool_permissions={server.server_id: ["delete"]} if permission_state == "stale-grants" else {}, + mcp_toolsets=["stale-toolset"] if permission_state == "stale-grants" else [], + mcp_access_groups=["stale-group"] if permission_state == "stale-grants" else [], + ) + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("denied-customer"), + value=LiteLLM_EndUserTable( + user_id="denied-customer", + blocked=False, + object_permission_id=permission.object_permission_id, + object_permission=None if permission_state == "unreadable" else permission, + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="denied-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + object_permission=None if admin else LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", mcp_servers=[server.server_id] + ), + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + + if surface == "manager": + assert await manager.get_allowed_mcp_servers(auth) == [] + elif surface == "mcp": + assert await operations._get_allowed_mcp_servers(auth, None) == [] + else: + request: Final = Request({"type": "http", "method": "POST", "path": "/mcp-rest/tools/call", "headers": []}) + with pytest.raises(HTTPException) as denied: + await rest_endpoints._resolve_allowed_mcp_servers_with_ip_filter(request, auth, server.server_id) + assert denied.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("customer", ["absent", "unrestricted", "empty-group", "limited"]) +@pytest.mark.parametrize("admin", [False, True]) +async def test_customer_without_explicit_deny_keeps_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + customer: Literal["absent", "unrestricted", "empty-group", "limited"], + admin: bool, +) -> None: + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("unrestricted-customer"), + value=LiteLLM_EndUserTable( + user_id="unrestricted-customer", + blocked=False, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="ordinary-customer", + mcp_servers=["another-server"] if customer == "limited" else [], + mcp_access_groups=["empty-group"] if customer == "empty-group" else [], + ), + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id=None if customer == "absent" else "unrestricted-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + + assert await manager.get_allowed_mcp_servers(auth) == [server.server_id]