This commit is contained in:
mrinal-berri 2026-11-06 02:26:22 +00:00 • committed by GitHub
commit f6243ecee5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 283 additions and 19 deletions

View file

@ -1135,13 +1135,16 @@ if MCP_AVAILABLE:
)
from litellm.proxy.proxy_server import prisma_client
access_groups: Final = set()
access_groups: Final = set[str]()
# Get from config-loaded servers
for server in global_mcp_server_manager.config_mcp_servers.values():
if server.access_groups:
access_groups.update(server.access_groups)
for server in global_mcp_server_manager.catalog.config_registrations.values():
access_groups.update(server.mcp_access_groups)
# Get from DB
if prisma_client is not None:
try:
@ -1243,11 +1246,17 @@ if MCP_AVAILABLE:
return "restricted"
def _redacted_server_tables(server_ids: Iterable[str]) -> list[LiteLLM_MCPServerTable]:
registered: Final = (global_mcp_server_manager.get_mcp_server_by_id(server_id) for server_id in server_ids)
selected_ids: Final = frozenset(server_ids)
registered: Final = (global_mcp_server_manager.get_mcp_server_by_id(server_id) for server_id in selected_ids)
tables: Final = [
global_mcp_server_manager.build_mcp_server_table(server) for server in registered if server is not None
]
return _redact_mcp_credentials_list(tables)
provider_tables: Final = [
server
for server_id, server in global_mcp_server_manager.catalog.provider_registrations().items()
if server_id in selected_ids
]
return _redact_mcp_credentials_list(tables + provider_tables)
async def _granted_mcp_server_ids(
object_permission_id: str | None,
@ -1301,7 +1310,7 @@ if MCP_AVAILABLE:
check_db_only=True,
)
team_server_ids: Final = await get_team_allowed_mcp_servers(team_obj)
team_server_ids: Final = await get_team_allowed_mcp_servers(team_obj, prisma_client=prisma_client)
allow_all_server_ids: Final = get_allow_all_keys_server_ids()
all_allowed_ids: Final = team_server_ids | allow_all_server_ids
@ -1379,7 +1388,7 @@ if MCP_AVAILABLE:
if view_providers
or getattr(server, "execution_mode", MCPExecutionMode.litellm) != MCPExecutionMode.provider
]
if view_all:
if view_all and len(visible_inventory) == len(inventory):
return visible_inventory
provider_servers: Final = [
@ -1387,13 +1396,21 @@ if MCP_AVAILABLE:
for server in visible_inventory
if getattr(server, "execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider
]
initial_servers: Final = visible_inventory if view_all else provider_servers
allowed_by_context: Final = [
await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context)
for auth_context in await build_effective_auth_contexts(user_api_key_dict)
]
aggregated: Final = {
server.server_id: server
for server in (*chain.from_iterable(reversed(allowed_by_context)), *provider_servers)
for server in (
*(
server
for server in chain.from_iterable(reversed(allowed_by_context))
if not view_all or server.execution_mode == MCPExecutionMode.provider
),
*initial_servers,
)
}
return list(aggregated.values())
@ -1528,7 +1545,9 @@ if MCP_AVAILABLE:
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict)
for server in redacted_mcp_servers:
server.connected_app_reachable = server.server_id in reachable_ids
server.connected_app_reachable = (
server.execution_mode == MCPExecutionMode.litellm and server.server_id in reachable_ids
)
# Annotate has_user_credential for BYOK servers (single batched query)
from litellm.proxy.proxy_server import prisma_client as _byok_prisma_client

View file

