feat(mcp): add replaces parameter to @tool decorator for override#41607
feat(mcp): add replaces parameter to @tool decorator for override#41607justinpark wants to merge 2 commits into
Conversation
justinpark
requested review from
eschutho,
geido,
kgabryje,
michael-s-molina,
rusackas and
villebro
as code owners
Code Review Agent Run #8e8041Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #41607 +/- ##
==========================================
+ Coverage 64.42% 64.52% +0.10%
==========================================
Files 2668 2673 +5
Lines 147182 147820 +638
Branches 33947 33949 +2
==========================================
+ Hits 94821 95387 +566
- Misses 50646 50714 +68
- Partials 1715 1719 +4
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![codeant-ai-for-open-source[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
| _CONTEXT_PATCH_TARGET = ( | ||
| "superset.core.mcp.core_mcp_injection.get_current_extension_context" | ||
| ) |
There was a problem hiding this comment.
Suggestion: Add an explicit type annotation for this module-level constant to satisfy the type-hint requirement for annotatable variables. [custom_rule]
Severity Level: Minor
Why it matters? 🤔
The new module-level constant is an annotatable variable and has no type hint, which matches the Python type-hint requirement for newly added code.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/test_tool_replaces.py
**Line:** 38:40
**Comment:**
*Custom Rule: Add an explicit type annotation for this module-level constant to satisfy the type-hint requirement for annotatable variables.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix| """``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() |
There was a problem hiding this comment.
Suggestion: Add a type annotation for this local mock variable since its type is known and can be declared. [custom_rule]
Severity Level: Minor
Why it matters? 🤔
This new local variable is unannotated even though it is a clearly annotatable Python variable introduced in the PR, so it fits the type-hint rule.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/test_tool_replaces.py
**Line:** 61:61
**Comment:**
*Custom Rule: Add a type annotation for this local mock variable since its type is known and can be declared.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix|
|
||
| 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() |
There was a problem hiding this comment.
Suggestion: Add a type annotation for this mock variable to comply with the rule requiring hints on annotatable variables. [custom_rule]
Severity Level: Minor
Why it matters? 🤔
This is another newly added, clearly annotatable local variable without a type hint, so the type-hint rule applies.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/test_tool_replaces.py
**Line:** 79:79
**Comment:**
*Custom Rule: Add a type annotation for this mock variable to comply with the rule requiring hints on annotatable variables.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
| can happen if registration order changes or in partial-init environments. | ||
| """ | ||
| try: | ||
| mcp.remove_tool(name) |
There was a problem hiding this comment.
Suggestion: The replacement-removal helper calls mcp.remove_tool, but this codebase removes tools through mcp.local_provider.remove_tool; if FastMCP does not expose remove_tool on the top-level object, replacement registration will raise AttributeError and fail before the new tool is added. Use the same removal path used elsewhere (local_provider.remove_tool) to avoid a runtime API mismatch. [api mismatch]
Severity Level: Critical 🚨
- ❌ Extensions replacing host tools crash during MCP tool registration.
- ❌ Replacement tools never register; host tools remain active.
- ⚠️ LLM clients cannot use optimized get_instance_info replacement.Steps of Reproduction ✅
1. Start the MCP service so `superset.mcp_service.app.mcp` is created as a `FastMCP`
instance via `create_mcp_app()` (superset/mcp_service/app.py:60-77, 600-637) and
`initialize_core_mcp_dependencies()` is called to inject concrete decorators
(superset/core/mcp/core_mcp_injection.py:54-76, 337-416;
superset/mcp_service/app.py:640-646).
2. Observe that all existing code paths that remove tools from the MCP registry use
`mcp.local_provider.remove_tool(tool_name)` (superset/mcp_service/app.py:20-36 and 39-43
in the 818-876 snippet, corresponding to file lines ~846 and ~859), and there are no other
references to `mcp.remove_tool` in the repository (verified by `Grep` for `remove_tool`).
3. In an extension package that uses the public stub decorator, import `tool` from
`superset_core.mcp.decorators` and define a replacement tool as shown in the stub
docstring (superset-core/src/superset_core/mcp/decorators.py:91-94):
`@tool(replaces="get_instance_info")` followed by `def my_custom_instance_info() -> dict:
...`. When the extension module is imported, `initialize_core_mcp_dependencies()` has
already replaced `superset_core.mcp.decorators.tool` with `create_tool_decorator`
(superset/core/mcp/core_mcp_injection.py:72-74, 337-416).
4. During decoration, `create_tool_decorator(..., replaces="get_instance_info")` calls
`_remove_tool_for_replacement(mcp, replaces)`
(superset/core/mcp/core_mcp_injection.py:78-80, 60-120), which in turn calls
`mcp.remove_tool(name)` at superset/core/mcp/core_mcp_injection.py:59 (35-74 snippet line
25). Because the codebase’s established removal API is `mcp.local_provider.remove_tool`
and there is no evidence of a top-level `FastMCP.remove_tool` method, this call will raise
`AttributeError` on the `FastMCP` instance, causing the decorator’s try-block to fail and
re-raise (superset/core/mcp/core_mcp_injection.py:68-88, 60-120). The replacement tool is
never registered and extension initialization fails.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/core/mcp/core_mcp_injection.py
**Line:** 59:59
**Comment:**
*Api Mismatch: The replacement-removal helper calls `mcp.remove_tool`, but this codebase removes tools through `mcp.local_provider.remove_tool`; if `FastMCP` does not expose `remove_tool` on the top-level object, replacement registration will raise `AttributeError` and fail before the new tool is added. Use the same removal path used elsewhere (`local_provider.remove_tool`) to avoid a runtime API mismatch.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
1 participant
SUMMARY
Added a replaces parameter to the @tool decorator in superset-core that allows extensions to cleanly replace existing host tools under their original name, without extension namespace prefixing
Problem
The original get_instance_info tool was taking ~280 seconds per call. The bottleneck was security filter subqueries (DashboardAccessFilter, ChartFilter, etc.) applied to every COUNT(*). Extensions had no clean way to swap out a host tool — the only option was monkey-patching module globals, which broke silently when re-exports obscured the target module.
Changes
TESTING INSTRUCTIONS
added specs
ADDITIONAL INFORMATION