You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/09/18 10:15:03 UTC
(superset) branch aminghadersohi/mcp-update-chart-response-size-guard created (now c1d444e8a75)
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a change to branch aminghadersohi/mcp-update-chart-response-size-guard
in repository https://gitbox.apache.org/repos/asf/superset.git
at c1d444e8a75 fix(mcp): make the committed-write guard cover every committing write tool
This branch includes the following new commits:
new c1d444e8a75 fix(mcp): make the committed-write guard cover every committing write tool
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: fix(mcp): make the committed-write guard cover every committing write tool
Posted by mi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch aminghadersohi/mcp-update-chart-response-size-guard
in repository https://gitbox.apache.org/repos/asf/superset.git
commit c1d444e8a7576b37ec00bcc1e42ecf134039347b
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Fri Sep 18 10:14:51 2026 +0000
fix(mcp): make the committed-write guard cover every committing write 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 <no...@anthropic.com>
---
superset/mcp_service/middleware.py | 172 ++++++++++++++-------
superset/mcp_service/utils/token_utils.py | 100 ++++++++++--
.../mcp_service/test_mcp_tool_registration.py | 82 ++++++++++
tests/unit_tests/mcp_service/test_middleware.py | 138 +++++++++++++++++
.../mcp_service/utils/test_token_utils.py | 45 ++++++
5 files changed, 464 insertions(+), 73 deletions(-)
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index 1557911bb7b..f9053943ce3 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -54,7 +54,9 @@ from superset.mcp_service.constants import (
DEFAULT_WARN_THRESHOLD_PCT,
)
from superset.mcp_service.utils.token_utils import (
+ COMMITTED_WRITE_SPECS,
COMMITTED_WRITE_TOOLS,
+ CommittedWriteSpec,
DATA_QUERY_TOOLS,
estimate_response_tokens,
format_size_limit_error,
@@ -88,6 +90,26 @@ _MINIMAL_FIELD_CHARS = 200
# backward-compatible alias, so both are kept.
_MINIMAL_ERROR_FIELDS = ("error_type", "error", "message", "error_code", "details")
+# Scalar fields kept when a committed write's identifying object (a nested
+# ``chart``/``dashboard``/``metric`` dict) has to be reduced to fit. These are
+# the names across the info models that answer "what was written" --
+# everything else on them is either unbounded or irrelevant to that question.
+# ``is_unsaved_state`` is here because update_chart defaults to
+# ``generate_preview=True`` and then persists nothing, so it is the caller's
+# only in-band way to tell a cached preview from a persisted write.
+_MINIMAL_IDENTITY_FIELDS = (
+ "id",
+ "uuid",
+ "url",
+ "slice_name",
+ "dashboard_title",
+ "metric_name",
+ "table_name",
+ "dataset_name",
+ "label",
+ "is_unsaved_state",
+)
+
def _sanitize_error_for_logging(error: Exception) -> str:
"""Sanitize error messages to prevent information disclosure in logs."""
@@ -1492,11 +1514,12 @@ class ResponseSizeGuardMiddleware(Middleware):
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
+ the protected identifying fields alone are 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.
"""
+ spec = COMMITTED_WRITE_SPECS[tool_name]
if (extracted := self._extract_payload_from_tool_result(response)) is not None:
payload = extracted
elif isinstance(response, dict):
@@ -1507,19 +1530,12 @@ class ResponseSizeGuardMiddleware(Middleware):
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."
+ f"re-read the {spec.resource} 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)
+ minimal = self._select_confirmation_fields(payload, spec)
+ minimal["_response_truncated"] = True
+ minimal["_truncation_notes"] = truncation_notes
+ self._shrink_minimal_response(minimal, spec)
logger.warning(
"Response for %s could not fit under the size limit after full "
"truncation (~%d tokens, limit %d); returning a minimal write "
@@ -1555,36 +1571,71 @@ class ResponseSizeGuardMiddleware(Middleware):
return self._rewrap_as_tool_result(minimal, response)
return minimal
- def _shrink_minimal_response(self, minimal: dict[str, Any]) -> None:
+ @staticmethod
+ def _select_confirmation_fields(
+ payload: dict[str, Any], spec: CommittedWriteSpec
+ ) -> dict[str, Any]:
+ """Keep only the payload keys that confirm the write happened.
+
+ Selecting from the payload's *own* keys rather than a fixed list is
+ what makes this work for every committed-write tool. A dashboard
+ response has no ``chart``, ``explore_url`` or ``success`` to copy, and
+ synthesizing them would put fields on the response that its model
+ never declares. Conversely, the fields worth keeping differ per tool
+ (``explore_url`` for charts, ``dashboard_url`` for dashboards,
+ ``changed_fields`` for patches) and are all already-bounded scalars,
+ so "every scalar the tool returned" names them without a per-tool
+ list.
+
+ Containers other than the identifying fields are dropped outright:
+ they are precisely what pushed the response over the limit, and none
+ of them answers "what was written".
+ """
+ minimal: dict[str, Any] = {
+ key: value
+ for key, value in payload.items()
+ if key in spec.identifying_fields
+ or key == "error"
+ or not isinstance(value, (list, dict))
+ }
+ if spec.reports_success:
+ # An unparseable payload yields nothing to copy, but a tool whose
+ # schema has ``success`` must still say the write succeeded.
+ minimal.setdefault("success", True)
+ return minimal
+
+ def _shrink_minimal_response(
+ self, minimal: dict[str, Any], spec: CommittedWriteSpec
+ ) -> None:
"""Force ``minimal`` under the token limit, degrading fields in place.
Every value here is copied from the *untruncated* payload, so a
"minimal" response is only actually small once each unbounded field
has been cut down:
- - ``chart`` is reduced to identifying scalars, which include
- ``is_unsaved_state``: update_chart defaults to
+ - each identifying field is a nested object (``chart``, ``dashboard``,
+ ``metric``) that is reduced to identifying scalars -- including
+ ``is_unsaved_state``, since update_chart defaults to
``generate_preview=True`` and then persists nothing, so that flag is
the caller's only in-band way to tell a cached preview from a
persisted write, and shrinking must not be what drops it;
- - those scalars (``slice_name``, ``url``) are themselves free-form
- strings, so they are clipped;
- - ``explore_url`` is a free-form string the reduction above does not
- reach, so it is clipped too;
- - ``error`` is not a string at all in the shapes the tools return --
- it is a nested error model that arrives here as a dict -- so it gets
- the same identifying-scalars treatment as ``chart`` rather than a
- plain clip (see ``_clip_error``).
-
- The remaining keys are bounded already: ``success`` and
- ``_response_truncated`` are booleans, ``schema_version`` and
- ``api_version`` are short fixed constants, and ``_truncation_notes``
- is fixed text. Reducing every unbounded field is what makes the result
- bounded by construction: identifying scalars plus fixed-text notes. A failed
+ - those scalars (``slice_name``, ``dashboard_title``, ``url``) are
+ themselves free-form strings, so they are clipped;
+ - the remaining top-level values are scalars by construction (see
+ ``_select_confirmation_fields``), but a scalar can still be a
+ free-form string -- ``explore_url``, ``dashboard_url``, ``message``
+ -- so every one of them is clipped too;
+ - ``error`` is not a string at all in most of the shapes the tools
+ return -- it is often a nested error model that arrives here as a
+ dict -- so it gets the same identifying-scalars treatment as the
+ identifying fields rather than a plain clip (see ``_clip_error``).
+
+ Reducing every unbounded field is what makes the result bounded by
+ construction: identifying scalars plus fixed-text notes. A failed
measurement counts as "too big" so the payload is degraded rather
- than optimistically returned, and the chart identity is never dropped
- just because the estimator errored -- surfacing which chart was
- written is the whole point of this fallback.
+ than optimistically returned, and the written object's identity is
+ never dropped just because the estimator errored -- surfacing what
+ was written is the whole point of this fallback.
With an extremely small ``token_limit`` even the fully clipped form
can exceed it. Returning it anyway is deliberate: this path exists so
@@ -1594,25 +1645,30 @@ class ResponseSizeGuardMiddleware(Middleware):
if _fits(minimal, self.token_limit):
return
- chart = minimal.get("chart")
- if isinstance(chart, dict):
- minimal["chart"] = {
- key: _clip_string(chart[key])
- for key in ("id", "uuid", "slice_name", "url", "is_unsaved_state")
- 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."
- )
+ for field in sorted(spec.identifying_fields):
+ if field not in minimal:
+ continue
+ value = minimal[field]
+ if isinstance(value, dict):
+ minimal[field] = {
+ key: _clip_string(value[key])
+ for key in _MINIMAL_IDENTITY_FIELDS
+ if key in value
+ }
+ minimal["_truncation_notes"].append(
+ f"'{field}' reduced to identifying fields only."
+ )
+ elif isinstance(value, list):
+ minimal[field] = []
+ minimal["_truncation_notes"].append(
+ f"'{field}' list cleared to fit the size limit."
+ )
- for key, clip in (("error", _clip_error), ("explore_url", _clip_string)):
- current = minimal.get(key)
- clipped = clip(current)
+ for key in list(minimal):
+ if key in spec.identifying_fields or key.startswith("_"):
+ continue
+ current = minimal[key]
+ clipped = _clip_error(current) if key == "error" else _clip_string(current)
if clipped is not current:
minimal[key] = clipped
minimal["_truncation_notes"].append(
@@ -1651,15 +1707,13 @@ class ResponseSizeGuardMiddleware(Middleware):
truncated (never for COMMITTED_WRITE_TOOLS).
"""
# Info tools and committed-write tools: field-level truncation
- # (strings, lists, dicts). Committed-write tools protect their
- # identifying field ('chart') so write confirmation survives even
- # the most aggressive truncation phase.
+ # (strings, lists, dicts). Committed-write tools protect their own
+ # identifying fields -- 'chart', 'dashboard', 'metric', or none at
+ # all where the identity is top-level scalars the phases never drop
+ # -- so write confirmation survives even the most aggressive phase.
if tool_name in INFO_TOOLS or tool_name in COMMITTED_WRITE_TOOLS:
- protected_keys = (
- frozenset({"chart"})
- if tool_name in COMMITTED_WRITE_TOOLS
- else frozenset()
- )
+ spec = COMMITTED_WRITE_SPECS.get(tool_name)
+ protected_keys = spec.identifying_fields if spec else frozenset()
truncated = self._try_truncate_info_response(
tool_name, response, estimated_tokens, protected_keys=protected_keys
)
diff --git a/superset/mcp_service/utils/token_utils.py b/superset/mcp_service/utils/token_utils.py
index e6b4bb66d6c..86a95dfc2ac 100644
--- a/superset/mcp_service/utils/token_utils.py
+++ b/superset/mcp_service/utils/token_utils.py
@@ -46,7 +46,7 @@ that risk.
from __future__ import annotations
import logging
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, NamedTuple, Union
from pydantic import BaseModel
from typing_extensions import TypeAlias
@@ -489,19 +489,91 @@ DATA_QUERY_TOOLS = frozenset(
}
)
-# Mutating tools whose transaction commits (via @transaction) before this
-# middleware ever inspects the response -- by the time an oversized response
-# is detected, the write already happened. Raising ToolError here would
-# report a completed write as a failure, and a retrying MCP client would
-# replay the mutation. These are truncated with the same field-level phases
-# as INFO_TOOLS (see ``_handle_oversized_response``), with the tool's
-# identifying field protected from the final "clear everything" phase so the
-# caller can always confirm what was written.
-COMMITTED_WRITE_TOOLS = frozenset(
- {
- "update_chart",
- }
-)
+
+class CommittedWriteSpec(NamedTuple):
+ """Per-tool facts the size guard needs to confirm a committed write.
+
+ ``resource`` is the noun used in the truncation note that tells the
+ caller what to re-read.
+
+ ``identifying_fields`` are the top-level keys that carry the write
+ confirmation *as a list or dict*, and so have to be protected from the
+ destructive truncation phases. Only containers need naming here: phases 4
+ and 5 summarize dicts and empty collections, and nothing in
+ ``truncate_oversized_response`` ever drops a top-level scalar (phase 1
+ only clips long strings). A response whose identity is scalars -- e.g.
+ ``delete_chart``'s ``deleted_id``, ``create_dataset``'s ``id`` and
+ ``table_name`` -- therefore needs no protection and correctly maps to an
+ empty set.
+
+ ``reports_success`` records whether the tool's response model actually
+ has a ``success`` field. It is consulted only when the payload could not
+ be parsed at all, to decide whether synthesizing ``success: True`` would
+ confirm the write or invent a field the schema does not have.
+
+ ``TestCommittedWriteSpecsMatchToolSchemas`` checks every entry against the
+ registered output schemas, so this table cannot drift from the models.
+ """
+
+ resource: str
+ identifying_fields: frozenset[str]
+ reports_success: bool
+
+
+def _spec(
+ resource: str,
+ *identifying_fields: str,
+ reports_success: bool = False,
+) -> CommittedWriteSpec:
+ """Build a ``CommittedWriteSpec`` with the field set spelled inline."""
+ return CommittedWriteSpec(
+ resource=resource,
+ identifying_fields=frozenset(identifying_fields),
+ reports_success=reports_success,
+ )
+
+
+# Mutating tools whose transaction commits (via @transaction, or an explicit
+# ``db.session.commit()``) before this middleware ever inspects the response --
+# by the time an oversized response is detected, the write already happened.
+# Raising ToolError here would report a completed write as a failure, and a
+# retrying MCP client would replay the mutation. These are truncated with the
+# same field-level phases as INFO_TOOLS (see ``_handle_oversized_response``),
+# with the tool's identifying fields protected from the final "clear
+# everything" phase so the caller can always confirm what was written.
+#
+# This is every mutating MCP tool except three, each excluded for a reason
+# that is about the invariant above rather than about response size:
+# - ``execute_sql`` commits against the *analytics* database, not Superset
+# metadata, and already degrades gracefully via DATA_QUERY_TOOLS row
+# truncation rather than reaching the hard-error path;
+# - ``generate_explore_link`` and ``update_chart_preview`` persist nothing
+# to the metadata database -- they only cache a form_data key -- so a
+# retry re-caches rather than replaying a mutation.
+COMMITTED_WRITE_SPECS: Dict[str, CommittedWriteSpec] = {
+ "add_chart_to_existing_dashboard": _spec("dashboard", "dashboard"),
+ "create_dataset": _spec("dataset"),
+ "create_theme": _spec("theme", reports_success=True),
+ "create_virtual_dataset": _spec("dataset"),
+ "delete_chart": _spec("chart", reports_success=True),
+ "delete_dashboard": _spec("dashboard", reports_success=True),
+ "duplicate_dashboard": _spec("dashboard", "dashboard"),
+ "generate_chart": _spec("chart", "chart", reports_success=True),
+ "generate_dashboard": _spec("dashboard", "dashboard"),
+ "manage_dashboard_certification": _spec("dashboard"),
+ "manage_dashboard_owners": _spec("dashboard"),
+ "manage_dashboard_roles": _spec("dashboard"),
+ "manage_native_filters": _spec("dashboard"),
+ "remove_chart_from_dashboard": _spec("dashboard", "dashboard"),
+ "restore_chart": _spec("chart", reports_success=True),
+ "restore_dashboard": _spec("dashboard", reports_success=True),
+ "save_sql_query": _spec("saved query"),
+ "update_chart": _spec("chart", "chart", reports_success=True),
+ "update_dashboard": _spec("dashboard", "dashboard"),
+ "update_dataset_metric": _spec("dataset", "metric"),
+}
+
+COMMITTED_WRITE_TOOLS = frozenset(COMMITTED_WRITE_SPECS)
# Tools whose oversized response is dominated by a single large string field
# with no row/page/limit parameter the caller could add to shrink it (e.g.
diff --git a/tests/unit_tests/mcp_service/test_mcp_tool_registration.py b/tests/unit_tests/mcp_service/test_mcp_tool_registration.py
index 4de4ad48d33..97650cfba43 100644
--- a/tests/unit_tests/mcp_service/test_mcp_tool_registration.py
+++ b/tests/unit_tests/mcp_service/test_mcp_tool_registration.py
@@ -25,6 +25,7 @@ from unittest.mock import MagicMock, patch
import pytest
from superset.mcp_service.app import get_default_instructions, init_fastmcp_server, mcp
+from superset.mcp_service.utils.token_utils import COMMITTED_WRITE_SPECS
# Patch target for the feature_flag_manager imported inside _apply_config_guards
_FFM_PATH = "superset.extensions.feature_flag_manager"
@@ -143,6 +144,19 @@ MUTATING_TOOLS = {
"update_dataset_metric",
}
+# Mutating tools deliberately outside COMMITTED_WRITE_SPECS. Each is excluded
+# for a reason about *what* it commits, not about how large its response is:
+# execute_sql commits against the analytics database rather than Superset
+# metadata (and degrades via row truncation), while the two preview tools
+# persist nothing to the metadata database -- they only cache a form_data key,
+# so a retry re-caches instead of replaying a mutation.
+NON_COMMITTING_MUTATING_TOOLS = {
+ "execute_sql",
+ "generate_explore_link",
+ "update_chart_preview",
+}
+
+
DESTRUCTIVE_TOOLS = {
"delete_chart",
"delete_dashboard",
@@ -640,3 +654,71 @@ def test_instructions_generated_after_disabled_tools_removed() -> None:
# get_default_instructions must have been called with the disabled set
assert len(captured) == 1
assert "execute_sql" in captured[0]
+
+
+def test_committed_write_specs_cover_every_committing_mutating_tool() -> None:
+ """Every mutating tool that commits before the size guard runs needs a spec.
+
+ The size guard only spares a tool from a hard ToolError -- and so only
+ stops a retrying client replaying an already-committed mutation -- when
+ the tool is in this table. Deriving the expectation from MUTATING_TOOLS
+ means a newly added write tool fails here instead of silently getting the
+ hard-block behaviour.
+ """
+ assert set(COMMITTED_WRITE_SPECS) == MUTATING_TOOLS - NON_COMMITTING_MUTATING_TOOLS
+
+
+def test_committed_write_specs_match_registered_output_schemas() -> None:
+ """Each spec must describe the response the tool actually returns.
+
+ The spec drives which fields survive truncation, so an identifying field
+ that is not on the response would protect nothing, and a wrong
+ ``reports_success`` would either drop the write confirmation or invent a
+ field the schema never declares.
+ """
+ registered = {tool.name: tool for tool in _run(mcp.list_tools())}
+ problems: dict[str, str] = {}
+
+ for name, spec in COMMITTED_WRITE_SPECS.items():
+ top_level = _top_level_schema_property_names(
+ registered[name].output_schema or {}
+ )
+ if missing := spec.identifying_fields - top_level:
+ problems[name] = (
+ f"identifying fields absent from response: {sorted(missing)}"
+ )
+ elif spec.reports_success is not ("success" in top_level):
+ problems[name] = (
+ f"reports_success={spec.reports_success} but 'success' in "
+ f"response schema is {'success' in top_level}"
+ )
+
+ assert not problems, problems
+
+
+def test_committed_write_identifying_fields_are_the_container_ones() -> None:
+ """Only list/dict fields need naming; scalars survive truncation anyway.
+
+ ``truncate_oversized_response`` only summarizes dicts and empties
+ collections -- no phase drops a top-level scalar. So a response whose
+ identity is scalars (delete_chart's ``deleted_id``, create_dataset's
+ ``id``/``table_name``) correctly maps to an empty set, and naming a scalar
+ here would be dead weight that hides that reasoning.
+ """
+ registered = {tool.name: tool for tool in _run(mcp.list_tools())}
+ scalar_types = {"string", "integer", "number", "boolean"}
+ named_scalars: dict[str, list[str]] = {}
+
+ for name, spec in COMMITTED_WRITE_SPECS.items():
+ schema = registered[name].output_schema or {}
+ properties = schema.get("properties", {})
+ for field in sorted(spec.identifying_fields):
+ declared = properties.get(field)
+ if not isinstance(declared, dict):
+ continue
+ # A scalar field declares a plain scalar "type" and nothing else;
+ # object/array fields use $ref, anyOf or an explicit container type.
+ if declared.get("type") in scalar_types and "$ref" not in declared:
+ named_scalars.setdefault(name, []).append(field)
+
+ assert not named_scalars, named_scalars
diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py
index 14e95433f9b..188d44c6ec6 100644
--- a/tests/unit_tests/mcp_service/test_middleware.py
+++ b/tests/unit_tests/mcp_service/test_middleware.py
@@ -802,6 +802,144 @@ class TestResponseSizeGuardMiddleware:
assert "was not rolled back by this size limit" in note
assert "committed" not in note
+ @pytest.mark.asyncio
+ async def test_update_dashboard_committed_write_is_not_hard_blocked(
+ self,
+ ) -> None:
+ """A committed update_dashboard write must never surface as ToolError.
+
+ update_dashboard commits (``db.session.commit()``) before it builds
+ UpdateDashboardResponse, so by the time the size guard runs the
+ dashboard is already written -- the same invariant that puts
+ update_chart on this path. Hard-blocking here would report a
+ completed write as a failure and let a retrying client replay it.
+ """
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "update_dashboard"
+ context.message.arguments = {"identifier": 42}
+
+ dashboard_state = {"dashboard_title": "Original"}
+
+ async def call_next_committing_write(_ctx: Any) -> dict[str, Any]:
+ # Mirrors the real tool: the commit happens here, before the
+ # middleware ever sees the (oversized) response.
+ dashboard_state["dashboard_title"] = "Q1 Revenue"
+ return {
+ "dashboard": {
+ "id": 42,
+ "uuid": "dash-uuid",
+ "dashboard_title": "Q1 Revenue",
+ "url": "/superset/dashboard/42/",
+ },
+ "dashboard_url": "/superset/dashboard/42/",
+ "changed_fields": ["dashboard_title"],
+ "error": None,
+ "position_json": {f"key_{i}": f"value_{i}" for i in range(200)},
+ }
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next_committing_write)
+
+ assert dashboard_state["dashboard_title"] == "Q1 Revenue"
+ assert isinstance(result, dict)
+ assert result["dashboard"]["id"] == 42
+ assert result.get("_response_truncated") is True
+
+ @pytest.mark.asyncio
+ async def test_update_dashboard_identifying_field_survives_nuclear_phase(
+ self,
+ ) -> None:
+ """The protected field must follow the tool, not a hardcoded 'chart'.
+
+ Phase 5 empties every unprotected dict, so protecting 'chart' on a
+ dashboard response would protect nothing and clear the very field
+ that says which dashboard was written.
+ """
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "update_dashboard"
+ context.message.arguments = {"identifier": 7}
+
+ large_response: dict[str, Any] = {
+ "dashboard": {
+ "id": 7,
+ "uuid": "dash-uuid",
+ "dashboard_title": "Wide Dashboard",
+ },
+ "dashboard_url": "/superset/dashboard/7/",
+ "changed_fields": ["css"],
+ "error": None,
+ }
+ # Dicts small enough to escape Phase 4's summarizer (<= 20 keys) and
+ # strings short enough to escape the Phase 1/3 clippers, but together
+ # far over budget -- so truncation has to reach Phase 5, the only
+ # phase that would empty the 'dashboard' dict.
+ for index in range(6):
+ large_response[f"filter_scope_{index}"] = {
+ f"key_{i}": "v" * 100 for i in range(20)
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ # The write confirmation survived the destructive phases intact.
+ assert result["dashboard"]["id"] == 7
+ assert result["dashboard"]["dashboard_title"] == "Wide Dashboard"
+
+ @pytest.mark.asyncio
+ async def test_minimal_response_adds_no_fields_the_schema_lacks(
+ self,
+ ) -> None:
+ """The minimal confirmation must not invent chart-shaped fields.
+
+ UpdateDashboardResponse declares no ``chart``, ``explore_url`` or
+ ``success``; synthesizing them would hand the caller a payload its
+ own output schema does not describe.
+ """
+ middleware = ResponseSizeGuardMiddleware(token_limit=500)
+
+ context = MagicMock()
+ context.message.name = "update_dashboard"
+ context.message.arguments = {"identifier": 7}
+
+ large_response = {
+ "dashboard": {"id": 7, "dashboard_title": "D" * 40000},
+ "dashboard_url": "/superset/dashboard/7/",
+ "changed_fields": ["css"],
+ "error": None,
+ }
+ call_next = AsyncMock(return_value=large_response)
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ # Force the minimal-response fallback to be the path under test.
+ patch(
+ "superset.mcp_service.middleware.estimate_response_tokens",
+ side_effect=[600, 600],
+ ),
+ ):
+ result = await middleware.on_call_tool(context, call_next)
+
+ assert isinstance(result, dict)
+ for absent in ("chart", "explore_url", "success"):
+ assert absent not in result
+ # ...while still confirming which dashboard was written, bounded.
+ assert result["dashboard"]["id"] == 7
+ assert estimate_token_count(utils_json.dumps(result)) <= 500
+ assert "re-read the dashboard" in result["_truncation_notes"][0]
+
@pytest.mark.asyncio
async def test_minimal_response_is_bounded_by_every_unbounded_field(
self,
diff --git a/tests/unit_tests/mcp_service/utils/test_token_utils.py b/tests/unit_tests/mcp_service/utils/test_token_utils.py
index 1bb0ca7b396..67c8eb29c42 100644
--- a/tests/unit_tests/mcp_service/utils/test_token_utils.py
+++ b/tests/unit_tests/mcp_service/utils/test_token_utils.py
@@ -36,6 +36,7 @@ from superset.mcp_service.utils.token_utils import (
_truncate_strings,
_truncate_strings_recursive,
CHARS_PER_TOKEN,
+ COMMITTED_WRITE_SPECS,
COMMITTED_WRITE_TOOLS,
estimate_response_tokens,
estimate_token_count,
@@ -457,15 +458,59 @@ class TestCommittedWriteToolsSet:
assert "update_chart" in COMMITTED_WRITE_TOOLS
def test_does_not_contain_read_only_tools(self) -> None:
+ """Read-only tools commit nothing, so an oversized response is safe
+ to hard-block: there is no completed write for a retry to replay."""
assert "get_chart_info" not in COMMITTED_WRITE_TOOLS
assert "list_charts" not in COMMITTED_WRITE_TOOLS
assert "execute_sql" not in COMMITTED_WRITE_TOOLS
+ def test_contains_update_dashboard(self) -> None:
+ """update_dashboard commits (db.session.commit(), before it builds
+ UpdateDashboardResponse) just as update_chart does, so it satisfies
+ the same invariant and must get the same protection -- otherwise an
+ oversized response hard-errors after the dashboard was already
+ written, and a retrying client replays the mutation."""
+ assert "update_dashboard" in COMMITTED_WRITE_TOOLS
+
+ def test_identifying_fields_are_per_tool(self) -> None:
+ """The protected field cannot be a single hardcoded name.
+
+ Protecting 'chart' on a dashboard response would protect nothing:
+ the field that carries the write confirmation differs per tool, so
+ each spec names its own.
+ """
+ assert COMMITTED_WRITE_SPECS["update_chart"].identifying_fields == frozenset(
+ {"chart"}
+ )
+ assert COMMITTED_WRITE_SPECS[
+ "update_dashboard"
+ ].identifying_fields == frozenset({"dashboard"})
+ assert COMMITTED_WRITE_SPECS[
+ "update_dataset_metric"
+ ].identifying_fields == frozenset({"metric"})
+
+ def test_scalar_identity_tools_protect_nothing(self) -> None:
+ """No truncation phase drops a top-level scalar, so a response whose
+ identity is scalars needs no protected field at all -- delete_chart's
+ deleted_id survives even the nuclear phase untouched."""
+ assert COMMITTED_WRITE_SPECS["delete_chart"].identifying_fields == frozenset()
+ assert COMMITTED_WRITE_SPECS["create_dataset"].identifying_fields == frozenset()
+
+ def test_reports_success_tracks_the_response_model(self) -> None:
+ """Consulted only when the payload is unparseable, to decide whether
+ synthesizing ``success`` confirms the write or invents a field the
+ schema never declares. GenerateChartResponse has one;
+ UpdateDashboardResponse does not."""
+ assert COMMITTED_WRITE_SPECS["update_chart"].reports_success is True
+ assert COMMITTED_WRITE_SPECS["update_dashboard"].reports_success is False
+
class TestStringFieldTruncationToolsMap:
"""Test the STRING_FIELD_TRUNCATION_TOOLS constant."""
def test_get_chart_sql_maps_to_sql_field(self) -> None:
+ """The map names the field to bisect: get_chart_sql's payload is
+ dominated by ``sql``, which is what truncation has to cut down."""
assert STRING_FIELD_TRUNCATION_TOOLS["get_chart_sql"] == "sql"