feat(mcp): implement selective column serialization for list tools (#36035)

tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py

@@ -112,12 +112,17 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is None
-            assert result.data.dashboard is not None
-            assert result.data.dashboard.id == 10
-            assert result.data.dashboard.dashboard_title == "Analytics Dashboard"
-            assert result.data.dashboard.chart_count == 2
-            assert "/superset/dashboard/10/" in result.data.dashboard_url
+            assert result.structured_content["error"] is None
+            assert result.structured_content["dashboard"] is not None
+            assert result.structured_content["dashboard"]["id"] == 10
+            assert (
+                result.structured_content["dashboard"]["dashboard_title"]
+                == "Analytics Dashboard"
+            )
+            assert result.structured_content["dashboard"]["chart_count"] == 2
+            assert (
+                "/superset/dashboard/10/" in result.structured_content["dashboard_url"]
+            )
 
     @patch("superset.db.session")
     @pytest.mark.asyncio
@@ -138,10 +143,10 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is not None
-            assert "Charts not found: [2]" in result.data.error
-            assert result.data.dashboard is None
-            assert result.data.dashboard_url is None
+            assert result.structured_content["error"] is not None
+            assert "Charts not found: [2]" in result.structured_content["error"]
+            assert result.structured_content["dashboard"] is None
+            assert result.structured_content["dashboard_url"] is None
 
     @patch("superset.commands.dashboard.create.CreateDashboardCommand")
     @patch("superset.db.session")
@@ -169,9 +174,9 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is None
-            assert result.data.dashboard.chart_count == 1
-            assert result.data.dashboard.published is True  # From mock
+            assert result.structured_content["error"] is None
+            assert result.structured_content["dashboard"]["chart_count"] == 1
+            assert result.structured_content["dashboard"]["published"] is True
 
     @patch("superset.commands.dashboard.create.CreateDashboardCommand")
     @patch("superset.db.session")
@@ -198,8 +203,8 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is None
-            assert result.data.dashboard.chart_count == 6
+            assert result.structured_content["error"] is None
+            assert result.structured_content["dashboard"]["chart_count"] == 6
 
             # Verify CreateDashboardCommand was called with proper layout
             mock_create_command.assert_called_once()
@@ -257,9 +262,9 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is not None
-            assert "Failed to create dashboard" in result.data.error
-            assert result.data.dashboard is None
+            assert result.structured_content["error"] is not None
+            assert "Failed to create dashboard" in result.structured_content["error"]
+            assert result.structured_content["dashboard"] is None
 
     @patch("superset.commands.dashboard.create.CreateDashboardCommand")
     @patch("superset.db.session")
@@ -287,8 +292,11 @@ class TestGenerateDashboard:
         async with Client(mcp_server) as client:
             result = await client.call_tool("generate_dashboard", {"request": request})
 
-            assert result.data.error is None
-            assert result.data.dashboard.dashboard_title == "Minimal Dashboard"
+            assert result.structured_content["error"] is None
+            assert (
+                result.structured_content["dashboard"]["dashboard_title"]
+                == "Minimal Dashboard"
+            )
 
             # Check that description was not included in call
             call_args = mock_create_command.call_args[0][0]
@@ -343,13 +351,15 @@ class TestAddChartToExistingDashboard:
                 "add_chart_to_existing_dashboard", {"request": request}
             )
 
-            assert result.data.error is None
-            assert result.data.dashboard is not None
-            assert result.data.dashboard.chart_count == 3
-            assert result.data.position is not None
-            assert "row" in result.data.position  # Should have row info
-            assert "chart_key" in result.data.position
-            assert "/superset/dashboard/1/" in result.data.dashboard_url
+            assert result.structured_content["error"] is None
+            assert result.structured_content["dashboard"] is not None
+            assert result.structured_content["dashboard"]["chart_count"] == 3
+            assert result.structured_content["position"] is not None
+            assert "row" in result.structured_content["position"]
+            assert "chart_key" in result.structured_content["position"]
+            assert (
+                "/superset/dashboard/1/" in result.structured_content["dashboard_url"]
+            )
 
     @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
     @pytest.mark.asyncio
@@ -364,8 +374,10 @@ class TestAddChartToExistingDashboard:
                 "add_chart_to_existing_dashboard", {"request": request}
             )
 
-            assert result.data.error is not None
-            assert "Dashboard with ID 999 not found" in result.data.error
+            assert result.structured_content["error"] is not None
+            assert (
+                "Dashboard with ID 999 not found" in result.structured_content["error"]
+            )
 
     @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
     @patch("superset.db.session")
@@ -384,8 +396,8 @@ class TestAddChartToExistingDashboard:
                 "add_chart_to_existing_dashboard", {"request": request}
             )
 
-            assert result.data.error is not None
-            assert "Chart with ID 999 not found" in result.data.error
+            assert result.structured_content["error"] is not None
+            assert "Chart with ID 999 not found" in result.structured_content["error"]
 
     @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
     @patch("superset.db.session")
@@ -407,8 +419,11 @@ class TestAddChartToExistingDashboard:
                 "add_chart_to_existing_dashboard", {"request": request}
             )
 
