sadpandajoe commented on code in PR #43842:
URL: https://github.com/apache/superset/pull/43842#discussion_r4129618950
##########
superset/core/mcp/core_mcp_injection.py:
##########
@@ -40,6 +41,21 @@
logger = logging.getLogger(__name__)
+# The concrete ``@tool``/``@prompt`` decorators register with the FastMCP
instance
+# that lives in ``superset.mcp_service.app``. That module is the memory-heavy
part
+# of MCP setup and is imported by ``initialize_core_mcp_host_tools`` before any
+# extension applies a decorator. Processes that never serve MCP (e.g. Celery
+# workers with ``CORE_MCP_HOST_TOOLS_ENABLED = False``) only swap the
decorators
+# and never import it. In those processes a decoration must stay a no-op rather
+# than importing the service app on demand -- otherwise applying ``@tool`` in
an
+# extension would still load the host-tool stack and defeat the memory savings.
+_MCP_SERVICE_APP_MODULE = "superset.mcp_service.app"
+
+
+def _mcp_host_tools_loaded() -> bool:
Review Comment:
`_mcp_host_tools_loaded()` only checks whether `superset.mcp_service.app` is
a key in `sys.modules`, and Python adds a module to `sys.modules` the moment it
starts importing it — before that module's own host-tool registration (further
down in `mcp_service/app.py`) has actually run. A caller relying on this helper
to mean "host tools are registered" could get a false positive while that
module is still mid-import. Could the name/docstring make clear this signals
"app module present in `sys.modules`" rather than "host tools finished
registering"?
##########
tests/unit_tests/initialization_test.py:
##########
@@ -162,6 +163,114 @@ def test_init_app_in_ctx_calls_sync_config_to_db(self,
mock_logger):
# Assert that sync_config_to_db was called on the app
mock_app.sync_config_to_db.assert_called_once()
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_host_tools")
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_decorators")
+
@patch("superset.core.api.core_api_injection.initialize_core_api_dependencies")
+ def test_init_core_dependencies_registers_host_tools_by_default(
+ self,
+ mock_init_api,
+ mock_init_decorators,
+ mock_init_host_tools,
+ ):
+ """Default config (flag absent): decorators and host tools both
register."""
+ mock_app = MagicMock()
+ mock_app.config = {}
+ app_initializer = SupersetAppInitializer(mock_app)
+
+ app_initializer.init_core_dependencies()
+
+ mock_init_api.assert_called_once()
+ # Decorators are always registered so extensions using @tool/@prompt at
+ # import time keep working in every process.
+ mock_init_decorators.assert_called_once()
+ # Host tools default to on.
+ mock_init_host_tools.assert_called_once()
+
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_host_tools")
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_decorators")
+
@patch("superset.core.api.core_api_injection.initialize_core_api_dependencies")
+ def test_init_core_dependencies_registers_host_tools_when_enabled(
+ self,
+ mock_init_api,
+ mock_init_decorators,
+ mock_init_host_tools,
+ ):
+ """Flag True: decorators and host tools both register (unchanged
behavior)."""
+ mock_app = MagicMock()
+ mock_app.config = {"CORE_MCP_HOST_TOOLS_ENABLED": True}
+ app_initializer = SupersetAppInitializer(mock_app)
+
+ app_initializer.init_core_dependencies()
+
+ mock_init_api.assert_called_once()
+ mock_init_decorators.assert_called_once()
+ mock_init_host_tools.assert_called_once()
+
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_host_tools")
+
@patch("superset.core.mcp.core_mcp_injection.initialize_core_mcp_decorators")
+
@patch("superset.core.api.core_api_injection.initialize_core_api_dependencies")
+ def test_init_core_dependencies_skips_host_tools_when_disabled(
+ self,
+ mock_init_api,
+ mock_init_decorators,
+ mock_init_host_tools,
+ ):
+ """Flag False: host-tool registration is skipped, but core init and the
+ decorator swap still run (so @tool/@prompt extensions don't break)."""
+ mock_app = MagicMock()
+ mock_app.config = {"CORE_MCP_HOST_TOOLS_ENABLED": False}
+ app_initializer = SupersetAppInitializer(mock_app)
+
+ app_initializer.init_core_dependencies()
+
+ # Core API deps and the decorator swap still run.
+ mock_init_api.assert_called_once()
+ mock_init_decorators.assert_called_once()
+ # The heavy MCP service app import is skipped.
+ mock_init_host_tools.assert_not_called()
+
+ def test_disabled_worker_decorators_do_not_import_service_app(self):
Review Comment:
This test calls `create_tool_decorator`/`create_prompt_decorator` directly
rather than through `superset_core.mcp.decorators.tool`/`.prompt`, and the
three tests above it mock `initialize_core_mcp_decorators()` out entirely. So
no test in this file exercises the actual assignment inside
`initialize_core_mcp_decorators()` (`superset_core.mcp.decorators.tool =
create_tool_decorator`, etc.) — a regression there (e.g. binding the wrong
function) would pass every current test. Could one test call the real
`initialize_core_mcp_decorators()` and then decorate through
`superset_core.mcp.decorators.tool`/`.prompt` to cover that wiring?
##########
superset/core/mcp/core_mcp_injection.py:
##########
@@ -404,3 +453,19 @@ def initialize_core_mcp_dependencies() -> None:
logger.info("MCP service app imported - host tools registered")
except Exception as e:
logger.error("Failed to register MCP host tools: %s", e)
+
+
+def initialize_core_mcp_dependencies() -> None:
+ """
+ Initialize MCP dependency injection by replacing abstract functions
+ in superset_core.api.mcp with concrete implementations, then registering
the
Review Comment:
This docstring says the decorator swap happens "in `superset_core.api.mcp`",
but the code replaces attributes on `superset_core.mcp.decorators` (see
`initialize_core_mcp_decorators` above and the assignments a few lines up). A
reader tracing where the swap happens from this docstring would look in the
wrong module.
```suggestion
in superset_core.mcp.decorators with concrete implementations, then
registering the
```
--
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]
