fix(mcp): sync per-statement row_count after execute_sql truncation

_bisect_row_limit already trims statements[*].data.rows in lockstep
with the top-level rows field to keep the response under the token
limit, but left each statement's row_count pointing at the
pre-truncation count, so a truncated execute_sql response could claim
more rows than statements[*].data actually contains. Track the
statement dicts (not just their row lists) and update row_count
alongside data.rows.

Also log a full traceback (logger.exception) instead of a plain
warning when data-query truncation hits an unexpected error, so a
real bug isn't indistinguishable from routine size-limit enforcement.

Adds regression tests for: total_rows surviving truncation distinct
from row_count, the no-recognised-row-field fallback path, and the
oversized-excel-export hard-failure path.

superset/mcp_service/middleware.py

@@ -52,10 +52,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
 
@@ -1247,6 +1249,25 @@ class ResponseSizeGuardMiddleware(Middleware):
             meta=original.meta if isinstance(original, ToolResult) else None,
         )
 
+    def _make_size_fn(
+        self, extracted: dict[str, Any] | None, original: Any
+    ) -> Callable[[dict[str, Any]], int]:
+        """Build a size function that measures the object actually returned.
+
+        When the payload was unwrapped from a ToolResult, re-wrapping adds
+        overhead (the ``content``/``meta`` envelope) on top of the bare
+        payload. Truncation must search for a fit against that final
+        wrapped size, not the unwrapped payload, or a candidate that looks
+        like it fits can still end up oversized once re-wrapped.
+        """
+        if extracted is None:
+            return estimate_response_tokens
+
+        def _size(data: dict[str, Any]) -> int:
+            return estimate_response_tokens(self._rewrap_as_tool_result(data, original))
+
+        return _size
+
     def _try_truncate_info_response(
         self,
         tool_name: str,
@@ -1274,11 +1295,14 @@ class ResponseSizeGuardMiddleware(Middleware):
             )
             truncation_target = response
 
+        size_fn = self._make_size_fn(extracted, response)
+
         try:
             truncated, was_truncated, notes = truncate_oversized_response(
                 truncation_target,
                 self.token_limit,
                 max_list_items=self.max_list_items,
+                size_fn=size_fn,
             )
         except (MemoryError, RecursionError) as trunc_error:
             logger.warning(
@@ -1292,7 +1316,20 @@ class ResponseSizeGuardMiddleware(Middleware):
         if not was_truncated:
             return None
 
-        truncated_tokens = estimate_response_tokens(truncated)
+        if isinstance(truncated, dict):
+            truncated["_response_truncated"] = True
+            truncated["_truncation_notes"] = notes
+
+        # Re-wrap into ToolResult if we unwrapped one *before* re-checking
+        # the size: the wrapper (content/meta envelope) adds overhead on top
+        # of the payload, so the limit must be enforced against the object
+        # that is actually returned to the caller, not the bare payload.
+        if extracted is not None and isinstance(truncated, dict):
+            final_response = self._rewrap_as_tool_result(truncated, response)
+        else:
+            final_response = truncated
+
+        truncated_tokens = estimate_response_tokens(final_response)
         if truncated_tokens > self.token_limit:
             return None
 
@@ -1321,15 +1358,163 @@ class ResponseSizeGuardMiddleware(Middleware):
         except Exception as log_error:  # noqa: BLE001
             logger.warning("Failed to log truncation event: %s", log_error)
 
+        return final_response
+
+    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
+        size_fn = self._make_size_fn(extracted, response)
+
+        try:
+            truncated, was_truncated, notes = truncate_query_result(
+                truncation_target,
+                self.token_limit,
+                tool_name=tool_name,
+                size_fn=size_fn,
+            )
+        except Exception as trunc_error:  # noqa: BLE001
+            # Deliberately broad: any failure here must fail safe into the
+            # hard size-limit error below rather than ship a broken
+            # truncation, so this stays a catch-all rather than narrowing to
+            # specific exception types. Logged with a full traceback (unlike
+            # the plain warning above) so a genuine bug (e.g. a KeyError from
+            # an unexpected response shape) is still visible instead of
+            # looking like routine size-limit enforcement.
+            logger.exception(
+                "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
+
         if isinstance(truncated, dict):
             truncated["_response_truncated"] = True
             truncated["_truncation_notes"] = notes
 
-        # Re-wrap into ToolResult if we unwrapped one
+        # Re-wrap into ToolResult if we unwrapped one *before* re-checking
+        # the size: the wrapper (content/meta envelope) adds overhead on top
+        # of the payload, so the limit must be enforced against the object
+        # that is actually returned to the caller, not the bare payload.
         if extracted is not None and isinstance(truncated, dict):
-            return self._rewrap_as_tool_result(truncated, response)
+            final_response = self._rewrap_as_tool_result(truncated, response)
+        else:
+            final_response = truncated
 
-        return truncated
+        # 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(final_response)
+        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)
+
+        return final_response
+
+    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,
@@ -1377,52 +1562,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,
+            return self._handle_oversized_response(
+                tool_name, response, estimated_tokens, params
             )
 
-            # 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,
-            )
-
-            raise ToolError(error_message)
-
         return response
 
 

