From 83fcacad086e4bb6938fd205348e6e7a802fe64d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Sat, 30 May 2026 09:20:24 +0530 Subject: [PATCH] fix(mcp): ignore stale ids on key save (#29128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * 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 * 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 * fix(mcp): normalize server aliases on key save Co-authored-by: Sameer Kankute * fix: remove unused stale MCP helper and capture normalized object_permission on key generation Co-authored-by: Yassin Kortam * 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 Co-authored-by: Sameer Kankute Co-authored-by: Yassin Kortam Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../key_management_endpoints.py | 19 +- .../object_permission_utils.py | 203 +++++++++++++- .../test_key_management_endpoints.py | 66 ++++- .../test_object_permission_utils.py | 265 +++++++++++++++++- 4 files changed, 535 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a67d8d934b..f49fb1f37b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index eb90d1b5ca..4c966b2541 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -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.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 38bf2d2c91..8fb242372e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index b36383dfd9..965580e875 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -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",