-            assert result.data.error is not None
-            assert "Chart 5 is already in dashboard 1" in result.data.error
+            assert result.structured_content["error"] is not None
+            assert (
+                "Chart 5 is already in dashboard 1"
+                in result.structured_content["error"]
+            )
 
     @patch("superset.commands.dashboard.update.UpdateDashboardCommand")
     @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@@ -437,9 +452,9 @@ class TestAddChartToExistingDashboard:
                 "add_chart_to_existing_dashboard", {"request": request}
             )
 
-            assert result.data.error is None
-            assert "row" in result.data.position  # Should have row info
-            assert result.data.position.get("row") == 0  # First row
+            assert result.structured_content["error"] is None
+            assert "row" in result.structured_content["position"]
+            assert result.structured_content["position"].get("row") == 0
 
             # Verify update was called with proper layout structure
             call_args = mock_update_command.call_args[0][1]

tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_tools.py

@@ -103,18 +103,17 @@ async def test_list_dashboards_basic(mock_list, mcp_server):
         result = await client.call_tool(
             "list_dashboards", {"request": request.model_dump()}
         )
-        dashboards = result.data.dashboards
+        dashboards = result.data["dashboards"]
         assert len(dashboards) == 1
-        assert dashboards[0].dashboard_title == "Test Dashboard"
-        assert dashboards[0].uuid == "test-dashboard-uuid-1"
-        assert dashboards[0].slug == "test-dashboard"
-        assert dashboards[0].published is True
+        assert dashboards[0]["dashboard_title"] == "Test Dashboard"
+        assert dashboards[0]["uuid"] == "test-dashboard-uuid-1"
+        assert dashboards[0]["slug"] == "test-dashboard"
+        assert dashboards[0]["published"] is True
 
-        # Verify UUID and slug are in default columns
-        assert "uuid" in result.data.columns_requested
-        assert "slug" in result.data.columns_requested
-        assert "uuid" in result.data.columns_loaded
-        assert "slug" in result.data.columns_loaded
+        assert "uuid" in result.data["columns_requested"]
+        assert "slug" in result.data["columns_requested"]
+        assert "uuid" in result.data["columns_loaded"]
+        assert "slug" in result.data["columns_loaded"]
 
 
 @patch("superset.daos.dashboard.DashboardDAO.list")
@@ -180,8 +179,8 @@ async def test_list_dashboards_with_filters(mock_list, mcp_server):
         result = await client.call_tool(
             "list_dashboards", {"request": request.model_dump()}
         )
-        assert result.data.count == 1
-        assert result.data.dashboards[0].dashboard_title == "Filtered Dashboard"
+        assert result.data["count"] == 1
+        assert result.data["dashboards"][0]["dashboard_title"] == "Filtered Dashboard"
 
 
 @patch("superset.daos.dashboard.DashboardDAO.list")
@@ -262,8 +261,8 @@ async def test_list_dashboards_with_search(mock_list, mcp_server):
         result = await client.call_tool(
             "list_dashboards", {"request": request.model_dump()}
         )
-        assert result.data.count == 1
-        assert result.data.dashboards[0].dashboard_title == "search_dashboard"
+        assert result.data["count"] == 1
+        assert result.data["dashboards"][0]["dashboard_title"] == "search_dashboard"
         args, kwargs = mock_list.call_args
         assert kwargs["search"] == "search_dashboard"
         assert "dashboard_title" in kwargs["search_columns"]
@@ -283,7 +282,7 @@ async def test_list_dashboards_with_simple_filters(mock_list, mcp_server):
         result = await client.call_tool(
             "list_dashboards", {"request": request.model_dump()}
         )
-        assert hasattr(result.data, "count")
+        assert "count" in result.data
 
 
 @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
@@ -503,16 +502,15 @@ async def test_list_dashboards_custom_uuid_slug_columns(mock_list, mcp_server):
         result = await client.call_tool(
             "list_dashboards", {"request": request.model_dump()}
         )
-        dashboards = result.data.dashboards
+        dashboards = result.data["dashboards"]
         assert len(dashboards) == 1
-        assert dashboards[0].uuid == "test-custom-uuid-123"
-        assert dashboards[0].slug == "custom-dashboard"
+        assert dashboards[0]["uuid"] == "test-custom-uuid-123"
+        assert dashboards[0]["slug"] == "custom-dashboard"
 
-        # Verify custom columns include UUID and slug
-        assert "uuid" in result.data.columns_requested
-        assert "slug" in result.data.columns_requested
-        assert "uuid" in result.data.columns_loaded
-        assert "slug" in result.data.columns_loaded
+        assert "uuid" in result.data["columns_requested"]
+        assert "slug" in result.data["columns_requested"]
+        assert "uuid" in result.data["columns_loaded"]
+        assert "slug" in result.data["columns_loaded"]
 
 
 class TestDashboardSortableColumns:
@@ -558,9 +556,8 @@ class TestDashboardSortableColumns:
             assert call_args["order_column"] == "dashboard_title"
             assert call_args["order_direction"] == "desc"
 
-            # Verify the result
-            assert result.data.count == 0
-            assert result.data.dashboards == []
+            assert result.data["count"] == 0
+            assert result.data["dashboards"] == []
 
     def test_sortable_columns_in_docstring(self):
         """Test that sortable columns are documented in tool docstring."""