superset/mcp_service/utils/token_utils.py

@@ -46,7 +46,7 @@ that risk.
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Union
+from typing import Any, Callable, Dict, List, Union
 
 from pydantic import BaseModel
 from typing_extensions import TypeAlias
@@ -58,6 +58,12 @@ logger = logging.getLogger(__name__)
 # Type alias for MCP tool responses (Pydantic models, dicts, lists, strings, bytes)
 ToolResponse: TypeAlias = Union[BaseModel, Dict[str, Any], List[Any], str, bytes]
 
+# Measures the token size of a candidate payload dict. Callers that re-wrap
+# the payload before returning it (e.g. back into a ToolResult) can supply a
+# size function that measures the *wrapped* size, so truncation converges on
+# a result that fits after re-wrapping, not just before it.
+SizeFn: TypeAlias = Callable[[Dict[str, Any]], int]
+
 # Fallback character-to-token ratio used when tiktoken is unavailable.
 # 3.0 is conservative for JSON content (the previous 3.5 under-counted
 # JSON-heavy payloads relative to Claude's actual tokenizer, which let
@@ -469,6 +475,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
@@ -592,11 +614,18 @@ def _replace_collections_with_summaries(data: Dict[str, Any], notes: List[str])
     return changed
 
 
-def _is_under_limit(data: Dict[str, Any], token_limit: int) -> bool:
-    """Check if the serialized data fits within the token limit."""
+def _default_size_fn(data: Dict[str, Any]) -> int:
+    """Estimate tokens for a bare payload dict (no re-wrap overhead)."""
     from superset.utils import json as utils_json
 