@ -7,6 +7,7 @@ import json
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Optional
@ -325,7 +326,15 @@ async def _resolve_mcp_server_identifiers_to_ids(
if identifier == registry_key or _mcp_server_identifier_matches(server, identifier):
resolved[identifier].add(server_id)
return resolved
for server in global_mcp_server_manager.catalog.config_registrations.values():
for identifier in identifiers:
if _mcp_server_identifier_matches(server, identifier):
resolved[identifier].add(server.server_id)
# An exact server_id grants only that server, matching runtime expand_permission_list,
# so another registration whose name or alias equals the ID is not also granted.
known_ids: Final = frozenset(chain.from_iterable(resolved.values()))
return {identifier: {identifier} if identifier in known_ids else ids for identifier, ids in resolved.items()}
_MCP_TOOL_PERMISSIONS_ADAPTER: Final = TypeAdapter(dict[str, list[str] | None])
@ -491,7 +500,7 @@ def _get_all_mcp_server_ids() -> set[str]:
global_mcp_server_manager,
)
return set(global_mcp_server_manager.get_registry().keys())
return global_mcp_server_manager.get_all_mcp_server_ids()
async def _existing_object_permission_mcp_servers(

View file

@ -3754,7 +3754,7 @@ class MCPRequestHandler:
if object_permissions is None:
return set(team_access_group_servers)
if SpecialMCPServerName.all_proxy_servers.value in (object_permissions.mcp_servers or []):
return set(global_mcp_server_manager.get_registry().keys())
return global_mcp_server_manager.get_all_mcp_server_ids()
legacy_access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups(
object_permissions.mcp_access_groups or [],
requires_fresh_policy=requires_fresh_policy or require_readable_policy,
@ -4758,8 +4758,21 @@ class MCPRequestHandler:
)
snapshot: Final = global_mcp_server_manager.catalog.current()
provider_registrations: Final = (
global_mcp_server_manager.catalog.provider_registrations()
if snapshot is not None
else global_mcp_server_manager.catalog.config_registrations
)
provider_ids: Final = {
server.server_id
for server in provider_registrations.values()
if set(access_groups).intersection(server.mcp_access_groups)
}
if snapshot is not None:
return list(MCPRequestHandler._get_config_server_ids_for_access_groups(snapshot.servers, access_groups))
return list(
MCPRequestHandler._get_config_server_ids_for_access_groups(snapshot.servers, access_groups)
| provider_ids
)
# Use the new helper for config-loaded servers
server_ids: Final = MCPRequestHandler._get_config_server_ids_for_access_groups(
@ -4772,7 +4785,7 @@ class MCPRequestHandler:
)
server_ids.update(db_server_ids)
return list(server_ids)
return list(server_ids | provider_ids)
except Exception as e:
if requires_fresh_policy:
raise

View file

@ -198,6 +198,16 @@ class CatalogSnapshots:
scoped: Final = self._operation.get()
return scoped[0] if scoped is not None and not scoped[1].is_set() else None
def provider_registrations(self) -> Mapping[str, LiteLLM_MCPServerTable]:
snapshot: Final = self.current()
if snapshot is not None:
return {
server_id: server
for server_id, server in snapshot.management_servers.items()
if server.execution_mode == MCPExecutionMode.provider
}
return {**self.config_registrations, **self._provider_native_servers}
def routing(self) -> MutableMapping[str, str]:
staged: Final = self._staged_routing.get()
if staged is not None and not staged[1].is_set():

View file

@ -3643,11 +3643,15 @@ class MCPServerManager:
Get all MCP server IDs
"""
all_servers: Final = list(self.get_registry().values())
return {server.server_id for server in all_servers}
return {server.server_id for server in all_servers} | self.catalog.provider_registrations().keys()
def get_allow_all_keys_server_ids(self) -> list[str]:
"""Return server IDs that bypass per-key restrictions."""
return [server.server_id for server in self.get_registry().values() if server.allow_all_keys is True]
return [server.server_id for server in self.get_registry().values() if server.allow_all_keys is True] + [
server.server_id
for server in self.catalog.provider_registrations().values()
if server.allow_all_keys is True
]
@staticmethod
def get_byom_submitted_servers_cache_key(user_id: str) -> str:
@ -3848,7 +3852,7 @@ class MCPServerManager:
else 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)
self.get_all_mcp_server_ids() if admin_unscoped else set(resolved_access.server_ids)
)
combined_servers: Final = set(resolved_server_ids)
verbose_logger.debug("Allowed MCP Servers for user api key auth: %s", combined_servers)
@ -7529,9 +7533,10 @@ class MCPServerManager:
if not identifiers:
return []
registry: Final = self.get_registry()
providers: Final = self.catalog.provider_registrations()
expanded: Final[set[str]] = set()
for identifier in identifiers:
if identifier in registry:
if identifier in registry or identifier in providers:
expanded.add(identifier)
continue
matches: list[str] = [
@ -7539,6 +7544,11 @@ class MCPServerManager:
for server_id, server in registry.items()
if server.alias == identifier or server.server_name == identifier or server.name == identifier
]
matches.extend(
server.server_id
for server in providers.values()
if identifier in (server.server_id, server.server_name, server.alias)
)
if matches:
expanded.update(matches)
else:
@ -7863,10 +7873,14 @@ class MCPServerManager:
allowed_server_ids: Final = await self.get_allowed_mcp_servers(user_api_key_auth)
list_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
providers: Final = self.catalog.provider_registrations()
for server_id in allowed_server_ids:
server = self.get_mcp_server_by_id(server_id)
if not server:
if server_id in providers:
list_mcp_servers.append(providers[server_id])
continue
verbose_logger.warning("MCP Server %s not found in registry", server_id)
continue

View file

@ -8053,7 +8053,10 @@ class TestConnectedAppViewAnnotation:
caller_auth = self._ui_session_auth()
admitted_auth = UserAPIKeyAuth(user_id="test_user_id")
admitted_auth.mcp_admitted_user_subject = True
mock_manager = self._mock_manager(self._servers(), ["server-1"])
provider: Final = LiteLLM_MCPServerTable(
server_id="native", execution_mode="provider", transport=None, url="https://native.example/mcp",
)
mock_manager = self._mock_manager([*self._servers(), provider], ["server-1", "native"])
reload_mock = AsyncMock(return_value=admitted_auth)
with (
@ -8077,7 +8080,7 @@ class TestConnectedAppViewAnnotation:
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
flags = {server.server_id: server.connected_app_reachable for server in result}
assert flags == {"server-1": True, "server-2": False}
assert flags == {"server-1": True, "server-2": False, "native": False}
reload_mock.assert_awaited_once_with("test_user_id", requires_fresh_policy=False)
mock_manager.get_allowed_mcp_servers.assert_awaited_once_with(admitted_auth)
@ -11825,3 +11828,43 @@ async def test_mcp_register_application_type_reaches_upstream_or_is_rejected(
assert posted["application_type"] == application_type
finally:
mgmt_endpoints._temporary_mcp_servers.pop(server.server_id, None)
@pytest.mark.asyncio
@pytest.mark.parametrize("view_mode", ["restricted", "view_all"])
@pytest.mark.parametrize("granted", [True, False])
async def test_provider_management_list_and_detail_share_explicit_grants(
monkeypatch: pytest.MonkeyPatch, view_mode: str, granted: bool,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.mcp_server import mcp_server_manager
manager: Final = MCPServerManager()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(mgmt_endpoints, "global_mcp_server_manager", manager)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
monkeypatch.setattr(proxy_server, "general_settings", {"user_mcp_management_mode": view_mode})
await manager.load_servers_from_config({"native_name": {
"server_id": "native-id", "alias": "native_alias", "execution_mode": "provider",
"url": "https://native.example/mcp", "allowed_tools": ["search"],
}})
auth: Final = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="native-grant", mcp_servers=["native_alias"] if granted else [],
),
)
servers: Final = await mgmt_endpoints._resolve_accessible_mcp_servers(auth)
assert [server.server_id for server in servers] == (["native-id"] if granted else [])
request: Final = Request({"type": "http", "headers": []})
if granted:
detail: Final = await mgmt_endpoints.fetch_mcp_server(request, "native_name", user_api_key_dict=auth)
assert detail.server_id == "native-id" and detail.allowed_tools == ["search"]
assert detail.last_health_check is None
assert [server.server_id for server in mgmt_endpoints._redacted_server_tables(iter(["native-id"]))] == ["native-id"]
else:
with pytest.raises(HTTPException) as denied:
await mgmt_endpoints.fetch_mcp_server(request, "native_alias", user_api_key_dict=auth)
assert denied.value.status_code == 403
assert manager.get_registry() == {}

View file

@ -267,6 +267,7 @@ def _make_mock_mcp_manager(*existing_ids: str, servers=None):
for server_id in existing_ids:
server_objs.setdefault(server_id, _make_mock_mcp_server(server_id))
mock_mgr.get_registry.return_value = server_objs
mock_mgr.get_all_mcp_server_ids.side_effect = lambda: set(server_objs)
mock_mgr.get_mcp_server_by_id.side_effect = lambda sid: server_objs.get(sid)
return mock_mgr
@ -1074,6 +1075,7 @@ async def test_resolve_team_all_proxy_sentinel_resolves_dynamically(mock_access_
}
mock_mgr = MagicMock()
mock_mgr.get_registry.return_value = registry
mock_mgr.get_all_mcp_server_ids.side_effect = lambda: set(registry)
team_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable)
team_perm.mcp_servers = [SpecialMCPServerName.all_proxy_servers.value]
@ -1596,6 +1598,21 @@ async def test_unambiguous_tool_permission_keys_persist_verbatim():
assert json.loads(upsert.record["mcp_tool_permissions"]) == tool_permissions
@pytest.mark.asyncio
async def test_exact_server_id_resolves_only_that_server_while_shared_alias_resolves_all():
"""An exact server_id grants only that server even when another row's alias equals it,
matching runtime expansion. A shared alias still resolves to every server using it."""
from litellm.proxy.management_helpers.object_permission_utils import _resolve_mcp_server_identifiers_to_ids
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
with patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", MCPServerManager()
):
resolved = await _resolve_mcp_server_identifiers_to_ids({"solo-id", "wiki"}, _make_ambiguity_prisma())
assert resolved == {"solo-id": {"solo-id"}, "wiki": {"wiki-a-id", "wiki-b-id"}}
@pytest.mark.asyncio
async def test_stored_ambiguous_tool_permission_key_is_grandfathered_until_changed():
"""A shared-alias entry already on the row may be re-sent unchanged so unrelated
@ -1636,3 +1653,40 @@ def test_object_permission_dict_mirrors_pydantic_model():
f"Only in Pydantic model: {sorted(pydantic_fields - typeddict_fields)}\n"
f"Only in TypedDict: {sorted(typeddict_fields - pydantic_fields)}"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("identifier", ["native-id", "native_name", "native_alias"])
async def test_provider_config_key_assignment_retains_server_and_tool_grants(
monkeypatch: pytest.MonkeyPatch, identifier: str,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.mcp_server import mcp_server_manager
manager: Final = mcp_server_manager.MCPServerManager()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(proxy_server, "prisma_client", None)
await manager.load_servers_from_config({"native_name": {
"server_id": "native-id", "alias": "native_alias", "execution_mode": "provider",
"url": "https://native.example/mcp", "allowed_tools": ["search"],
}})
permissions: Final[ObjectPermissionDict] = {
"mcp_servers": [identifier], "mcp_tool_permissions": {identifier: ["search"]},
}
team: Final = LiteLLM_TeamTableCachedObj(
team_id="native-team", object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="team-grant", mcp_servers=["native-id"],
),
)
assert await validate_key_mcp_servers_against_team(permissions, team) == {
"mcp_servers": [identifier], "mcp_tool_permissions": {identifier: ["search"]},
}
assert await validate_key_mcp_servers_against_team(
permissions, team, existing_key_object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="key-grant", mcp_servers=[identifier],
),
) == {"mcp_servers": [identifier], "mcp_tool_permissions": {identifier: ["search"]}}
with pytest.raises(HTTPException) as denied:
await validate_key_mcp_servers_against_team(permissions, None)
assert denied.value.status_code == 403
assert manager.get_registry() == {}

