aminghadersohi commented on code in PR #44386:
URL: https://github.com/apache/superset/pull/44386#discussion_r4037868041
##########
superset/mcp_service/middleware.py:
##########
@@ -1357,6 +1388,193 @@ def _try_truncate_data_query_response(
return truncated
+ def _try_truncate_string_field_response(
+ self,
+ tool_name: str,
+ response: Any,
+ estimated_tokens: int,
+ field: str,
+ ) -> Any | None:
+ """Attempt to truncate a response by bisecting one oversized string
field.
+
+ Returns the truncated response if successful, None otherwise.
+ """
+ extracted = self._extract_payload_from_tool_result(response)
+ if extracted is None and isinstance(response, ToolResult):
+ # A ToolResult whose payload can't be parsed is opaque: truncating
+ # it would model_dump() the wrapper itself and hand FastMCP a
+ # plain dict, which then fails in to_mcp_result(). Decline instead
+ # and let the caller fall through to its own fallback.
+ logger.warning(
+ "Cannot truncate %s: ToolResult payload is not a JSON object",
+ tool_name,
+ )
+ return None
+
+ truncation_target = extracted if extracted is not None else response
+
+ try:
+ truncated, was_truncated, notes = truncate_string_field_response(
+ truncation_target, self.token_limit, field
+ )
+ except Exception as trunc_error: # noqa: BLE001
+ logger.warning(
+ "String field truncation failed for %s due to %s: %s",
+ tool_name,
+ type(trunc_error).__name__,
+ trunc_error,
+ )
+ return None
+
+ if not was_truncated:
+ return None
+
+ truncated_tokens = estimate_response_tokens(truncated)
+ if truncated_tokens > self.token_limit:
+ return None
+
+ logger.warning(
+ "Response 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",
+ dashboard_id=None,
+ duration_ms=None,
+ slice_id=None,
+ referrer=None,
+ 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)
+
+ if extracted is not None and isinstance(truncated, dict):
+ return self._rewrap_as_tool_result(truncated, response)
+
+ return truncated
+
+ def _minimal_committed_write_response(
+ self,
+ tool_name: str,
+ response: Any,
+ estimated_tokens: int,
+ ) -> Any:
+ """Build a guaranteed-small success response for a committed write.
+
+ Last-resort fallback for COMMITTED_WRITE_TOOLS: reached only when
+ even the nuclear phase of ``truncate_oversized_response`` can't bring
+ the response under budget (in practice this should not happen, since
+ the protected identifying field alone is tiny). The underlying
+ mutation already committed by the time this middleware runs, so this
+ path must never raise -- it keeps only what confirms the write
+ succeeded and drops everything else.
+ """
+ if (extracted := self._extract_payload_from_tool_result(response)) is
not None:
+ payload = extracted
+ elif isinstance(response, dict):
+ payload = response
+ else:
+ payload = {}
+ truncation_notes = [
+ f"Response for {tool_name} exceeded the size limit even after "
+ "truncation; non-essential fields were dropped. The tool call "
+ "itself completed and was not rolled back by this size limit -- "
+ "re-read the chart to see its full state."
+ ]
+ minimal = {
+ "chart": payload.get("chart"),
+ "error": payload.get("error"),
+ "success": payload.get("success", True),
+ "explore_url": payload.get("explore_url"),
+ "schema_version": payload.get("schema_version"),
+ "api_version": payload.get("api_version"),
+ "_response_truncated": True,
+ "_truncation_notes": truncation_notes,
+ }
+ self._shrink_minimal_response(minimal)
+ logger.warning(
+ "Response for %s could not fit under the size limit after full "
+ "truncation (~%d tokens, limit %d); returning a minimal write "
+ "confirmation instead of blocking a completed write.",
+ tool_name,
+ estimated_tokens,
+ self.token_limit,
+ )
+ try:
+ user_id = get_user_id()
+ event_logger.log(
+ user_id=user_id,
+ action="mcp_response_truncated",
+ dashboard_id=None,
+ duration_ms=None,
+ slice_id=None,
+ referrer=None,
+ curated_payload={
+ "tool": tool_name,
+ "original_tokens": estimated_tokens,
+ "token_limit": self.token_limit,
+ "truncation_notes": truncation_notes,
+ },
+ )
+ except Exception as log_error: # noqa: BLE001
+ logger.warning("Failed to log truncation event: %s", log_error)
+
+ # Rewrap whenever the tool returned a ToolResult, including the case
+ # where its payload could not be parsed: returning a bare dict there
+ # would blow up in FastMCP's ``result.to_mcp_result()`` and surface
+ # the completed write as an internal error after all.
+ if isinstance(response, ToolResult):
+ return self._rewrap_as_tool_result(minimal, response)
+ return minimal
+
+ def _shrink_minimal_response(self, minimal: dict[str, Any]) -> None:
+ """Force ``minimal`` under the token limit, degrading ``chart`` in
place.
+
+ ``chart`` is copied from the *untruncated* payload, so when it is
+ itself the oversized field the "minimal" response is not actually
+ small. Reduce it to identifying scalars, which is bounded by
+ construction, rather than handing back something the transport will
+ reject.
+
+ Only one measurement is taken, and a failed measurement counts as
+ "too big": the reduced form is small enough that there is nothing to
+ re-check, and the chart identity is never dropped just because the
+ estimator errored -- surfacing which chart was written is the whole
+ point of this fallback.
+ """
+ if _fits(minimal, self.token_limit):
+ return
+
+ chart = minimal.get("chart")
+ if isinstance(chart, dict):
+ minimal["chart"] = {
+ key: chart[key]
+ for key in ("id", "uuid", "slice_name", "url")
+ if key in chart
+ }
+ minimal["_truncation_notes"].append(
+ "Chart details reduced to identifying fields only."
+ )
+ else:
+ minimal["chart"] = None
+ minimal["_truncation_notes"].append(
+ "Chart details omitted entirely to fit the size limit."
Review Comment:
Confirmed and fixed in 28096e0041 — thanks, this was a real residual gap.
Reproduced it directly against the fallback path. `_shrink_minimal_response`
reduces `chart` to identifying fields, but that reduction reaches nothing else:
`error` and `explore_url` are copied verbatim from the *untruncated* payload,
and the retained `slice_name`/`url` are themselves free-form strings. With a
200-token limit:
| oversized field | result before | after |
| --- | --- | --- |
| `explore_url` (40k chars) | 10142 tokens | 242 |
| `error` (40k chars) | 10127 tokens | 192 |
| `chart.slice_name` (40k chars) | 20126 tokens | 230 |
So the "reduced form is small enough that there is nothing to re-check"
claim in the old docstring was simply not true.
The fix clips every free-form string to a fixed cap, which makes the
confirmation bounded by construction (identifying scalars plus fixed-text
notes) rather than bounded by assumption, and re-checks afterward.
One deliberate limit worth stating: with an extremely small `token_limit`
even the fully clipped form can exceed it, because a write confirmation has a
non-zero floor. This path exists precisely so a completed write is never
reported as a failure, so it logs a warning and returns anyway instead of
degrading further and losing the confirmation.
Regression test: `test_minimal_response_is_bounded_by_every_unbounded_field`
— fails before the change, passes after.
--
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]
