You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ju...@apache.org on 2026/08/04 00:22:36 UTC
(superset) 01/02: fix(mcp): resolve real tool name and capture entity IDs through call_tool proxy and request wrapper
This is an automated email from the ASF dual-hosted git repository.
justinpark pushed a commit to branch 6.2
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 309aec3dab11ee77c12d0a44480f6595032c1c6c
Author: justinpark <ju...@gmail.com>
AuthorDate: Mon Aug 3 17:21:24 2026 -0700
fix(mcp): resolve real tool name and capture entity IDs through call_tool proxy and request wrapper
_resolve_tool_name() was dead code (never invoked, and _CALL_TOOL_PROXY
was undefined), so tool search proxy calls always logged tool="call_tool".
Separately, every MCP tool takes a single `request` argument, so real
arguments arrive nested as {"request": {...}} (and one layer deeper when
routed through the call_tool proxy), which the previous flat
params.get("dashboard_id") lookup could never see through. Also extend
the create-tool output backfill to cover create_virtual_dataset's
dataset_id, which generate_chart/generate_dashboard already had for
chart/dashboard ids.
---
superset/mcp_service/middleware.py | 107 ++++++++++----
.../mcp_service/test_middleware_logging.py | 155 ++++++++++++++++++++-
2 files changed, 230 insertions(+), 32 deletions(-)
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index 105b0afff4e..ecc2aebe375 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -130,6 +130,9 @@ class LoggingMiddleware(Middleware):
in on_message().
"""
+ _CALL_TOOL_PROXY = "call_tool"
+ _DATASET_CREATE_TOOLS = frozenset({"create_virtual_dataset"})
+
def _is_error_response(self, result: ToolResult) -> bool:
"""Check if a tool result contains an error schema response.
@@ -142,6 +145,35 @@ class LoggingMiddleware(Middleware):
except (AttributeError, IndexError):
return False
+ @staticmethod
+ def _unwrap_params_for_id_extraction(params: Any) -> dict[str, Any]:
+ """Unwrap tool arguments down to the level where entity IDs live.
+
+ Real tool arguments arrive wrapped in up to two extra layers that
+ the naive top-level ``params.get("dashboard_id")`` lookup can't see
+ through:
+
+ 1. When the MCP client calls tools via the ``call_tool`` search
+ proxy, arguments are ``{"name": ..., "arguments": {...}}``
+ instead of the real tool's arguments directly.
+ 2. Every MCP tool in this service takes a single ``request``
+ pydantic argument (e.g. ``def get_chart_info(request, ctx)``),
+ so the real fields are nested one level under ``{"request": {...}}``.
+
+ Without unwrapping both layers, dashboard_id/chart_id/dataset_id
+ extraction silently returns None for virtually all real traffic.
+ """
+ if not isinstance(params, dict):
+ return {}
+ effective = params
+ if isinstance(effective.get("name"), str) and isinstance(
+ effective.get("arguments"), dict
+ ):
+ effective = effective["arguments"]
+ if isinstance(effective.get("request"), dict):
+ effective = effective["request"]
+ return effective
+
def _extract_context_info(
self, context: MiddlewareContext
) -> tuple[
@@ -162,31 +194,39 @@ class LoggingMiddleware(Middleware):
user_id = get_user_id()
except (RuntimeError, AttributeError):
user_id = None
- if isinstance(params, dict):
- dashboard_id = params.get("dashboard_id")
- slice_id = params.get("chart_id") or params.get("slice_id")
- dataset_id = params.get("dataset_id")
+ effective_params = self._unwrap_params_for_id_extraction(params)
+ dashboard_id = effective_params.get("dashboard_id")
+ slice_id = effective_params.get("chart_id") or effective_params.get("slice_id")
+ dataset_id = effective_params.get("dataset_id")
return agent_id, user_id, dashboard_id, slice_id, dataset_id, params
- def _extract_output_ids(self, result: ToolResult) -> tuple[int | None, int | None]:
- """Extract dashboard/chart IDs created by the tool from its response.
-
- Create-style tools (generate_chart, generate_dashboard) don't take
- chart_id/dashboard_id as input, so _extract_context_info never sees
- them and every retry logs slice_id/dashboard_id=None even on the
- attempt that actually persisted the object. Look at the response
- body instead, since that's the only place the new ID appears.
- Supports both flat ("chart_id"/"dashboard_id") and nested
- ("chart"/"dashboard" objects with an "id" field) response shapes.
+ def _extract_output_ids(
+ self, tool_name: str | None, result: ToolResult
+ ) -> tuple[int | None, int | None, int | None]:
+ """Extract dashboard/chart/dataset IDs created by the tool from its response.
+
+ Create-style tools (generate_chart, generate_dashboard,
+ create_virtual_dataset) don't take the new object's ID as input, so
+ _extract_context_info never sees it and every retry logs
+ slice_id/dashboard_id/dataset_id=None even on the attempt that
+ actually persisted the object. Look at the response body instead,
+ since that's the only place the new ID appears. Supports both flat
+ ("chart_id"/"dashboard_id") and nested ("chart"/"dashboard" objects
+ with an "id" field) response shapes.
+
+ dataset_id extraction is gated on tool_name (_DATASET_CREATE_TOOLS)
+ because create_virtual_dataset's response has a flat "id" field
+ with no distinguishing wrapper key, unlike chart/dashboard -- a
+ blind "id" lookup would misattribute IDs from unrelated tools.
"""
from superset.utils.json import loads as json_loads
try:
data = json_loads(result.content[0].text)
except (AttributeError, IndexError, ValueError, TypeError):
- return None, None
+ return None, None, None
if not isinstance(data, dict):
- return None, None
+ return None, None, None
slice_id = None
chart = data.get("chart")
@@ -202,7 +242,11 @@ class LoggingMiddleware(Middleware):
if dashboard_id is None:
dashboard_id = data.get("dashboard_id")
- return dashboard_id, slice_id
+ dataset_id = None
+ if tool_name in self._DATASET_CREATE_TOOLS:
+ dataset_id = data.get("id")
+
+ return dashboard_id, slice_id, dataset_id
@staticmethod
def _resolve_tool_name(tool_name: str | None, params: Any) -> str | None:
@@ -228,26 +272,32 @@ class LoggingMiddleware(Middleware):
def _backfill_output_ids(
self,
+ tool_name: str | None,
success: bool,
result: Any,
dashboard_id: int | None,
slice_id: int | None,
- ) -> tuple[int | None, int | None]:
+ dataset_id: int | None,
+ ) -> tuple[int | None, int | None, int | None]:
"""Fill in missing ids from a create tool's response on success.
- Create-style tools (generate_chart, generate_dashboard) don't take
- the new object's ID as input, so it's missing from params. On a
- successful call, pull it from the response instead so retried
- creates are distinguishable.
+ Create-style tools (generate_chart, generate_dashboard,
+ create_virtual_dataset) don't take the new object's ID as input, so
+ it's missing from params. On a successful call, pull it from the
+ response instead so retried creates are distinguishable.
"""
if not success or not isinstance(result, ToolResult):
- return dashboard_id, slice_id
- output_dashboard_id, output_slice_id = self._extract_output_ids(result)
+ return dashboard_id, slice_id, dataset_id
+ output_dashboard_id, output_slice_id, output_dataset_id = (
+ self._extract_output_ids(tool_name, result)
+ )
if dashboard_id is None:
dashboard_id = output_dashboard_id
if slice_id is None:
slice_id = output_slice_id
- return dashboard_id, slice_id
+ if dataset_id is None:
+ dataset_id = output_dataset_id
+ return dashboard_id, slice_id, dataset_id
@staticmethod
def _build_call_tool_payload(
@@ -289,8 +339,8 @@ class LoggingMiddleware(Middleware):
start_time: float,
) -> None:
duration_ms = int((time.time() - start_time) * 1000)
- dashboard_id, slice_id = self._backfill_output_ids(
- success, result, dashboard_id, slice_id
+ dashboard_id, slice_id, dataset_id = self._backfill_output_ids(
+ tool_name, success, result, dashboard_id, slice_id, dataset_id
)
payload = self._build_call_tool_payload(
tool_name=tool_name,
@@ -337,6 +387,9 @@ class LoggingMiddleware(Middleware):
self._extract_context_info(context)
)
tool_name = getattr(context.message, "name", None)
+ resolved_tool_name = self._resolve_tool_name(tool_name, params)
+ if resolved_tool_name:
+ tool_name = resolved_tool_name
start_time = time.time()
success = False
diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py
index 1665bffaf52..46ae4a771b9 100644
--- a/tests/unit_tests/mcp_service/test_middleware_logging.py
+++ b/tests/unit_tests/mcp_service/test_middleware_logging.py
@@ -253,6 +253,45 @@ class TestLoggingMiddlewareOnCallTool:
call_kwargs = mock_event_logger.log.call_args[1]
assert call_kwargs["slice_id"] == 111
+ @patch("superset.mcp_service.middleware.event_logger")
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=42)
+ @pytest.mark.asyncio
+ async def test_on_call_tool_resolves_name_and_dataset_id_through_proxy(
+ self, mock_get_user_id, mock_event_logger
+ ) -> None:
+ """End-to-end regression test for the reported bug: a
+ create_virtual_dataset call made through the ``call_tool`` search
+ proxy must log the real tool name (not "call_tool") and the
+ created dataset_id (backfilled from the response), instead of
+ tool="call_tool" and dataset_id=None."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(
+ name="call_tool",
+ params={
+ "name": "create_virtual_dataset",
+ "arguments": {
+ "request": {
+ "database_id": 108,
+ "schema": "jitney",
+ "table_name": "Airchat Monthly Usage Summary",
+ "sql": "SELECT 1",
+ }
+ },
+ },
+ )
+ response_text = '{"id": 42, "dataset_name": "Airchat Monthly Usage Summary"}'
+ original_result = ToolResult(
+ content=[mt.TextContent(type="text", text=response_text)]
+ )
+ call_next = AsyncMock(return_value=original_result)
+
+ await middleware.on_call_tool(ctx, call_next)
+
+ call_kwargs = mock_event_logger.log.call_args[1]
+ assert call_kwargs["curated_payload"]["tool"] == "create_virtual_dataset"
+ assert call_kwargs["curated_payload"]["dataset_id"] == 42
+ assert call_kwargs["curated_payload"]["success"] is True
+
class TestLoggingMiddlewareOnMessage:
"""Tests for LoggingMiddleware.on_message()."""
@@ -359,6 +398,72 @@ class TestExtractContextInfo:
assert params == {"dashboard_id": 7}
assert dashboard_id == 7
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=1)
+ def test_extract_unwraps_request_wrapper(self, mock_get_user_id) -> None:
+ """Every MCP tool in this service takes a single ``request``
+ pydantic argument, so real arguments arrive as
+ ``{"request": {"dashboard_id": ..., "chart_id": ...}}`` (e.g.
+ add_chart_to_existing_dashboard). Without unwrapping this layer,
+ dashboard_id/chart_id extraction silently returns None."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(
+ name="add_chart_to_existing_dashboard",
+ params={"request": {"dashboard_id": 10, "chart_id": 20}},
+ )
+
+ _, _, dashboard_id, slice_id, _, params = middleware._extract_context_info(ctx)
+
+ assert dashboard_id == 10
+ assert slice_id == 20
+ # The raw params (used for the logged payload) stay unwrapped.
+ assert params == {"request": {"dashboard_id": 10, "chart_id": 20}}
+
+ @patch("superset.mcp_service.middleware.get_user_id", return_value=1)
+ def test_extract_unwraps_call_tool_proxy_and_request(
+ self, mock_get_user_id
+ ) -> None:
+ """Reproduces the reported bug: when the client calls tools through
+ the ``call_tool`` search proxy, arguments arrive as
+ ``{"name": ..., "arguments": {"request": {...}}}`` -- two layers of
+ wrapping. Both must be unwrapped to see dataset_id."""
+ middleware = LoggingMiddleware()
+ ctx = _make_context(
+ name="call_tool",
+ params={
+ "name": "create_virtual_dataset",
+ "arguments": {
+ "request": {
+ "database_id": 108,
+ "dataset_id": 42,
+ }
+ },
+ },
+ )
+
+ _, _, _, _, dataset_id, _ = middleware._extract_context_info(ctx)
+
+ assert dataset_id == 42
+
+
+class TestResolveToolName:
+ """Tests for LoggingMiddleware._resolve_tool_name()."""
+
+ def test_resolves_real_tool_name_from_call_tool_proxy(self) -> None:
+ resolved = LoggingMiddleware._resolve_tool_name(
+ "call_tool", {"name": "create_virtual_dataset", "arguments": {}}
+ )
+ assert resolved == "create_virtual_dataset"
+
+ def test_returns_none_when_not_call_tool_proxy(self) -> None:
+ resolved = LoggingMiddleware._resolve_tool_name(
+ "get_chart_info", {"identifier": 123}
+ )
+ assert resolved is None
+
+ def test_returns_none_when_name_missing(self) -> None:
+ resolved = LoggingMiddleware._resolve_tool_name("call_tool", {})
+ assert resolved is None
+
class TestIsErrorResponse:
"""Tests for LoggingMiddleware._is_error_response()."""
@@ -434,17 +539,24 @@ class TestExtractOutputIds:
def test_returns_none_for_non_json_body(self) -> None:
"""A malformed/non-JSON response body must not raise -- the
- defensive try/except should fall back to (None, None)."""
+ defensive try/except should fall back to (None, None, None)."""
middleware = LoggingMiddleware()
result = ToolResult(
content=[mt.TextContent(type="text", text="not valid json {{{")]
)
- assert middleware._extract_output_ids(result) == (None, None)
+ assert middleware._extract_output_ids("generate_chart", result) == (
+ None,
+ None,
+ None,
+ )
def test_returns_none_for_empty_content(self) -> None:
"""A ToolResult with no content items must not raise."""
middleware = LoggingMiddleware()
- assert middleware._extract_output_ids(ToolResult(content=[])) == (
+ assert middleware._extract_output_ids(
+ "generate_chart", ToolResult(content=[])
+ ) == (
+ None,
None,
None,
)
@@ -454,13 +566,46 @@ class TestExtractOutputIds:
list) must not raise and must yield no IDs."""
middleware = LoggingMiddleware()
result = ToolResult(content=[mt.TextContent(type="text", text="[1, 2, 3]")])
- assert middleware._extract_output_ids(result) == (None, None)
+ assert middleware._extract_output_ids("generate_chart", result) == (
+ None,
+ None,
+ None,
+ )
def test_extracts_both_ids_from_flat_response(self) -> None:
middleware = LoggingMiddleware()
response_text = '{"chart_id": 123, "dashboard_id": 456}'
result = ToolResult(content=[mt.TextContent(type="text", text=response_text)])
- assert middleware._extract_output_ids(result) == (456, 123)
+ assert middleware._extract_output_ids("generate_chart", result) == (
+ 456,
+ 123,
+ None,
+ )
+
+ def test_extracts_dataset_id_for_create_virtual_dataset(self) -> None:
+ """create_virtual_dataset's response is a flat {"id": ...} with no
+ wrapper key, unlike chart/dashboard -- extraction must be gated on
+ tool_name so unrelated tools' "id" fields aren't misattributed."""
+ middleware = LoggingMiddleware()
+ response_text = '{"id": 789, "dataset_name": "my_dataset"}'
+ result = ToolResult(content=[mt.TextContent(type="text", text=response_text)])
+ assert middleware._extract_output_ids("create_virtual_dataset", result) == (
+ None,
+ None,
+ 789,
+ )
+
+ def test_does_not_extract_dataset_id_for_unrelated_tool(self) -> None:
+ """A generic "id" field on an unrelated tool's response must not be
+ misread as a dataset_id."""
+ middleware = LoggingMiddleware()
+ response_text = '{"id": 789}'
+ result = ToolResult(content=[mt.TextContent(type="text", text=response_text)])
+ assert middleware._extract_output_ids("get_chart_info", result) == (
+ None,
+ None,
+ None,
+ )
class TestMiddlewareChainOrder: