Repository navigation
fix(mcp): never block a committed update_chart write as ToolError - #44386
aminghadersohi merged 12 commits into
Conversation
update_chart persists its write before the response-size guard checks the serialized result, so an oversized response was reported to MCP clients as a failed call even though the chart had already been updated -- a common outcome that let retrying agents replay the same mutation. Make committed-write tools (update_chart) truncation-eligible like the existing INFO_TOOLS, protecting the tool's identifying field so a write confirmation always survives even the most aggressive truncation phase, and add a guaranteed-small fallback response for the (practically unreachable) case where truncation still can't fit -- this path never raises ToolError, since the mutation already committed. Also generalize the CSV-export bisection helper to any named string field and use it for get_chart_sql, whose oversized responses are dominated by one large SQL string with no row/limit parameter to shrink it.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #44386 +/- ##
==========================================
+ Coverage 80.57% 80.62% +0.05%
==========================================
Files 2940 2942 +2
Lines 175316 175780 +464
Branches 40696 40795 +99
==========================================
+ Hits 141253 141730 +477
+ Misses 31398 31388 -10
+ Partials 2665 2662 -3
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:
|
…nable
Follow-up hardening on the response-size guard:
- Phase 4's dict summarizer ignored protected_keys, so a ChartInfo payload
(more than the 20-key threshold) was replaced with a {_truncated,_message}
marker before Phase 5 could protect it -- destroying the write confirmation
in exactly the case the guard exists for.
- The minimal committed-write fallback returned a bare dict when a ToolResult's
payload could not be parsed, which fails in FastMCP's to_mcp_result() and
surfaces a completed write as an internal error anyway. Rewrap on
isinstance(response, ToolResult) instead of on successful extraction.
- That fallback also claimed the write "already committed successfully", but
update_chart defaults to generate_preview=True and persists nothing. Reworded
to state only what is true: the call was not rolled back by the size limit.
- The fallback copied chart from the untruncated payload and never re-checked
size, so the "guaranteed-small" response could itself exceed the limit.
It now degrades chart to identifying scalars.
- A bisected SQL prefix stayed syntactically valid, so a statement cut before
its WHERE/LIMIT clause could be executed unfiltered and unbounded. The kept
prefix now carries an in-band marker, reserved inside the token budget.
Second review pass on the response-size guard: - The truncation marker was a `--` line comment, so a cut statement stayed valid and executable -- exactly the full-table-scan hazard it was meant to prevent. Use an unterminated /* block comment, which is a parse error in every dialect. - The three _try_truncate_* helpers gated rewrapping on successful payload extraction, so a ToolResult whose text is not a JSON object had its own wrapper fields model_dump()ed and returned as a plain dict, failing in FastMCP's to_mcp_result(). They now decline to truncate an opaque ToolResult and let the caller fall through to its own fallback. - test_untruncated_sql_gets_no_marker returned at the under-limit guard and never reached the branch it was named for. Split into a pass-through test and a direct _bisect_string_length test that does cover it.
The unterminated /* block comment was not the universal parse error it claimed to be: SQLite treats a block comment as terminated at end of input and runs the statement, so a get_chart_sql response truncated before its WHERE clause was still executable against it -- the exact hazard the marker exists to close. Verified against sqlite3 directly. Use an unterminated string literal instead, which is a tokenizer error in every dialect including SQLite, and add a test that asserts the truncated SQL fails to parse across sqlite/mysql/postgres/trino/bigquery.
The bisect point is arbitrary, so it can land inside a comment or a string literal, and the previous marker only handled a cut in normal state. A cut landing inside a block comment swallowed the whole marker and SQLite ran the truncated statement -- verified against sqlite3 directly. The marker now ends any open line comment, closes an open block comment, and opens an unterminated string literal. Tested against a real sqlite3 engine (not just sqlglot, which rejects markers SQLite happily executes) for cuts landing in normal, line-comment, block-comment and string-literal state, across sqlite/mysql/postgres/trino/bigquery. Documented the marker as defense in depth: _response_truncated and _truncation_notes remain the authoritative signal that the SQL is partial.
The rationale comment said the marker "opens an unterminated string literal" in every cut state. That is wrong for a cut landing mid-string: there the quote terminates the in-progress literal and the */ is swallowed into it, so rejection rests entirely on the trailing bare identifiers. Trimming the marker to what the comment described would silently reintroduce the bug. Also widen the test's dialect list to cover duckdb, snowflake and tsql, which the comment previously claimed without guarding.
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
_shrink_minimal_response measured the response once, before reducing chart to its identifying fields, and never again. That reduction only reaches chart: error and explore_url are copied verbatim from the untruncated payload, and the identifying slice_name/url are themselves free-form strings. Any one of them could leave the 'minimal' response far over budget -- a 40k-character error field produced a ~10k-token confirmation against a 200-token limit. Clip every free-form string to a fixed cap so the confirmation is bounded by construction, and warn if even that still exceeds an extremely small limit rather than silently returning it. Also keep chart.is_unsaved_state among the retained identifying fields. update_chart defaults to generate_preview=True, which caches a preview and persists nothing, so that flag is the caller's only in-band way to tell a preview from a persisted write -- shrinking should not be what removes it.
The minimal write confirmation clipped `error` as if it were a string, but `update_chart` returns `GenerateChartResponse`, whose `error` is a `ChartGenerationError` -- a nested model that reaches the size guard as a dict once the ToolResult payload is parsed. `_clip_string` returns dicts unchanged, so the only shape the tools emit was never bounded: an error carrying a large `query_info`, `details` or `validation_errors` still pushed the "minimal" confirmation over the limit, which the commit's own post-reduction re-check duly logged. Reduce a structured error to clipped identifying scalars the same way `chart` is reduced, dropping the unbounded containers, and correct the docstring's account of which fields are bounded and why.
![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 #26bf41
Actionable Suggestions - 3
-
superset/mcp_service/utils/token_utils.py - 1
- Committed-write set incomplete · Line 492-504
-
tests/unit_tests/mcp_service/utils/test_token_utils.py - 2
- Missing test docstring · Line 459-459
- Missing test docstring · Line 468-468
Additional Suggestions - 6
-
superset/mcp_service/middleware.py - 3
-
Misleading error-fields comment · Line 84-87This comment asserts everything not kept on `MCPBaseError`/subclasses is an unbounded container, but `timestamp` (fixed-size datetime) and `help_url` (short string) are also dropped and are neither - verified against `error_schemas.py`. A maintainer auditing the field list against the schema hits this contradiction; reword so the drop rule for non-container fields is stated too.
-
Duplicated opaque-payload guard · Line 1240-1249The opaque-ToolResult guard added at 1240-1249 is duplicated verbatim in `_try_truncate_data_query_response` (1334-1343) and `_try_truncate_string_field_response` (1417-1426). Three identical copies of decline logic (extract, isinstance check, warning, return None) will drift independently — a future fix to one copy resurfaces the FastMCP `to_mcp_result()` failure on the missed path. Please hoist into one helper used by all three.
-
Tool-schema coupling in middleware · Line 1512-1521`_minimal_committed_write_response` hardcodes update_chart's payload shape (chart/error/explore_url/schema_version/api_version) inside the generic ResponseSizeGuardMiddleware, while COMMITTED_WRITE_TOOLS (token_utils.py:500) is an extensible registry. A second committed-write tool with a different schema would get a fabricated confirmation: chart=None, success defaulting to True. Worth making the kept fields tool-driven before the set grows.
-
-
superset/mcp_service/utils/token_utils.py - 3
-
Duplicates _is_under_limit helper · Line 1114-1114`_is_under_limit` (token_utils.py:692-696) already wraps exactly `estimate_token_count(utils_json.dumps(data)) <= token_limit` and is used by `truncate_oversized_response`; this new function re-inlines the expression. Reuse the helper so the budget check lives in one place; the local `utils_json` import at line 1105 then becomes unused and can be removed.
-
Response-normalization block duplicated · Line 1107-1112Lines 1107-1112 repeat the `model_dump()`/`dict`/else normalization already in `truncate_query_result` (lines 1052-1057) and `truncate_oversized_response` (line 734). A shared `_response_to_data()` helper would keep ToolResponse-handling rules in one place; the wrapper-handling subtleties documented at middleware.py:1417-1426 show how easily these copies drift.
-
Marker registry bypassed · Line 1077-1079The `sql` path in `truncate_string_field_response` sources its in-band marker from `_STRING_FIELD_TRUNCATION_MARKERS`, but this `csv_data` fallback hardcodes no marker. Behavior is identical today, yet a future `csv_data` entry in the registry would be silently ignored here. Consider consulting the registry for consistency: `suffix_marker=_STRING_FIELD_TRUNCATION_MARKERS.get("csv_data", "")`.
-
Review Details
-
Files reviewed - 4 · Commit Range:
14ead58..92a3ad3- 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 an incremental AI Review. -
/review full- 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
… tool
COMMITTED_WRITE_TOOLS stated a category invariant -- mutating tools whose
transaction commits before the size guard inspects the response -- but held
only update_chart. update_dashboard commits at update_dashboard.py:360 and
was absent, so an oversized UpdateDashboardResponse still hard-errored after
the dashboard was written, and a retrying client would replay the mutation.
Membership alone was not enough: _handle_oversized_response hardcoded
frozenset({"chart"}) as the protected field, and the minimal confirmation
copied a fixed chart-shaped key list. Both are now per-tool, driven by
COMMITTED_WRITE_SPECS: each entry names the resource noun, the container
fields that carry the write confirmation, and whether the response model
actually declares `success`. Responses whose identity is top-level scalars
(delete_chart's deleted_id, create_dataset's id) need no protected field --
no truncation phase drops a top-level scalar.
The minimal confirmation now selects from the payload's own keys instead of
a fixed list, so a dashboard response no longer gains a `chart`,
`explore_url` or `success` its schema never declares.
Tests assert the table covers every mutating tool except the three that
commit nothing to the metadata database, and check every spec against the
registered output schemas so the table cannot drift from the models.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
Resolves a semantic conflict the drift guard caught in CI: master added the mutating tool `apply_dashboard_filters`, which commits a dashboard permalink before the size guard inspects the response, so it belongs in COMMITTED_WRITE_SPECS. Its identity (dashboard_id, dashboard_url, permalink_key) is top-level scalars, which no truncation phase drops, so it needs no protected container field. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Code Review Agent Run #4fe4d2
Actionable Suggestions - 1
-
superset/mcp_service/middleware.py - 1
- Identity fields miss restore keys · Line 100-111
Additional Suggestions - 2
-
tests/unit_tests/mcp_service/test_mcp_tool_registration.py - 1
-
Ambiguous if/elif condition binding · Line 686-690The if/elif chain mixes condition styles: line 686 binds the walrus and tests the resulting set's truthiness, while line 690 parenthesizes the test and applies `is not`, which reads like an identity check against a bound name. Naming both conditions (`missing`, `success_in_response`) makes each branch's test unambiguous. Behavior is correct today; this is clarity-only.
-
-
superset/mcp_service/middleware.py - 1
-
changed_fields list dropped · Line 1594-1600`_select_confirmation_fields` keeps identifying containers plus top-level scalars, but `changed_fields` on `UpdateDashboardResponse` is `List[str]` (dashboard/schemas.py:996), so the minimal path drops it while the normal path returns it. The docstring's claim that `changed_fields` is kept does not hold. Consider retaining short string lists so write confirmations stay consistent.
-
Filtered by Review Rules
Bito filtered these suggestions based on rules created automatically for your feedback. Manage rules.
-
superset/mcp_service/utils/token_utils.py - 1
- Misleading docstring test reference · Line 514-515
Review Details
-
Files reviewed - 5 · Commit Range:
92a3ad3..c15c11f- superset/mcp_service/middleware.py
- superset/mcp_service/utils/token_utils.py
- tests/unit_tests/mcp_service/test_mcp_tool_registration.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 an incremental AI Review. -
/review full- 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
rebenitez1802
left a comment
rebenitez1802
left a comment
There was a problem hiding this comment.
Approve: careful, well-iterated fix — the write-already-committed problem is correctly modeled, the chart field is protected through both destructive truncation phases, the minimal fallback is bounded by construction, and the SQL truncation marker is genuinely robust (and correctly framed as defense-in-depth behind _response_truncated). No blocking issues. Two test-quality nits below, neither gating.
🟢 Low — Minimal-fallback tests pass by swallowing StopIteration, not by the mocked signal
The five tests at tests/unit_tests/mcp_service/test_middleware.py:761,806,868,921,997 patch estimate_response_tokens with side_effect=[600, 600], but the committed-write fallback path calls it up to four times (on_call_tool → _try_truncate_info_response post-check → two _fits() calls in _shrink_minimal_response). The 3rd/4th calls exhaust the list and raise StopIteration, which _fits's except Exception (middleware.py:1822) catches and treats as "does not fit" — so the shrink is driven by exception-swallowing, and the inline comment "Force every truncation phase to report still over budget" doesn't describe those calls. The assertions themselves are sound (the boundedness checks use the unpatched estimate_token_count), so behavior is verified correctly — but the setup is brittle: narrowing _fits to except (TypeError, ValueError) (a plausible cleanup, since a broad except is discouraged here) would make all five tests error out with StopIteration despite correct middleware. Fix: replace side_effect=[600, 600] with return_value=600 so every call returns a real over-budget value and the tests don't depend on _fits catching StopIteration; update the comment accordingly.
🟢 Low — Two of the three new opaque-ToolResult guard branches are untested
The PR adds the same if extracted is None and isinstance(response, ToolResult): return None guard to _try_truncate_info_response, _try_truncate_data_query_response (middleware.py:1334), and _try_truncate_string_field_response (middleware.py:1417). Only the info-tool branch is exercised: test_opaque_tool_result_is_blocked_not_returned_as_dict (get_dashboard_info) and test_committed_write_fallback_rewraps_unparseable_tool_result (update_chart) both route through _try_truncate_info_response. The get_chart_sql/string-field test and the data-query tests all pass parseable payloads, so the guards at :1334 and :1417 never execute — a regression that reintroduced model_dump() of the wrapper there would ship undetected. Fix: add one test feeding an unparseable ToolResult through a STRING_FIELD_TRUNCATION_TOOLS tool (get_chart_sql) and one through a DATA_QUERY_TOOLS tool, asserting a ToolError is raised (mirroring the existing info-tool test).
rebenitez1802
left a comment
rebenitez1802
left a comment
There was a problem hiding this comment.
Approve — the response-size-guard rework is correct and well-tested. The committed-write invariant (transaction commits before the guard inspects the response) holds for all 21 COMMITTED_WRITE_SPECS entries, verified against the actual command/tool code; the minimal fallback is provably bounded whichever field was oversized; the SQL truncation marker holds up across cut-states and dialects; and the truncated response bypasses output-schema validation (via the CallToolResult/meta path), so it can't be re-rejected. No security boundary is affected — the guard runs post-RBAC and post-execution and only clips data already bound for the same caller.
Two non-blocking nits:
🟢 Minimal-response tests drive the shrink via a swallowed StopIteration, not the mocked estimate. In the four plain-dict committed-write tests, estimate_response_tokens is patched with side_effect=[600, 600], but both values are consumed by on_call_tool and the info-truncation re-check before _shrink_minimal_response runs. Its _fits calls then hit the exhausted iterator and raise StopIteration, which _fits's except Exception: return False swallows — so the shrink runs because the iterator ran out, not because the mock reported over-budget, and the _fits early-return-True branch is never covered. Not a production defect (the code is correct, and each test still independently asserts estimate_token_count(...) <= 500), but brittle: narrowing that except would break the tests with no behavior change. Consider a callable side_effect that returns 600 for any call, plus one case where the minimal payload already fits.
🟢 success: True is synthesized for an unparseable payload. _select_confirmation_fields does minimal.setdefault("success", True) when reports_success, so an unparseable payload ({}) yields success: True. Essentially unreachable in practice (these tools always serialize to a JSON object), and harmless even if reached since the write did commit under the caller's authorization — worth just a one-line comment noting the assumption.
|
@rebenitez1802 Thanks for both approval reviews! Addressed the three nits in 5763594: over-budget mocks now return 600 on every call, with corrected comments and an already-fits preservation test; opaque ToolResult coverage now includes get_chart_sql and execute_sql, asserting ToolError without dumping the wrapper; and a one-line comment documents the success fallback assumption. All 239 affected tests and branch-scoped pre-commit checks pass. No review re-requested. |
Code Review Agent Run #eb75dbActionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
|
@rebenitez1802 Following up on both review bodies (21:19 review, 21:33 review): the swallowed- The non-exhausting over-budget mocks, already-fits preservation case, opaque Deliberate-break evidence (all temporary mutations restored):
Both review bodies are addressed; GitHub has no body-level resolve operation or associated unresolved thread, so this linked reply records closure. I have not re-requested review or dismissed either approval. Acceptance criteria remain disclosed: AC 3 is PARTIAL (oversized-response truncation is handled, but |
|
Reply to Bito's 2026-09-17 review body, “Marker registry bypassed”: valid consistency finding, addressed in 83f1ebc.
Affected tests: 251 passed. Staged/branch-scoped pre-commit and explicit mypy passed. This is body-only feedback with no associated unresolved thread; this linked reply records it as addressed. |
|
Reply to Bito's 2026-09-18 review body, “changed_fields list dropped”: valid, addressed in 83f1ebc.
Affected tests: 251 passed. Staged/branch-scoped pre-commit and explicit mypy passed. This is body-only feedback with no associated unresolved thread; this linked reply records it as addressed. The separate restore/delete identity finding remains disproved and closed; it is not being reopened. AC 3 remains PARTIAL and AC 6 remains post-deploy verification. |
Code Review Agent Run #47bd8aActionable Suggestions - 0Additional Suggestions - 4
Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
|
@aminghadersohi Good fix, and the regression test earns its place. I checked out Not a blocker, but the invariant this PR establishes does not hold when
|
SUMMARY
update_chartpersists its chart write and only afterward does the MCPresponse-size guard estimate the serialized result. When that response
exceeds the configured token limit, the guard raised
ToolError, so asuccessful chart update was reported to the calling agent as a failed tool
call -- a common outcome for charts with a non-trivial
form_data. Agentsroutinely retry on tool failure, replaying a call that already went
through.
The guard's truncation was previously limited to two closed sets:
INFO_TOOLS(read-only single-object tools) andDATA_QUERY_TOOLS(row-based results). Mutating tools were excluded from both, so any
response over budget was hard-blocked regardless of whether the underlying
write had already happened.
This PR:
COMMITTED_WRITE_TOOLSset (currentlyupdate_chart) that istruncated the same way as
INFO_TOOLS, but with the tool's identifyingfield (
chart) protected from the two destructive truncation phases(large-dict summarization and the final "clear everything" phase), so a
caller can always confirm what was written.
COMMITTED_WRITE_TOOLS,used only if truncation still can't fit the response under budget. Unlike
every other tool category, this path never raises
ToolError. It alsodegrades the identifying field itself (to just
id/uuid/slice_name/url) if that field alone is what's oversized, and correctly re-wraps aToolResultwhose payload couldn't be parsed as JSON instead of handingFastMCP a bare dict it can't encode.
named-string-field truncator, and uses it for
get_chart_sql, whoseoversized responses are dominated by one large rendered-SQL string with
no
limit/row_limitparameter that could shrink it. Previously suchresponses were hard-blocked with an unactionable
Reduction needed: ~0%message; now the SQL is bisected to the largest prefix that fits.
executable SQL -- e.g. a statement cut before its
WHERE/LIMITclausewould silently scan far more data than the original query. The kept
prefix now gets an in-band marker designed to make the truncated
statement a syntax error regardless of where the cut landed (mid-comment,
mid-string-literal, or mid-expression), verified against both
sqlite3and
sqlglotacross several SQL dialects._response_truncated/_truncation_notesremain the authoritative signal that the SQL ispartial; the marker is defense in depth.
BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
N/A (backend-only change, no UI)
TESTING INSTRUCTIONS
pytest tests/unit_tests/mcp_service/test_middleware.py tests/unit_tests/mcp_service/utils/test_token_utils.pytest_update_chart_write_committed_before_size_check_still_succeedsreproduces the original bug (asserts the mutation happens, then fails
against pre-fix code because
ToolErrorwas raised) and now asserts theguard returns a success response with a truncation marker instead.
re-wrapping, the
chartfield surviving both destructive truncationphases, the
get_chart_sqlstring-field bisection, and the SQL markerrejecting every cut position across multiple real and simulated SQL
engines.
tests/unit_tests/mcp_service/suite passes unchanged (4313 passed,2 skipped).
ADDITIONAL INFORMATION