scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

refactor(mcp): resolve all-proxy-mcps through one grant parser

Browse source 
Route the all-proxy-mcps / no-mcp-servers sentinels through a single
tagged-union parser (AllServers | NoServers | ExplicitServers) that both
the read chokepoint (expand_permission_list) and the write path
(_resolve_team_allowed_mcp_servers, validate_key_mcp_servers_against_team)
match on with assert_never, so a future sentinel becomes a type error
rather than a silently half-wired feature.

Fixes the two gaps the sentinel left in the write path: creating a key
under an all-proxy-mcps team no longer 403s, and the sentinel is no
longer silently stripped on save. A key may carry all-proxy-mcps only
via a teamless proxy-admin assignment; a key in a team inherits the team
grant at request time and is rejected if it tries to self-set it, so the
grant cannot be persisted into a standalone escalation.
This commit is contained in:
ryan-crabbe-berri 2026-06-25 15:39:11 -07:00
parent fb34108c7a
commit bca7b2694c
7 changed files with 568 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -8,6 +8,10 @@ from starlette.types import Scope
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy._experimental.mcp_server.permission_grant import (
NoServers,
parse_mcp_server_grant,
)
from litellm.proxy._types import (
LiteLLM_TeamTable,
ProxyException,
@ -646,9 +650,8 @@ class MCPRequestHandler:
# The key explicitly opted out of every MCP server. This overrides
# team inheritance and additive grants (mirrors no-default-models).
if (
SpecialMCPServerNames.no_mcp_servers.value
in allowed_mcp_servers_for_key
if isinstance(
parse_mcp_server_grant(allowed_mcp_servers_for_key), NoServers
):
return []
@ -1070,8 +1073,9 @@ class MCPRequestHandler:
# Sentinel opt-out: surface it unexpanded so the caller can short-circuit
# to zero servers instead of inheriting the team.
if SpecialMCPServerNames.no_mcp_servers.value in (
key_object_permission.mcp_servers or []
if isinstance(
parse_mcp_server_grant(key_object_permission.mcp_servers or []),
NoServers,
):
return [SpecialMCPServerNames.no_mcp_servers.value]

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

@ -16,6 +16,8 @@ import time
from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast
from urllib.parse import urlparse
from typing_extensions import assert_never
import anyio
from fastapi import HTTPException
from httpx import HTTPStatusError
@ -49,6 +51,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
from litellm.proxy._experimental.mcp_server.permission_grant import (
AllServers,
ExplicitServers,
NoServers,
parse_mcp_server_grant,
)
from litellm.proxy._experimental.mcp_server.elicitation_handler import (
MCP_ELICITATION_AVAILABLE,
)
@ -1367,9 +1375,9 @@ class MCPServerManager:
key_object_permission = (
user_api_key_auth.object_permission if user_api_key_auth else None
)
if key_object_permission is not None and (
SpecialMCPServerNames.no_mcp_servers.value
in (key_object_permission.mcp_servers or [])
if key_object_permission is not None and isinstance(
parse_mcp_server_grant(key_object_permission.mcp_servers or []),
NoServers,
):
return []
@ -4132,6 +4140,30 @@ class MCPServerManager:
if server.available_on_public_internet or server.server_id in public_ids
]
def _resolve_permission_identifier(
self, identifier: str, registry: Dict[str, MCPServer]
) -> Tuple[str, ...]:
if identifier in registry:
return (identifier,)
matches = tuple(
server_id
for server_id, server in registry.items()
if server.alias == identifier
or server.server_name == identifier
or server.name == identifier
)
if matches:
return matches
# %r quotes and escapes control chars so an admin-controlled identifier
# with newlines cannot forge log lines.
verbose_logger.debug(
"MCP permission entry %r does not resolve to any known "
"server (config + DB union). Passing through — the "
"downstream access check will deny it if it's stale.",
identifier,
)
return (identifier,)
def expand_permission_list(self, identifiers: List[str]) -> List[str]:
"""
Expand a permission list of server_ids/names/aliases into concrete
@ -4147,38 +4179,31 @@ class MCPServerManager:
The ``all-proxy-mcps`` sentinel expands to every server in the live
registry, so an entity granted it stays in sync as servers are added
or removed without re-editing its permission list.
or removed without re-editing its permission list. The
``no-mcp-servers`` sentinel surfaces unexpanded so opt-out-aware callers
can short-circuit before inheriting any other scope.
"""
if not identifiers:
return []
registry = self.get_registry()
if SpecialMCPServerNames.all_proxy_mcp_servers.value in identifiers:
return list(registry.keys())
expanded: Set[str] = set()
for identifier in identifiers:
if identifier in registry:
expanded.add(identifier)
continue
matches: List[str] = [
server_id
for server_id, server in registry.items()
if server.alias == identifier
or server.server_name == identifier
or server.name == identifier
]
if matches:
expanded.update(matches)
else:
# %r quotes and escapes control chars so an admin-controlled
# identifier with newlines cannot forge log lines.
verbose_logger.debug(
"MCP permission entry %r does not resolve to any known "
"server (config + DB union). Passing through — the "
"downstream access check will deny it if it's stale.",
identifier,
grant = parse_mcp_server_grant(identifiers)
match grant:
case AllServers():
return list(registry.keys())
case NoServers():
return [SpecialMCPServerNames.no_mcp_servers.value]
case ExplicitServers(identifiers=explicit):
return list(
frozenset(
server_id
for identifier in explicit
for server_id in self._resolve_permission_identifier(
identifier, registry
)
)
)
expanded.add(identifier)
return list(expanded)
case _:
assert_never(grant)
def expand_tool_permissions(
self,

54
litellm/proxy/_experimental/mcp_server/permission_grant.py Normal file
View file

@ -0,0 +1,54 @@
"""Parse raw ``mcp_servers`` permission lists into a tagged grant union.
This module is pure: it only knows about the sentinel strings via the
``SpecialMCPServerNames`` enum and never touches the server manager, the DB, or
the proxy server. Downstream resolvers turn an ``ExplicitServers`` set of
identifiers into concrete deployments.
"""
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Union
from litellm.proxy._types import SpecialMCPServerNames
MCP_GRANT_SENTINELS: frozenset[str] = frozenset(
{
SpecialMCPServerNames.no_mcp_servers.value,
SpecialMCPServerNames.all_proxy_mcp_servers.value,
}
)
@dataclass(frozen=True, slots=True)
class AllServers:
"""Grant every MCP server on the proxy, including ones added later."""
@dataclass(frozen=True, slots=True)
class NoServers:
"""Block all MCP servers."""
@dataclass(frozen=True, slots=True)
class ExplicitServers:
"""Grant a specific, unresolved set of server_ids / aliases / names."""
identifiers: frozenset[str]
MCPServerGrant = Union[AllServers, NoServers, ExplicitServers]
def parse_mcp_server_grant(raw: Iterable[str]) -> MCPServerGrant:
"""Resolve precedence once, most-restrictive first.
``no-mcp-servers`` blocks everything and beats ``all-proxy-mcps``; absent any
sentinel the entries are explicit identifiers.
"""
values = frozenset(raw)
if SpecialMCPServerNames.no_mcp_servers.value in values:
return NoServers()
if SpecialMCPServerNames.all_proxy_mcp_servers.value in values:
return AllServers()
return ExplicitServers(values)

189
litellm/proxy/management_helpers/object_permission_utils.py
View file

@ -7,11 +7,18 @@ import json
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
from fastapi import HTTPException, status
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import SpecialMCPServerNames
from litellm.proxy._experimental.mcp_server.permission_grant import (
MCP_GRANT_SENTINELS,
AllServers,
ExplicitServers,
NoServers,
parse_mcp_server_grant,
)
from litellm.proxy.utils import PrismaClient
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.table_repositories import MCPServerRepository
@ -286,12 +293,15 @@ def _rewrite_object_permission_mcp_servers(
if not isinstance(mcp_servers, list):
return
normalized_servers: List[str] = []
for identifier in mcp_servers:
if identifier == SpecialMCPServerNames.no_mcp_servers.value:
normalized_servers.append(SpecialMCPServerNames.no_mcp_servers.value)
continue
normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, [])))
normalized_servers = [
resolved
for identifier in mcp_servers
for resolved in (
[identifier]
if identifier in MCP_GRANT_SENTINELS
else sorted(identifier_to_server_ids.get(identifier, []))
)
]
object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers)
@ -355,30 +365,55 @@ async def _resolve_team_allowed_mcp_servers(
- Direct mcp_servers list
- Servers from mcp_access_groups
- Server IDs referenced in mcp_tool_permissions keys
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
direct_servers: List[str] = team_object_permission.mcp_servers or []
access_group_servers: List[
str
] = await MCPRequestHandler._get_mcp_servers_from_access_groups(
team_object_permission.mcp_access_groups or []
)
raw_tool_perms = team_object_permission.mcp_tool_permissions or {}
if isinstance(raw_tool_perms, str):
raw_tool_perms = json.loads(raw_tool_perms)
tool_perm_servers: List[str] = list(raw_tool_perms.keys())
raw_servers = set(direct_servers + access_group_servers + tool_perm_servers)
resolved_servers = await _resolve_mcp_server_identifiers_to_ids(
identifiers=raw_servers,
prisma_client=prisma_client,
)
unresolved_servers = {
server_id for server_id in raw_servers if not resolved_servers.get(server_id)
}
return _flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers
The grant on the direct list decides everything: ``all-proxy-mcps`` resolves
to every registered server, ``no-mcp-servers`` blocks all of them.
"""
grant = parse_mcp_server_grant(team_object_permission.mcp_servers or [])
match grant:
case NoServers():
return set()
case AllServers():
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
return set(global_mcp_server_manager.get_registry().keys())
case ExplicitServers(identifiers):
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
access_group_servers = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
team_object_permission.mcp_access_groups or []
)
)
raw_tool_perms_value = team_object_permission.mcp_tool_permissions or {}
raw_tool_perms = (
json.loads(raw_tool_perms_value)
if isinstance(raw_tool_perms_value, str)
else raw_tool_perms_value
)
raw_servers = (
set(identifiers)
| set(access_group_servers)
| set(raw_tool_perms.keys())
)
resolved_servers = await _resolve_mcp_server_identifiers_to_ids(
identifiers=raw_servers,
prisma_client=prisma_client,
)
unresolved_servers = {
server_id
for server_id in raw_servers
if not resolved_servers.get(server_id)
}
return (
_flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers
)
case _:
assert_never(grant)
def _get_allow_all_keys_server_ids() -> Set[str]:
@ -426,17 +461,21 @@ def _extract_requested_mcp_server_ids(
if not object_permission or not isinstance(object_permission, dict):
return set()
server_ids: Set[str] = set()
mcp_servers = object_permission.get("mcp_servers")
if isinstance(mcp_servers, list):
server_ids.update(mcp_servers)
server_ids.discard(SpecialMCPServerNames.no_mcp_servers.value)
explicit_servers = (
frozenset(s for s in mcp_servers if s not in MCP_GRANT_SENTINELS)
if isinstance(mcp_servers, list)
else frozenset()
)
mcp_tool_permissions = object_permission.get("mcp_tool_permissions")
if isinstance(mcp_tool_permissions, dict):
server_ids.update(mcp_tool_permissions.keys())
tool_permission_servers = (
frozenset(mcp_tool_permissions.keys())
if isinstance(mcp_tool_permissions, dict)
else frozenset()
)
return server_ids
return set(explicit_servers | tool_permission_servers)
def _extract_requested_mcp_access_groups(
@ -493,6 +532,34 @@ async def validate_key_mcp_servers_against_team(
requested_toolsets = _extract_requested_mcp_toolsets(object_permission)
key_mcp_servers = (
object_permission.get("mcp_servers")
if isinstance(object_permission, dict)
else None
)
key_grant = (
parse_mcp_server_grant(key_mcp_servers)
if isinstance(key_mcp_servers, list)
else ExplicitServers(frozenset())
)
team_grant = (
parse_mcp_server_grant(team_obj.object_permission.mcp_servers or [])
if team_obj is not None and team_obj.object_permission is not None
else ExplicitServers(frozenset())
)
if isinstance(key_grant, AllServers) and not teamless_admin_assignment:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": (
"all-proxy-mcps cannot be granted directly on a key. A key in a "
"team inherits the team's MCP grant at request time; only a proxy "
"admin may assign all-proxy-mcps to a teamless key."
)
},
)
# Nothing to validate
if not requested_servers and not requested_access_groups and not requested_toolsets:
return object_permission
@ -532,30 +599,32 @@ async def validate_key_mcp_servers_against_team(
identifier_to_server_ids
)
allowed_servers = all_allowed_servers
if teamless_admin_assignment:
allowed_servers = all_allowed_servers | active_requested_servers
disallowed_servers = active_requested_servers - allowed_servers
if disallowed_servers:
if team_obj is not None:
team_id = team_obj.team_id
detail = (
f"Key requests MCP servers not allowed by team '{team_id}': "
f"{sorted(disallowed_servers)}. "
f"Team allows: {sorted(team_allowed_servers)}. "
f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}."
)
else:
detail = (
f"Key is not in a team. Only globally available (allow_all_keys) MCP servers "
f"can be assigned: {sorted(allow_all_keys_servers)}. "
f"Disallowed servers: {sorted(disallowed_servers)}."
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": detail},
if not isinstance(team_grant, AllServers):
allowed_servers = (
all_allowed_servers | active_requested_servers
if teamless_admin_assignment
else all_allowed_servers
)
disallowed_servers = active_requested_servers - allowed_servers
if disallowed_servers:
if team_obj is not None:
team_id = team_obj.team_id
detail = (
f"Key requests MCP servers not allowed by team '{team_id}': "
f"{sorted(disallowed_servers)}. "
f"Team allows: {sorted(team_allowed_servers)}. "
f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}."
)
else:
detail = (
f"Key is not in a team. Only globally available (allow_all_keys) MCP servers "
f"can be assigned: {sorted(allow_all_keys_servers)}. "
f"Disallowed servers: {sorted(disallowed_servers)}."
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": detail},
)
# Validate requested access groups (must be subset of team's access groups)
if requested_access_groups:

49
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -160,6 +160,55 @@ class TestMCPServerManager:
# blanket "return everything".
assert manager.expand_permission_list(["srv-a"]) == ["srv-a"]
async def test_expand_permission_list_resolution_branches(self):
"""Pin every non-all-proxy branch of expand_permission_list: empty,
concrete id passthrough, alias/server_name expansion, and unknown
passthrough. Guards the ExplicitServers resolution loop."""
manager = MCPServerManager()
await manager.add_server(
LiteLLM_MCPServerTable(
server_id="srv-x",
alias="my-alias",
server_name="my-server-name",
description="",
url=None,
transport=MCPTransport.stdio,
command="python",
args=["-m", "server"],
env={},
created_at=datetime.now(),
updated_at=datetime.now(),
)
)
assert manager.expand_permission_list([]) == []
assert manager.expand_permission_list(["srv-x"]) == ["srv-x"]
assert manager.expand_permission_list(["my-alias"]) == ["srv-x"]
assert manager.expand_permission_list(["my-server-name"]) == ["srv-x"]
assert manager.expand_permission_list(["ghost"]) == ["ghost"]
async def test_expand_permission_list_no_mcp_servers_passthrough(self):
"""The no-mcp-servers sentinel is not a registered server, so today it
passes through unchanged. Callers (key opt-out, team expansion) rely on
this literal surviving expand_permission_list, so pin it."""
manager = MCPServerManager()
await manager.add_server(
LiteLLM_MCPServerTable(
server_id="srv-x",
alias="srv-x",
description="",
url=None,
transport=MCPTransport.stdio,
command="python",
args=["-m", "server"],
env={},
created_at=datetime.now(),
updated_at=datetime.now(),
)
)
sentinel = SpecialMCPServerNames.no_mcp_servers.value
assert manager.expand_permission_list([sentinel]) == [sentinel]
async def test_create_mcp_client_stdio(self):
"""Test creating MCP client for stdio transport"""
manager = MCPServerManager()

63
tests/test_litellm/proxy/_experimental/mcp_server/test_permission_grant.py Normal file
View file

@ -0,0 +1,63 @@
import dataclasses
import pytest
from litellm.proxy._experimental.mcp_server.permission_grant import (
AllServers,
ExplicitServers,
MCPServerGrant,
NoServers,
parse_mcp_server_grant,
)
from litellm.proxy._types import SpecialMCPServerNames
NO_MCP = SpecialMCPServerNames.no_mcp_servers.value
ALL_PROXY = SpecialMCPServerNames.all_proxy_mcp_servers.value
def test_concrete_only_returns_explicit_servers() -> None:
grant = parse_mcp_server_grant(["srv-a", "srv-b", "alias-c"])
assert grant == ExplicitServers(frozenset({"srv-a", "srv-b", "alias-c"}))
def test_all_proxy_sentinel_returns_all_servers() -> None:
assert parse_mcp_server_grant([ALL_PROXY]) == AllServers()
def test_all_proxy_sentinel_dominates_concrete_ids() -> None:
assert parse_mcp_server_grant([ALL_PROXY, "srv-a"]) == AllServers()
def test_no_mcp_sentinel_returns_no_servers() -> None:
assert parse_mcp_server_grant([NO_MCP]) == NoServers()
def test_block_all_wins_over_all_proxy_regardless_of_order() -> None:
assert parse_mcp_server_grant([NO_MCP, ALL_PROXY]) == NoServers()
assert parse_mcp_server_grant([ALL_PROXY, NO_MCP]) == NoServers()
def test_empty_list_returns_empty_explicit_servers() -> None:
assert parse_mcp_server_grant([]) == ExplicitServers(frozenset())
def test_grants_are_frozen() -> None:
grant = ExplicitServers(frozenset({"srv-a"}))
with pytest.raises(dataclasses.FrozenInstanceError):
grant.identifiers = frozenset({"other"}) # pyright: ignore[reportAttributeAccessIssue]
def test_grants_equal_by_value() -> None:
assert AllServers() == AllServers()
assert NoServers() == NoServers()
assert ExplicitServers(frozenset({"a"})) == ExplicitServers(frozenset({"a"}))
assert AllServers() != NoServers()
def test_union_alias_admits_each_variant() -> None:
grants: tuple[MCPServerGrant, ...] = (
AllServers(),
NoServers(),
ExplicitServers(frozenset({"a"})),
)
assert len(grants) == 3

210
tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py
View file

@ -9,7 +9,7 @@ sys.path.insert(0, os.path.abspath("../../../.."))
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerNames
from litellm.proxy.management_helpers.object_permission_utils import (
_extract_requested_mcp_access_groups,
_extract_requested_mcp_server_ids,
@ -890,3 +890,211 @@ async def test_validate_search_tools_raises_when_not_subset():
team_obj=_make_team_obj_search(search_tools=["t1"]),
)
assert exc.value.status_code == 403
# ---- Tests for the all-proxy-mcps grant ----
ALL_PROXY = SpecialMCPServerNames.all_proxy_mcp_servers.value
NO_MCP = SpecialMCPServerNames.no_mcp_servers.value
def test_extract_requested_mcp_server_ids_excludes_all_proxy_sentinel():
obj_perm = {"mcp_servers": [ALL_PROXY, "server-1"]}
assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"}
def test_rewrite_object_permission_mcp_servers_preserves_all_proxy_sentinel():
obj_perm = {"mcp_servers": [ALL_PROXY, "alias-1"]}
_rewrite_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}})
assert obj_perm["mcp_servers"] == [ALL_PROXY, "server-1"]
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1", "server-2"),
)
async def test_resolve_team_allowed_mcp_servers_all_proxy_expands_to_registry():
mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable)
mock_perm.mcp_servers = [ALL_PROXY]
mock_perm.mcp_access_groups = []
mock_perm.mcp_tool_permissions = {}
result = await _resolve_team_allowed_mcp_servers(mock_perm)
assert result == {"server-1", "server-2"}
@pytest.mark.asyncio
async def test_resolve_team_allowed_mcp_servers_no_mcp_servers_blocks_all():
mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable)
mock_perm.mcp_servers = [NO_MCP]
mock_perm.mcp_access_groups = ["group-a"]
mock_perm.mcp_tool_permissions = {"server-x": ["tool1"]}
result = await _resolve_team_allowed_mcp_servers(mock_perm)
assert result == set()
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1", "server-2"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_team_all_proxy_grants_any_real_server(
mock_access_groups, mock_allow_all
):
"""A team scoped to all-proxy-mcps lets a key request any real server without a 403."""
team_obj = _make_team_obj(mcp_servers=[ALL_PROXY])
await validate_key_mcp_servers_against_team(
object_permission={"mcp_servers": ["server-1"]},
team_obj=team_obj,
) # Must not raise
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("registry-server"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._resolve_mcp_server_identifiers_to_ids",
new_callable=AsyncMock,
return_value={"db-only": {"db-only"}},
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_team_all_proxy_allows_server_outside_registry(
mock_access_groups, mock_allow_all, mock_resolve
):
"""An all-proxy team must allow a key's server that resolves via the DB even when
it is absent from the in-memory registry keys; the validator short-circuits on the
team's all-proxy grant instead of enumerating the registry."""
team_obj = _make_team_obj(mcp_servers=[ALL_PROXY])
await validate_key_mcp_servers_against_team(
object_permission={"mcp_servers": ["db-only"]},
team_obj=team_obj,
) # Must not raise
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_key_all_proxy_under_explicit_team_raises(
mock_access_groups, mock_allow_all
):
"""A key may not self-grant all-proxy-mcps when its team only allows specific servers."""
team_obj = _make_team_obj(mcp_servers=["server-1"])
with pytest.raises(HTTPException) as exc_info:
await validate_key_mcp_servers_against_team(
object_permission={"mcp_servers": [ALL_PROXY]},
team_obj=team_obj,
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_key_all_proxy_in_team_raises_even_when_team_all_proxy(
mock_access_groups, mock_allow_all
):
"""A key in a team may not set all-proxy-mcps directly; it inherits the team's
grant at request time. Persisting the literal on the key would let it keep
all-server access if the team's grant were later revoked, so it is rejected."""
team_obj = _make_team_obj(mcp_servers=[ALL_PROXY])
with pytest.raises(HTTPException) as exc_info:
await validate_key_mcp_servers_against_team(
object_permission={"mcp_servers": [ALL_PROXY]},
team_obj=team_obj,
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_teamless_admin_can_assign_all_proxy(
mock_access_groups, mock_allow_all
):
"""A proxy admin assigning all-proxy-mcps to a teamless key is allowed and preserved."""
obj_perm = {"mcp_servers": [ALL_PROXY]}
await validate_key_mcp_servers_against_team(
object_permission=obj_perm,
team_obj=None,
is_proxy_admin=True,
)
assert obj_perm["mcp_servers"] == [ALL_PROXY]
@pytest.mark.asyncio
@patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
new=_make_mock_mcp_manager("server-1"),
)
@patch(
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
return_value=set(),
)
@patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
)
async def test_validate_teamless_non_admin_all_proxy_raises(
mock_access_groups, mock_allow_all
):
"""A non-admin teamless key cannot self-grant all-proxy-mcps; it is rejected, not dropped."""
with pytest.raises(HTTPException) as exc_info:
await validate_key_mcp_servers_against_team(
object_permission={"mcp_servers": [ALL_PROXY]},
team_obj=None,
is_proxy_admin=False,
)
assert exc_info.value.status_code == 403