-    return estimate_token_count(utils_json.dumps(data)) <= token_limit
+    return estimate_token_count(utils_json.dumps(data))
+
+
+def _is_under_limit(
+    data: Dict[str, Any], token_limit: int, size_fn: SizeFn = _default_size_fn
+) -> bool:
+    """Check if the serialized data fits within the token limit."""
+    return size_fn(data) <= token_limit
 
 
 def truncate_oversized_response(
@@ -604,6 +633,7 @@ def truncate_oversized_response(
     token_limit: int,
     # Configurable via MCP_RESPONSE_SIZE_CONFIG["max_list_items"]
     max_list_items: int = DEFAULT_MAX_LIST_ITEMS,
+    size_fn: SizeFn = _default_size_fn,
 ) -> tuple[ToolResponse, bool, list[str]]:
     """
     Dynamically truncate large fields in a response to fit within the token limit.
@@ -619,6 +649,11 @@ def truncate_oversized_response(
         response: The tool response (Pydantic model, dict, or other).
         token_limit: Maximum estimated tokens allowed.
         max_list_items: Maximum items to keep in list fields during Phase 2.
+        size_fn: Measures the token size of a candidate payload dict. Defaults
+            to measuring the bare dict; callers that re-wrap the payload
+            before returning it (e.g. back into a ToolResult) should pass a
+            function that measures the *wrapped* size, so truncation
+            converges on a result that still fits after re-wrapping.
 
     Returns:
         A tuple of (possibly-truncated response, was_truncated, list of notes).
@@ -637,24 +672,24 @@ def truncate_oversized_response(
 
     # Phase 1: Truncate long string fields
     was_truncated |= _truncate_strings(data, notes)
-    if _is_under_limit(data, token_limit):
+    if _is_under_limit(data, token_limit, size_fn):
         return data, was_truncated, notes
 
     # Phase 2: Truncate large list fields
     was_truncated |= _truncate_lists(data, notes, max_list_items)
-    if _is_under_limit(data, token_limit):
+    if _is_under_limit(data, token_limit, size_fn):
         return data, was_truncated, notes
 
     # Phase 3: Recursively truncate strings inside nested structures
     # (e.g. charts[i].description, native_filters[i].config, etc.)
     was_truncated |= _truncate_strings_recursive(data, notes)
-    if _is_under_limit(data, token_limit):
+    if _is_under_limit(data, token_limit, size_fn):
         return data, was_truncated, notes
 
     # Phase 4: Aggressively reduce lists and summarize large dicts
     was_truncated |= _truncate_lists(data, notes, max_items=10)
     was_truncated |= _summarize_large_dicts(data, notes)
-    if _is_under_limit(data, token_limit):
+    if _is_under_limit(data, token_limit, size_fn):
         return data, was_truncated, notes
 
     # Phase 5: Nuclear — replace all collections with empty values
@@ -663,6 +698,338 @@ def truncate_oversized_response(
     return data, was_truncated, notes
 
 
+def _linked_statement_entries(data: Dict[str, Any]) -> list[Dict[str, Any]]:
+    """Find per-statement dicts with row data duplicated from an
+    ``execute_sql`` response's top-level ``rows`` field.
+
+    ``ExecuteSqlResponse.rows`` is a copy of the last data-bearing
+    statement's ``data.rows`` (see ``_convert_to_response`` in
+    ``sql_lab/tool/execute_sql.py``), and every data-bearing statement's
+    rows are additionally serialised under ``statements[*].data.rows``. If
+    only the top-level ``rows`` field is bisected, these nested copies keep
+    the full untruncated payload, which can leave the response oversized
+    even after "truncation". Returns the statement dicts themselves (not
+    just their row lists) so the caller can shrink ``data.rows`` in
+    lockstep with the top-level field *and* keep each statement's
+    ``row_count`` in sync, rather than leaving it pointing at the
+    pre-truncation count.
+    """
+    statements = data.get("statements")
+    if not isinstance(statements, list):
+        return []
+
+    linked = []
+    for statement in statements:
+        if not isinstance(statement, dict):
+            continue
+        statement_data = statement.get("data")
+        if isinstance(statement_data, dict) and isinstance(
+            statement_data.get("rows"), list
+        ):
+            linked.append(statement)
+    return linked
+
+
+def _bisect_row_limit(
+    data: Dict[str, Any],
+    row_field: str,
+    original_rows: List[Any],
+    token_limit: int,
+    size_fn: SizeFn = _default_size_fn,
+) -> 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).
+
+    When ``data`` is an ``execute_sql`` response, per-statement row lists
+    under ``statements[*].data.rows`` duplicate ``data[row_field]`` (see
+    ``_linked_statement_entries``) and are shrunk to the same length at
+    each step — along with each statement's ``row_count``, so it doesn't
+    keep reporting the pre-truncation count — so the size measured during
+    the search, and the response ultimately returned, reflects the
+    truncated nested data too, not just the top-level field.
+
+    ``size_fn`` measures each candidate; pass a function that re-wraps the
+    payload (e.g. back into a ToolResult) before measuring so the search
+    converges on a prefix that still fits the limit after re-wrapping.
+    """
+    linked_statements = _linked_statement_entries(data)
+    linked_rows = [list(stmt["data"]["rows"]) for stmt in linked_statements]
+
+    def _apply(count: int) -> None:
+        data[row_field] = original_rows[:count]
+        for original_linked, statement in zip(
+            linked_rows, linked_statements, strict=False
+        ):
+            kept_linked = original_linked[: min(count, len(original_linked))]
+            statement["data"]["rows"] = kept_linked
+            statement["row_count"] = len(kept_linked)
+
+    lo, hi = 0, len(original_rows)
+    while lo < hi:
+        mid = (lo + hi + 1) // 2
+        _apply(mid)
+        if size_fn(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
+
+    _apply(kept)
+    return kept
+
+
+def _bisect_csv_row_limit(
+    data: Dict[str, Any],
+    field: str,
+    original_value: str,
+    token_limit: int,
+    size_fn: SizeFn = _default_size_fn,
+) -> int:
+    """Binary-search for the largest whole-record prefix of a CSV string.
+
+    Cuts on CSV record boundaries (via the ``csv`` module, which
+    understands quoted fields and embedded delimiters/newlines) rather than
+    an arbitrary character offset, so the result is always parseable CSV.
+    The header row (record 0) is always kept when present, since
+    ``get_chart_data`` always writes one via ``csv.DictWriter.writeheader()``.
+
+    Mutates ``data[field]`` during the search and leaves it at the final
+    kept CSV text on return.
+
+    ``size_fn`` measures each candidate; pass a function that re-wraps the
+    payload before measuring so the search converges on a prefix that still
+    fits the limit after re-wrapping.
+
+    Returns:
+        The number of records kept, including the header row.
+    """
+    import csv
+    import io
+
+    records = list(csv.reader(io.StringIO(original_value)))
+    total = len(records)
+    if total <= 1:
+        # No data rows to trim (just a header, or unparseable) — nothing to
+        # cut on a record boundary; leave the field untouched.
+        return total
+
+    def _render(count: int) -> str:
+        buf = io.StringIO()
+        csv.writer(buf).writerows(records[:count])
+        return buf.getvalue()
+
+    lo, hi = 1, total  # always keep at least the header row
+    while lo < hi:
+        mid = (lo + hi + 1) // 2
+        data[field] = _render(mid)
+        if size_fn(data) <= token_limit:
+            lo = mid
+        else:
+            hi = mid - 1
+
+    data[field] = _render(lo)
+    return lo
+
+
+# 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,
+    size_fn: SizeFn = _default_size_fn,
+) -> 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, size_fn)
+
+    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,
+    size_fn: SizeFn = _default_size_fn,
+) -> list[str] | None:
+    """Try to trim 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``. Truncation cuts
+    on CSV record boundaries (see ``_bisect_csv_row_limit``) rather than an
+    arbitrary character offset, so the result is always parseable CSV; the
+    header row is always preserved. ``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.
+    """
+    import csv
+    import io
+
+    csv_data = data.get("csv_data")
+    if not isinstance(csv_data, str) or not csv_data:
+        return None
+
+    total_records = sum(1 for _ in csv.reader(io.StringIO(csv_data)))
+    original_data_rows = max(total_records - 1, 0)
+    if original_data_rows == 0:
+        # Just a header (or unparseable) — nothing to cut on a record
+        # boundary.
+        return None
+
+    # Same reservation trick as ``_truncate_rows_field``, keyed on the
+    # data-row count rather than the raw character count.
+    placeholder_note = (
+        f"CSV content truncated: {original_data_rows} of {original_data_rows} "
+        f"data rows returned (limit ~{token_limit:,} tokens). {advice}"
+    )
+    data["_response_truncated"] = True
+    data["_truncation_notes"] = [placeholder_note]
+
+    kept_records = _bisect_csv_row_limit(
+        data, "csv_data", csv_data, token_limit, size_fn
+    )
+    if (kept_data_rows := max(kept_records - 1, 0)) < original_data_rows:
+        notes = [
+            f"CSV content truncated: {kept_data_rows} of {original_data_rows} "
+            f"data rows returned (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,
+    size_fn: SizeFn = _default_size_fn,
+) -> 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.
+
+    For ``execute_sql``, per-statement row lists under
+    ``statements[*].data.rows`` (and their ``row_count``) are shrunk
+    alongside the top-level field (see ``_linked_statement_entries``) so a
+    duplicate, untruncated copy of the same rows can't leave the response
+    oversized, or its metadata inconsistent, after "truncation".
+
+    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``).
+        size_fn: Measures the token size of a candidate payload dict.
+            Defaults to measuring the bare dict; callers that re-wrap the
+            payload before returning it (e.g. back into a ToolResult) should
+            pass a function that measures the *wrapped* size, so the binary
+            search converges on a result that still fits after re-wrapping.
+
+    Returns:
+        A tuple of (possibly-truncated response, was_truncated, list of notes).
+    """
+    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, size_fn=size_fn)
+
+    if size_fn(data) <= token_limit:
+        return data, False, []
+
+    notes = _truncate_rows_field(data, row_field, token_limit, advice, size_fn)
+    if notes is None:
+        notes = _truncate_csv_data_field(data, token_limit, advice, size_fn)
+
+    return data, notes is not None, notes or []
+
+
 def format_size_limit_error(
     tool_name: str,
     params: Dict[str, Any] | None,

tests/unit_tests/mcp_service/test_middleware.py

@@ -380,6 +380,335 @@ 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_truncation_preserves_total_rows(self) -> None:
+        """``total_rows`` (the true pre-truncation count) must survive.
+
+        ``row_count`` is overwritten to reflect the truncated row list, but
+        ``total_rows`` is the signal callers rely on to detect that a
+        response is partial — it must keep describing the original,
+        untruncated dataset.
+        """
+        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,
+            "total_rows": 5000,
+            "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 result["_response_truncated"] is True
+        assert result["row_count"] == len(result["data"])
+        assert result["row_count"] < 200
+        assert result["total_rows"] == 5000
+
+    @pytest.mark.asyncio
+    async def test_get_chart_data_excel_export_raises_tool_error(self) -> None:
+        """Oversized excel exports must hard-fail, not ship a corrupt file.
+
+        ``excel_data`` is base64-encoded binary and is deliberately left
+        untouched by truncation (cutting it would produce an unreadable
+        file), and there are no other truncatable fields (``data=[]``), so
+        ``truncate_query_result`` reports ``was_truncated=False`` and the
+        response must fall through to the hard size-limit error.
+        """
+        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",
+            "data": [],
+            "row_count": 200,
+            "summary": "",
+            "excel_data": "QUJDREVGRw==" * 5000,
+            "format": "excel",
+        }
+        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_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 +1082,78 @@ 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_data_query_truncation_fits_after_rewrap(self) -> None:
+        """The re-wrapped ToolResult must itself fit under the token limit.
+
+        Regression test: truncation used to size candidates against the bare
+        unwrapped payload, then re-wrap into a ToolResult afterward. The
+        ``content``/``meta`` envelope added by re-wrapping can push a
+        candidate that "fit" back over the limit. The binary search must
+        size candidates against the final wrapped object it will return.
+        """
+        from superset.mcp_service.utils.token_utils import estimate_response_tokens
+
+        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 estimate_response_tokens(result) <= middleware.token_limit
+
     @pytest.mark.asyncio
     async def test_meta_preserved_after_truncation(self) -> None:
         """Should preserve the original ToolResult meta through truncation."""

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,290 @@ 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_csv_truncation_keeps_header_and_whole_records(self) -> None:
+        """Truncated CSV must stay parseable: header intact, no partial rows.
+
+        Regression test: naive character-offset truncation can cut a CSV
+        payload mid-record (e.g. inside a quoted field), producing invalid
+        CSV. Truncation must cut on record boundaries instead.
+        """
+        import csv
+        import io
+
+        header = "col_0,col_1,col_2"
+        # Include a quoted field containing a comma and an embedded newline,
+        # which would corrupt a naive character-offset cut.
+        rows = [
+            f'{i},"value, with comma","line1\nline2 for row {i}"' for i in range(500)
+        ]
+        csv_data = header + "\n" + "\n".join(rows) + "\n"
+        response: dict[str, Any] = {
+            "chart_id": 1,
+            "data": [],
+            "csv_data": csv_data,
+            "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 estimate_response_tokens(result) <= 500
+
+        parsed_rows = list(csv.reader(io.StringIO(result["csv_data"])))
+        assert parsed_rows, "truncated CSV must remain parseable"
+        assert parsed_rows[0] == header.split(",")
+        # Every parsed record must be a complete, well-formed row (3 fields),
+        # proving no record was cut mid-field.
+        assert all(len(row) == 3 for row in parsed_rows)
+        assert len(parsed_rows) - 1 < len(rows)
+        assert any("data rows" 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)
+
+    def test_execute_sql_statements_rows_truncated_alongside_top_level(self) -> None:
+        """Real ``ExecuteSqlResponse`` shapes duplicate rows under ``statements``.
+
+        Regression test: ``ExecuteSqlResponse.rows`` is a copy of the last
+        data-bearing statement's ``data.rows``, and the SAME rows are also
+        serialised under ``statements[0].data.rows``. If only the top-level
+        field is bisected, the nested copy stays at full size and the
+        response can remain oversized after "truncation".
+        """
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        rows = [row] * 200
+        response: dict[str, Any] = {
+            "success": True,
+            "rows": rows,
+            "row_count": 200,
+            "statements": [
+                {
+                    "original_sql": "SELECT * FROM t",
+                    "executed_sql": "SELECT * FROM t",
+                    "row_count": 200,
+                    "data": {
+                        "rows": rows,
+                        "columns": [
+                            {"name": f"col_{i}", "type": "TEXT"} for i in range(10)
+                        ],
+                    },
+                }
+            ],
+        }
+        result, was_truncated, notes = truncate_query_result(
+            response, 500, tool_name="execute_sql"
+        )
+        assert was_truncated is True
+        assert isinstance(result, dict)
+        # The overall payload — including the nested statements copy — must
+        # actually fit; this is what the middleware's post-truncation size
+        # check relies on to avoid falling back to a hard error.
+        assert estimate_response_tokens(result) <= 500
+        assert len(result["rows"]) < 200
+        assert len(result["statements"][0]["data"]["rows"]) == len(result["rows"])
+        assert notes
+
+    def test_execute_sql_statements_row_count_stays_in_sync(self) -> None:
+        """Each statement's ``row_count`` must track its trimmed ``data.rows``.
+
+        Regression test: bisection already shrinks the nested
+        ``statements[*].data.rows`` list in lockstep with the top-level
+        field (see the sibling ``..._rows_truncated_alongside_top_level``
+        test), but the per-statement ``row_count`` field is a separate
+        value copied at response-build time and must be updated too, or a
+        client would see a statement whose declared row count doesn't match
+        the row list actually returned.
+        """
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        rows = [row] * 200
+        response: dict[str, Any] = {
+            "success": True,
+            "rows": rows,
+            "row_count": 200,
+            "statements": [
+                {
+                    "original_sql": "SELECT * FROM t",
+                    "executed_sql": "SELECT * FROM t",
+                    "row_count": 200,
+                    "data": {
+                        "rows": rows,
+                        "columns": [
+                            {"name": f"col_{i}", "type": "TEXT"} for i in range(10)
+                        ],
+                    },
+                }
+            ],
+        }
+        result, was_truncated, _ = truncate_query_result(
+            response, 500, tool_name="execute_sql"
+        )
+        assert was_truncated is True
+        statement = result["statements"][0]
+        assert statement["row_count"] == len(statement["data"]["rows"])
+        assert statement["row_count"] < 200
+
+    def test_falls_back_to_generic_truncation_when_no_row_field(self) -> None:
+        """Responses without a recognised row/data field truncate generically.
+
+        A DML-final ``execute_sql`` response (or any data-query response
+        with no ``rows``/``data`` list) has no row field to bisect on, so
+        ``truncate_query_result`` delegates to the generic
+        ``truncate_oversized_response`` field-truncation phases instead of
+        raising or returning the payload untouched.
+        """
+        response: dict[str, Any] = {
+            "success": True,
+            "affected_rows": 5,
+            "multi_statement_warning": "x" * 5000,
+        }
+        result, was_truncated, notes = truncate_query_result(response, 200)
+        assert was_truncated is True
+        assert notes
+        assert len(result["multi_statement_warning"]) < 5000
+        assert result["affected_rows"] == 5
+
+    def test_execute_sql_multiple_statements_each_capped(self) -> None:
+        """Every statement's rows are capped, not just the last (top-level) one."""
+        row = {f"col_{i}": f"value_{i}" for i in range(10)}
+        first_stmt_rows = [row] * 150
+        last_stmt_rows = [row] * 150
+        response: dict[str, Any] = {
+            "success": True,
+            "rows": last_stmt_rows,
+            "row_count": 150,
+            "multi_statement_warning": "2 data-bearing statements",
+            "statements": [
+                {
+                    "original_sql": "SELECT * FROM a",
+                    "executed_sql": "SELECT * FROM a",
+                    "row_count": 150,
+                    "data": {
+                        "rows": first_stmt_rows,
+                        "columns": [
+                            {"name": f"col_{i}", "type": "TEXT"} for i in range(10)
+                        ],
+                    },
+                },
+                {
+                    "original_sql": "SELECT * FROM b",
+                    "executed_sql": "SELECT * FROM b",
+                    "row_count": 150,
+                    "data": {
+                        "rows": last_stmt_rows,
+                        "columns": [
+                            {"name": f"col_{i}", "type": "TEXT"} for i in range(10)
+                        ],
+                    },
+                },
+            ],
+        }
+        result, was_truncated, notes = truncate_query_result(
+            response, 2000, tool_name="execute_sql"
+        )
+        assert was_truncated is True
+        assert isinstance(result, dict)
+        assert estimate_response_tokens(result) <= 2000
+        assert len(result["statements"][0]["data"]["rows"]) < 150
+        assert len(result["statements"][1]["data"]["rows"]) < 150
+        assert notes