aminghadersohi commented on code in PR #41607: URL: https://github.com/apache/superset/pull/41607#discussion_r3524909458
########## superset/core/mcp/core_mcp_injection.py: ########## @@ -38,6 +38,32 @@ logger = logging.getLogger(__name__) +def _resolve_tool_name(base_name: str, replaces: Optional[str]) -> tuple[str, str]: + """Return (tool_name, context_type) for registration. + + When *replaces* is set the tool is replacing a host tool, so the name is + used as-is and context_type is always "host". + """ + if replaces is not None: + return replaces, "host" + return _get_prefixed_id_with_context(base_name) + + +def _remove_tool_for_replacement(mcp: Any, name: str) -> None: + """Remove *name* from the MCP registry before registering a replacement. + + Logs a warning (rather than raising) when the tool is not found, which + can happen if registration order changes or in partial-init environments. Review Comment: MEDIUM: resolution is "last decorator registered wins" with no collision detection — two tools both declaring `replaces="X"`, or `replaces` targeting another extension's (not-yet-loaded) tool rather than a host tool, silently depend on import/registration order. Nothing here actually restricts `replaces` to host tools despite the docstring's claim, and the log line only records the *argument* value, not whether a removal actually occurred — operators can't distinguish "replaced the host tool" from "silently overwrote another extension's replacement." ########## tests/unit_tests/mcp_service/test_tool_replaces.py: ########## @@ -0,0 +1,163 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for the ``replaces`` parameter on the MCP ``@tool`` decorator. + +``replaces`` lets an extension swap out a host tool while keeping the same +name visible to the LLM: the original tool is removed from the registry and +the decorated function is registered under the replaced name, bypassing the +usual extension-namespace prefixing. +""" + +import logging +from typing import Any, Callable +from unittest.mock import MagicMock, patch + +import pytest + +from superset.core.mcp.core_mcp_injection import ( + _remove_tool_for_replacement, + _resolve_tool_name, + create_tool_decorator, +) +from superset.mcp_service.app import mcp as real_mcp + +_CONTEXT_PATCH_TARGET = ( + "superset.core.mcp.core_mcp_injection.get_current_extension_context" +) + + +# --------------------------------------------------------------------------- +# _resolve_tool_name +# --------------------------------------------------------------------------- + + +def test_resolve_tool_name_without_replaces_uses_prefixed_id() -> None: + """With no ``replaces``, the existing prefixing behavior is unchanged.""" + with patch(_CONTEXT_PATCH_TARGET, return_value=None): + tool_name, context_type = _resolve_tool_name("my_tool", None) + + assert tool_name == "my_tool" + assert context_type == "host" + + +def test_resolve_tool_name_with_replaces_bypasses_extension_prefix() -> None: + """``replaces`` is used verbatim even inside an active extension context, + so the decorated function overrides the host tool name rather than being + registered under an extension-namespaced name.""" + mock_context = MagicMock() + mock_context.manifest.publisher = "acme" + mock_context.manifest.name = "analytics" + + with patch(_CONTEXT_PATCH_TARGET, return_value=mock_context): + tool_name, context_type = _resolve_tool_name("my_tool", "get_instance_info") + + assert tool_name == "get_instance_info" + assert context_type == "host" + + +# --------------------------------------------------------------------------- +# _remove_tool_for_replacement +# --------------------------------------------------------------------------- + + +def test_remove_tool_for_replacement_removes_existing_tool() -> None: + """The named tool is removed from the registry via ``mcp.remove_tool``.""" + mock_mcp = MagicMock() + + _remove_tool_for_replacement(mock_mcp, "get_instance_info") + + mock_mcp.remove_tool.assert_called_once_with("get_instance_info") + + +# --------------------------------------------------------------------------- +# create_tool_decorator(replaces=...) +# --------------------------------------------------------------------------- + + +def _custom_instance_info() -> dict[str, Any]: + """Custom replacement tool.""" + return {} + + +def _brand_new_tool() -> dict[str, Any]: + """A normal, non-replacing tool.""" + return {} + + +def test_create_tool_decorator_with_replaces_removes_old_tool_first() -> None: + """Decorating with ``replaces`` removes the original tool before the + replacement is registered.""" + decorator: Callable[..., Any] = create_tool_decorator(replaces="get_instance_info") + + with ( + patch.object(real_mcp, "remove_tool") as mock_remove, + patch.object(real_mcp, "add_tool") as mock_add, + ): + decorator(_custom_instance_info) + + mock_remove.assert_called_once_with("get_instance_info") + mock_add.assert_called_once() + + +def test_create_tool_decorator_with_replaces_registers_under_replaced_name() -> None: + """The new tool is registered under the ``replaces`` name, not the + decorated function's own name.""" + decorator: Callable[..., Any] = create_tool_decorator(replaces="get_instance_info") + + with ( + patch.object(real_mcp, "remove_tool"), + patch.object(real_mcp, "add_tool") as mock_add, + ): + decorator(_custom_instance_info) + + registered_tool = mock_add.call_args.args[0] + assert registered_tool.name == "get_instance_info" + + +def test_create_tool_decorator_without_replaces_does_not_remove_anything() -> None: + """Without ``replaces``, no removal call is made — only additive + registration, matching pre-existing behavior.""" + decorator: Callable[..., Any] = create_tool_decorator() + + with ( + patch.object(real_mcp, "remove_tool") as mock_remove, + patch.object(real_mcp, "add_tool"), + ): + decorator(_brand_new_tool) + + mock_remove.assert_not_called() + + +def test_create_tool_decorator_replaces_survives_missing_original_tool( + caplog: pytest.LogCaptureFixture, +) -> None: + """If the tool being replaced was never registered, ``create_tool_decorator`` + still registers the replacement instead of failing the whole decoration.""" + decorator: Callable[..., Any] = create_tool_decorator(replaces="missing") + + with ( + patch.object( + real_mcp, "remove_tool", side_effect=KeyError("missing") Review Comment: MEDIUM: this injects `side_effect=KeyError(...)` on the mocked `remove_tool`, but the real `FastMCP.remove_tool()` raises `fastmcp.exceptions.NotFoundError` for a missing tool (see the BLOCKER on core_mcp_injection.py:59). This test doesn't exercise the actual failure path and gives false confidence that the "missing tool" case is handled in production. ########## superset/core/mcp/core_mcp_injection.py: ########## @@ -38,6 +38,32 @@ logger = logging.getLogger(__name__) +def _resolve_tool_name(base_name: str, replaces: Optional[str]) -> tuple[str, str]: + """Return (tool_name, context_type) for registration. + + When *replaces* is set the tool is replacing a host tool, so the name is + used as-is and context_type is always "host". + """ + if replaces is not None: + return replaces, "host" + return _get_prefixed_id_with_context(base_name) + + +def _remove_tool_for_replacement(mcp: Any, name: str) -> None: + """Remove *name* from the MCP registry before registering a replacement. + + Logs a warning (rather than raising) when the tool is not found, which + can happen if registration order changes or in partial-init environments. + """ + try: + mcp.remove_tool(name) Review Comment: BLOCKER: `mcp.remove_tool()` is FastMCP 3.4.2's *deprecated* top-level API — it catches the underlying `local_provider.remove_tool()`'s `KeyError` internally and re-raises it as `fastmcp.exceptions.NotFoundError` (verified against the fastmcp 3.4.2 wheel). The `except KeyError` here never fires, so replacing a non-existent tool raises uncaught, which `create_tool_decorator`'s outer `except Exception: raise` turns into a hard failure of the whole decoration — not the "log a warning and continue" behavior the docstring/tests claim. `test_create_tool_decorator_replaces_survives_missing_original_tool` only passes because it mocks `remove_tool` with a raw `KeyError` side effect instead of exercising the real fastmcp path. Fix: call `mcp.local_provider.remove_tool(name)` directly (same pattern already used in `superset/mcp_service/app.py:849,862`) — it raises the plain `KeyError` this code expects, and sidesteps the deprecation warning on the public wrapper entirely. (Corrects the codeant-ai bot comment on this line, which claims `remove_tool` doesn't exist at all — it does exist, just deprecated with a different exception type than assumed here.) ########## superset/core/mcp/core_mcp_injection.py: ########## @@ -105,8 +134,9 @@ def decorator(func: F) -> F: tool_description = description or func.__doc__ or f"Tool: {base_tool_name}" tool_tags = tags or [] - # Get prefixed ID based on ambient context - tool_name, context_type = _get_prefixed_id_with_context(base_tool_name) + tool_name, context_type = _resolve_tool_name(base_tool_name, replaces) + if replaces is not None: + _remove_tool_for_replacement(mcp, replaces) Review Comment: HIGH: no check that the replacement's `protect`/`class_permission_name`/`method_permission_name` are at least as restrictive as the tool being replaced. An extension can silently swap a permissioned/protected host tool for one with `protect=False` or no RBAC at all — the only trace is a routine `logger.info` line below. Given `replaces` is explicitly meant for cross-package host-tool overrides (per the PR description), this is a meaningful blind spot: worth at least a warning-level log when the new tool's protection is weaker than what it replaces. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
