feat(mcp): allow skipping MCP host-tool registration in non-MCP processes - #43842
feat(mcp): allow skipping MCP host-tool registration in non-MCP processes#43842eschutho wants to merge 1 commit into
Conversation
…sses init_core_dependencies() calls initialize_core_mcp_dependencies() in every process that builds the app, including Celery workers that never serve MCP. Importing the MCP service app to register host tools has a real per-process memory cost that can OOM high-concurrency prefork worker pools on boot. Add a CORE_MCP_HOST_TOOLS_ENABLED config flag (default True, no behavior change) and gate the call on it, so deployments can turn host-tool registration off for processes that do not serve MCP. Signed-off-by: Elizabeth Thompson <elizabeth@preset.io> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Code Review Agent Run #31e528Actionable Suggestions - 0Additional Suggestions - 2
Review 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 |
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
| if self.config.get("CORE_MCP_HOST_TOOLS_ENABLED", True): | ||
| initialize_core_mcp_dependencies() |
There was a problem hiding this comment.
Suggestion: With this flag false, the abstract tool and prompt decorators remain active, so extensions using them raise NotImplementedError and are skipped during extension loading. [api mismatch]
Assessment: 🟠 Major · 🔁 Occurrence: Sometimes
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/initialization/__init__.py
**Line:** 565:566
**Comment:**
*Api Mismatch: With this flag false, the abstract `tool` and `prompt` decorators remain active, so extensions using them raise `NotImplementedError` and are skipped during extension loading.
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
There was a problem hiding this comment.
Agreed—the disabled path leaves the Core decorators as raising stubs, so an extension entrypoint using @tool or @prompt can fail before its non-MCP contributions, including Celery tasks, are registered. Should decorator replacement be kept independent of host-tool registration?
There was a problem hiding this comment.
Yes. Decorator replacement should be independent of host-tool registration.
CORE_MCP_HOST_TOOLS_ENABLED=False should prevent importing/registering the MCP host-tool stack, but it must still install the functional @tool and @prompt decorators before extensions load. Otherwise an extension that declares MCP metadata can raise NotImplementedError and prevent unrelated contributions, such as Celery tasks, from being registered.
A minimal structure would be:
initialize_core_mcp_decorators()
if self.config.get("CORE_MCP_HOST_TOOLS_ENABLED", True):
initialize_core_mcp_host_tools()This likely requires splitting initialize_core_mcp_dependencies() into a lightweight decorator-initialization step and a host-tool registration step. The decorator path should avoid importing the MCP service application, preserving the worker memory benefit while maintaining extension-loading compatibility.
|
The flagged issue is correct. When To resolve this, you should ensure that the decorators themselves are either conditionally registered or provide a safe fallback when the MCP stack is disabled. A common approach is to make the decorators no-ops or return a dummy implementation when the feature is disabled, rather than raising an error. Would you like me to fetch all comments and validate the rest of the PR to implement a comprehensive fix? superset/initialization/init.py |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #43842 +/- ##
===========================================
+ Coverage 64.16% 75.89% +11.73%
===========================================
Files 2591 2791 +200
Lines 138162 158943 +20781
Branches 32048 35283 +3235
===========================================
+ Hits 88647 120634 +31987
+ Misses 47986 35808 -12178
- Partials 1529 2501 +972
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:
|
aminghadersohi
left a comment
aminghadersohi
left a comment
There was a problem hiding this comment.
LGTM — approving.
The diagnosis and the shape of the fix both look right, and defaulting to True means this is a genuine no-op until a deployment opts out.
One thing I verified explicitly, because it looks wrong at a glance and I suspect it'll come up again in review: the from superset.core.mcp.core_mcp_injection import initialize_core_mcp_dependencies at the top of init_core_dependencies sits outside the new if, so at first read the import cost appears to be paid regardless of the flag — which would defeat the stated purpose. It isn't. core_mcp_injection's module-level imports are just logging, typing, and superset.extensions.context; the expensive work (import superset_core.mcp.decorators, from fastmcp.tools import Tool, and from superset.mcp_service import app) all happens inside initialize_core_mcp_dependencies(). So gating the call does skip the heavy imports, and the memory win is real. Might be worth a one-line comment saying so, purely to save the next reader the same detour.
One question rather than a change request: gating the call also skips the two decorator replacements —
superset_core.mcp.decorators.tool = create_tool_decorator
superset_core.mcp.decorators.prompt = create_prompt_decoratorso on a process with the flag off, superset_core.mcp.decorators.tool stays the abstract version. If any extension that gets loaded in a Celery worker applies @tool/@prompt at import time, it'll now hit the abstract decorator instead of the concrete one. I couldn't check what the abstract implementation does from here (superset_core isn't in my environment). If it's an inert stub, this is a non-issue; if it raises, it'd be worth either leaving the decorator swap ungated or noting the constraint in the config docstring. You'll know the worker-side extension set better than I do.
Two trivia, neither blocking:
- The new constant lands between the feature-flag lifecycle comment block (
# - stable: .../# - deprecated: ...) andDEFAULT_FEATURE_FLAGS, so it reads as though it belongs to the feature-flag section. It isn't a feature flag — which is the right call, since this is a process-startup concern rather than a runtime toggle — so a few lines earlier or later would make that clearer. - [x] Changes UIis ticked in the description, but this is backend-only.
| # (the web app and the standalone MCP service). Deployments can disable | ||
| # this in processes that never serve MCP -- e.g. Celery workers -- to | ||
| # avoid importing the MCP stack where it is unused. | ||
| if self.config.get("CORE_MCP_HOST_TOOLS_ENABLED", True): |
There was a problem hiding this comment.
This opt-out still imports superset.core.mcp.core_mcp_injection before evaluating the flag. That module imports mcp.types, whose package initializer imports MCP client/server modules, so high-concurrency workers still load much of the stack this setting is meant to avoid. Could the injection import move inside the enabled branch?
| # (the web app and the standalone MCP service). Deployments can disable | ||
| # this in processes that never serve MCP -- e.g. Celery workers -- to | ||
| # avoid importing the MCP stack where it is unused. | ||
| if self.config.get("CORE_MCP_HOST_TOOLS_ENABLED", True): |
There was a problem hiding this comment.
There is no regression coverage for the process-level opt-out. Could a focused initialization test set CORE_MCP_HOST_TOOLS_ENABLED=False and assert MCP registration/import is skipped while core API initialization still runs, alongside the default-enabled case?
3 participants
SUMMARY
init_core_dependencies()callsinitialize_core_mcp_dependencies()in every process thatbuilds the Flask app — including Celery workers, which never serve MCP. Registering the MCP host
tools imports the MCP service app, which carries a real per-process memory cost. On high-concurrency
prefork worker pools this import can push workers over their memory limit and OOM them at boot.
This adds a
CORE_MCP_HOST_TOOLS_ENABLEDconfig flag (defaultTrue, so no behavior change) andgates the
initialize_core_mcp_dependencies()call on it. Deployments can set it toFalseforprocesses that never serve MCP (e.g. Celery workers) to avoid importing the MCP stack where it is
unused.
CHANGES
superset/config.py: newCORE_MCP_HOST_TOOLS_ENABLED = Trueflag with docstring.superset/initialization/__init__.py: gateinitialize_core_mcp_dependencies()onself.config.get("CORE_MCP_HOST_TOOLS_ENABLED", True).TESTING INSTRUCTIONS
True): the web app and standalone MCP service register host tools exactly as before.CORE_MCP_HOST_TOOLS_ENABLED = Falsein a worker's config: the worker boots without importingthe MCP host-tool stack; web / MCP processes are unaffected.
ADDITIONAL INFORMATION