diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a323c277b52..67d816f5a4b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1210,6 +1210,23 @@ if MCP_AVAILABLE: def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]: return ((server.server_name or server.alias or server.server_id).lower(), server.server_id) + def _provider_registration_for_identifier( + servers: Sequence[LiteLLM_MCPServerTable], identifier: str + ) -> LiteLLM_MCPServerTable | None: + """An exact server_id wins over every name or alias, so another row's alias cannot shadow it.""" + exact: Final = next((server for server in servers if server.server_id == identifier), None) + if exact is not None: + return exact if exact.execution_mode == MCPExecutionMode.provider else None + return next( + ( + server + for server in servers + if server.execution_mode == MCPExecutionMode.provider + and identifier in (server.server_name, server.alias) + ), + None, + ) + def _get_user_mcp_management_mode() -> UserMCPManagementMode: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -1979,25 +1996,16 @@ if MCP_AVAILABLE: is_admin_view: Final = user_api_key_has_admin_view(user_api_key_dict) is_restricted_virtual_key: Final = _is_restricted_virtual_key_request(user_api_key_dict) snapshot: Final = global_mcp_server_manager.catalog.current() - config_match: Final = ( - next( - ( - server - for server in snapshot.management_servers.values() - if server.is_config - and server.execution_mode == MCPExecutionMode.provider - and server_id in (server.server_id, server.server_name, server.alias) - ), - None, - ) + provider_match: Final = ( + _provider_registration_for_identifier(tuple(snapshot.management_servers.values()), server_id) if snapshot is not None else None ) resolved: Final = ( - ResolvedMCPServer(table=config_match, runtime=None, source="registry") - if config_match is not None + ResolvedMCPServer(table=provider_match, runtime=None, source="registry") + if provider_match is not None and provider_match.is_config else await resolve_mcp_server( - server_id, + provider_match.server_id if provider_match is not None else server_id, manager=global_mcp_server_manager, db_lookup=lambda sid: get_mcp_server( get_prisma_client_or_throw("Database not connected. Connect a database to your proxy"), diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index ebcaa606831..be4d4a0f6b2 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -1286,6 +1286,27 @@ def _validate_mcp_protocol_write( raise HTTPException(status_code=400, detail=str(exc)) from exc +_TOOL_METADATA_FIELDS: Final = ("tool_name_to_description", "tool_name_to_display_name") + + +def _tool_metadata_for_new_allowlist( + current: Mapping[str, object], patch: Mapping[str, object] +) -> Mapping[str, object]: + allowed: Final = patch.get("allowed_tools") + if not isinstance(allowed, Sequence) or isinstance(allowed, str): + return {} + kept: Final = frozenset(TypeAdapter(tuple[str, ...]).validate_python(allowed)) + return { + field: { + tool: value + for tool, value in TypeAdapter(dict[str, object]).validate_python(stored).items() + if tool in kept + } + for field in _TOOL_METADATA_FIELDS + if field not in patch and isinstance(stored := current.get(field), Mapping) + } + + def _validate_provider_registration_update( stored: "prisma_db_models.LiteLLM_MCPServerTable", data: UpdateMCPServerRequest, @@ -1294,7 +1315,8 @@ def _validate_provider_registration_update( try: patch: Final = data.provider_registration_patch(fields_set) registration: Final = LiteLLM_MCPServerTable.model_validate(stored.model_dump()) - merged: Final = {**registration.model_dump(include=set(NewMCPServerRequest.model_fields)), **patch} + current: Final = registration.model_dump(include=set(NewMCPServerRequest.model_fields)) + merged: Final = {**current, **_tool_metadata_for_new_allowlist(current, patch), **patch} return NewMCPServerRequest.model_validate(merged) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 6b2a1a3ffc4..47dcd94ea44 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -218,6 +218,34 @@ async def test_mcp_publication_list_and_detail_derive_current_status( assert record.mcp_info == original_metadata +def _provider_row(server_id: str, **fields: object) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, execution_mode="provider", transport=None, + url=f"https://{server_id}.example/mcp", allowed_tools=["search"], **fields, + ) + + +@pytest.mark.parametrize( + "identifier,expected", + ( + ("native-id", "native-id"), + ("docs", "native-id"), + ("docs_alias", "native-id"), + ("gateway-id", None), + ("unknown", None), + ), +) +def test_provider_detail_lookup_prefers_exact_ids_over_other_rows_aliases(identifier: str, expected: str | None) -> None: + rows: Final = ( + _provider_row("shadow-id", alias="native-id"), + _provider_row("gateway-shadow", alias="gateway-id"), + _provider_row("native-id", server_name="docs", alias="docs_alias"), + LiteLLM_MCPServerTable(server_id="gateway-id", transport=MCPTransport.http, url="https://gw.example/mcp"), + ) + match: Final = mgmt_endpoints._provider_registration_for_identifier(rows, identifier) + assert (match.server_id if match is not None else None) == expected + + @pytest.mark.parametrize("approval_status", ("pending_review", "rejected", "draft", "active")) @pytest.mark.parametrize("strict", (False, True)) def test_mcp_publication_projection_excludes_unregistered_lifecycle_records( diff --git a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py index 83b14feb1f6..ea104f581dd 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py +++ b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py @@ -1405,6 +1405,32 @@ def test_provider_partial_update_preserves_tools_and_empty_list_stays_enforced() assert cleared.mcp_info["tool_allowlist_enforced"] is True +@pytest.mark.parametrize( + "allowed_tools,descriptions,display_names", + [ + ([], {}, {}), + (["fetch"], {"fetch": "Fetch a page"}, {"fetch": "Fetch"}), + ], +) +def test_provider_allowlist_change_drops_metadata_for_removed_tools( + allowed_tools: list[str], descriptions: dict[str, str], display_names: dict[str, str] +) -> None: + from litellm.proxy.mcp_server.db import _validate_provider_registration_update + + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, + allowed_tools=["search", "fetch"], mcp_info={"tool_allowlist_enforced": True}, + tool_name_to_description={"search": "Search docs", "fetch": "Fetch a page"}, + tool_name_to_display_name={"search": "Search", "fetch": "Fetch"}, + ) + updated = _validate_provider_registration_update( + stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=allowed_tools) + ) + assert (updated.allowed_tools, updated.tool_name_to_description, updated.tool_name_to_display_name) == ( + allowed_tools, descriptions, display_names, + ) + + @pytest.mark.parametrize("field", ["description", "instructions", "source_url", "tool_name_to_description"]) def test_provider_update_clears_explicit_null_metadata_without_writing_omitted_fields(field: str) -> None: from litellm.proxy.mcp_server.db import _prepare_provider_registration_update