mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 14:23:44 +00:00
fix(mcp): ignore stale ids on key save (#29128)
* fix(mcp): ignore stale server ids during key permission validation Prevent virtual key save failures when object_permission still includes MCP server IDs that were deleted from the registry. The validator now drops stale IDs before team scope checks and adds coverage for stale-ID and active unauthorized-ID behavior. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): use DB find_many to detect stale server IDs, add tool-permission test Addresses greptile feedback on PR #29128: P1 – Race-condition / authoritative-signal concern: Extract _get_stale_mcp_server_ids() helper that performs a single DB find_many (LiteLLM_MCPServerTable) when prisma_client is available, making server-existence checks authoritative even during in-memory registry warm-up. Falls back to the in-memory registry only when no DB client is present (config-only deployments / unit tests). Pass prisma_client through both call sites in key_management_endpoints. P2 – Missing test coverage for mcp_tool_permissions: Add test_validate_stale_ids_in_mcp_tool_permissions_silently_dropped to confirm stale server IDs that appear only as keys in mcp_tool_permissions are also stripped without raising 403. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): also check in-memory registry when DB is present to avoid misclassifying config-file servers as stale Config-file MCP servers are never written to LiteLLM_MCPServerTable, so a DB-only query would classify them as stale and silently drop their IDs from the authorization check, allowing unauthorized key access. Fix: treat a server ID as stale only when it is absent from BOTH the DB and the in-memory registry (which holds config-file servers). Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): normalize server aliases on key save Co-authored-by: Sameer Kankute <Sameerlite@users.noreply.github.com> * fix: remove unused stale MCP helper and capture normalized object_permission on key generation Co-authored-by: Yassin Kortam <yassin@berri.ai> * test(mcp): keep validate_key_mcp_servers_against_team stub in sync with return contract validate_key_mcp_servers_against_team now returns the (normalized) object_permission, and the key-generation helper assigns that return value back into the request data. The two key-generation tests stubbed the function with a bare AsyncMock, whose default MagicMock return value clobbered object_permission and skipped permission-record creation. Make the stubs pass object_permission through unchanged. * fix(mcp): preserve provided object_permission fields on key update The key-update path reconstructs data.object_permission from a full model_dump(), which marks every field as set. Downstream model_dump(exclude_unset=True) then emits models/blocked_tools/ search_tools as None, and those are non-nullable array columns, so the Prisma write fails whenever the UI submits object_permission with only a subset of fields populated (e.g. a TPM/RPM-only edit). Build the dict with exclude_unset=True so the normalized object_permission keeps the caller's original field set. --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Sameer Kankute <Sameerlite@users.noreply.github.com> Co-authored-by: Yassin Kortam <yassin@berri.ai> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
co-authored by
Cursor
Sameer Kankute
Yassin Kortam
mateo-berri
parent
1d9095f914
commit
83fcacad08
@@ -833,10 +833,13 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
||||
data_json.pop("tags")
|
||||
|
||||
# Validate MCP servers in object_permission are within team scope
|
||||
await validate_key_mcp_servers_against_team(
|
||||
normalized_object_permission = await validate_key_mcp_servers_against_team(
|
||||
object_permission=data_json.get("object_permission"),
|
||||
team_obj=team_table,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if normalized_object_permission is not None:
|
||||
data_json["object_permission"] = normalized_object_permission
|
||||
await validate_key_search_tools_against_team(
|
||||
object_permission=data_json.get("object_permission"),
|
||||
team_obj=team_table,
|
||||
@@ -2112,7 +2115,7 @@ async def _validate_mcp_servers_for_key_update(
|
||||
existing_key_row: Any,
|
||||
prisma_client: Any,
|
||||
user_api_key_cache: Any,
|
||||
) -> None:
|
||||
) -> Optional[dict]:
|
||||
"""Validate MCP servers in object_permission against the effective team."""
|
||||
effective_team_obj = team_obj
|
||||
# If team_id isn't being changed, resolve the existing key's team
|
||||
@@ -2126,18 +2129,20 @@ async def _validate_mcp_servers_for_key_update(
|
||||
object_permission_dict: Optional[dict] = None
|
||||
if data.object_permission is not None:
|
||||
object_permission_dict = (
|
||||
data.object_permission.model_dump()
|
||||
data.object_permission.model_dump(exclude_unset=True)
|
||||
if hasattr(data.object_permission, "model_dump")
|
||||
else dict(data.object_permission) # type: ignore[arg-type]
|
||||
)
|
||||
await validate_key_mcp_servers_against_team(
|
||||
normalized_object_permission = await validate_key_mcp_servers_against_team(
|
||||
object_permission=object_permission_dict,
|
||||
team_obj=effective_team_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
await validate_key_search_tools_against_team(
|
||||
object_permission=object_permission_dict,
|
||||
team_obj=effective_team_obj,
|
||||
)
|
||||
return normalized_object_permission
|
||||
|
||||
|
||||
async def _validate_update_key_data(
|
||||
@@ -2350,13 +2355,17 @@ async def _validate_update_key_data(
|
||||
|
||||
# Validate MCP servers in object_permission against the effective team
|
||||
if data.object_permission is not None:
|
||||
await _validate_mcp_servers_for_key_update(
|
||||
normalized_object_permission = await _validate_mcp_servers_for_key_update(
|
||||
data=data,
|
||||
team_obj=team_obj,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
if normalized_object_permission is not None:
|
||||
data.object_permission = LiteLLM_ObjectPermissionBase(
|
||||
**normalized_object_permission
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
||||
@@ -4,7 +4,7 @@ organizations, teams, and keys.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
@@ -192,8 +192,155 @@ async def _set_object_permission(
|
||||
return data_json
|
||||
|
||||
|
||||
def _dedupe_preserving_order(values: List[str]) -> List[str]:
|
||||
seen: Set[str] = set()
|
||||
result: List[str] = []
|
||||
for value in values:
|
||||
if value in seen:
|
||||
continue
|
||||
seen.add(value)
|
||||
result.append(value)
|
||||
return result
|
||||
|
||||
|
||||
def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool:
|
||||
return identifier in {
|
||||
getattr(server, "server_id", None),
|
||||
getattr(server, "alias", None),
|
||||
getattr(server, "server_name", None),
|
||||
getattr(server, "name", None),
|
||||
}
|
||||
|
||||
|
||||
async def _get_db_mcp_servers_by_identifiers(
|
||||
identifiers: Set[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
) -> List[Any]:
|
||||
if prisma_client is None or not identifiers:
|
||||
return []
|
||||
|
||||
identifier_list = list(identifiers)
|
||||
return await prisma_client.db.litellm_mcpservertable.find_many(
|
||||
where={
|
||||
"OR": [
|
||||
{"server_id": {"in": identifier_list}},
|
||||
{"alias": {"in": identifier_list}},
|
||||
{"server_name": {"in": identifier_list}},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_mcp_server_identifiers_to_ids(
|
||||
identifiers: Set[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
) -> Dict[str, Set[str]]:
|
||||
"""
|
||||
Resolve MCP permission entries written as server_id, alias, or server_name
|
||||
to canonical server IDs.
|
||||
|
||||
DB rows are authoritative when available; the in-memory registry is still
|
||||
consulted for config-file servers, which are not persisted in the MCP table.
|
||||
"""
|
||||
if not identifiers:
|
||||
return {}
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
resolved: Dict[str, Set[str]] = {identifier: set() for identifier in identifiers}
|
||||
|
||||
for server in await _get_db_mcp_servers_by_identifiers(
|
||||
identifiers=identifiers,
|
||||
prisma_client=prisma_client,
|
||||
):
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if not server_id:
|
||||
continue
|
||||
for identifier in identifiers:
|
||||
if _mcp_server_identifier_matches(server, identifier):
|
||||
resolved[identifier].add(server_id)
|
||||
|
||||
for registry_key, server in global_mcp_server_manager.get_registry().items():
|
||||
server_id = getattr(server, "server_id", None) or registry_key
|
||||
if not server_id:
|
||||
continue
|
||||
for identifier in identifiers:
|
||||
if identifier == registry_key or _mcp_server_identifier_matches(
|
||||
server, identifier
|
||||
):
|
||||
resolved[identifier].add(server_id)
|
||||
|
||||
return resolved
|
||||
|
||||
|
||||
def _rewrite_object_permission_mcp_servers(
|
||||
object_permission: dict,
|
||||
identifier_to_server_ids: Dict[str, Set[str]],
|
||||
) -> None:
|
||||
mcp_servers = object_permission.get("mcp_servers")
|
||||
if not isinstance(mcp_servers, list):
|
||||
return
|
||||
|
||||
normalized_servers: List[str] = []
|
||||
for identifier in mcp_servers:
|
||||
normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, [])))
|
||||
object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers)
|
||||
|
||||
|
||||
def _rewrite_object_permission_mcp_tool_permissions(
|
||||
object_permission: dict,
|
||||
identifier_to_server_ids: Dict[str, Set[str]],
|
||||
) -> None:
|
||||
mcp_tool_permissions = object_permission.get("mcp_tool_permissions")
|
||||
if not isinstance(mcp_tool_permissions, dict):
|
||||
return
|
||||
|
||||
normalized_tool_permissions: Dict[str, List[str]] = {}
|
||||
for identifier, tools in mcp_tool_permissions.items():
|
||||
if not isinstance(tools, list):
|
||||
tools = []
|
||||
for server_id in sorted(identifier_to_server_ids.get(identifier, [])):
|
||||
normalized_tool_permissions.setdefault(server_id, [])
|
||||
normalized_tool_permissions[server_id].extend(tools)
|
||||
|
||||
object_permission["mcp_tool_permissions"] = {
|
||||
server_id: _dedupe_preserving_order(tools)
|
||||
for server_id, tools in normalized_tool_permissions.items()
|
||||
}
|
||||
|
||||
|
||||
def _rewrite_object_permission_mcp_identifiers(
|
||||
object_permission: Optional[dict],
|
||||
identifier_to_server_ids: Dict[str, Set[str]],
|
||||
) -> None:
|
||||
if not object_permission or not isinstance(object_permission, dict):
|
||||
return
|
||||
|
||||
_rewrite_object_permission_mcp_servers(
|
||||
object_permission=object_permission,
|
||||
identifier_to_server_ids=identifier_to_server_ids,
|
||||
)
|
||||
_rewrite_object_permission_mcp_tool_permissions(
|
||||
object_permission=object_permission,
|
||||
identifier_to_server_ids=identifier_to_server_ids,
|
||||
)
|
||||
|
||||
|
||||
def _flatten_resolved_mcp_server_ids(
|
||||
identifier_to_server_ids: Dict[str, Set[str]],
|
||||
) -> Set[str]:
|
||||
return {
|
||||
server_id
|
||||
for server_ids in identifier_to_server_ids.values()
|
||||
for server_id in server_ids
|
||||
}
|
||||
|
||||
|
||||
async def _resolve_team_allowed_mcp_servers(
|
||||
team_object_permission: "LiteLLM_ObjectPermissionTable",
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Resolve the full set of MCP server IDs a team has access to.
|
||||
@@ -217,7 +364,15 @@ async def _resolve_team_allowed_mcp_servers(
|
||||
if isinstance(raw_tool_perms, str):
|
||||
raw_tool_perms = json.loads(raw_tool_perms)
|
||||
tool_perm_servers: List[str] = list(raw_tool_perms.keys())
|
||||
return set(direct_servers + access_group_servers + tool_perm_servers)
|
||||
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
|
||||
|
||||
|
||||
def _get_allow_all_keys_server_ids() -> Set[str]:
|
||||
@@ -231,6 +386,7 @@ def _get_allow_all_keys_server_ids() -> Set[str]:
|
||||
|
||||
async def _get_team_allowed_mcp_servers(
|
||||
team_obj: Optional["LiteLLM_TeamTableCachedObj"],
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Get the full set of MCP server IDs a team allows.
|
||||
@@ -245,7 +401,10 @@ async def _get_team_allowed_mcp_servers(
|
||||
if team_object_permission is None:
|
||||
return set()
|
||||
|
||||
return await _resolve_team_allowed_mcp_servers(team_object_permission)
|
||||
return await _resolve_team_allowed_mcp_servers(
|
||||
team_object_permission=team_object_permission,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
|
||||
def _extract_requested_mcp_server_ids(
|
||||
@@ -302,7 +461,8 @@ def _extract_requested_mcp_toolsets(
|
||||
async def validate_key_mcp_servers_against_team(
|
||||
object_permission: Optional[dict],
|
||||
team_obj: Optional["LiteLLM_TeamTableCachedObj"],
|
||||
):
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
Validate that MCP servers requested on a key are within the allowed scope.
|
||||
|
||||
@@ -322,17 +482,44 @@ async def validate_key_mcp_servers_against_team(
|
||||
|
||||
# Nothing to validate
|
||||
if not requested_servers and not requested_access_groups and not requested_toolsets:
|
||||
return
|
||||
return object_permission
|
||||
|
||||
allow_all_keys_servers = _get_allow_all_keys_server_ids()
|
||||
team_allowed_servers = await _get_team_allowed_mcp_servers(team_obj)
|
||||
team_allowed_servers = await _get_team_allowed_mcp_servers(
|
||||
team_obj=team_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# Combined allowed set = team servers + allow_all_keys servers
|
||||
all_allowed_servers = team_allowed_servers | allow_all_keys_servers
|
||||
|
||||
# Validate requested server IDs
|
||||
if requested_servers:
|
||||
disallowed_servers = requested_servers - all_allowed_servers
|
||||
# Normalize aliases/names before authorization. Only entries that do not
|
||||
# resolve to a server in the DB or config registry are treated as stale.
|
||||
identifier_to_server_ids = await _resolve_mcp_server_identifiers_to_ids(
|
||||
identifiers=requested_servers,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
stale_identifiers = {
|
||||
identifier
|
||||
for identifier in requested_servers
|
||||
if not identifier_to_server_ids.get(identifier)
|
||||
}
|
||||
if stale_identifiers:
|
||||
verbose_proxy_logger.warning(
|
||||
"validate_key_mcp_servers_against_team: ignoring stale MCP server "
|
||||
f"identifiers (no longer in registry or DB): {sorted(stale_identifiers)}"
|
||||
)
|
||||
_rewrite_object_permission_mcp_identifiers(
|
||||
object_permission=object_permission,
|
||||
identifier_to_server_ids=identifier_to_server_ids,
|
||||
)
|
||||
active_requested_servers = _flatten_resolved_mcp_server_ids(
|
||||
identifier_to_server_ids
|
||||
)
|
||||
|
||||
disallowed_servers = active_requested_servers - all_allowed_servers
|
||||
if disallowed_servers:
|
||||
if team_obj is not None:
|
||||
team_id = team_obj.team_id
|
||||
@@ -404,6 +591,8 @@ async def validate_key_mcp_servers_against_team(
|
||||
},
|
||||
)
|
||||
|
||||
return object_permission
|
||||
|
||||
|
||||
def _extract_requested_search_tools(object_permission: Optional[dict]) -> List[str]:
|
||||
"""Return search_tool_name values from a key's object_permission dict."""
|
||||
|
||||
@@ -611,7 +611,9 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
|
||||
AsyncMock(),
|
||||
AsyncMock(
|
||||
side_effect=lambda object_permission=None, **kwargs: object_permission
|
||||
),
|
||||
)
|
||||
|
||||
from litellm.proxy._types import (
|
||||
@@ -859,6 +861,64 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update_object_permission_does_not_add_null_fields():
|
||||
"""
|
||||
Updating a key with an object_permission that only sets a subset of fields
|
||||
must not normalize the unset list fields to ``None``.
|
||||
|
||||
The UI always submits object_permission with empty MCP/vector lists, even for
|
||||
a TPM/RPM-only edit. ``models``/``blocked_tools``/``search_tools`` are
|
||||
non-nullable array columns, so emitting them as ``None`` makes the downstream
|
||||
Prisma write fail. The normalized object_permission must keep the same field
|
||||
set the caller provided.
|
||||
"""
|
||||
data = UpdateKeyRequest(
|
||||
key="sk-test-key",
|
||||
tpm_limit=123,
|
||||
rpm_limit=456,
|
||||
object_permission={
|
||||
"vector_stores": [],
|
||||
"mcp_servers": [],
|
||||
"mcp_access_groups": [],
|
||||
"mcp_toolsets": [],
|
||||
"agents": [],
|
||||
"agent_access_groups": [],
|
||||
},
|
||||
)
|
||||
provided_fields = set(data.object_permission.model_fields_set)
|
||||
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "admin_user"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
existing_key_row.organization_id = None
|
||||
existing_key_row.project_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin_user",
|
||||
)
|
||||
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
|
||||
normalized = data.object_permission.model_dump(exclude_unset=True)
|
||||
assert set(normalized.keys()) == provided_fields
|
||||
assert "models" not in normalized
|
||||
assert "blocked_tools" not in normalized
|
||||
assert "search_tools" not in normalized
|
||||
assert "mcp_tool_permissions" not in normalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_info_returns_object_permission(monkeypatch):
|
||||
"""
|
||||
@@ -3006,7 +3066,9 @@ async def test_generate_key_with_object_permission():
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
|
||||
new_callable=AsyncMock,
|
||||
new=AsyncMock(
|
||||
side_effect=lambda object_permission=None, **kwargs: object_permission
|
||||
),
|
||||
),
|
||||
):
|
||||
# Execute
|
||||
|
||||
@@ -152,6 +152,34 @@ def _make_team_obj(
|
||||
return mock_team
|
||||
|
||||
|
||||
def _make_mock_mcp_server(
|
||||
server_id: str,
|
||||
alias=None,
|
||||
server_name=None,
|
||||
name=None,
|
||||
):
|
||||
mock_server = MagicMock()
|
||||
mock_server.server_id = server_id
|
||||
mock_server.alias = alias
|
||||
mock_server.server_name = server_name
|
||||
mock_server.name = name or server_name or alias or server_id
|
||||
return mock_server
|
||||
|
||||
|
||||
def _make_mock_mcp_manager(*existing_ids: str, servers=None):
|
||||
"""
|
||||
Return a mock global_mcp_server_manager with a registry containing every
|
||||
explicit server plus simple server objects for every ID in *existing_ids.
|
||||
"""
|
||||
mock_mgr = MagicMock()
|
||||
server_objs = {server.server_id: server for server in (servers or [])}
|
||||
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_mcp_server_by_id.side_effect = lambda sid: server_objs.get(sid)
|
||||
return mock_mgr
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
@@ -171,6 +199,10 @@ async def test_validate_no_object_permission(mock_access_groups, mock_allow_all)
|
||||
|
||||
|
||||
@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(),
|
||||
@@ -192,6 +224,10 @@ async def test_validate_key_servers_within_team_scope(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("server-1", "server-outside"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value=set(),
|
||||
@@ -204,7 +240,7 @@ async def test_validate_key_servers_within_team_scope(
|
||||
async def test_validate_key_servers_outside_team_scope_raises(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
"""Key requests servers NOT in the team's scope — should raise 403."""
|
||||
"""Key requests a server that exists but is NOT in the team's scope — should raise 403."""
|
||||
team_obj = _make_team_obj(mcp_servers=["server-1"])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_key_mcp_servers_against_team(
|
||||
@@ -216,6 +252,10 @@ async def test_validate_key_servers_outside_team_scope_raises(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("server-1", "global-server"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value={"global-server"},
|
||||
@@ -237,6 +277,10 @@ async def test_validate_allow_all_keys_servers_always_allowed(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("global-server"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value={"global-server"},
|
||||
@@ -248,7 +292,6 @@ async def test_validate_allow_all_keys_servers_always_allowed(
|
||||
)
|
||||
async def test_validate_no_team_only_allow_all_keys(mock_access_groups, mock_allow_all):
|
||||
"""Key without a team can only use allow_all_keys servers."""
|
||||
# This should pass — requesting a global server without a team
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission={"mcp_servers": ["global-server"]},
|
||||
team_obj=None,
|
||||
@@ -256,6 +299,10 @@ async def test_validate_no_team_only_allow_all_keys(mock_access_groups, mock_all
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("private-server"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value={"global-server"},
|
||||
@@ -268,7 +315,7 @@ async def test_validate_no_team_only_allow_all_keys(mock_access_groups, mock_all
|
||||
async def test_validate_no_team_non_global_server_raises(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
"""Key without a team requesting a non-global server — should raise 403."""
|
||||
"""Key without a team requesting an existing non-global server — should raise 403."""
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission={"mcp_servers": ["private-server"]},
|
||||
@@ -279,6 +326,10 @@ async def test_validate_no_team_non_global_server_raises(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("some-server"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value=set(),
|
||||
@@ -302,6 +353,10 @@ async def test_validate_team_no_mcp_config_blocks_all(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager("server-outside"),
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
return_value=set(),
|
||||
@@ -314,7 +369,7 @@ async def test_validate_team_no_mcp_config_blocks_all(
|
||||
async def test_validate_tool_permissions_validated_against_team(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
"""Server IDs in mcp_tool_permissions should also be validated."""
|
||||
"""Server IDs in mcp_tool_permissions should also be validated when they exist."""
|
||||
team_obj = _make_team_obj(mcp_servers=["server-1"])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_key_mcp_servers_against_team(
|
||||
@@ -325,6 +380,208 @@ async def test_validate_tool_permissions_validated_against_team(
|
||||
assert "server-outside" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager(), # empty registry — all IDs are stale
|
||||
)
|
||||
@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_stale_mcp_server_ids_are_silently_dropped(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
"""
|
||||
Stale MCP server IDs (servers deleted and no longer in the registry) must not
|
||||
block a key save with a 403. They are silently stripped instead.
|
||||
|
||||
Scenario: key/team were configured with S1+S2, those servers were deleted and
|
||||
replaced with S3+S4. The UI form still holds S1+S2 in its local state. Saving
|
||||
should succeed, not raise a 403.
|
||||
"""
|
||||
team_obj = _make_team_obj(mcp_servers=["s3", "s4"])
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission={"mcp_servers": ["s1-stale", "s2-stale"]},
|
||||
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(), # empty registry — all IDs are stale
|
||||
)
|
||||
@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_stale_ids_in_mcp_tool_permissions_silently_dropped(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
"""
|
||||
Stale server IDs referenced only as keys in mcp_tool_permissions (not in
|
||||
mcp_servers) must also be silently stripped rather than raising a 403.
|
||||
"""
|
||||
team_obj = _make_team_obj(mcp_servers=["s3", "s4"])
|
||||
object_permission = {"mcp_tool_permissions": {"s1-stale": ["tool1"]}}
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission=object_permission,
|
||||
team_obj=team_obj,
|
||||
) # Must not raise
|
||||
assert object_permission["mcp_tool_permissions"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager(), # empty registry — all IDs are stale
|
||||
)
|
||||
@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_stale_mcp_server_ids_are_removed_from_object_permission(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
team_obj = _make_team_obj(mcp_servers=["s3", "s4"])
|
||||
object_permission = {"mcp_servers": ["s1-stale", "s2-stale"]}
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission=object_permission,
|
||||
team_obj=team_obj,
|
||||
)
|
||||
assert object_permission["mcp_servers"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager(
|
||||
"team-server",
|
||||
servers=[
|
||||
_make_mock_mcp_server(
|
||||
"private-server-id",
|
||||
alias="private-alias",
|
||||
server_name="Private Server",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
@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_mcp_server_alias_outside_team_scope_raises(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
team_obj = _make_team_obj(mcp_servers=["team-server"])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission={"mcp_servers": ["private-alias"]},
|
||||
team_obj=team_obj,
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "private-server-id" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager(
|
||||
servers=[
|
||||
_make_mock_mcp_server(
|
||||
"allowed-server-id",
|
||||
alias="allowed-alias",
|
||||
server_name="Allowed Server",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
@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_mcp_server_alias_is_normalized_before_save(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
team_obj = _make_team_obj(mcp_servers=["allowed-server-id"])
|
||||
object_permission = {
|
||||
"mcp_servers": ["allowed-alias"],
|
||||
"mcp_tool_permissions": {"Allowed Server": ["tool1"], "stale-id": ["tool2"]},
|
||||
}
|
||||
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission=object_permission,
|
||||
team_obj=team_obj,
|
||||
)
|
||||
|
||||
assert object_permission["mcp_servers"] == ["allowed-server-id"]
|
||||
assert object_permission["mcp_tool_permissions"] == {"allowed-server-id": ["tool1"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
new=_make_mock_mcp_manager(),
|
||||
)
|
||||
@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_db_mcp_server_alias_outside_team_scope_raises_when_registry_empty(
|
||||
mock_access_groups, mock_allow_all
|
||||
):
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_db_server = MagicMock()
|
||||
mock_db_server.server_id = "private-server-id"
|
||||
mock_db_server.alias = "private-alias"
|
||||
mock_db_server.server_name = "Private Server"
|
||||
mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(
|
||||
return_value=[mock_db_server]
|
||||
)
|
||||
|
||||
team_obj = _make_team_obj(mcp_servers=[])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_key_mcp_servers_against_team(
|
||||
object_permission={"mcp_servers": ["private-alias"]},
|
||||
team_obj=team_obj,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "private-server-id" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids",
|
||||
|
||||
Reference in New Issue
Block a user