View file

@ -9759,7 +9759,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
# registry may be a list of bare server_ids (MagicMock servers) OR a dict of
# {server_id: server_obj} for tests that need real alias/name resolution (config servers).
reg = registry if isinstance(registry, dict) else {s: MagicMock() for s in registry}
reg = registry if isinstance(registry, dict) else {s: MagicMock(server_id=s) for s in registry}
cms.append(patch.object(global_mcp_server_manager, "get_registry", return_value=reg))
with contextlib.ExitStack() as es:
for cm in cms:
@ -13252,3 +13252,48 @@ async def test_admin_dashboard_session_keeps_the_whole_mcp_registry(
assert set(await global_mcp_server_manager.get_allowed_mcp_servers(dashboard_session)) == set(
global_mcp_server_manager.get_registry()
)
@pytest.mark.asyncio
@pytest.mark.parametrize("key_tools,team_tools,use_toolset,expected", [
(["search", "fetch"], ["search"], False, ["search"]),
(["fetch"], ["search"], False, []),
(None, ["search"], False, ["search"]),
([], None, False, []),
(None, ["search"], True, ["search"]),
])
async def test_provider_tool_grants_use_existing_key_team_and_toolset_resolution(
monkeypatch: pytest.MonkeyPatch,
key_tools: list[str] | None, team_tools: list[str] | None, use_toolset: bool, expected: list[str],
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.mcp_server import mcp_server_manager
manager: Final = mcp_server_manager.MCPServerManager()
cache: Final = DualCache()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
await manager.load_servers_from_config({"native_name": {
"server_id": "native-id", "alias": "native_alias", "execution_mode": "provider",
"url": "https://native.example/mcp", "allowed_tools": ["search", "fetch"],
}})
team: Final = LiteLLM_TeamTable(
team_id="native-team", object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="native-team-grant", mcp_servers=["native_name"],
mcp_tool_permissions={"native_name": team_tools} if team_tools is not None else {},
),
)
await cache.async_set_cache(key="team_id:native-team", value=team)
await cache.async_set_cache(key="toolset_perms:native-set", value={"native-id": ["search"]})
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace())
auth: Final = UserAPIKeyAuth(
team_id="native-team", object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="native-key-grant",
mcp_servers=[] if use_toolset else ["native_alias"],
mcp_tool_permissions={"native_alias": key_tools} if key_tools is not None else {},
mcp_toolsets=["native-set"] if use_toolset else [],
),
)
assert await MCPRequestHandler.get_allowed_tools_for_server("native-id", auth) == expected
assert manager.get_registry() == {}

