From 5e3ca00bb2adc553fee1c72dc181bee7f132ae23 Mon Sep 17 00:00:00 2001
From: alexandrusoare <alexandrusoare98@gmail.com>
Date: Thu, 30 Apr 2026 15:24:45 +0300
Subject: [PATCH] cover edge case

---
 superset/mcp_service/middleware.py            |  3 ++
 .../mcp_service/test_middleware_logging.py    | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index a7767a03478..a90c8448321 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -241,6 +241,7 @@ class LoggingMiddleware(Middleware):
         tool_name = getattr(context.message, "name", None)
 
         mcp_call_id = uuid.uuid4().hex[:12]
+        context.mcp_call_id = mcp_call_id
         start_time = time.time()
         success = False
         try:
@@ -396,8 +397,10 @@ class StructuredContentStripperMiddleware(Middleware):
             # unhandled exception — including ToolError from
             # GlobalErrorHandlerMiddleware, ValueError, TypeError, etc. —
             # will cause encoding failures on the wire.
+            mcp_call_id = getattr(context, "mcp_call_id", None)
             return ToolResult(
                 content=[mt.TextContent(type="text", text=f"Error: {e}")],
+                meta={"mcp_call_id": mcp_call_id} if mcp_call_id else None,
             )
         if isinstance(result, ToolResult) and result.structured_content is not None:
             result = ToolResult(content=result.content, meta=result.meta)
diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py
index 1da6f7d4575..ca68689557b 100644
--- a/tests/unit_tests/mcp_service/test_middleware_logging.py
+++ b/tests/unit_tests/mcp_service/test_middleware_logging.py
@@ -426,3 +426,35 @@ class TestMiddlewareChainOrder:
             "them. Ensure StructuredContentStripper is outermost "
             "(first added) in build_middleware_list()."
         )
+
+    @patch("superset.mcp_service.middleware.event_logger")
+    @patch("superset.mcp_service.middleware.get_user_id", return_value=42)
+    @pytest.mark.asyncio
+    async def test_real_middleware_chain_error_result_has_mcp_call_id(
+        self, mock_get_user_id, mock_event_logger
+    ):
+        """When a tool raises, the error ToolResult from
+        StructuredContentStripper still carries mcp_call_id in meta."""
+        from functools import partial
+
+        from fastmcp.tools.tool import ToolResult
+
+        from superset.mcp_service.server import build_middleware_list
+
+        middleware_list = build_middleware_list()
+
+        async def failing_tool(context: Any) -> Any:
+            raise ValueError("chart not found")
+
+        chain = failing_tool
+        for mw in reversed(middleware_list):
+            chain = partial(mw, call_next=chain)
+
+        ctx = _make_context(name="get_chart_info")
+        result = await chain(ctx)
+
+        assert isinstance(result, ToolResult)
+        assert result.content[0].text.startswith("Error:")
+        assert result.meta is not None
+        assert "mcp_call_id" in result.meta
+        assert len(result.meta["mcp_call_id"]) == 12