fix(mcp): truncate query-tool responses instead of hard-failing

[aminghadersohi]

SUMMARY

execute_sql, query_dataset, and get_chart_data raised a hard "Response too large" error whenever their result exceeded the MCP response token limit, even when the response was only a few rows over. This extends ResponseSizeGuardMiddleware to truncate these three tools the same way info tools (get_chart_info, get_dataset_info, etc.) are already truncated, instead of hard-failing:

• Binary-search the largest row-count prefix that keeps the serialized response under the token limit, returning the truncated result instead of erroring.
• Handle CSV chart-data exports, where the payload lives in the scalar csv_data field (data=[]) rather than a row list, by truncating that field directly. excel_data is base64-encoded binary and is left untouched, since truncating it would produce a corrupt file.
• Re-check the truncated response's size before returning it (mirroring the existing info-tool path), so a single row/CSV blob that alone exceeds the limit falls back to the informative size-limit error instead of silently shipping an over-budget response. The bisection reserves headroom for the truncation-note metadata up front so this fallback only triggers when truncation genuinely can't fit.
• Tailor the truncation-note wording per tool (limit/row_limit parameter vs. SQL LIMIT clause) instead of always suggesting a SQL LIMIT.

BEFORE/AFTER

Before: execute_sql/query_dataset/get_chart_data responses that exceeded the token limit always raised ToolError: Response too large, even if truncating a handful of rows would have brought them under the limit.

After: these tools truncate the row set (or CSV content) and return a partial result with _response_truncated / _truncation_notes set, falling back to the hard error only when truncation genuinely cannot bring the response under the limit.

TESTING INSTRUCTIONS

• pytest tests/unit_tests/mcp_service/ (3086 passed)
• ruff check / ruff format --diff / mypy on the changed files

ADDITIONAL INFORMATION

• Has associated tests
• Confirm this PR does not include user-facing docs strings changes requiring docs/ updates

[netlify]

✅ Deploy Preview for superset-docs-preview ready!

Name	Link
🔨 Latest commit	ce8e066
🔍 Latest deploy log	https://app.netlify.com/projects/superset-docs-preview/deploys/6a5f9298f0829c00089c0e52
😎 Deploy Preview	https://deploy-preview-42244--superset-docs-preview.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.
🤖 Make changes	Run an agent on this branch

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 126 lines in your changes missing coverage. Please review.
✅ Project coverage is 65.20%. Comparing base (f751716) to head (ce8e066).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
superset/mcp_service/utils/token_utils.py	0.00%	87 Missing ⚠️
superset/mcp_service/middleware.py	0.00%	39 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #42244      +/-   ##
==========================================
- Coverage   65.27%   65.20%   -0.08%     
==========================================
  Files        2769     2768       -1     
  Lines      156343   156431      +88     
  Branches    35788    35801      +13     
==========================================
- Hits       102054   101994      -60     
- Misses      52322    52469     +147     
- Partials     1967     1968       +1     

Flag	Coverage Δ
hive	38.61% <0.00%> (-0.06%)	⬇️
mysql	57.84% <0.00%> (-0.09%)	⬇️
postgres	57.89% <0.00%> (-0.09%)	⬇️
presto	40.54% <0.00%> (-0.07%)	⬇️
python	59.29% <0.00%> (-0.10%)	⬇️
sqlite	57.50% <0.00%> (-0.09%)	⬇️
unit	100.00% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[bito-code-review]

The flagged issue is correct. In _try_truncate_data_query_response, the token estimation is performed on the truncated payload alone, ignoring the overhead of the ToolResult wrapper that is added back before returning. This can result in a returned object that still exceeds the token_limit.

To resolve this, you should re-estimate the tokens on the final object after rewrapping, or account for the wrapper overhead during the check. Here is the corrected logic for the re-check:

        # ... (after truncation)
        if not was_truncated:
            return None

        # Re-wrap to check total size including wrapper overhead
        final_result = self._rewrap_as_tool_result(truncated, response) if extracted is not None and isinstance(truncated, dict) else truncated
        truncated_tokens = estimate_response_tokens(final_result)
        
        if truncated_tokens > self.token_limit:
            return None
        
        # ... (proceed with logging and return final_result)

Would you like me to check the other comments on this PR and implement fixes for them as well?

superset/mcp_service/middleware.py