View file

@ -1052,3 +1052,60 @@ async def test_provider_config_changes_invalidate_cached_inventory_without_datab
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == []
await manager.load_servers_from_config({})
assert await manager.get_all_mcp_servers_unfiltered() == []
@pytest.mark.asyncio
async def test_provider_permission_identifiers_and_groups_follow_config_reload(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.mcp_server import mcp_server_manager
from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
manager: Final = mcp_server_manager.MCPServerManager()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(proxy_server, "prisma_client", None)
config: Final = {"native_name": {
"server_id": "native-id", "alias": "native_alias", "execution_mode": "provider",
"url": "https://native.example/mcp", "allowed_tools": ["search"],
"access_groups": ["research"], "allow_all_keys": True,
}}
await manager.load_servers_from_config(config)
async with manager.catalog.operation():
assert manager.get_all_mcp_server_ids() == {"native-id"}
assert manager.get_allow_all_keys_server_ids() == ["native-id"]
assert manager.expand_permission_list(["native_name", "native_alias"]) == ["native-id"]
assert manager.expand_tool_permissions({"native_alias": ["search"]}) == {"native-id": ["search"]}
assert await MCPRequestHandler.get_mcp_servers_from_access_groups(["research"]) == ["native-id"]
assert await manager.catalog.list() == {}
await manager.load_servers_from_config({})
assert manager.get_all_mcp_server_ids() == {"native-id"}
assert manager.get_all_mcp_server_ids() == set()
assert manager.get_allow_all_keys_server_ids() == []
assert await MCPRequestHandler.get_mcp_servers_from_access_groups(["research"]) == []
@pytest.mark.asyncio
async def test_exact_provider_id_grant_ignores_another_registrations_matching_alias(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.models.mcp_server import LiteLLM_MCPServerTable
from litellm.proxy import proxy_server
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
granted: Final = LiteLLM_MCPServerTable(
server_id="native-id", execution_mode="provider", transport=None,
url="https://granted.example/mcp", allowed_tools=["search"],
)
other: Final = LiteLLM_MCPServerTable(
server_id="other-id", alias="native-id", execution_mode="provider", transport=None,
url="https://other.example/mcp", allowed_tools=["delete"],
)
tables: Final = SimpleNamespace(
litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[granted, other])),
litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))),
)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
monkeypatch.setattr(proxy_server, "general_settings", {})
manager: Final = MCPServerManager()
async with manager.catalog.operation():
assert len(await manager.get_all_mcp_servers_unfiltered()) == 2
assert manager.expand_permission_list(["native-id"]) == ["native-id"]