fix(mcp): add dynamic response truncation for oversized info tool responses (#39107)

(cherry picked from commit 83ad1eca26)
This commit is contained in:
Amin Ghadersohi
2026-04-06 12:36:03 -04:00
committed by Michael S. Molina
parent bc8690187f
commit 4764eea1dc
4 changed files with 643 additions and 4 deletions

View File

@@ -936,6 +936,72 @@ class ResponseSizeGuardMiddleware(Middleware):
excluded_tools = [excluded_tools]
self.excluded_tools = set(excluded_tools or [])
def _try_truncate_info_response(
self,
tool_name: str,
response: Any,
estimated_tokens: int,
) -> Any | None:
"""Attempt to dynamically truncate an info tool response to fit the limit.
Returns the truncated response if successful, None otherwise.
"""
from superset.mcp_service.utils.token_utils import (
estimate_response_tokens,
truncate_oversized_response,
)
try:
truncated, was_truncated, notes = truncate_oversized_response(
response, self.token_limit
)
except (MemoryError, RecursionError) as trunc_error:
logger.warning(
"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). Fields: %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",
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 isinstance(truncated, dict):
truncated["_response_truncated"] = True
truncated["_truncation_notes"] = notes
return truncated
async def on_call_tool(
self,
context: MiddlewareContext,
@@ -984,9 +1050,18 @@ class ResponseSizeGuardMiddleware(Middleware):
# Block if over limit
if estimated_tokens > self.token_limit:
# Extract params for smart suggestions
params = getattr(context.message, "params", {}) or {}
# For info tools, try dynamic truncation before blocking
from superset.mcp_service.utils.token_utils import INFO_TOOLS
if tool_name in INFO_TOOLS:
truncated = self._try_truncate_info_response(
tool_name, response, estimated_tokens
)
if truncated is not None:
return truncated
# Log the blocked response
logger.error(
"Response blocked for %s: ~%d tokens exceeds limit of %d",
@@ -1011,9 +1086,6 @@ class ResponseSizeGuardMiddleware(Middleware):
except Exception as log_error: # noqa: BLE001
logger.warning("Failed to log size exceeded event: %s", log_error)
# Generate helpful error message with suggestions
# Avoid passing the full `response` (which may be huge) into the formatter
# to prevent large-memory operations during error formatting.
error_message = format_size_limit_error(
tool_name=tool_name,
params=params,

View File

@@ -382,6 +382,210 @@ def _get_tool_specific_suggestions(
return suggestions
# Tools eligible for dynamic response truncation instead of hard blocking.
# These tools return single objects (not paginated lists) where truncation
# is preferable to returning an error.
INFO_TOOLS = frozenset(
{
"get_chart_info",
"get_dataset_info",
"get_dashboard_info",
"get_instance_info",
}
)
# Maximum character length for string fields before truncation
_MAX_STRING_CHARS = 500
# Maximum items to keep in list fields before truncation
_MAX_LIST_ITEMS = 30
# Maximum keys to keep when summarizing large dict fields
_MAX_DICT_KEYS = 20
def _truncate_strings(
data: Dict[str, Any], notes: List[str], max_chars: int = _MAX_STRING_CHARS
) -> bool:
"""Truncate string fields exceeding max_chars at the top level only."""
changed = False
for key, value in data.items():
if isinstance(value, str) and len(value) > max_chars:
original_len = len(value)
data[key] = value[:max_chars] + f"... [truncated from {original_len} chars]"
notes.append(f"Field '{key}' truncated from {original_len} chars")
changed = True
return changed
def _truncate_strings_recursive(
data: Any,
notes: List[str],
max_chars: int = _MAX_STRING_CHARS,
path: str = "",
_depth: int = 0,
) -> bool:
"""Recursively truncate strings throughout the entire data tree.
Walks nested dicts and list items to catch strings like
``charts[0].description`` that top-level truncation misses.
Depth is capped at 10 to avoid runaway recursion.
"""
if _depth > 10:
return False
changed = False
if isinstance(data, dict):
for key, value in data.items():
field_path = f"{path}.{key}" if path else key
if isinstance(value, str) and len(value) > max_chars:
original_len = len(value)
data[key] = (
value[:max_chars] + f"... [truncated from {original_len} chars]"
)
notes.append(
f"Field '{field_path}' truncated from {original_len} chars"
)
changed = True
elif isinstance(value, (dict, list)):
changed |= _truncate_strings_recursive(
value, notes, max_chars, field_path, _depth + 1
)
elif isinstance(data, list):
for i, item in enumerate(data):
if isinstance(item, (dict, list)):
changed |= _truncate_strings_recursive(
item, notes, max_chars, f"{path}[{i}]", _depth + 1
)
return changed
def _truncate_lists(data: Dict[str, Any], notes: List[str], max_items: int) -> bool:
"""Truncate list fields exceeding max_items. Returns True if any truncated.
Does NOT append marker objects into the list to preserve the element type
contract (e.g. ``List[TableColumnInfo]`` stays homogeneous). Truncation
metadata is communicated through the *notes* list and top-level response
fields ``_response_truncated`` / ``_truncation_notes``.
"""
changed = False
for key, value in data.items():
if isinstance(value, list) and len(value) > max_items:
original_len = len(value)
data[key] = value[:max_items]
notes.append(
f"Field '{key}' truncated from {original_len} to {max_items} items"
)
changed = True
return changed
def _summarize_large_dicts(
data: Dict[str, Any], notes: List[str], max_keys: int = _MAX_DICT_KEYS
) -> bool:
"""Replace large dict fields with key summaries. Returns True if any changed."""
changed = False
for key, value in data.items():
if isinstance(value, dict) and len(value) > max_keys:
keys_list = list(value.keys())[:max_keys]
data[key] = {
"_truncated": True,
"_message": (
f"Dict with {len(value)} keys truncated. "
f"Keys: {', '.join(str(k) for k in keys_list)}..."
),
}
notes.append(f"Field '{key}' dict summarized ({len(value)} keys)")
changed = True
return changed
def _replace_collections_with_summaries(data: Dict[str, Any], notes: List[str]) -> bool:
"""Replace all non-empty list/dict fields with empty/minimal values.
Lists are emptied (preserving the list type) rather than replaced with
marker objects to avoid breaking typed list contracts.
"""
changed = False
for key, value in list(data.items()):
if not isinstance(value, (list, dict)) or not value:
continue
count = len(value)
if isinstance(value, list):
data[key] = []
notes.append(f"Field '{key}' list ({count} items) cleared to fit limit")
else:
data[key] = {}
notes.append(f"Field '{key}' dict ({count} keys) cleared to fit limit")
changed = True
return changed
def _is_under_limit(data: Dict[str, Any], token_limit: int) -> bool:
"""Check if the serialized data fits within the token limit."""
from superset.utils import json as utils_json
return estimate_token_count(utils_json.dumps(data)) <= token_limit
def truncate_oversized_response(
response: ToolResponse,
token_limit: int,
) -> tuple[ToolResponse, bool, list[str]]:
"""
Dynamically truncate large fields in a response to fit within the token limit.
Applies five progressive phases of truncation:
1. Truncate long top-level string fields
2. Truncate large list fields to _MAX_LIST_ITEMS
3. Recursively truncate strings in nested structures (list items, nested dicts)
4. Aggressively reduce lists to 10 items and summarize large dicts
5. Replace all collections with empty values
Args:
response: The tool response (Pydantic model, dict, or other).
token_limit: Maximum estimated tokens allowed.
Returns:
A tuple of (possibly-truncated response, was_truncated, list of notes).
"""
notes: list[str] = []
# Convert to a mutable dict for manipulation
if hasattr(response, "model_dump"):
data = response.model_dump()
elif isinstance(response, dict):
data = dict(response)
else:
return response, False, notes
was_truncated = False
# Phase 1: Truncate long string fields
was_truncated |= _truncate_strings(data, notes)
if _is_under_limit(data, token_limit):
return data, was_truncated, notes
# Phase 2: Truncate large list fields
was_truncated |= _truncate_lists(data, notes, _MAX_LIST_ITEMS)
if _is_under_limit(data, token_limit):
return data, was_truncated, notes
# Phase 3: Recursively truncate strings inside nested structures
# (e.g. charts[i].description, native_filters[i].config, etc.)
was_truncated |= _truncate_strings_recursive(data, notes)
if _is_under_limit(data, token_limit):
return data, was_truncated, notes
# Phase 4: Aggressively reduce lists and summarize large dicts
was_truncated |= _truncate_lists(data, notes, max_items=10)
was_truncated |= _summarize_large_dicts(data, notes)
if _is_under_limit(data, token_limit):
return data, was_truncated, notes
# Phase 5: Nuclear — replace all collections with empty values
was_truncated |= _replace_collections_with_summaries(data, notes)
return data, was_truncated, notes
def format_size_limit_error(
tool_name: str,
params: Dict[str, Any] | None,

View File

@@ -207,6 +207,106 @@ class TestResponseSizeGuardMiddleware:
call_args = mock_event_logger.log.call_args
assert call_args.kwargs["action"] == "mcp_response_size_exceeded"
@pytest.mark.asyncio
async def test_truncates_info_tool_instead_of_blocking(self) -> None:
"""Should truncate info tool responses instead of blocking them."""
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "get_dataset_info"
context.message.params = {}
# Large info tool response with a big description
large_response = {
"id": 1,
"table_name": "test",
"description": "x" * 50000,
}
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)
# Should return truncated response, not raise ToolError
assert isinstance(result, dict)
assert result["id"] == 1
assert result["_response_truncated"] is True
assert "[truncated" in result["description"]
@pytest.mark.asyncio
async def test_truncates_chart_info_with_large_form_data(self) -> None:
"""Should truncate get_chart_info with large form_data."""
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "get_chart_info"
context.message.params = {}
large_response = {
"id": 1,
"slice_name": "My Chart",
"form_data": {f"key_{i}": f"value_{i}" for i in range(100)},
}
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)
assert result["id"] == 1
assert result["_response_truncated"] is True
@pytest.mark.asyncio
async def test_still_blocks_non_info_tools(self) -> None:
"""Should still block non-info tools that exceed limit."""
middleware = ResponseSizeGuardMiddleware(token_limit=100)
context = MagicMock()
context.message.name = "list_charts" # Not an info tool
context.message.params = {}
large_response = {"data": "x" * 10000}
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"),
pytest.raises(ToolError),
):
await middleware.on_call_tool(context, call_next)
@pytest.mark.asyncio
async def test_logs_truncation_event(self) -> None:
"""Should log mcp_response_truncated event on successful truncation."""
middleware = ResponseSizeGuardMiddleware(token_limit=500)
context = MagicMock()
context.message.name = "get_dashboard_info"
context.message.params = {}
large_response = {
"id": 1,
"description": "x" * 50000,
}
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") as mock_event_logger,
):
await middleware.on_call_tool(context, call_next)
# Should log truncation event (not size_exceeded)
mock_event_logger.log.assert_called()
call_args = mock_event_logger.log.call_args
assert call_args.kwargs["action"] == "mcp_response_truncated"
class TestCreateResponseSizeGuardMiddleware:
"""Test create_response_size_guard_middleware factory function."""

View File

@@ -24,6 +24,11 @@ from typing import Any, List
from pydantic import BaseModel
from superset.mcp_service.utils.token_utils import (
_replace_collections_with_summaries,
_summarize_large_dicts,
_truncate_lists,
_truncate_strings,
_truncate_strings_recursive,
CHARS_PER_TOKEN,
estimate_response_tokens,
estimate_token_count,
@@ -31,6 +36,8 @@ from superset.mcp_service.utils.token_utils import (
format_size_limit_error,
generate_size_reduction_suggestions,
get_response_size_bytes,
INFO_TOOLS,
truncate_oversized_response,
)
@@ -374,3 +381,259 @@ class TestCalculatedSuggestions:
# Should mention ~66% reduction needed (int truncation of 66.6%)
combined = " ".join(suggestions)
assert "66%" in combined
class TestInfoToolsSet:
"""Test the INFO_TOOLS constant."""
def test_info_tools_contains_expected_tools(self) -> None:
"""Should contain all info tools."""
assert "get_chart_info" in INFO_TOOLS
assert "get_dataset_info" in INFO_TOOLS
assert "get_dashboard_info" in INFO_TOOLS
assert "get_instance_info" in INFO_TOOLS
def test_info_tools_does_not_contain_list_tools(self) -> None:
"""Should not contain list or write tools."""
assert "list_charts" not in INFO_TOOLS
assert "execute_sql" not in INFO_TOOLS
assert "generate_chart" not in INFO_TOOLS
class TestTruncateStrings:
"""Test _truncate_strings helper."""
def test_truncates_long_strings(self) -> None:
"""Should truncate strings exceeding max_chars."""
data: dict[str, Any] = {"description": "x" * 1000, "name": "short"}
notes: list[str] = []
changed = _truncate_strings(data, notes, max_chars=500)
assert changed is True
assert len(data["description"]) < 1000
assert "[truncated from 1000 chars]" in data["description"]
assert data["name"] == "short"
assert len(notes) == 1
def test_does_not_truncate_short_strings(self) -> None:
"""Should not truncate strings within limit."""
data: dict[str, Any] = {"name": "hello", "id": 123}
notes: list[str] = []
changed = _truncate_strings(data, notes, max_chars=500)
assert changed is False
assert data["name"] == "hello"
assert len(notes) == 0
class TestTruncateStringsRecursive:
"""Test _truncate_strings_recursive helper."""
def test_truncates_nested_strings_in_list_items(self) -> None:
"""Should truncate strings inside list items (e.g. charts[i].description)."""
data: dict[str, Any] = {
"id": 1,
"charts": [
{"id": 1, "description": "x" * 1000},
{"id": 2, "description": "short"},
],
}
notes: list[str] = []
changed = _truncate_strings_recursive(data, notes, max_chars=500)
assert changed is True
assert "[truncated" in data["charts"][0]["description"]
assert data["charts"][1]["description"] == "short"
assert len(notes) == 1
assert "charts[0].description" in notes[0]
def test_truncates_nested_strings_in_dicts(self) -> None:
"""Should truncate strings inside nested dicts."""
data: dict[str, Any] = {
"filter_state": {
"dataMask": {"some_filter": "y" * 2000},
},
}
notes: list[str] = []
changed = _truncate_strings_recursive(data, notes, max_chars=500)
assert changed is True
assert "[truncated" in data["filter_state"]["dataMask"]["some_filter"]
def test_respects_depth_limit(self) -> None:
"""Should stop recursing at depth 10."""
# Build a deeply nested structure (15 levels)
data: dict[str, Any] = {"level": "x" * 1000}
current = data
for _ in range(15):
current["nested"] = {"level": "x" * 1000}
current = current["nested"]
notes: list[str] = []
_truncate_strings_recursive(data, notes, max_chars=500)
# Should truncate levels 0-10 but stop before 15
assert len(notes) <= 11
def test_handles_empty_structures(self) -> None:
"""Should handle empty dicts and lists gracefully."""
data: dict[str, Any] = {"items": [], "meta": {}, "name": "ok"}
notes: list[str] = []
changed = _truncate_strings_recursive(data, notes, max_chars=500)
assert changed is False
def test_dashboard_with_many_charts_edge_case(self) -> None:
"""Simulate a dashboard with 30 charts each having long descriptions."""
data: dict[str, Any] = {
"id": 1,
"dashboard_title": "Big Dashboard",
"charts": [
{"id": i, "slice_name": f"Chart {i}", "description": "d" * 2000}
for i in range(30)
],
}
notes: list[str] = []
changed = _truncate_strings_recursive(data, notes, max_chars=500)
assert changed is True
# All 30 chart descriptions should be truncated
assert len(notes) == 30
for chart in data["charts"]:
assert len(chart["description"]) < 2000
assert "[truncated" in chart["description"]
class TestTruncateLists:
"""Test _truncate_lists helper."""
def test_truncates_long_lists(self) -> None:
"""Should truncate lists exceeding max_items without inline markers."""
data: dict[str, Any] = {
"columns": [{"name": f"col_{i}"} for i in range(50)],
"tags": [1, 2],
}
notes: list[str] = []
changed = _truncate_lists(data, notes, max_items=10)
assert changed is True
# Exactly 10 items — no marker appended (preserves type contract)
assert len(data["columns"]) == 10
assert all(isinstance(c, dict) and "name" in c for c in data["columns"])
assert data["tags"] == [1, 2] # Not truncated
assert len(notes) == 1
assert "50" in notes[0]
def test_does_not_truncate_short_lists(self) -> None:
"""Should not truncate lists within limit."""
data: dict[str, Any] = {"items": [1, 2, 3]}
notes: list[str] = []
changed = _truncate_lists(data, notes, max_items=10)
assert changed is False
class TestSummarizeLargeDicts:
"""Test _summarize_large_dicts helper."""
def test_summarizes_large_dicts(self) -> None:
"""Should replace large dicts with key summaries."""
big_dict = {f"key_{i}": f"value_{i}" for i in range(30)}
data: dict[str, Any] = {"form_data": big_dict, "id": 1}
notes: list[str] = []
changed = _summarize_large_dicts(data, notes, max_keys=20)
assert changed is True
assert data["form_data"]["_truncated"] is True
assert "30 keys" in data["form_data"]["_message"]
assert data["id"] == 1
def test_does_not_summarize_small_dicts(self) -> None:
"""Should not summarize dicts within limit."""
data: dict[str, Any] = {"params": {"a": 1, "b": 2}}
notes: list[str] = []
changed = _summarize_large_dicts(data, notes, max_keys=20)
assert changed is False
class TestReplaceCollectionsWithSummaries:
"""Test _replace_collections_with_summaries helper."""
def test_replaces_lists_and_dicts(self) -> None:
"""Should clear non-empty collections to reduce size."""
data: dict[str, Any] = {
"columns": [1, 2, 3],
"params": {"a": 1},
"name": "test",
"empty": [],
}
notes: list[str] = []
changed = _replace_collections_with_summaries(data, notes)
assert changed is True
# Lists become empty lists (preserves type)
assert data["columns"] == []
# Dicts become empty dicts (preserves type)
assert data["params"] == {}
# Scalars unchanged
assert data["name"] == "test"
# Empty collections unchanged
assert data["empty"] == []
assert len(notes) == 2
class TestTruncateOversizedResponse:
"""Test truncate_oversized_response function."""
def test_no_truncation_needed(self) -> None:
"""Should return original data when under limit."""
response = {"id": 1, "name": "test"}
result, was_truncated, notes = truncate_oversized_response(response, 10000)
assert was_truncated is False
assert notes == []
def test_truncates_large_string_fields(self) -> None:
"""Should truncate long strings to fit."""
response = {
"id": 1,
"description": "x" * 50000, # Very large description
}
result, was_truncated, notes = truncate_oversized_response(response, 500)
assert was_truncated is True
assert isinstance(result, dict)
assert "[truncated" in result["description"]
assert any("description" in n for n in notes)
def test_truncates_large_lists(self) -> None:
"""Should truncate lists when strings alone are not enough."""
response = {
"id": 1,
"columns": [{"name": f"col_{i}", "type": "VARCHAR"} for i in range(200)],
}
result, was_truncated, notes = truncate_oversized_response(response, 500)
assert was_truncated is True
assert isinstance(result, dict)
# Should have been truncated
assert len(result["columns"]) < 200
def test_handles_pydantic_model(self) -> None:
"""Should handle Pydantic model input."""
class FakeInfo(BaseModel):
id: int = 1
description: str = "x" * 5000
response = FakeInfo()
result, was_truncated, notes = truncate_oversized_response(response, 200)
assert was_truncated is True
assert isinstance(result, dict)
def test_returns_non_dict_unchanged(self) -> None:
"""Should return non-dict/model responses unchanged."""
result, was_truncated, notes = truncate_oversized_response("just a string", 100)
assert was_truncated is False
assert result == "just a string"
def test_progressive_truncation(self) -> None:
"""Should progressively apply truncation phases."""
# Build a response that's quite large
response = {
"id": 1,
"description": "x" * 2000,
"css": "y" * 2000,
"columns": [{"name": f"col_{i}"} for i in range(100)],
"form_data": {f"key_{i}": f"val_{i}" for i in range(50)},
}
result, was_truncated, notes = truncate_oversized_response(response, 300)
assert was_truncated is True
assert isinstance(result, dict)
assert result["id"] == 1 # Scalar fields preserved
assert len(notes) > 0