fix(mcp): truncate query-tool responses instead of hard-failing#42244
fix(mcp): truncate query-tool responses instead of hard-failing#42244aminghadersohi wants to merge 2 commits into
Conversation
execute_sql, query_dataset, and get_chart_data raised a hard "Response too large" error whenever their result set exceeded the MCP response token limit, even though the majority of oversized responses are just a few rows too many. Extend ResponseSizeGuardMiddleware to truncate these three tools the same way info tools are already truncated: - Binary-search the largest row-count prefix that keeps the serialized response under the token limit, and return that instead of erroring. - Also handle CSV chart-data exports, where the payload lives in the scalar csv_data field (data=[]) rather than a row list. 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 now 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 SQL.
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #42244 +/- ##
==========================================
+ Coverage 56.01% 65.16% +9.15%
==========================================
Files 2768 2768
Lines 156280 156225 -55
Branches 35774 35743 -31
==========================================
+ Hits 87534 101806 +14272
+ Misses 67917 52453 -15464
- Partials 829 1966 +1137
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
|
The flagged issue is correct. In 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 |
![bito-code-review[bot]](https://avatars.githubusercontent.com/in/1061978?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
AI Code Review powered by
| if extracted is not None and isinstance(truncated, dict): | ||
| return self._rewrap_as_tool_result(truncated, response) | ||
|
|
||
| return truncated |
There was a problem hiding this comment.
The _try_truncate_data_query_response method is missing the _response_truncated and _truncation_notes metadata markers that _try_truncate_info_response sets at line 1326-1328. Callers of data-query tools (execute_sql, query_dataset, get_chart_data) will not know truncation occurred, making debugging harder and output inconsistent with info-tool responses. (CWE-252)
Code Review Run #f8f8ed
Should Bito avoid suggestions like this for future reviews? (Manage Rules)
- Yes, avoid them
2 participants
SUMMARY
execute_sql,query_dataset, andget_chart_dataraised 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 extendsResponseSizeGuardMiddlewareto truncate these three tools the same way info tools (get_chart_info,get_dataset_info, etc.) are already truncated, instead of hard-failing:csv_datafield (data=[]) rather than a row list, by truncating that field directly.excel_datais base64-encoded binary and is left untouched, since truncating it would produce a corrupt file.limit/row_limitparameter vs. SQLLIMITclause) instead of always suggesting a SQLLIMIT.BEFORE/AFTER
Before:
execute_sql/query_dataset/get_chart_dataresponses that exceeded the token limit always raisedToolError: 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_notesset, 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/mypyon the changed filesADDITIONAL INFORMATION
docs/updates