aminghadersohi opened a new pull request, #44386:
URL: https://github.com/apache/superset/pull/44386
### SUMMARY
`update_chart` persists its chart write and only afterward does the MCP
response-size guard estimate the serialized result. When that response
exceeds the configured token limit, the guard raised `ToolError`, so a
successful chart update was reported to the calling agent as a failed tool
call -- a common outcome for charts with a non-trivial `form_data`. Agents
routinely retry on tool failure, replaying an already-committed mutation.
The guard's truncation was previously limited to two closed sets:
`INFO_TOOLS` (read-only single-object tools) and `DATA_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 committed.
This PR:
- Adds a `COMMITTED_WRITE_TOOLS` set (currently `update_chart`) that is
truncated the same way as `INFO_TOOLS`, but with the tool's identifying
field (`chart`) protected from the final "clear everything" truncation
phase, so a caller can always confirm what was written.
- Adds a last-resort minimal-success fallback for `COMMITTED_WRITE_TOOLS`
that is used only if truncation still can't fit the response under
budget. Unlike every other tool category, this path never raises
`ToolError`, because the mutation has already committed by the time the
guard runs.
- Generalizes the existing CSV-export bisection helper into a reusable
named-string-field truncator, and uses it for `get_chart_sql`, whose
oversized responses are dominated by one large rendered-SQL string with
no `limit`/`row_limit` parameter that could shrink it. Previously such
responses were hard-blocked with an unactionable `Reduction needed: ~0%`
message; now the SQL is bisected to the largest prefix that fits, and the
remaining hard-block path (reached only if bisection alone can't help)
gives an actionable suggestion instead.
### 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.py`
- New regression test
`TestResponseSizeGuardMiddleware::test_update_chart_write_committed_before_size_check_still_succeeds`
reproduces the bug (asserts the mutation happens, then fails against
pre-fix code because `ToolError` was raised) and now asserts the guard
returns a success response with a truncation marker instead.
- New tests cover the minimal-fallback path
(`test_committed_write_falls_back_to_minimal_response_not_error`), the
`get_chart_sql` string-field bisection
(`test_truncates_get_chart_sql_by_bisecting_sql_field`), and the
`protected_keys` mechanism at the `token_utils` layer.
- Full `tests/unit_tests/mcp_service/` suite passes unchanged (4304 passed,
2 skipped).
### ADDITIONAL INFORMATION
- [ ] Has associated issue:
- [ ] Required feature flags:
- [ ] Changes UI
- [ ] Includes DB Migration (follow approval process in
[SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [ ] Introduces new feature or API
- [ ] Removes existing feature or API
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
