scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 24

fix(proxy): enforce customer MCP denial before open grants

Browse source
This commit is contained in:
lewis 2026-10-17 21:11:21 -04:00
parent 7266452d14
commit 0f52da9232
3 changed files with 151 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

10
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

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

111
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

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