mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 4814b00604 into 408f79b5c3
This commit is contained in:
commit
e75ed9cb6e
9 changed files with 283 additions and 19 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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() == {}
|
||||
|
|
|
|||
|
|
@ -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() == {}
|
||||
|
|
|
|||
|
|
@ -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() == {}
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue