diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py index 363a8e8c6ad9..c74494e34b47 100644 --- a/superset/mcp_service/middleware.py +++ b/superset/mcp_service/middleware.py @@ -246,7 +246,7 @@ def _extract_context_info( dashboard_id = None slice_id = None dataset_id = None - params = getattr(context.message, "params", {}) or {} + params = getattr(context.message, "arguments", {}) or {} if hasattr(context, "metadata") and context.metadata: agent_id = context.metadata.get("agent_id") if not agent_id and hasattr(context, "session") and context.session: @@ -1152,7 +1152,7 @@ async def on_call_tool( ) if estimated_tokens > self.token_limit: - params = getattr(context.message, "params", {}) or {} + params = getattr(context.message, "arguments", {}) or {} return self._handle_oversized_response( tool_name, response, estimated_tokens, params ) diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py index 622487109116..3f0957fd9b82 100644 --- a/tests/unit_tests/mcp_service/test_middleware.py +++ b/tests/unit_tests/mcp_service/test_middleware.py @@ -95,7 +95,7 @@ async def test_allows_small_response(self) -> None: # Create mock context context = MagicMock() context.message.name = "list_charts" - context.message.params = {} + context.message.arguments = {} # Create mock call_next that returns small response small_response = {"charts": [{"id": 1, "name": "test"}]} @@ -118,7 +118,7 @@ async def test_blocks_large_response(self) -> None: # Create mock context context = MagicMock() context.message.name = "list_charts" - context.message.params = {"page_size": 100} + context.message.arguments = {"page_size": 100} # Create large response large_response = { @@ -148,7 +148,7 @@ async def test_skips_excluded_tools(self) -> None: # Create mock context for excluded tool context = MagicMock() context.message.name = "health_check" - context.message.params = {} + context.message.arguments = {} # Create response that would exceed limit large_response = {"data": "x" * 10000} @@ -173,7 +173,7 @@ async def test_logs_warning_at_threshold(self) -> None: context = MagicMock() context.message.name = "list_charts" - context.message.params = {} + context.message.arguments = {} response = {"data": "approaching the limit"} call_next = AsyncMock(return_value=response) @@ -201,7 +201,7 @@ async def test_error_includes_suggestions(self) -> None: context = MagicMock() context.message.name = "list_charts" - context.message.params = {"page_size": 100} + context.message.arguments = {"page_size": 100} large_response = {"charts": [{"id": i} for i in range(1000)]} call_next = AsyncMock(return_value=large_response) @@ -226,7 +226,7 @@ async def test_logs_size_exceeded_event(self) -> None: context = MagicMock() context.message.name = "list_charts" - context.message.params = {} + context.message.arguments = {} large_response = {"data": "x" * 10000} call_next = AsyncMock(return_value=large_response) @@ -250,7 +250,7 @@ async def test_truncates_info_tool_instead_of_blocking(self) -> None: context = MagicMock() context.message.name = "get_dataset_info" - context.message.params = {} + context.message.arguments = {} # Large info tool response with a big description large_response = { @@ -279,7 +279,7 @@ async def test_truncates_chart_info_with_large_form_data(self) -> None: context = MagicMock() context.message.name = "get_chart_info" - context.message.params = {} + context.message.arguments = {} large_response = { "id": 1, @@ -305,7 +305,7 @@ async def test_still_blocks_non_info_tools(self) -> None: context = MagicMock() context.message.name = "list_charts" # Not an info tool - context.message.params = {} + context.message.arguments = {} large_response = {"data": "x" * 10000} call_next = AsyncMock(return_value=large_response) @@ -324,7 +324,7 @@ async def test_logs_truncation_event(self) -> None: context = MagicMock() context.message.name = "get_dashboard_info" - context.message.params = {} + context.message.arguments = {} large_response = { "id": 1, @@ -357,7 +357,7 @@ async def test_truncates_dashboard_info_with_custom_max_list_items(self) -> None context = MagicMock() context.message.name = "get_dashboard_info" - context.message.params = {} + context.message.arguments = {} large_response = { "id": 1, @@ -387,7 +387,7 @@ async def test_truncates_execute_sql_rows_instead_of_blocking(self) -> None: context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -417,7 +417,7 @@ async def test_truncates_query_dataset_data_field(self) -> None: context = MagicMock() context.message.name = "query_dataset" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -446,7 +446,7 @@ async def test_truncates_get_chart_data_rows(self) -> None: context = MagicMock() context.message.name = "get_chart_data" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -476,7 +476,7 @@ async def test_data_query_truncation_updates_row_count(self) -> None: context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -502,7 +502,7 @@ async def test_data_query_truncation_note_mentions_limit_clause(self) -> None: context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -529,7 +529,7 @@ async def test_data_query_truncation_logs_truncation_event(self) -> None: context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_response = { @@ -557,7 +557,7 @@ async def test_truncates_get_chart_data_csv_export(self) -> None: context = MagicMock() context.message.name = "get_chart_data" - context.message.params = {} + context.message.arguments = {} large_response: dict[str, Any] = { "chart_id": 1, @@ -595,7 +595,7 @@ async def test_data_query_blocks_when_single_row_still_exceeds_limit(self) -> No context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} huge_row = {"col": "x" * 5000} large_response = { @@ -619,7 +619,7 @@ async def test_data_query_under_limit_passes_through(self) -> None: context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} small_response = { "status": "success", @@ -949,7 +949,7 @@ async def test_info_tool_result_is_truncated_and_rewrapped(self) -> None: middleware = ResponseSizeGuardMiddleware(token_limit=500) context = MagicMock() context.message.name = "get_dataset_info" - context.message.params = {} + context.message.arguments = {} large_payload = {"id": 1, "table_name": "test", "description": "x" * 50000} tool_result = self._make_tool_result(large_payload) @@ -975,7 +975,7 @@ async def test_small_tool_result_passes_through_unchanged(self) -> None: middleware = ResponseSizeGuardMiddleware(token_limit=25000) context = MagicMock() context.message.name = "get_chart_info" - context.message.params = {} + context.message.arguments = {} small_payload = {"id": 1, "name": "My Chart"} tool_result = self._make_tool_result(small_payload) @@ -995,7 +995,7 @@ async def test_large_non_info_tool_result_is_blocked(self) -> None: middleware = ResponseSizeGuardMiddleware(token_limit=100) context = MagicMock() context.message.name = "list_charts" - context.message.params = {} + context.message.arguments = {} large_payload = { "charts": [{"id": i, "name": f"chart_{i}"} for i in range(500)] @@ -1026,7 +1026,7 @@ async def test_data_query_tool_result_is_truncated_and_rewrapped(self) -> None: middleware = ResponseSizeGuardMiddleware(token_limit=500) context = MagicMock() context.message.name = "execute_sql" - context.message.params = {} + context.message.arguments = {} row = {f"col_{i}": f"value_{i}" for i in range(10)} large_payload = { @@ -1058,7 +1058,7 @@ async def test_meta_preserved_after_truncation(self) -> None: middleware = ResponseSizeGuardMiddleware(token_limit=500) context = MagicMock() context.message.name = "get_dashboard_info" - context.message.params = {} + context.message.arguments = {} meta = {"request_id": "abc-123"} large_payload = {"id": 1, "title": "My Dashboard", "description": "x" * 50000} @@ -1093,7 +1093,7 @@ class ChartInfo(BaseModel): context = MagicMock() context.message.name = "get_chart_info" - context.message.params = {} + context.message.arguments = {} response = ChartInfo(id=1, name="Test Chart") call_next = AsyncMock(return_value=response) @@ -1113,7 +1113,7 @@ async def test_list_response(self) -> None: context = MagicMock() context.message.name = "list_charts" - context.message.params = {} + context.message.arguments = {} response = [{"id": 1}, {"id": 2}, {"id": 3}] call_next = AsyncMock(return_value=response) @@ -1133,7 +1133,7 @@ async def test_string_response(self) -> None: context = MagicMock() context.message.name = "health_check" - context.message.params = {} + context.message.arguments = {} response = "OK" call_next = AsyncMock(return_value=response) diff --git a/tests/unit_tests/mcp_service/test_middleware_logging.py b/tests/unit_tests/mcp_service/test_middleware_logging.py index d838a478a1d2..eed0334648ee 100644 --- a/tests/unit_tests/mcp_service/test_middleware_logging.py +++ b/tests/unit_tests/mcp_service/test_middleware_logging.py @@ -49,7 +49,7 @@ def _make_context( ctx.method = method message = MagicMock() message.name = name - message.params = params or {} + message.arguments = params or {} ctx.message = message if metadata is not None: ctx.metadata = metadata @@ -503,6 +503,33 @@ def test_extract_slice_id_from_slice_id(self, mock_get_user_id) -> None: assert slice_id == 66 + @patch("superset.mcp_service.middleware.get_user_id", return_value=1) + def test_extract_reads_arguments_on_real_call_tool_request_params( + self, mock_get_user_id + ) -> None: + """Regression test: the real MCP ``CallToolRequestParams`` object + exposes tool arguments as ``.arguments``, not ``.params`` -- a + ``MagicMock``-based context would auto-vivify a ``.params`` + attribute and hide a mismatch. Using the real SDK type here + ensures params/dashboard_id/etc. are actually populated instead + of silently logging as empty.""" + middleware = LoggingMiddleware() + message = mt.CallToolRequestParams( + name="get_dashboard_info", + arguments={"dashboard_id": 7}, + ) + ctx = MagicMock() + ctx.message = message + ctx.metadata = None + ctx.session = None + + agent_id, user_id, dashboard_id, slice_id, dataset_id, params = ( + middleware._extract_context_info(ctx) + ) + + assert params == {"dashboard_id": 7} + assert dashboard_id == 7 + class TestIsErrorResponse: """Tests for LoggingMiddleware._is_error_response()."""