This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 2317d9cf918 fix(mcp): truncate query-tool responses instead of
hard-failing (#42244)
2317d9cf918 is described below
commit 2317d9cf918248e2b6b8d9887742ea9aacd59533
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Thu Jul 23 14:38:49 2026 -0400
fix(mcp): truncate query-tool responses instead of hard-failing (#42244)
---
superset/mcp_service/middleware.py | 181 ++++++++++---
superset/mcp_service/utils/token_utils.py | 244 +++++++++++++++++
tests/unit_tests/mcp_service/test_middleware.py | 295 +++++++++++++++++++++
.../mcp_service/utils/test_token_utils.py | 103 +++++++
4 files changed, 781 insertions(+), 42 deletions(-)
diff --git a/superset/mcp_service/middleware.py
b/superset/mcp_service/middleware.py
index fdb5681eaa0..5faebbb19c7 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -51,10 +51,12 @@ from superset.mcp_service.constants import (
DEFAULT_WARN_THRESHOLD_PCT,
)
from superset.mcp_service.utils.token_utils import (
+ DATA_QUERY_TOOLS,
estimate_response_tokens,
format_size_limit_error,
INFO_TOOLS,
truncate_oversized_response,
+ truncate_query_result,
)
from superset.utils.core import get_user_id
@@ -840,6 +842,141 @@ class ResponseSizeGuardMiddleware(Middleware):
return truncated
+ def _try_truncate_data_query_response(
+ self,
+ tool_name: str,
+ response: Any,
+ estimated_tokens: int,
+ ) -> Any | None:
+ """Attempt to truncate a data-query tool response by dropping tail
rows.
+
+ Returns the truncated response if successful, None otherwise.
+ """
+ extracted = self._extract_payload_from_tool_result(response)
+ truncation_target = extracted if extracted is not None else response
+
+ try:
+ truncated, was_truncated, notes = truncate_query_result(
+ truncation_target, self.token_limit, tool_name=tool_name
+ )
+ except Exception as trunc_error: # noqa: BLE001
+ logger.warning(
+ "Query result truncation failed for %s due to %s: %s",
+ tool_name,
+ type(trunc_error).__name__,
+ trunc_error,
+ )
+ return None
+
+ if not was_truncated:
+ return None
+
+ # Mirror the info-tool path: if truncation couldn't bring the
+ # response back under the limit (e.g. a single row/scalar field
+ # alone exceeds it), fall back to the hard size-limit error instead
+ # of shipping an over-budget response.
+ truncated_tokens = estimate_response_tokens(truncated)
+ if truncated_tokens > self.token_limit:
+ return None
+
+ logger.warning(
+ "Query result for %s truncated from ~%d to ~%d tokens (limit: %d).
%s",
+ tool_name,
+ estimated_tokens,
+ truncated_tokens,
+ self.token_limit,
+ "; ".join(notes),
+ )
+
+ try:
+ user_id = get_user_id()
+ event_logger.log(
+ user_id=user_id,
+ action="mcp_response_truncated",
+ curated_payload={
+ "tool": tool_name,
+ "original_tokens": estimated_tokens,
+ "truncated_tokens": truncated_tokens,
+ "token_limit": self.token_limit,
+ "truncation_notes": notes,
+ },
+ )
+ except Exception as log_error: # noqa: BLE001
+ logger.warning("Failed to log truncation event: %s", log_error)
+
+ if extracted is not None and isinstance(truncated, dict):
+ return self._rewrap_as_tool_result(truncated, response)
+
+ return truncated
+
+ def _handle_oversized_response(
+ self,
+ tool_name: str,
+ response: Any,
+ estimated_tokens: int,
+ params: dict[str, Any],
+ ) -> Any:
+ """Attempt truncation for known tool categories; block everything else.
+
+ For info tools (``INFO_TOOLS``) and data-query tools
+ (``DATA_QUERY_TOOLS``), tries dynamic truncation first and returns
+ the truncated result if successful. Falls through to a hard
+ ``ToolError`` for all other tools, or when truncation cannot reduce
+ the response to fit the limit.
+
+ Raises:
+ ToolError: When the response exceeds the limit and cannot be
+ truncated.
+ """
+ # Info tools: field-level truncation (strings, lists, dicts).
+ if tool_name in INFO_TOOLS:
+ truncated = self._try_truncate_info_response(
+ tool_name, response, estimated_tokens
+ )
+ if truncated is not None:
+ return truncated
+
+ # Data-query tools: row-level truncation.
+ if tool_name in DATA_QUERY_TOOLS:
+ truncated = self._try_truncate_data_query_response(
+ tool_name, response, estimated_tokens
+ )
+ if truncated is not None:
+ return truncated
+
+ # Log the blocked response (user-caused: requested too much data)
+ logger.warning(
+ "Response blocked for %s: ~%d tokens exceeds limit of %d",
+ tool_name,
+ estimated_tokens,
+ self.token_limit,
+ )
+
+ try:
+ user_id = get_user_id()
+ event_logger.log(
+ user_id=user_id,
+ action="mcp_response_size_exceeded",
+ curated_payload={
+ "tool": tool_name,
+ "estimated_tokens": estimated_tokens,
+ "token_limit": self.token_limit,
+ "params": _sanitize_params(params),
+ },
+ )
+ except Exception as log_error: # noqa: BLE001
+ logger.warning("Failed to log size exceeded event: %s", log_error)
+
+ raise ToolError(
+ format_size_limit_error(
+ tool_name=tool_name,
+ params=params,
+ estimated_tokens=estimated_tokens,
+ token_limit=self.token_limit,
+ response=None,
+ )
+ )
+
async def on_call_tool(
self,
context: MiddlewareContext,
@@ -886,52 +1023,12 @@ class ResponseSizeGuardMiddleware(Middleware):
self.token_limit,
)
- # Block if over limit
if estimated_tokens > self.token_limit:
params = getattr(context.message, "params", {}) or {}
-
- # For info tools, try dynamic truncation before blocking
- if tool_name in INFO_TOOLS:
- truncated = self._try_truncate_info_response(
- tool_name, response, estimated_tokens
- )
- if truncated is not None:
- return truncated
-
- # Log the blocked response (user-caused: requested too much data)
- logger.warning(
- "Response blocked for %s: ~%d tokens exceeds limit of %d",
- tool_name,
- estimated_tokens,
- self.token_limit,
- )
-
- # Log to event logger for monitoring
- try:
- user_id = get_user_id()
- event_logger.log(
- user_id=user_id,
- action="mcp_response_size_exceeded",
- curated_payload={
- "tool": tool_name,
- "estimated_tokens": estimated_tokens,
- "token_limit": self.token_limit,
- "params": _sanitize_params(params),
- },
- )
- except Exception as log_error: # noqa: BLE001
- logger.warning("Failed to log size exceeded event: %s",
log_error)
-
- error_message = format_size_limit_error(
- tool_name=tool_name,
- params=params,
- estimated_tokens=estimated_tokens,
- token_limit=self.token_limit,
- response=None,
+ return self._handle_oversized_response(
+ tool_name, response, estimated_tokens, params
)
- raise ToolError(error_message)
-
return response
diff --git a/superset/mcp_service/utils/token_utils.py
b/superset/mcp_service/utils/token_utils.py
index c8e482dffd7..5d546e2e635 100644
--- a/superset/mcp_service/utils/token_utils.py
+++ b/superset/mcp_service/utils/token_utils.py
@@ -469,6 +469,22 @@ INFO_TOOLS = frozenset(
}
)
+# Data-query tools that return tabular results. When the response exceeds the
+# token limit, these tools are truncated by dropping tail rows rather than
+# raising a hard ToolError. A truncation note is appended so the caller
+# knows the result is partial and how to narrow the query.
+DATA_QUERY_TOOLS = frozenset(
+ {
+ "execute_sql",
+ "query_dataset",
+ "get_chart_data",
+ }
+)
+
+# Data field names used by the three query tools (in priority order).
+# ``rows`` is used by execute_sql; ``data`` by query_dataset and
get_chart_data.
+_DATA_ROW_FIELDS = ("rows", "data")
+
# Maximum character length for string fields before truncation
_MAX_STRING_CHARS = 500
# Maximum keys to keep when summarizing large dict fields
@@ -663,6 +679,234 @@ def truncate_oversized_response(
return data, was_truncated, notes
+def _bisect_row_limit(
+ data: Dict[str, Any],
+ row_field: str,
+ original_rows: List[Any],
+ token_limit: int,
+) -> int:
+ """Binary-search for the largest row prefix that keeps data under limit.
+
+ Mutates ``data[row_field]`` during the search and leaves it at the final
+ kept count on return. Returns the number of rows kept (>= 1 if the
+ original list was non-empty).
+ """
+ from superset.utils import json as utils_json
+
+ lo, hi = 0, len(original_rows)
+ while lo < hi:
+ mid = (lo + hi + 1) // 2
+ data[row_field] = original_rows[:mid]
+ if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ lo = mid
+ else:
+ hi = mid - 1
+
+ kept = lo
+ if kept == 0 and original_rows:
+ # Even a single row is too large — keep it anyway so the caller gets
+ # at least some data rather than an empty list.
+ kept = 1
+
+ data[row_field] = original_rows[:kept]
+ return kept
+
+
+def _bisect_string_length(
+ data: Dict[str, Any],
+ field: str,
+ original_value: str,
+ token_limit: int,
+) -> int:
+ """Binary-search for the largest string prefix that keeps data under limit.
+
+ Mutates ``data[field]`` during the search and leaves it at the final
+ kept length on return.
+ """
+ from superset.utils import json as utils_json
+
+ lo, hi = 0, len(original_value)
+ while lo < hi:
+ mid = (lo + hi + 1) // 2
+ data[field] = original_value[:mid]
+ if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ lo = mid
+ else:
+ hi = mid - 1
+
+ kept = lo
+ if kept == 0 and original_value:
+ kept = 1
+
+ data[field] = original_value[:kept]
+ return kept
+
+
+# Row-truncation advice, keyed by tool name. ``get_chart_data`` and
+# ``query_dataset`` cap output via a row/page-size parameter, not a SQL
+# clause, so the generic "Add a LIMIT clause" wording is wrong for them.
+_ROW_LIMIT_ADVICE = {
+ "get_chart_data": "Use the 'limit' parameter to restrict rows returned.",
+ "query_dataset": "Use the 'row_limit' parameter to restrict rows
returned.",
+}
+_DEFAULT_ROW_LIMIT_ADVICE = (
+ "Add a LIMIT clause or reduce selected columns to get all rows."
+)
+
+
+def _truncate_rows_field(
+ data: Dict[str, Any],
+ row_field: str,
+ token_limit: int,
+ advice: str,
+) -> list[str] | None:
+ """Try to bisect ``data[row_field]`` down to fit the limit.
+
+ Returns the truncation notes on success, or ``None`` if no rows were
+ (or could be) dropped, in which case ``data`` is left unmodified.
+ """
+ original_rows: list[Any] = data[row_field]
+ original_count = len(original_rows)
+ if not original_rows:
+ return None
+
+ # Reserve space for the truncation-note metadata *before* bisecting so
+ # the search already accounts for it, rather than measuring fit on bare
+ # row data and finding out only afterward that appending the note
+ # pushed the payload back over the limit. The placeholder uses
+ # original_count for both halves of "X of Y rows", which is the
+ # longest the real note can ever be (kept <= original_count).
+ placeholder_note = (
+ f"Result truncated: {original_count} of {original_count} rows returned
"
+ f"(limit ~{token_limit:,} tokens). {advice}"
+ )
+ data["_response_truncated"] = True
+ data["_truncation_notes"] = [placeholder_note]
+
+ kept = _bisect_row_limit(data, row_field, original_rows, token_limit)
+
+ if kept < original_count:
+ notes = [
+ f"Result truncated: {kept} of {original_count} rows returned "
+ f"(limit ~{token_limit:,} tokens). {advice}"
+ ]
+ data["_truncation_notes"] = notes
+ if "row_count" in data:
+ data["row_count"] = kept
+ return notes
+
+ # Reserving headroom turned out not to be needed after all — no rows
+ # were actually dropped, so undo the placeholder metadata.
+ del data["_response_truncated"]
+ del data["_truncation_notes"]
+ return None
+
+
+def _truncate_csv_data_field(
+ data: Dict[str, Any],
+ token_limit: int,
+ advice: str,
+) -> list[str] | None:
+ """Try to bisect the scalar ``csv_data`` field down to fit the limit.
+
+ Used when there are no rows to trim — e.g. a CSV export where ``data``
+ is empty and the actual payload lives in ``csv_data``. ``excel_data``
+ is base64-encoded binary and is intentionally left alone: cutting it
+ would produce a corrupt file, so oversized Excel exports still fall
+ through to the hard size-limit error.
+
+ Returns the truncation notes on success, or ``None`` if nothing could
+ be trimmed, in which case ``data`` is left unmodified.
+ """
+ csv_data = data.get("csv_data")
+ if not isinstance(csv_data, str) or not csv_data:
+ return None
+
+ original_len = len(csv_data)
+ # Same reservation trick as ``_truncate_rows_field``, keyed on the
+ # character count rather than a row count.
+ placeholder_note = (
+ f"CSV content truncated: kept {original_len:,} of {original_len:,} "
+ f"characters (limit ~{token_limit:,} tokens). {advice}"
+ )
+ data["_response_truncated"] = True
+ data["_truncation_notes"] = [placeholder_note]
+
+ kept_len = _bisect_string_length(data, "csv_data", csv_data, token_limit)
+ if kept_len < original_len:
+ notes = [
+ f"CSV content truncated: kept {kept_len:,} of {original_len:,} "
+ f"characters (limit ~{token_limit:,} tokens). {advice}"
+ ]
+ data["_truncation_notes"] = notes
+ return notes
+
+ del data["_response_truncated"]
+ del data["_truncation_notes"]
+ return None
+
+
+def truncate_query_result(
+ response: ToolResponse,
+ token_limit: int,
+ tool_name: str | None = None,
+) -> tuple[ToolResponse, bool, list[str]]:
+ """Truncate a data-query tool response to fit within the token limit.
+
+ Unlike ``truncate_oversized_response`` (which targets info-tool dict
+ fields), this function targets the rows/data list directly: it performs
+ a binary search to find the largest prefix of rows that keeps the
+ serialised payload under the limit, then records a truncation note so
+ the caller knows the result is partial.
+
+ Handles both dict responses and Pydantic models (e.g. ExecuteSqlResponse,
+ QueryDatasetResponse, GetChartDataResponse) that contain a ``rows`` or
+ ``data`` field. Note that summary/column-level statistics (e.g.
+ ``summary``, per-column ``null_count``/``unique_count``, ``insights``)
+ describe the pre-truncation dataset and are not recomputed here; callers
+ rely on ``_truncation_notes`` and the true ``total_rows`` to catch this.
+
+ Args:
+ response: The tool response containing tabular row data.
+ token_limit: Maximum estimated tokens allowed.
+ tool_name: Name of the calling tool, used to tailor the truncation
+ advice (e.g. ``limit`` vs. ``row_limit`` vs. SQL ``LIMIT``).
+
+ Returns:
+ A tuple of (possibly-truncated response, was_truncated, list of notes).
+ """
+ from superset.utils import json as utils_json
+
+ advice = _ROW_LIMIT_ADVICE.get(tool_name or "", _DEFAULT_ROW_LIMIT_ADVICE)
+
+ if hasattr(response, "model_dump"):
+ data = response.model_dump()
+ elif isinstance(response, dict):
+ data = dict(response)
+ else:
+ return response, False, []
+
+ # Find which field holds the row list.
+ row_field: str | None = None
+ for field in _DATA_ROW_FIELDS:
+ if field in data and isinstance(data[field], list):
+ row_field = field
+ break
+
+ if row_field is None:
+ # No recognised row field — fall back to generic field truncation.
+ return truncate_oversized_response(response, token_limit)
+
+ if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ return data, False, []
+
+ notes = _truncate_rows_field(data, row_field, token_limit, advice)
+ if notes is None:
+ notes = _truncate_csv_data_field(data, token_limit, advice)
+
+ return data, notes is not None, notes or []
+
+
def format_size_limit_error(
tool_name: str,
params: Dict[str, Any] | None,
diff --git a/tests/unit_tests/mcp_service/test_middleware.py
b/tests/unit_tests/mcp_service/test_middleware.py
index 818450e049c..62248710911 100644
--- a/tests/unit_tests/mcp_service/test_middleware.py
+++ b/tests/unit_tests/mcp_service/test_middleware.py
@@ -380,6 +380,263 @@ class TestResponseSizeGuardMiddleware:
# native_filters (48 items) fits under the custom cap, untouched
assert len(result["native_filters"]) == 48
+ @pytest.mark.asyncio
+ async def test_truncates_execute_sql_rows_instead_of_blocking(self) ->
None:
+ """execute_sql should truncate rows, not raise ToolError."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "status": "success",
+ "rows": [row] * 200,
+ "columns": [{"name": f"col_{i}", "type": "STRING"} for i in
range(10)],
+ "row_count": 200,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ assert result["_response_truncated"] is True
+ assert len(result["rows"]) < 200
+ assert isinstance(result.get("_truncation_notes"), list)
+ assert result["_truncation_notes"]
+
+ @pytest.mark.asyncio
+ async def test_truncates_query_dataset_data_field(self) -> None:
+ """query_dataset should truncate the 'data' list, not raise
ToolError."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "query_dataset"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "dataset_id": 1,
+ "dataset_name": "test",
+ "data": [row] * 200,
+ "row_count": 200,
+ "summary": "",
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ assert result["_response_truncated"] is True
+ assert len(result["data"]) < 200
+
+ @pytest.mark.asyncio
+ async def test_truncates_get_chart_data_rows(self) -> None:
+ """get_chart_data should truncate rows, not raise ToolError."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "get_chart_data"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "chart_id": 1,
+ "chart_name": "test chart",
+ "chart_type": "table",
+ "data": [row] * 200,
+ "row_count": 200,
+ "summary": "",
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ assert result["_response_truncated"] is True
+ assert len(result["data"]) < 200
+
+ @pytest.mark.asyncio
+ async def test_data_query_truncation_updates_row_count(self) -> None:
+ """row_count should reflect the truncated count, not the original."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "status": "success",
+ "rows": [row] * 200,
+ "row_count": 200,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert result["_response_truncated"] is True
+ assert result["row_count"] == len(result["rows"])
+
+ @pytest.mark.asyncio
+ async def test_data_query_truncation_note_mentions_limit_clause(self) ->
None:
+ """Truncation note must tell the caller to add a LIMIT clause."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "status": "success",
+ "rows": [row] * 200,
+ "row_count": 200,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ notes = result.get("_truncation_notes", [])
+ assert notes, "Should have at least one truncation note"
+ assert any("LIMIT" in note for note in notes)
+
+ @pytest.mark.asyncio
+ async def test_data_query_truncation_logs_truncation_event(self) -> None:
+ """Should log mcp_response_truncated (not size_exceeded) for query
tools."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_response = {
+ "status": "success",
+ "rows": [row] * 200,
+ "row_count": 200,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger") as
mock_event_logger,
+ ):
+ await middleware.on_call_tool(context, call_next)
+
+ mock_event_logger.log.assert_called()
+ assert (
+ mock_event_logger.log.call_args.kwargs["action"] ==
"mcp_response_truncated"
+ )
+
+ @pytest.mark.asyncio
+ async def test_truncates_get_chart_data_csv_export(self) -> None:
+ """CSV exports (data=[], payload in csv_data) should be truncated
too."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "get_chart_data"
+ context.message.params = {}
+
+ large_response: dict[str, Any] = {
+ "chart_id": 1,
+ "chart_name": "test chart",
+ "chart_type": "table",
+ "columns": [],
+ "data": [],
+ "row_count": 200,
+ "summary": "",
+ "csv_data": "col_0,col_1\n" + ("value,value\n" * 2000),
+ "format": "csv",
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ assert result["_response_truncated"] is True
+ assert len(result["csv_data"]) < len(large_response["csv_data"])
+
+ @pytest.mark.asyncio
+ async def test_data_query_blocks_when_single_row_still_exceeds_limit(self)
-> None:
+ """Should raise ToolError, not ship an over-budget response.
+
+ ``_bisect_row_limit`` always keeps at least one row when the
+ original list is non-empty, even if that one row alone exceeds the
+ token limit. The middleware must re-check the truncated size and
+ fall back to the hard error rather than treating this as success.
+ """
+ middleware = ResponseSizeGuardMiddleware(token_limit=50)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ huge_row = {"col": "x" * 5000}
+ large_response = {
+ "status": "success",
+ "rows": [huge_row] * 3,
+ "row_count": 3,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ pytest.raises(ToolError),
+ ):
+ await middleware.on_call_tool(context, call_next)
+
+ @pytest.mark.asyncio
+ async def test_data_query_under_limit_passes_through(self) -> None:
+ """Small query results should pass through unchanged."""
+ middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ small_response = {
+ "status": "success",
+ "rows": [{"a": 1, "b": 2}] * 3,
+ "row_count": 3,
+ }
+ call_next = AsyncMock(return_value=small_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert result == small_response
+ assert "_response_truncated" not in result
+
class TestCreateResponseSizeGuardMiddleware:
"""Test create_response_size_guard_middleware factory function."""
@@ -753,6 +1010,44 @@ class TestToolResultWrapping:
):
await middleware.on_call_tool(context, call_next)
+ @pytest.mark.asyncio
+ async def test_data_query_tool_result_is_truncated_and_rewrapped(self) ->
None:
+ """Truncate a ToolResult-wrapped execute_sql response and re-wrap it.
+
+ Regression test for the production path: FastMCP always wraps tool
+ return values in ToolResult before middleware sees them, so
+ data-query truncation must be exercised through that wrapper, not
+ just against a plain dict.
+ """
+ from fastmcp.tools.tool import ToolResult
+
+ from superset.utils import json
+
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ context = MagicMock()
+ context.message.name = "execute_sql"
+ context.message.params = {}
+
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ large_payload = {
+ "status": "success",
+ "rows": [row] * 200,
+ "row_count": 200,
+ }
+ tool_result = self._make_tool_result(large_payload)
+ call_next = AsyncMock(return_value=tool_result)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id",
return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, ToolResult)
+ reparsed = json.loads(result.content[0].text)
+ assert reparsed["_response_truncated"] is True
+ assert len(reparsed["rows"]) < 200
+
@pytest.mark.asyncio
async def test_meta_preserved_after_truncation(self) -> None:
"""Should preserve the original ToolResult meta through truncation."""
diff --git a/tests/unit_tests/mcp_service/utils/test_token_utils.py
b/tests/unit_tests/mcp_service/utils/test_token_utils.py
index f42401385a6..d183bd770b7 100644
--- a/tests/unit_tests/mcp_service/utils/test_token_utils.py
+++ b/tests/unit_tests/mcp_service/utils/test_token_utils.py
@@ -40,6 +40,7 @@ from superset.mcp_service.utils.token_utils import (
get_response_size_bytes,
INFO_TOOLS,
truncate_oversized_response,
+ truncate_query_result,
)
@@ -749,3 +750,105 @@ class TestTruncateOversizedResponse:
assert isinstance(result, dict)
assert len(result["charts"]) == 5
assert any("form_data" in n for n in notes)
+
+
+class TestTruncateQueryResult:
+ """Tests for ``truncate_query_result`` (data-query row/scalar
truncation)."""
+
+ def _rows_response(self, row_field: str, count: int = 200) -> dict[str,
Any]:
+ row = {f"col_{i}": f"value_{i}" for i in range(10)}
+ return {
+ "status": "success",
+ row_field: [row] * count,
+ "row_count": count,
+ }
+
+ def test_no_truncation_needed(self) -> None:
+ response = self._rows_response("rows", count=3)
+ result, was_truncated, notes = truncate_query_result(response, 25000)
+ assert was_truncated is False
+ assert notes == []
+ assert result == response
+
+ def test_truncated_result_fits_under_limit(self) -> None:
+ """The final payload (rows + note metadata) must itself fit.
+
+ Regression test: the note is built from the kept row count, but
+ that note text also consumes tokens. The bisection must reserve
+ room for it up front rather than measuring fit on bare rows and
+ appending the note afterward, which could push the final payload
+ back over the limit.
+ """
+ response = self._rows_response("rows")
+ result, was_truncated, notes = truncate_query_result(response, 500)
+ assert was_truncated is True
+ assert isinstance(result, dict)
+ assert estimate_response_tokens(result) <= 500
+ assert result["row_count"] == len(result["rows"])
+ assert result["row_count"] < 200
+
+ def test_single_oversized_row_is_kept_anyway(self) -> None:
+ """A single row that alone exceeds the limit is still returned.
+
+ ``truncate_query_result`` always keeps >=1 row when the original
+ list is non-empty; it is the caller's (middleware) job to reject
+ a still-oversized result rather than ship it silently.
+ """
+ response = {
+ "status": "success",
+ "rows": [{"col": "x" * 5000}] * 3,
+ "row_count": 3,
+ }
+ result, was_truncated, notes = truncate_query_result(response, 50)
+ assert was_truncated is True
+ assert isinstance(result, dict)
+ assert len(result["rows"]) == 1
+ assert notes
+
+ def test_truncates_csv_scalar_field_when_rows_empty(self) -> None:
+ """CSV exports carry their payload in ``csv_data`` with ``data=[]``."""
+ response: dict[str, Any] = {
+ "chart_id": 1,
+ "data": [],
+ "csv_data": "col_0,col_1\n" + ("value,value\n" * 2000),
+ "format": "csv",
+ }
+ result, was_truncated, notes = truncate_query_result(
+ response, 500, tool_name="get_chart_data"
+ )
+ assert was_truncated is True
+ assert isinstance(result, dict)
+ assert len(result["csv_data"]) < len(response["csv_data"])
+ assert estimate_response_tokens(result) <= 500
+ assert any("CSV" in n for n in notes)
+
+ def test_does_not_truncate_excel_binary_field(self) -> None:
+ """excel_data is base64 binary — truncating it would corrupt the
file."""
+ response = {
+ "chart_id": 1,
+ "data": [],
+ "excel_data": "QUJDREVGRw==" * 5000,
+ "format": "excel",
+ }
+ result, was_truncated, notes = truncate_query_result(response, 500)
+ assert was_truncated is False
+ assert notes == []
+ assert isinstance(result, dict)
+ assert result["excel_data"] == response["excel_data"]
+
+ def test_get_chart_data_advice_mentions_limit_param(self) -> None:
+ response = self._rows_response("data")
+ _, _, notes = truncate_query_result(response, 500,
tool_name="get_chart_data")
+ assert any("'limit' parameter" in n for n in notes)
+ assert not any("LIMIT clause" in n for n in notes)
+
+ def test_query_dataset_advice_mentions_row_limit_param(self) -> None:
+ response = self._rows_response("data")
+ _, _, notes = truncate_query_result(response, 500,
tool_name="query_dataset")
+ assert any("'row_limit' parameter" in n for n in notes)
+ assert not any("LIMIT clause" in n for n in notes)
+
+ def test_execute_sql_advice_mentions_limit_clause(self) -> None:
+ response = self._rows_response("rows")
+ _, _, notes = truncate_query_result(response, 500,
tool_name="execute_sql")
+ assert any("LIMIT clause" in n for n in notes)
