feat(mcp): return form_data and form_data_key in generate_chart and generate_explore_link responses (#36539)

superset/mcp_service/chart/schemas.py

@@ -1155,6 +1155,16 @@ class GenerateChartResponse(BaseModel):
         default_factory=dict, description="Related API endpoints for data/updates"
     )
 
+    # Form data for rendering charts in external clients (chatbot rendering)
+    form_data: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Complete form_data configuration for rendering the chart",
+    )
+    form_data_key: str | None = Field(
+        None,
+        description="Cache key for the form_data, used in explore URLs",
+    )
+
     # Performance and accessibility
     performance: PerformanceMetadata | None = Field(
         None, description="Performance metrics"

superset/mcp_service/chart/tool/generate_chart.py

@@ -20,6 +20,7 @@ MCP tool: generate_chart (simplified schema)
 
 import logging
 import time
+from urllib.parse import parse_qs, urlparse
 
 from fastmcp import Context
 from superset_core.mcp import tool
@@ -284,6 +285,43 @@ async def generate_chart(  # noqa: C901
                 raise
             # Update explore URL to use saved chart
             explore_url = f"{get_superset_base_url()}/explore/?slice_id={chart.id}"
+
+            # Generate form_data_key for saved charts (needed for chatbot rendering)
+            try:
+                from superset.commands.explore.form_data.parameters import (
+                    CommandParameters,
+                )
+                from superset.mcp_service.commands.create_form_data import (
+                    MCPCreateFormDataCommand,
+                )
+                from superset.utils.core import DatasourceType
+
+                # Add datasource to form_data for the cache
+                form_data_with_datasource = {
+                    **form_data,
+                    "datasource": f"{dataset.id}__table",
+                }
+
+                cmd_params = CommandParameters(
+                    datasource_type=DatasourceType.TABLE,
+                    datasource_id=dataset.id,
+                    chart_id=chart.id,
+                    tab_id=None,
+                    form_data=json.dumps(form_data_with_datasource),
+                )
+                form_data_key = MCPCreateFormDataCommand(cmd_params).run()
+                await ctx.debug(
+                    "Generated form_data_key for saved chart: form_data_key=%s"
+                    % (form_data_key,)
+                )
+            except Exception as fdk_error:
+                logger.warning(
+                    "Failed to generate form_data_key for saved chart: %s", fdk_error
+                )
+                await ctx.warning(
+                    "Failed to generate form_data_key: error=%s" % (str(fdk_error),)
+                )
+                # form_data_key remains None but chart is still valid
         else:
             await ctx.report_progress(2, 5, "Generating temporary chart preview")
             # Generate explore link with cached form_data for preview-only mode
@@ -292,9 +330,13 @@ async def generate_chart(  # noqa: C901
             explore_url = generate_explore_link(request.dataset_id, form_data)
             await ctx.debug("Generated explore link: explore_url=%s" % (explore_url,))
 
-            # Extract form_data_key from the explore URL
-            if explore_url and "form_data_key=" in explore_url:
-                form_data_key = explore_url.split("form_data_key=")[1].split("&")[0]
+            # Extract form_data_key from the explore URL using proper URL parsing
+            if explore_url:
+                parsed = urlparse(explore_url)
+                query_params = parse_qs(parsed.query)
+                form_data_key_list = query_params.get("form_data_key", [])
+                if form_data_key_list:
+                    form_data_key = form_data_key_list[0]
 
         # Generate semantic analysis
         capabilities = analyze_chart_capabilities(chart, request.config)
@@ -405,25 +447,37 @@ async def generate_chart(  # noqa: C901
 
         # Return enhanced data while maintaining backward compatibility
         await ctx.report_progress(4, 5, "Building response")
+
+        # Build chart info using serialize_chart_object for saved charts
+        chart_info = None
+        if request.save_chart and chart:
+            from superset.mcp_service.chart.schemas import serialize_chart_object
+
+            chart_info = serialize_chart_object(chart)
+            if chart_info:
+                # Override the URL with explore_url
+                chart_info.url = explore_url
+
+        # Safely serialize chart_info - handle both Pydantic models and dicts
+        chart_data = None
+        if chart_info:
+            if hasattr(chart_info, "model_dump"):
+                chart_data = chart_info.model_dump()
+            elif isinstance(chart_info, dict):
+                chart_data = chart_info
+            else:
+                chart_data = chart_info  # Pass through as-is
+
         result = {
-            "chart": {
-                "id": chart.id if chart else None,
-                "slice_name": chart.slice_name
-                if chart
-                else generate_chart_name(request.config),
-                "viz_type": chart.viz_type if chart else form_data.get("viz_type"),
-                "url": explore_url,
-                "uuid": str(chart.uuid) if chart and chart.uuid else None,
-                "saved": request.save_chart,
-            }
-            if request.save_chart and chart
-            else None,
+            "chart": chart_data,
             "error": None,
             # Enhanced fields for better LLM integration
             "previews": previews,
             "capabilities": capabilities.model_dump() if capabilities else None,
             "semantics": semantics.model_dump() if semantics else None,
             "explore_url": explore_url,
+            # Form data fields - REQUIRED for chatbot/external client rendering
+            "form_data": form_data,
             "form_data_key": form_data_key,
             "api_endpoints": {
                 "data": f"{get_superset_base_url()}/api/v1/chart/{chart.id}/data/"

superset/mcp_service/explore/tool/generate_explore_link.py

@@ -23,6 +23,7 @@ chart configuration.
 """
 
 from typing import Any, Dict
+from urllib.parse import parse_qs, urlparse
 
 from fastmcp import Context
 from superset_core.mcp import tool
@@ -106,14 +107,26 @@ async def generate_explore_link(
         # Generate explore link using shared utilities
         explore_url = generate_url(dataset_id=request.dataset_id, form_data=form_data)
 
+        # Extract form_data_key from the explore URL using proper URL parsing
+        form_data_key = None
+        if explore_url:
+            parsed = urlparse(explore_url)
+            query_params = parse_qs(parsed.query)
+            form_data_key_list = query_params.get("form_data_key", [])
+            if form_data_key_list:
+                form_data_key = form_data_key_list[0]
+
         await ctx.report_progress(3, 3, "URL generation complete")
         await ctx.info(
-            "Explore link generated successfully: url_length=%s, dataset_id=%s"
-            % (len(explore_url), request.dataset_id)
+            "Explore link generated successfully: url_length=%s, dataset_id=%s, "
+            "form_data_key=%s"
+            % (len(explore_url or ""), request.dataset_id, form_data_key)
         )
 
         return {
             "url": explore_url,
+            "form_data": form_data,
+            "form_data_key": form_data_key,
             "error": None,
         }
 
@@ -124,5 +137,7 @@ async def generate_explore_link(
         )
         return {
             "url": "",
+            "form_data": {},
+            "form_data_key": None,
             "error": f"Failed to generate explore link: {str(e)}",
         }

tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py

@@ -589,3 +589,97 @@ class TestGenerateExploreLink:
                 expected_url = f"http://localhost:9001/explore/?datasource_type=table&datasource_id={dataset_id}"
                 assert result.data["error"] is None
                 assert result.data["url"] == expected_url
+
+    @pytest.mark.asyncio
+    async def test_generate_explore_link_tool_exception_handling(self, mcp_server):
+        """Test that tool-level exceptions are properly handled and return error."""
+        import sys
+
+        # Get the actual module object from sys.modules (not via __init__.py which
+        # returns the function)
+        explore_module = sys.modules[
+            "superset.mcp_service.explore.tool.generate_explore_link"
+        ]
+
+        original_func = explore_module.map_config_to_form_data
+
+        def raise_error(*args, **kwargs):
+            raise ValueError("Invalid config structure")
+
+        explore_module.map_config_to_form_data = raise_error
+        try:
+            config = TableChartConfig(
+                chart_type="table", columns=[ColumnRef(name="test_col")]
+            )
+            request = GenerateExploreLinkRequest(dataset_id="1", config=config)
+
+            async with Client(mcp_server) as client:
+                result = await client.call_tool(
+                    "generate_explore_link", {"request": request.model_dump()}
+                )
+
+                # Should return error response with empty URL
+                assert result.data["url"] == ""
+                assert result.data["form_data"] == {}
+                assert result.data["form_data_key"] is None
+                assert "Invalid config structure" in result.data["error"]
+        finally:
+            # Restore original function
+            explore_module.map_config_to_form_data = original_func
+
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    @patch(
+        "superset.mcp_service.commands.create_form_data.MCPCreateFormDataCommand.run"
+    )
+    @pytest.mark.asyncio
+    async def test_generate_explore_link_returns_form_data_key(
+        self, mock_create_form_data, mock_find_dataset, mcp_server
+    ):
+        """Test that form_data_key is properly extracted from URL."""
+        mock_create_form_data.return_value = "extracted_form_key_xyz"
+        mock_find_dataset.return_value = _mock_dataset(id=1)
+
+        config = TableChartConfig(
+            chart_type="table", columns=[ColumnRef(name="region")]
+        )
+        request = GenerateExploreLinkRequest(dataset_id="1", config=config)
+
+        async with Client(mcp_server) as client:
+            result = await client.call_tool(
+                "generate_explore_link", {"request": request.model_dump()}
+            )
+
+            assert result.data["error"] is None
+            assert result.data["form_data_key"] == "extracted_form_key_xyz"
+            assert "form_data_key=extracted_form_key_xyz" in result.data["url"]
+
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    @patch(
+        "superset.mcp_service.commands.create_form_data.MCPCreateFormDataCommand.run"
+    )
+    @pytest.mark.asyncio
+    async def test_generate_explore_link_returns_form_data(
+        self, mock_create_form_data, mock_find_dataset, mcp_server
+    ):
+        """Test that form_data dict is returned for external rendering."""
+        mock_create_form_data.return_value = "form_data_test_key"
+        mock_find_dataset.return_value = _mock_dataset(id=1)
+
+        config = XYChartConfig(
+            chart_type="xy",
+            x=ColumnRef(name="date"),
+            y=[ColumnRef(name="sales", aggregate="SUM")],
+            kind="line",
+        )
+        request = GenerateExploreLinkRequest(dataset_id="1", config=config)
+
+        async with Client(mcp_server) as client:
+            result = await client.call_tool(
+                "generate_explore_link", {"request": request.model_dump()}
+            )
+
+            assert result.data["error"] is None
+            assert "form_data" in result.data
+            assert isinstance(result.data["form_data"], dict)
+            assert result.data["form_data"].get("viz_type") == "echarts_timeseries_line"
+            assert result.data["form_data"].get("x_axis") == "date"