# Re-wrap to check total size including wrapper overhead
        final_result = self._rewrap_as_tool_result(truncated, response) if extracted is not None and isinstance(truncated, dict) else truncated
        truncated_tokens = estimate_response_tokens(final_result)
        
        if truncated_tokens > self.token_limit:
            return None

[bito-code-review]

Code Review Agent Run #f8f8ed

Actionable Suggestions - 1

• superset/mcp_service/middleware.py - 1

  • CWE-252: Missing truncation metadata · Line 1398-1401

Review Details

• Files reviewed - 4 · Commit Range: 9d4e317..2d6ea36

  • superset/mcp_service/middleware.py
  • superset/mcp_service/utils/token_utils.py
  • tests/unit_tests/mcp_service/test_middleware.py
  • tests/unit_tests/mcp_service/utils/test_token_utils.py
• Files skipped - 0
• Tools

  • MyPy (Static Code Analysis) - ✔︎ Successful
  • Astral Ruff (Static Code Analysis) - ✔︎ Successful
  • Whispers (Secret Scanner) - ✔︎ Successful
  • Detect-secrets (Secret Scanner) - ✔︎ Successful

---

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

• /review - Manually triggers a full AI review.

• /pause - Pauses automatic reviews on this pull request.

• /resume - Resumes automatic reviews.

• /resolve - Marks all Bito-posted review comments as resolved.

• /abort - Cancels all in-progress reviews.

Refer to the documentation for additional commands.

Configuration

This repository uses Superset You can customize the agent settings here or contact your Bito workspace admin at evan@preset.io.

Documentation & Help

• Customize agent settings
• Review rules
• General documentation
• FAQ

AI Code Review powered by

[rebenitez1802]

Request changes: sound approach and the size-guard/fail-safe design is solid, but for the single-SELECT execute_sql case (the most common one) truncation is defeated by a duplicated row list, so the headline fix mostly doesn't fire there.

🔴 High — execute_sql row list is serialized twice, so row truncation can't shrink it
truncate_query_result picks row_field="rows" and _bisect_row_limit only mutates data["rows"], but in execute_sql._convert_to_response the top-level rows = last_data_stmt.data.rows is the same data also serialized under statements[*].data.rows (superset/mcp_service/sql_lab/tool/execute_sql.py:359-361). The bisection measures json.dumps(data), which still contains the full statements copy. Result for an oversized single SELECT: either the untouched statements copy alone still exceeds the limit → middleware re-check fails → hard ToolError (feature never engages), or it just fits → a partial payload ships where rows/row_count say "1 of N" while statements[0].data.rows + statements[0].row_count still hold the full set (metadata lies). Fix: when row_field == "rows" and a statements array is present, trim statements[*].data.rows in lockstep and update each statements[*].row_count — or bisect on the last data-bearing statement and derive the top-level slice from it.

🟡 Medium — row_field is None fallback truncates silently (no _response_truncated/_truncation_notes)
In truncate_query_result (superset/mcp_service/utils/token_utils.py:896-898) the no-row-field branch delegates to truncate_oversized_response, which never sets the truncation markers — the info path relies on the middleware to add them (middleware.py:1326-1329), but _try_truncate_data_query_response does not (middleware.py:1398). Scenario: a DML-final multi-statement execute_sql where top-level rows is None but statements[] is large → the client gets a truncated statements array with no marker/note and treats partial data as complete. Fix: mirror the info path — set _response_truncated/_truncation_notes on the dict after a successful truncation in _try_truncate_data_query_response (or inside the fallback branch of truncate_query_result).

🟡 Medium — Test gaps around the partial-data signal and the excel/fallback safety nets
Three untested paths could each hide a real regression: (1) no fixture carries total_rows, yet the docstring says callers "rely on the true total_rows" to detect partial data — nothing asserts it's preserved while row_count is overwritten; (2) the row_field is None → truncate_oversized_response fallback is never exercised; (3) the excel-export → hard ToolError guarantee is only unit-tested, never through on_call_tool. Fix: add a middleware test with row_count=200, total_rows=5000 asserting total_rows survives; a truncate_query_result test with only a huge scalar field; and an on_call_tool get_chart_data excel test asserting pytest.raises(ToolError).

🟢 Low — CSV truncation cuts mid-row
_bisect_string_length on csv_data (superset/mcp_service/utils/token_utils.py:835) slices at an arbitrary character offset, so the last kept line is a partial/unterminated record that a strict parser can reject. excel_data is correctly left alone (verified). Fix: after computing kept_len, back off to the last \n at or before it (keeping the header) so only whole rows remain.

🟢 Low — Broad except Exception masks truncation regressions
_try_truncate_data_query_response catches bare Exception (middleware.py:1349-1360) where the info path catches only (MemoryError, RecursionError). It's fail-safe (any error → hard ToolError, never an over-budget response), but a real KeyError/TypeError in truncate_query_result would be silently swallowed and every oversized query would degrade to a hard error at WARNING level, hiding the bug. Fix: narrow the catch or log at logger.exception.

---

Cleared (checked, no issue): binary searches converge correctly with no off-by-one/infinite-loop; the fail-open path is closed (post-truncation re-check falls back to hard error); placeholder-headroom reservation is sound (real note ≤ placeholder); estimated_tokens is measured on the same extracted payload that gets re-wrapped; DATA_QUERY_TOOLS names match the registered tools and are disjoint from INFO_TOOLS; no schema-validation bypass; mutation operates on fresh copies. Full suite passes locally (153 passed, with real tiktoken).

The High item is the one to resolve before merge — without it, the PR's core promise doesn't hold for the most common execute_sql path.

_{🤖 Team review via Claude Code}

[aminghadersohi]

Thanks for the thorough review — went through each point:

🔴 High — execute_sql row list serialized twice. Verified this empirically before touching anything: _linked_statement_row_lists (added in the prior ea1e5e3 fix, for the wrapper-overhead bug) already trims statements[*].data.rows in lockstep with the top-level rows field during bisection, so the "headline fix doesn't fire" concern doesn't hold — a simulated oversized single-SELECT execute_sql response truncates correctly and the final wrapped size comes in under the limit.

What was real: statements[*].row_count (a separate field copied at response-build time) wasn't updated when data.rows got trimmed — so a truncated response could show data.rows with 8 entries while row_count still said 2000. Fixed in 00c5e79: renamed _linked_statement_row_lists → _linked_statement_entries to track the statement dicts themselves (not just their row lists), and row_count is now set alongside data.rows on every bisection step. Added test_execute_sql_statements_row_count_stays_in_sync.

🟡 Medium — missing truncation markers on the row_field is None fallback. Already fixed in 0d9bc32 (pushed before this review) — _try_truncate_data_query_response now sets _response_truncated/_truncation_notes explicitly, mirroring the info-tool path, so this covers the fallback branch too (the row/CSV path already set them internally). Added test_falls_back_to_generic_truncation_when_no_row_field to cover that branch directly.

🟡 Medium — test gaps. Added:

• test_data_query_truncation_preserves_total_rows — row_count=200, total_rows=5000, asserts total_rows survives while row_count is overwritten to the truncated count.
• test_falls_back_to_generic_truncation_when_no_row_field — exercises the truncate_oversized_response fallback directly.
• test_get_chart_data_excel_export_raises_tool_error — oversized excel export through on_call_tool, asserts pytest.raises(ToolError).

🟢 Low — CSV truncation cuts mid-row. Checked this one and don't think it's live: the function actually responsible for csv_data (_bisect_csv_row_limit) already cuts on CSV record boundaries via the csv module (respects quoted fields/embedded newlines), not an arbitrary character offset — there's an existing regression test for exactly this (test_csv_truncation_keeps_header_and_whole_records) with a quoted-comma-and-newline field that would fail on a naive character cut. _bisect_string_length doesn't exist in this file; I couldn't find the referenced line, so I suspect this was a mix-up with the generic _truncate_strings phase (which does do a raw character cut, but only runs when there's no recognised row/data field at all — and every current data-query tool that emits csv_data also emits data=[], which routes it through the record-boundary-safe path instead). Let me know if you had a specific reachable path in mind that I'm missing.

🟢 Low — broad except Exception. Kept the broad catch (it's intentionally fail-safe — narrowing it wouldn't change behavior since any exception already falls through to the hard ToolError), but switched logger.warning → logger.exception in 00c5e79 so a genuine bug gets a full traceback instead of blending in with routine size-limit enforcement.

Rebased onto master (no conflicts) and pushed. Full commit list: 72df1f7 (original fix) → 984a491 (wrapper-overhead re-check) → 0d9bc32 (fallback truncation markers) → 00c5e79 (this round).

[rebenitez1802]

Looks like this just need a rebase, LGTM