You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by vi...@apache.org on 2026/09/22 19:54:25 UTC
(superset) branch master updated: fix(mcp): drop tiktoken, size-guard responses by exact byte length (#44508)
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 7024401aef2 fix(mcp): drop tiktoken, size-guard responses by exact byte length (#44508)
7024401aef2 is described below
commit 7024401aef2e35efe30ebd74a54c6a72dd0f7929
Author: Ville Brofeldt <33...@users.noreply.github.com>
AuthorDate: Tue Sep 22 12:54:18 2026 -0700
fix(mcp): drop tiktoken, size-guard responses by exact byte length (#44508)
Co-authored-by: rusackas <ev...@preset.io>
Co-authored-by: Claude Fable 5.1 <no...@anthropic.com>
Co-authored-by: Amin Ghadersohi <am...@gmail.com>
---
UPDATING.md | 19 +-
docs/admin_docs/configuration/mcp-server.mdx | 14 +-
pyproject.toml | 4 -
requirements/development.txt | 5 -
superset/mcp_service/__main__.py | 4 +-
superset/mcp_service/constants.py | 2 +-
superset/mcp_service/dashboard/schemas.py | 2 +-
superset/mcp_service/mcp_config.py | 23 +-
superset/mcp_service/middleware.py | 208 +++++++-------
superset/mcp_service/semantic_layer/schemas.py | 18 +-
.../{token_utils.py => response_size_utils.py} | 319 +++++++--------------
.../semantic_layer/tool/test_list_metrics.py | 62 ++--
.../mcp_service/test_mcp_tool_registration.py | 2 +-
tests/unit_tests/mcp_service/test_middleware.py | 227 +++++++++------
..._token_utils.py => test_response_size_utils.py} | 265 ++++++++---------
15 files changed, 548 insertions(+), 626 deletions(-)
diff --git a/UPDATING.md b/UPDATING.md
index bc8e4626418..364f3d2c33f 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -24,6 +24,23 @@ assists people when migrating to a new version.
## Next
+### MCP response size guard: byte limit instead of estimated token count
+
+The MCP response-size guard no longer estimates LLM token counts (it
+previously used `tiktoken`'s `cl100k_base` encoding, with a character-based
+fallback). An MCP server has no way to know which client or tokenizer is
+actually consuming a response, so token estimation was replaced with the
+exact serialized UTF-8 byte length of the response, which is deterministic
+and tokenizer-agnostic. This also removes the `tiktoken` dependency
+entirely, including the unannounced network request it could make to
+download its vocabulary on a cold cache.
+
+`MCP_RESPONSE_SIZE_CONFIG["token_limit"]` is renamed to
+`MCP_RESPONSE_SIZE_CONFIG["max_bytes"]`, and its default changes from
+25,000 (estimated tokens) to 50,000 (exact bytes). Any deployment that has
+set `token_limit` in `superset_config.py` must rename the key to `max_bytes`
+and adjust the value for byte semantics.
+
### Scheduled report and alert retry admission
Run `superset db upgrade` before starting workers with this version. The migration
@@ -104,7 +121,7 @@ for the chosen metrics, or explicitly request `include_compatible_dimensions=tru
with `page_size` at most 8 for every scope, including built-in `dataset_id`
requests. Non-embedded requests retain the 500-metric ceiling.
This fixed embedding cap is independent of the operator's
-`MCP_RESPONSE_SIZE_CONFIG['token_limit']` (25,000 by default); it does not guarantee
+`MCP_RESPONSE_SIZE_CONFIG['max_bytes']` (50,000 by default); it does not guarantee
that every payload fits a configured response limit.
### Default Docker image is now batteries-included; the minimal image moves to `-lean`
diff --git a/docs/admin_docs/configuration/mcp-server.mdx b/docs/admin_docs/configuration/mcp-server.mdx
index 78900705c87..d07c06e54c2 100644
--- a/docs/admin_docs/configuration/mcp-server.mdx
+++ b/docs/admin_docs/configuration/mcp-server.mdx
@@ -663,14 +663,14 @@ Limits response sizes to prevent exceeding LLM context windows:
```python
MCP_RESPONSE_SIZE_CONFIG = {
"enabled": True,
- "token_limit": 25000,
+ "max_bytes": 50000,
"warn_threshold_pct": 80,
"max_list_items": 100,
"excluded_tools": [
"health_check",
- "get_chart_preview",
"generate_explore_link",
"open_sql_lab_with_context",
+ "search_tools",
],
}
```
@@ -678,9 +678,9 @@ MCP_RESPONSE_SIZE_CONFIG = {
| Key | Default | Description |
| -------------------- | --------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `enabled` | `True` | Enable response size checking |
-| `token_limit` | `25000` | Maximum estimated token count per response |
+| `max_bytes` | `50000` | Maximum serialized response size, in bytes |
| `warn_threshold_pct` | `80` | Warn when response exceeds this percentage of the limit |
-| `max_list_items` | `100` | Cap on list-field length (e.g. `charts`, `native_filters`) applied to the `get_*_info` tools before falling back to more aggressive truncation. Raised from a hardcoded 30 in earlier versions; the higher default only keeps more data before the same token-budget fallback kicks in, so it's not a breaking change, but tenants that tuned workflows around the old 30-item cap should lower this value explicitly. |
+| `max_list_items` | `100` | Cap on list-field length (e.g. `charts`, `native_filters`) applied to the `get_*_info` tools before falling back to more aggressive truncation. Raised from a hardcoded 30 in earlier versions; the higher default only keeps more data before the same size-budget fallback kicks in, so it's not a breaking change, but tenants that tuned workflows around the old 30-item cap should lower this value explicitly. |
| `excluded_tools` | See above | Tools exempt from size checking (e.g., tools that return URLs, not data) |
Write tools that cannot fit after truncation return a minimal confirmation rather
@@ -855,7 +855,7 @@ Every MCP request passes through a middleware stack before reaching the tool fun
| `ToolResultCompatibilityMiddleware` | Gates structured results and sanitizes errors that escape the inner error handler | Enabled |
| `LoggingMiddleware` | Logs each tool call with user, parameters, and duration | Enabled |
| `GlobalErrorHandlerMiddleware` | Catches unhandled exceptions and sanitizes sensitive data before it reaches the client | Enabled |
-| `ResponseSizeGuardMiddleware` | Estimates token count, warns at 80% of limit, blocks at limit | Enabled (configurable via `MCP_RESPONSE_SIZE_CONFIG`) |
+| `ResponseSizeGuardMiddleware` | Measures serialized response size, warns at 80% of limit, blocks at limit | Enabled (configurable via `MCP_RESPONSE_SIZE_CONFIG`) |
| `ResponseCachingMiddleware` | Caches read-heavy tool responses (in-memory or Redis) | Disabled (enable via `MCP_CACHE_CONFIG`) |
Additional middleware classes (`RateLimitMiddleware`, `FieldPermissionsMiddleware`, `PrivateToolMiddleware`) are implemented in `superset/mcp_service/middleware.py` but are not added to the default pipeline. They are available for operators who want to layer them in via a custom startup path.
@@ -987,9 +987,9 @@ Every tool with a side effect (create/update/delete/execute) is always excluded
### Response too large
-- If a tool call returns an error about exceeding token limits, the response size guard is blocking an oversized result
+- If a tool call returns an error about exceeding size limits, the response size guard is blocking an oversized result
- Reduce `page_size` or `limit` parameters, use `select_columns` to exclude large fields, or add filters to narrow results
-- To adjust the threshold, change `token_limit` in `MCP_RESPONSE_SIZE_CONFIG`
+- To adjust the threshold, change `max_bytes` in `MCP_RESPONSE_SIZE_CONFIG`
- To disable the guard entirely, set `MCP_RESPONSE_SIZE_CONFIG = {"enabled": False}`
---
diff --git a/pyproject.toml b/pyproject.toml
index 44c95d23f18..137d3805677 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -202,10 +202,6 @@ fastmcp = [
# the conflict (#43876). Capped to match what fastmcp 3.x actually
# supports until this extra moves to fastmcp 4.x, which accepts mcp 2.x.
"mcp>=1.29.1,<2.0",
- # tiktoken backs the response-size-guard token estimator. Without
- # it, the middleware falls back to a coarser character-based
- # heuristic that under-counts JSON-heavy MCP responses.
- "tiktoken>=0.14.0,<1.0",
]
# sqlalchemy-firebird >=2.0.0 unconditionally requires SQLAlchemy 2.0 on
# Python >=3.8 (which covers Superset's >=3.11 floor), with no dual-compat
diff --git a/requirements/development.txt b/requirements/development.txt
index f12f69fdad0..238846fb24b 100644
--- a/requirements/development.txt
+++ b/requirements/development.txt
@@ -964,8 +964,6 @@ referencing==0.36.2
# jsonschema
# jsonschema-path
# jsonschema-specifications
-regex==2026.4.4
- # via tiktoken
requests==2.33.0
# via
# -c requirements/base-constraint.txt
@@ -979,7 +977,6 @@ requests==2.33.0
# requests-cache
# requests-oauthlib
# shillelagh
- # tiktoken
# trino
requests-cache==1.2.1
# via
@@ -1124,8 +1121,6 @@ tabulate==0.10.0
# apache-superset
testcontainers==4.15.0
# via -r requirements/development.in
-tiktoken==0.14.0
- # via apache-superset
tomli-w==1.2.0
# via apache-superset-extensions-cli
tomlkit==0.13.3
diff --git a/superset/mcp_service/__main__.py b/superset/mcp_service/__main__.py
index b67ed5f81fe..230d990ad3a 100644
--- a/superset/mcp_service/__main__.py
+++ b/superset/mcp_service/__main__.py
@@ -87,8 +87,8 @@ def _add_default_middlewares(
# Response size guard is innermost (added last), then response caching.
if size_guard := create_response_size_guard_middleware():
mcp.add_middleware(size_guard)
- limit = size_guard.token_limit
- sys.stderr.write(f"[MCP] Response size guard enabled (token_limit={limit})\n")
+ limit = size_guard.max_bytes
+ sys.stderr.write(f"[MCP] Response size guard enabled (max_bytes={limit})\n")
if caching_middleware := create_response_caching_middleware():
mcp.add_middleware(caching_middleware)
diff --git a/superset/mcp_service/constants.py b/superset/mcp_service/constants.py
index 407bca2e35b..4e6b8a9c240 100644
--- a/superset/mcp_service/constants.py
+++ b/superset/mcp_service/constants.py
@@ -26,7 +26,7 @@ DEFAULT_PAGE_SIZE = 10 # Default number of items per page
MAX_PAGE_SIZE = 100 # Maximum allowed page_size to prevent oversized responses
# Response size guard defaults
-DEFAULT_TOKEN_LIMIT = 25_000 # ~25k tokens prevents overwhelming LLM context windows
+DEFAULT_MAX_RESPONSE_BYTES = 50_000 # ~50KB preserves the former 25K-token guard
DEFAULT_WARN_THRESHOLD_PCT = 80 # Log warnings above 80% of limit
# Phase 2 list-field truncation cap; matches MAX_PAGE_SIZE
DEFAULT_MAX_LIST_ITEMS = 100
diff --git a/superset/mcp_service/dashboard/schemas.py b/superset/mcp_service/dashboard/schemas.py
index 3cee9b6bc51..498f3c9a215 100644
--- a/superset/mcp_service/dashboard/schemas.py
+++ b/superset/mcp_service/dashboard/schemas.py
@@ -460,7 +460,7 @@ class DashboardInfo(BaseModel):
description=(
"Charts on this dashboard. May be capped below chart_count "
"(cap: MCP_RESPONSE_SIZE_CONFIG['max_list_items']) when the full "
- "response would exceed the token budget. "
+ "response would exceed the size budget. "
"Compare len(charts) to chart_count to detect this. For "
"dashboards with more charts than the cap, call list_charts "
"with filters=[{'col': 'dashboards', 'opr': 'eq', "
diff --git a/superset/mcp_service/mcp_config.py b/superset/mcp_service/mcp_config.py
index 28feed30afc..502ffad1506 100644
--- a/superset/mcp_service/mcp_config.py
+++ b/superset/mcp_service/mcp_config.py
@@ -28,7 +28,7 @@ from superset.constants import CHANGE_ME_GUEST_TOKEN_JWT_SECRET
from superset.mcp_service.composite_token_verifier import CompositeTokenVerifier
from superset.mcp_service.constants import (
DEFAULT_MAX_LIST_ITEMS,
- DEFAULT_TOKEN_LIMIT,
+ DEFAULT_MAX_RESPONSE_BYTES,
DEFAULT_WARN_THRESHOLD_PCT,
)
from superset.mcp_service.guest_token_verifier import GuestTokenVerifier
@@ -399,12 +399,13 @@ MCP_CACHE_CONFIG: dict[str, Any] = {
# Overview:
# ---------
# The Response Size Guard prevents oversized responses from overwhelming LLM
-# clients (e.g., Claude Desktop). When a tool response exceeds the token limit,
+# clients (e.g., Claude Desktop). When a tool response exceeds the byte limit,
# it returns a helpful error with suggestions for reducing the response size.
#
# How it works:
# -------------
-# 1. After a tool executes, the middleware estimates the response's token count
+# 1. After a tool executes, the middleware measures the response's serialized
+# UTF-8 byte size
# 2. If the response exceeds the configured limit, it blocks the response
# 3. Instead, it returns an error message with smart suggestions:
# - Reduce page_size/limit
@@ -415,24 +416,26 @@ MCP_CACHE_CONFIG: dict[str, Any] = {
# Configuration:
# --------------
# - enabled: Toggle the guard on/off (default: True)
-# - token_limit: Maximum estimated tokens per response (default: 25,000)
+# - max_bytes: Maximum serialized response size in bytes (default: 50,000)
# - excluded_tools: Tools to skip checking (e.g., streaming tools)
# - warn_threshold_pct: Log warnings above this % of limit (default: 80%)
# - max_list_items: Cap applied to list fields (e.g. ``charts``,
# ``native_filters``) during Phase 2 of dynamic truncation for the "info"
# tools (get_chart_info, get_dataset_info, get_dashboard_info,
-# get_instance_info) when a response exceeds token_limit (default: 100).
+# get_instance_info) when a response exceeds max_bytes (default: 100).
# Operators with tenants that have unusually large dashboards (hundreds of
# charts/filters) can raise this value to return more complete responses.
#
-# Token Estimation:
-# -----------------
-# Uses character-based heuristic (~3.5 chars per token for JSON).
-# This is intentionally conservative to avoid underestimating.
+# Size Measurement:
+# ------------------
+# Uses the exact serialized UTF-8 byte length of the response. This is not an
+# LLM token estimate: an MCP server cannot know which client (Claude, GPT,
+# Gemini, a local model) or tokenizer is consuming a given response, so byte
+# size is used as a deterministic, tokenizer-agnostic proxy for response size.
# =============================================================================
MCP_RESPONSE_SIZE_CONFIG: dict[str, Any] = {
"enabled": True, # Enabled by default to protect LLM clients
- "token_limit": DEFAULT_TOKEN_LIMIT,
+ "max_bytes": DEFAULT_MAX_RESPONSE_BYTES,
"warn_threshold_pct": DEFAULT_WARN_THRESHOLD_PCT,
"max_list_items": DEFAULT_MAX_LIST_ITEMS,
"excluded_tools": [ # Tools to skip size checking
diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py
index beaefccc559..59893fd9fdf 100644
--- a/superset/mcp_service/middleware.py
+++ b/superset/mcp_service/middleware.py
@@ -50,17 +50,18 @@ from superset.mcp_service.auth import (
)
from superset.mcp_service.constants import (
DEFAULT_MAX_LIST_ITEMS,
- DEFAULT_TOKEN_LIMIT,
+ DEFAULT_MAX_RESPONSE_BYTES,
DEFAULT_WARN_THRESHOLD_PCT,
)
-from superset.mcp_service.utils.token_utils import (
+from superset.mcp_service.utils.response_size_utils import (
COMMITTED_WRITE_SPECS,
COMMITTED_WRITE_TOOLS,
CommittedWriteSpec,
DATA_QUERY_TOOLS,
- estimate_response_tokens,
format_size_limit_error,
+ get_response_size_bytes,
INFO_TOOLS,
+ string_clip_chars,
STRING_FIELD_TRUNCATION_TOOLS,
truncate_oversized_response,
truncate_query_result,
@@ -80,6 +81,7 @@ _METRIC_TOOL_NAME_RE = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.\-]{0,127}")
# confirmation (see ``_shrink_minimal_response``). Generous enough to keep a
# chart name or a short error readable, small enough that the whole confirmation
# stays bounded no matter how large the fields were in the original payload.
+# Small byte budgets lower it further (see ``string_clip_chars``).
_MINIMAL_FIELD_CHARS = 200
# Bound both list overhead and total string content in write confirmations.
@@ -1157,7 +1159,7 @@ class ResponseSizeGuardMiddleware(Middleware):
"""
Middleware that prevents oversized responses from overwhelming LLM clients.
- When a tool response exceeds the configured token limit, this middleware
+ When a tool response exceeds the configured byte limit, this middleware
intercepts it and returns a helpful error message with suggestions for
reducing the response size.
@@ -1166,7 +1168,7 @@ class ResponseSizeGuardMiddleware(Middleware):
Configuration via MCP_RESPONSE_SIZE_CONFIG in superset_config.py:
- enabled: Toggle the guard on/off (default: True)
- - token_limit: Maximum estimated tokens per response (default: 25,000)
+ - max_bytes: Maximum serialized response size in bytes (default: 50,000)
- warn_threshold_pct: Log warnings above this % of limit (default: 80%)
- max_list_items: Cap for list fields during dynamic truncation (default: 100)
- excluded_tools: Tools to skip checking
@@ -1174,14 +1176,14 @@ class ResponseSizeGuardMiddleware(Middleware):
def __init__(
self,
- token_limit: int = DEFAULT_TOKEN_LIMIT,
+ max_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
warn_threshold_pct: int = DEFAULT_WARN_THRESHOLD_PCT,
excluded_tools: list[str] | str | None = None,
max_list_items: int = DEFAULT_MAX_LIST_ITEMS,
) -> None:
- self.token_limit = token_limit
+ self.max_bytes = max_bytes
self.warn_threshold_pct = warn_threshold_pct
- self.warn_threshold = int(token_limit * warn_threshold_pct / 100)
+ self.warn_threshold = int(max_bytes * warn_threshold_pct / 100)
if isinstance(excluded_tools, str):
excluded_tools = [excluded_tools]
self.excluded_tools = set(excluded_tools or [])
@@ -1244,7 +1246,7 @@ class ResponseSizeGuardMiddleware(Middleware):
self,
tool_name: str,
response: Any,
- estimated_tokens: int,
+ actual_bytes: int,
protected_keys: frozenset[str] = frozenset(),
) -> Any | None:
"""Attempt to dynamically truncate an info tool response to fit the limit.
@@ -1286,7 +1288,7 @@ class ResponseSizeGuardMiddleware(Middleware):
try:
truncated, was_truncated, notes = truncate_oversized_response(
truncation_target,
- self.token_limit,
+ self.max_bytes,
max_list_items=self.max_list_items,
protected_keys=protected_keys,
)
@@ -1302,16 +1304,16 @@ class ResponseSizeGuardMiddleware(Middleware):
if not was_truncated:
return None
- truncated_tokens = estimate_response_tokens(truncated)
- if truncated_tokens > self.token_limit:
+ truncated_bytes = get_response_size_bytes(truncated)
+ if truncated_bytes > self.max_bytes:
return None
logger.warning(
- "Response for %s truncated from ~%d to ~%d tokens (limit: %d). Fields: %s",
+ "Response for %s truncated from %d to %d bytes (limit: %d). Fields: %s",
tool_name,
- estimated_tokens,
- truncated_tokens,
- self.token_limit,
+ actual_bytes,
+ truncated_bytes,
+ self.max_bytes,
"; ".join(notes),
)
@@ -1326,9 +1328,9 @@ class ResponseSizeGuardMiddleware(Middleware):
referrer=None,
curated_payload={
"tool": tool_name,
- "original_tokens": estimated_tokens,
- "truncated_tokens": truncated_tokens,
- "token_limit": self.token_limit,
+ "original_bytes": actual_bytes,
+ "truncated_bytes": truncated_bytes,
+ "max_bytes": self.max_bytes,
"truncation_notes": notes,
},
)
@@ -1349,7 +1351,7 @@ class ResponseSizeGuardMiddleware(Middleware):
self,
tool_name: str,
response: Any,
- estimated_tokens: int,
+ actual_bytes: int,
) -> Any | None:
"""Attempt to truncate a data-query tool response by dropping tail rows.
@@ -1371,7 +1373,7 @@ class ResponseSizeGuardMiddleware(Middleware):
try:
truncated, was_truncated, notes = truncate_query_result(
- truncation_target, self.token_limit, tool_name=tool_name
+ truncation_target, self.max_bytes, tool_name=tool_name
)
except Exception as trunc_error: # noqa: BLE001
logger.warning(
@@ -1389,16 +1391,16 @@ class ResponseSizeGuardMiddleware(Middleware):
# response back under the limit (e.g. a single row/scalar field
# alone exceeds it), fall back to the hard size-limit error instead
# of shipping an over-budget response.
- truncated_tokens = estimate_response_tokens(truncated)
- if truncated_tokens > self.token_limit:
+ truncated_bytes = get_response_size_bytes(truncated)
+ if truncated_bytes > self.max_bytes:
return None
logger.warning(
- "Query result for %s truncated from ~%d to ~%d tokens (limit: %d). %s",
+ "Query result for %s truncated from %d to %d bytes (limit: %d). %s",
tool_name,
- estimated_tokens,
- truncated_tokens,
- self.token_limit,
+ actual_bytes,
+ truncated_bytes,
+ self.max_bytes,
"; ".join(notes),
)
@@ -1413,9 +1415,9 @@ class ResponseSizeGuardMiddleware(Middleware):
referrer=None,
curated_payload={
"tool": tool_name,
- "original_tokens": estimated_tokens,
- "truncated_tokens": truncated_tokens,
- "token_limit": self.token_limit,
+ "original_bytes": actual_bytes,
+ "truncated_bytes": truncated_bytes,
+ "max_bytes": self.max_bytes,
"truncation_notes": notes,
},
)
@@ -1431,7 +1433,7 @@ class ResponseSizeGuardMiddleware(Middleware):
self,
tool_name: str,
response: Any,
- estimated_tokens: int,
+ actual_bytes: int,
field: str,
) -> Any | None:
"""Attempt to truncate a response by bisecting one oversized string field.
@@ -1454,7 +1456,7 @@ class ResponseSizeGuardMiddleware(Middleware):
try:
truncated, was_truncated, notes = truncate_string_field_response(
- truncation_target, self.token_limit, field
+ truncation_target, self.max_bytes, field
)
except Exception as trunc_error: # noqa: BLE001
logger.warning(
@@ -1468,16 +1470,16 @@ class ResponseSizeGuardMiddleware(Middleware):
if not was_truncated:
return None
- truncated_tokens = estimate_response_tokens(truncated)
- if truncated_tokens > self.token_limit:
+ truncated_bytes = get_response_size_bytes(truncated)
+ if truncated_bytes > self.max_bytes:
return None
logger.warning(
- "Response for %s truncated from ~%d to ~%d tokens (limit: %d). %s",
+ "Response for %s truncated from %d to %d bytes (limit: %d). %s",
tool_name,
- estimated_tokens,
- truncated_tokens,
- self.token_limit,
+ actual_bytes,
+ truncated_bytes,
+ self.max_bytes,
"; ".join(notes),
)
@@ -1492,9 +1494,9 @@ class ResponseSizeGuardMiddleware(Middleware):
referrer=None,
curated_payload={
"tool": tool_name,
- "original_tokens": estimated_tokens,
- "truncated_tokens": truncated_tokens,
- "token_limit": self.token_limit,
+ "original_bytes": actual_bytes,
+ "truncated_bytes": truncated_bytes,
+ "max_bytes": self.max_bytes,
"truncation_notes": notes,
},
)
@@ -1510,7 +1512,7 @@ class ResponseSizeGuardMiddleware(Middleware):
self,
tool_name: str,
response: Any,
- estimated_tokens: int,
+ actual_bytes: int,
) -> Any:
"""Build a guaranteed-small success response for a committed write.
@@ -1541,11 +1543,11 @@ class ResponseSizeGuardMiddleware(Middleware):
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 "
+ "truncation (%d bytes, limit %d); returning a minimal write "
"confirmation instead of blocking a completed write.",
tool_name,
- estimated_tokens,
- self.token_limit,
+ actual_bytes,
+ self.max_bytes,
)
try:
user_id = get_user_id()
@@ -1558,8 +1560,8 @@ class ResponseSizeGuardMiddleware(Middleware):
referrer=None,
curated_payload={
"tool": tool_name,
- "original_tokens": estimated_tokens,
- "token_limit": self.token_limit,
+ "original_bytes": actual_bytes,
+ "max_bytes": self.max_bytes,
"truncation_notes": truncation_notes,
},
)
@@ -1617,7 +1619,7 @@ class ResponseSizeGuardMiddleware(Middleware):
def _shrink_minimal_response(
self, minimal: dict[str, Any], spec: CommittedWriteSpec
) -> None:
- """Force ``minimal`` under the token limit, degrading fields in place.
+ """Force ``minimal`` under the byte 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
@@ -1647,21 +1649,25 @@ class ResponseSizeGuardMiddleware(Middleware):
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
+ With an extremely small ``max_bytes`` even the fully clipped form
can exceed it. Returning it anyway is deliberate: this path exists so
a completed write is never reported as a failure, and there is
nothing further to give up without losing that confirmation.
"""
- if _fits(minimal, self.token_limit):
+ if _fits(minimal, self.max_bytes):
return
+ # Clip proportionally under small budgets so the clipped fields do
+ # not by themselves exceed the limit they are being shrunk to fit.
+ max_chars = string_clip_chars(self.max_bytes, _MINIMAL_FIELD_CHARS)
+
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])
+ key: _clip_string(value[key], max_chars)
for key in _MINIMAL_IDENTITY_FIELDS
if key in value
}
@@ -1678,26 +1684,30 @@ class ResponseSizeGuardMiddleware(Middleware):
if key in spec.identifying_fields or key.startswith("_"):
continue
current = minimal[key]
- clipped = _clip_error(current) if key == "error" else _clip_string(current)
+ clipped = (
+ _clip_error(current, max_chars)
+ if key == "error"
+ else _clip_string(current, max_chars)
+ )
if clipped is not current:
minimal[key] = clipped
minimal["_truncation_notes"].append(
f"'{key}' was reduced to fit the size limit."
)
- if not _fits(minimal, self.token_limit):
+ if not _fits(minimal, self.max_bytes):
logger.warning(
- "Minimal write confirmation still estimates over the token "
+ "Minimal write confirmation still estimates over the byte "
"limit (%d) after full reduction; returning it anyway rather "
"than reporting a completed write as a failure.",
- self.token_limit,
+ self.max_bytes,
)
def _handle_oversized_response(
self,
tool_name: str,
response: Any,
- estimated_tokens: int,
+ actual_bytes: int,
params: dict[str, Any],
) -> Any:
"""Attempt truncation for known tool categories; block everything else.
@@ -1725,7 +1735,7 @@ class ResponseSizeGuardMiddleware(Middleware):
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
+ tool_name, response, actual_bytes, protected_keys=protected_keys
)
if truncated is not None:
return truncated
@@ -1733,7 +1743,7 @@ class ResponseSizeGuardMiddleware(Middleware):
# Data-query tools: row-level truncation.
if tool_name in DATA_QUERY_TOOLS:
truncated = self._try_truncate_data_query_response(
- tool_name, response, estimated_tokens
+ tool_name, response, actual_bytes
)
if truncated is not None:
return truncated
@@ -1744,7 +1754,7 @@ class ResponseSizeGuardMiddleware(Middleware):
truncated = self._try_truncate_string_field_response(
tool_name,
response,
- estimated_tokens,
+ actual_bytes,
STRING_FIELD_TRUNCATION_TOOLS[tool_name],
)
if truncated is not None:
@@ -1754,15 +1764,15 @@ class ResponseSizeGuardMiddleware(Middleware):
# The mutation already committed -- never report it as a failed
# call, no matter how badly truncation underperformed.
return self._minimal_committed_write_response(
- tool_name, response, estimated_tokens
+ tool_name, response, actual_bytes
)
# Log the blocked response (user-caused: requested too much data)
logger.warning(
- "Response blocked for %s: ~%d tokens exceeds limit of %d",
+ "Response blocked for %s: %d bytes exceeds limit of %d",
tool_name,
- estimated_tokens,
- self.token_limit,
+ actual_bytes,
+ self.max_bytes,
)
try:
@@ -1776,8 +1786,8 @@ class ResponseSizeGuardMiddleware(Middleware):
referrer=None,
curated_payload={
"tool": tool_name,
- "estimated_tokens": estimated_tokens,
- "token_limit": self.token_limit,
+ "actual_bytes": actual_bytes,
+ "max_bytes": self.max_bytes,
"params": _sanitize_params(params),
},
)
@@ -1788,8 +1798,8 @@ class ResponseSizeGuardMiddleware(Middleware):
format_size_limit_error(
tool_name=tool_name,
params=params,
- estimated_tokens=estimated_tokens,
- token_limit=self.token_limit,
+ actual_bytes=actual_bytes,
+ max_bytes=self.max_bytes,
response=None,
)
)
@@ -1809,41 +1819,31 @@ class ResponseSizeGuardMiddleware(Middleware):
# Execute the tool
response = await call_next(context)
- # When the response is a ToolResult, estimate tokens on the actual
- # payload inside content[0].text rather than on the ToolResult
- # wrapper (which would double-serialize the JSON string).
+ # When the response is a ToolResult, measure the size of the actual
+ # payload inside content[0].text rather than the ToolResult wrapper
+ # (which would double-serialize the JSON string).
extracted = self._extract_payload_from_tool_result(response)
estimation_target = extracted if extracted is not None else response
- try:
- estimated_tokens = estimate_response_tokens(estimation_target)
- except MemoryError as me:
- logger.warning(
- "MemoryError while estimating tokens for %s: %s", tool_name, me
- )
- # Treat as over limit to avoid further serialization
- estimated_tokens = self.token_limit + 1
- except Exception as e: # noqa: BLE001
- logger.warning(
- "Failed to estimate response tokens for %s: %s", tool_name, e
- )
- # Conservative fallback: block rather than risk OOM
- estimated_tokens = self.token_limit + 1
+ # Never raises: a response that cannot be serialized measures as
+ # UNMEASURABLE_RESPONSE_BYTES, which exceeds any limit and so takes
+ # the oversized path below rather than slipping through unmeasured.
+ actual_bytes = get_response_size_bytes(estimation_target)
# Log warning if approaching limit
- if estimated_tokens > self.warn_threshold:
+ if actual_bytes > self.warn_threshold:
logger.warning(
- "Response size warning for %s: ~%d tokens (%.0f%% of %d limit)",
+ "Response size warning for %s: %d bytes (%.0f%% of %d limit)",
tool_name,
- estimated_tokens,
- (estimated_tokens / self.token_limit * 100) if self.token_limit else 0,
- self.token_limit,
+ actual_bytes,
+ (actual_bytes / self.max_bytes * 100) if self.max_bytes else 0,
+ self.max_bytes,
)
- if estimated_tokens > self.token_limit:
+ if actual_bytes > self.max_bytes:
params = getattr(context.message, "arguments", {}) or {}
return self._handle_oversized_response(
- tool_name, response, estimated_tokens, params
+ tool_name, response, actual_bytes, params
)
return response
@@ -1860,7 +1860,7 @@ def _clip_string(value: Any, max_chars: int = _MINIMAL_FIELD_CHARS) -> Any:
return value
-def _clip_error(value: Any) -> Any:
+def _clip_error(value: Any, max_chars: int = _MINIMAL_FIELD_CHARS) -> Any:
"""Bound an ``error`` field of either shape it can arrive in.
Tool responses type ``error`` as a nested model (``ChartGenerationError``),
@@ -1875,24 +1875,22 @@ def _clip_error(value: Any) -> Any:
"""
if isinstance(value, dict):
reduced = {
- key: _clip_string(value[key])
+ key: _clip_string(value[key], max_chars)
for key in _MINIMAL_ERROR_FIELDS
if key in value
}
return value if reduced == value else reduced
- return _clip_string(value)
+ return _clip_string(value, max_chars)
-def _fits(payload: Any, token_limit: int) -> bool:
- """Best-effort check that ``payload`` estimates under ``token_limit``.
+def _fits(payload: Any, max_bytes: int) -> bool:
+ """Check that ``payload`` measures under ``max_bytes``.
- Treats an estimation failure as "does not fit" so callers degrade the
- payload further rather than optimistically returning something oversized.
+ A measurement failure reads as "does not fit" (the helper reports it as
+ ``UNMEASURABLE_RESPONSE_BYTES``), so callers degrade the payload further
+ rather than optimistically returning something oversized.
"""
- try:
- return estimate_response_tokens(payload) <= token_limit
- except Exception: # noqa: BLE001
- return False
+ return get_response_size_bytes(payload) <= max_bytes
def _safe_int_config(config: dict[str, Any], key: str, default: int) -> int:
@@ -1946,7 +1944,7 @@ def create_response_size_guard_middleware() -> ResponseSizeGuardMiddleware | Non
)
middleware = ResponseSizeGuardMiddleware(
- token_limit=_safe_int_config(config, "token_limit", DEFAULT_TOKEN_LIMIT),
+ max_bytes=_safe_int_config(config, "max_bytes", DEFAULT_MAX_RESPONSE_BYTES),
warn_threshold_pct=_safe_int_config(
config, "warn_threshold_pct", DEFAULT_WARN_THRESHOLD_PCT
),
@@ -1955,8 +1953,8 @@ def create_response_size_guard_middleware() -> ResponseSizeGuardMiddleware | Non
)
logger.info(
- "Created ResponseSizeGuardMiddleware with token_limit=%d",
- middleware.token_limit,
+ "Created ResponseSizeGuardMiddleware with max_bytes=%d",
+ middleware.max_bytes,
)
return middleware
diff --git a/superset/mcp_service/semantic_layer/schemas.py b/superset/mcp_service/semantic_layer/schemas.py
index bc54f652f72..7339afbe5fd 100644
--- a/superset/mcp_service/semantic_layer/schemas.py
+++ b/superset/mcp_service/semantic_layer/schemas.py
@@ -89,10 +89,16 @@ class MetricInfo(BaseModel):
# ---------------------------------------------------------------------------
-# Measured with 40 dimensions per metric: 20x40 = 42,336 tokens;
-# 10x40 = 21,186; 8x40 = 16,956 (18,284 fallback), against the default 25,000 limit.
-# The fixed page cap is not a guarantee for every payload or operator token limit.
-EMBEDDED_DIMENSIONS_MAX_PAGE_SIZE: int = 8
+# Measured with get_response_size_bytes on this schema, 40 dimensions per
+# metric: with short descriptive text on every metric and dimension,
+# page_size=4 serializes to ~39 KB, 5 to ~48 KB, 6 to ~58 KB and 8 to ~77 KB
+# (about ~27 / ~34 / ~41 / ~55 KB with names only), against the
+# response-size guard's default MCP_RESPONSE_SIZE_CONFIG['max_bytes'] of
+# 50,000 bytes. 4 is the largest page size with real margin in both variants;
+# 8 (the old cap) already exceeds the default even with no descriptions at
+# all. The fixed page cap is not a guarantee for every payload or
+# operator-configured limit.
+EMBEDDED_DIMENSIONS_MAX_PAGE_SIZE: int = 4
class ListMetricsRequest(BaseModel):
@@ -133,9 +139,9 @@ class ListMetricsRequest(BaseModel):
raise ValueError(
"Embedded compatible dimensions require "
f"page_size <= {EMBEDDED_DIMENSIONS_MAX_PAGE_SIZE}: each "
- "metric's dimension list can consume roughly 1–2k tokens or more, "
+ "metric's dimension list can consume several KB or more, "
"and the MCP response guard uses "
- "MCP_RESPONSE_SIZE_CONFIG['token_limit'] (~25k by default). "
+ "MCP_RESPONSE_SIZE_CONFIG['max_bytes'] (~50k by default). "
"This fixed page cap does not guarantee that every response fits. "
"Reduce page_size or use include_compatible_dimensions=false "
"and get_compatible_dimensions for the chosen metric."
diff --git a/superset/mcp_service/utils/token_utils.py b/superset/mcp_service/utils/response_size_utils.py
similarity index 81%
rename from superset/mcp_service/utils/token_utils.py
rename to superset/mcp_service/utils/response_size_utils.py
index 332cbdc7769..96907b8ef7f 100644
--- a/superset/mcp_service/utils/token_utils.py
+++ b/superset/mcp_service/utils/response_size_utils.py
@@ -16,36 +16,23 @@
# under the License.
"""
-Token counting and response size utilities for MCP service.
+Response size utilities for the MCP service.
-This module provides utilities to estimate token counts and generate smart
+This module provides utilities to measure response size and generate smart
suggestions when responses exceed configured limits. This prevents large
responses from overwhelming LLM clients like Claude Desktop.
-Token counting strategy:
-
-1. ``tiktoken`` with the ``cl100k_base`` encoding when the package is
- installed (it is shipped as part of the ``fastmcp`` extra). This is a
- real BPE tokenizer trained on a similar vocabulary to Claude's; for
- English and JSON-heavy MCP payloads it tracks Claude's tokenizer
- within roughly ±10%, which is far more accurate than the legacy
- character heuristic.
-2. A character-based fallback (``CHARS_PER_TOKEN``) when tiktoken is not
- importable. The fallback uses a slightly more conservative ratio than
- before (3.0 chars/token instead of 3.5) so that JSON-heavy responses
- are not under-counted, which previously let oversized payloads slip
- past the response-size guard.
-
-The exact-Claude tokenizer is only available via Anthropic's network
-``count_tokens`` API; calling it from a synchronous middleware on every
-tool result is too slow and adds an external dependency on every
-response. ``tiktoken`` is the closest approximation we can ship without
-that risk.
+Response size is measured as the exact serialized UTF-8 byte length, not an
+estimated LLM token count. An MCP server has no way to know which client
+(Claude, GPT, Gemini, a local model) or tokenizer is actually consuming a
+given response, so any token estimate would be a guess about a vocabulary
+that cannot be known. Byte length is deterministic and tokenizer-agnostic.
"""
from __future__ import annotations
import logging
+import sys
from typing import Any, Dict, List, NamedTuple, Union
from pydantic import BaseModel
@@ -58,128 +45,20 @@ logger = logging.getLogger(__name__)
# Type alias for MCP tool responses (Pydantic models, dicts, lists, strings, bytes)
ToolResponse: TypeAlias = Union[BaseModel, Dict[str, Any], List[Any], str, bytes]
-# Fallback character-to-token ratio used when tiktoken is unavailable.
-# 3.0 is conservative for JSON content (the previous 3.5 under-counted
-# JSON-heavy payloads relative to Claude's actual tokenizer, which let
-# oversized responses slip past the response-size guard).
-CHARS_PER_TOKEN = 3.0
-
-# Encoding used when tiktoken is available. cl100k_base is OpenAI's
-# tokenizer for GPT-3.5/4; it is BPE-based with a vocabulary similar to
-# Claude's and tracks Claude's token counts within roughly ±10% for
-# English and JSON-heavy MCP responses.
-_TIKTOKEN_ENCODING_NAME = "cl100k_base"
-
-
-def _load_tiktoken_encoding() -> Any:
- """Return a tiktoken encoding instance, or None if tiktoken is unavailable.
-
- Imported lazily so the module can be used in environments without
- tiktoken installed. The encoding is small (~1 MB) so we cache it on
- first use.
- """
- try:
- import tiktoken
- except ImportError:
- logger.info(
- "tiktoken not installed; falling back to char-based token "
- "estimation (CHARS_PER_TOKEN=%s). Install the 'fastmcp' extra "
- "for accurate counts.",
- CHARS_PER_TOKEN,
- )
- return None
-
- try:
- return tiktoken.get_encoding(_TIKTOKEN_ENCODING_NAME)
- except (KeyError, ValueError) as exc:
- # tiktoken installed but the requested encoding is missing — this
- # only happens on partial installs. Treat as no tokenizer rather
- # than crashing on every tool call.
- logger.warning(
- "tiktoken encoding '%s' unavailable: %s; falling back to "
- "char-based token estimation",
- _TIKTOKEN_ENCODING_NAME,
- exc,
- )
- return None
-
-
-# Cached encoding instance (None if tiktoken not importable).
-_ENCODING = _load_tiktoken_encoding()
-
-
-def estimate_token_count(text: str | bytes) -> int:
- """
- Estimate the token count for a given text.
-
- Uses tiktoken's ``cl100k_base`` encoding when available for
- Claude-aligned accuracy (within ~10%), falling back to a
- character-based heuristic otherwise.
-
- Args:
- text: The text to estimate tokens for (string or bytes)
-
- Returns:
- Estimated number of tokens
- """
- if isinstance(text, bytes):
- text = text.decode("utf-8", errors="replace")
-
- if not text:
- return 0
-
- if _ENCODING is not None:
- try:
- return len(_ENCODING.encode(text))
- except (ValueError, UnicodeError) as exc:
- # Defensive: if tiktoken chokes on a specific input, fall
- # back to the char heuristic for this call rather than
- # raising — the response size guard must never fail-open.
- logger.warning("tiktoken encode failed (%s); using fallback", exc)
-
- return max(1, int(len(text) / CHARS_PER_TOKEN))
-
-
-def estimate_response_tokens(response: ToolResponse) -> int:
- """
- Estimate token count for an MCP tool response.
-
- Handles various response types including Pydantic models, dicts, and strings.
-
- Args:
- response: The response object to estimate
-
- Returns:
- Estimated number of tokens
- """
- try:
- from superset.utils import json
-
- # Convert response to JSON string for accurate estimation
- if hasattr(response, "model_dump"):
- # Pydantic model
- response_str = json.dumps(response.model_dump())
- elif isinstance(response, (dict, list)):
- response_str = json.dumps(response)
- elif isinstance(response, bytes):
- # Delegate to estimate_token_count which handles decoding safely
- return estimate_token_count(response)
- elif isinstance(response, str):
- response_str = response
- else:
- response_str = str(response)
-
- return estimate_token_count(response_str)
- except Exception as e: # noqa: BLE001
- logger.warning("Failed to estimate response tokens: %s", e)
- # Return a high estimate to be safe (conservative fallback)
- return 100000
+# Reported by ``get_response_size_bytes`` when a response cannot be serialized
+# and so cannot be measured. It exceeds any configurable ``max_bytes``, so an
+# unmeasurable response is always treated as oversized rather than allowed
+# through on a guess that happens to sit under the operator's limit.
+UNMEASURABLE_RESPONSE_BYTES = sys.maxsize
def get_response_size_bytes(response: ToolResponse) -> int:
"""
Get the size of a response in bytes.
+ Never raises: a response that cannot be serialized measures as
+ ``UNMEASURABLE_RESPONSE_BYTES`` so that callers treat it as oversized.
+
Args:
response: The response object
@@ -203,9 +82,9 @@ def get_response_size_bytes(response: ToolResponse) -> int:
return len(response_str.encode("utf-8"))
except Exception as e: # noqa: BLE001
logger.warning("Failed to get response size: %s", e)
- # Return a conservative large value to avoid allowing oversized responses
- # to bypass size checks (returning 0 would underestimate)
- return 1_000_000 # 1MB fallback
+ # A fixed fallback would read as "fits" under any limit configured
+ # above it, so report the size as unknown-and-oversized instead.
+ return UNMEASURABLE_RESPONSE_BYTES
def extract_query_params(params: Dict[str, Any] | None) -> Dict[str, Any]:
@@ -244,8 +123,8 @@ def extract_query_params(params: Dict[str, Any] | None) -> Dict[str, Any]:
def generate_size_reduction_suggestions(
tool_name: str,
params: Dict[str, Any] | None,
- estimated_tokens: int,
- token_limit: int,
+ actual_bytes: int,
+ max_bytes: int,
response: ToolResponse | None = None,
) -> List[str]:
"""
@@ -256,8 +135,8 @@ def generate_size_reduction_suggestions(
Args:
tool_name: Name of the MCP tool
params: The tool parameters
- estimated_tokens: Estimated token count of the response
- token_limit: Configured token limit
+ actual_bytes: Actual size of the response in bytes
+ max_bytes: Configured byte limit
response: Optional response object for additional analysis
Returns:
@@ -265,10 +144,8 @@ def generate_size_reduction_suggestions(
"""
suggestions = []
query_params = extract_query_params(params)
- reduction_needed = estimated_tokens - token_limit
- reduction_pct = (
- int((reduction_needed / estimated_tokens) * 100) if estimated_tokens else 0
- )
+ reduction_needed = actual_bytes - max_bytes
+ reduction_pct = int((reduction_needed / actual_bytes) * 100) if actual_bytes else 0
# Suggestion 1: Reduce page_size or limit
raw_page_size = query_params.get("page_size") or query_params.get("limit")
@@ -280,9 +157,7 @@ def generate_size_reduction_suggestions(
# Calculate suggested new limit based on reduction needed
suggested_limit = max(
1,
- int(current_page_size * (token_limit / estimated_tokens))
- if estimated_tokens
- else 1,
+ int(current_page_size * (max_bytes / actual_bytes)) if actual_bytes else 1,
)
suggestions.append(
f"Reduce page_size/limit from {current_page_size} to {suggested_limit} "
@@ -478,7 +353,7 @@ INFO_TOOLS = frozenset(
)
# Data-query tools that return tabular results. When the response exceeds the
-# token limit, these tools are truncated by dropping tail rows rather than
+# byte limit, these tools are truncated by dropping tail rows rather than
# raising a hard ToolError. A truncation note is appended so the caller
# knows the result is partial and how to narrow the query.
DATA_QUERY_TOOLS = frozenset(
@@ -626,10 +501,30 @@ _DATA_ROW_FIELDS = ("rows", "data")
# Maximum character length for string fields before truncation
_MAX_STRING_CHARS = 500
+# Floor for the budget-derived clip length (see ``string_clip_chars``), so a
+# clipped string stays recognizable even under a very small byte budget.
+_MIN_STRING_CHARS = 32
+# Above `_CLIP_BUDGET_DIVISOR * ceiling` bytes, string_clip_chars leaves the
+# ceiling unchanged; smaller budgets shrink the clip length proportionally.
+_CLIP_BUDGET_DIVISOR = 4
# Maximum keys to keep when summarizing large dict fields
_MAX_DICT_KEYS = 20
+def string_clip_chars(max_bytes: int, ceiling: int = _MAX_STRING_CHARS) -> int:
+ """Derive the per-string clip length from the response byte budget.
+
+ A fixed clip length only works while the budget dwarfs it: a single
+ clipped field plus its truncation marker must not exhaust the whole
+ budget by itself, or string truncation can never bring a response under
+ the limit. Above ``_CLIP_BUDGET_DIVISOR * ceiling`` bytes the ceiling
+ applies unchanged (the default 50 KB budget keeps the full
+ ``_MAX_STRING_CHARS``); smaller budgets shrink the clip length
+ proportionally, down to ``_MIN_STRING_CHARS``.
+ """
+ return min(ceiling, max(_MIN_STRING_CHARS, max_bytes // _CLIP_BUDGET_DIVISOR))
+
+
def _truncate_strings(
data: Dict[str, Any], notes: List[str], max_chars: int = _MAX_STRING_CHARS
) -> bool:
@@ -767,22 +662,20 @@ def _replace_collections_with_summaries(
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 _is_under_limit(data: Dict[str, Any], max_bytes: int) -> bool:
+ """Check if the serialized data fits within the byte limit."""
+ return get_response_size_bytes(data) <= max_bytes
def truncate_oversized_response(
response: ToolResponse,
- token_limit: int,
+ max_bytes: int,
# Configurable via MCP_RESPONSE_SIZE_CONFIG["max_list_items"]
max_list_items: int = DEFAULT_MAX_LIST_ITEMS,
protected_keys: frozenset[str] = frozenset(),
) -> tuple[ToolResponse, bool, list[str]]:
"""
- Dynamically truncate large fields in a response to fit within the token limit.
+ Dynamically truncate large fields in a response to fit within the byte limit.
Applies five progressive phases of truncation:
1. Truncate long top-level string fields
@@ -793,14 +686,14 @@ def truncate_oversized_response(
Args:
response: The tool response (Pydantic model, dict, or other).
- token_limit: Maximum estimated tokens allowed.
+ max_bytes: Maximum serialized response size in bytes.
max_list_items: Maximum items to keep in list fields during Phase 2.
protected_keys: Top-level keys the destructive phases (4 and 5) must
never summarize or clear, even if the response is still over
budget afterward. Used for committed-write tools so their
identifying field (e.g. ``chart``) always survives. Callers that
pass this must be prepared for the returned response to still
- exceed ``token_limit``.
+ exceed ``max_bytes``.
Returns:
A tuple of (possibly-truncated response, was_truncated, list of notes).
@@ -816,27 +709,30 @@ def truncate_oversized_response(
return response, False, notes
was_truncated = False
+ # Clip length scales down with small budgets so the string phases can
+ # actually converge instead of leaving one clipped field over the limit.
+ max_chars = string_clip_chars(max_bytes)
# Phase 1: Truncate long string fields
- was_truncated |= _truncate_strings(data, notes)
- if _is_under_limit(data, token_limit):
+ was_truncated |= _truncate_strings(data, notes, max_chars)
+ if _is_under_limit(data, max_bytes):
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):
+ if _is_under_limit(data, max_bytes):
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):
+ was_truncated |= _truncate_strings_recursive(data, notes, max_chars)
+ if _is_under_limit(data, max_bytes):
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, protected_keys=protected_keys)
- if _is_under_limit(data, token_limit):
+ if _is_under_limit(data, max_bytes):
return data, was_truncated, notes
# Phase 5: Nuclear — replace all collections with empty values
@@ -851,7 +747,7 @@ def _bisect_row_limit(
data: Dict[str, Any],
row_field: str,
original_rows: List[Any],
- token_limit: int,
+ max_bytes: int,
) -> int:
"""Binary-search for the largest row prefix that keeps data under limit.
@@ -859,13 +755,11 @@ def _bisect_row_limit(
kept count on return. Returns the number of rows kept (>= 1 if the
original list was non-empty).
"""
- from superset.utils import json as utils_json
-
lo, hi = 0, len(original_rows)
while lo < hi:
mid = (lo + hi + 1) // 2
data[row_field] = original_rows[:mid]
- if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ if get_response_size_bytes(data) <= max_bytes:
lo = mid
else:
hi = mid - 1
@@ -884,7 +778,7 @@ def _bisect_string_length(
data: Dict[str, Any],
field: str,
original_value: str,
- token_limit: int,
+ max_bytes: int,
suffix: str = "",
) -> int:
"""Binary-search for the largest string prefix that keeps data under limit.
@@ -895,13 +789,11 @@ def _bisect_string_length(
pushing the payload back over the limit afterward. It is only appended
when the value is actually shortened.
"""
- from superset.utils import json as utils_json
-
lo, hi = 0, len(original_value)
while lo < hi:
mid = (lo + hi + 1) // 2
data[field] = original_value[:mid] + suffix
- if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ if get_response_size_bytes(data) <= max_bytes:
lo = mid
else:
hi = mid - 1
@@ -931,7 +823,7 @@ _DEFAULT_ROW_LIMIT_ADVICE = (
def _truncate_rows_field(
data: Dict[str, Any],
row_field: str,
- token_limit: int,
+ max_bytes: int,
advice: str,
) -> list[str] | None:
"""Try to bisect ``data[row_field]`` down to fit the limit.
@@ -952,17 +844,17 @@ def _truncate_rows_field(
# longest the real note can ever be (kept <= original_count).
placeholder_note = (
f"Result truncated: {original_count} of {original_count} rows returned "
- f"(limit ~{token_limit:,} tokens). {advice}"
+ f"(limit ~{max_bytes:,} bytes). {advice}"
)
data["_response_truncated"] = True
data["_truncation_notes"] = [placeholder_note]
- kept = _bisect_row_limit(data, row_field, original_rows, token_limit)
+ kept = _bisect_row_limit(data, row_field, original_rows, max_bytes)
if kept < original_count:
notes = [
f"Result truncated: {kept} of {original_count} rows returned "
- f"(limit ~{token_limit:,} tokens). {advice}"
+ f"(limit ~{max_bytes:,} bytes). {advice}"
]
data["_truncation_notes"] = notes
if "row_count" in data:
@@ -977,11 +869,9 @@ def _truncate_rows_field(
def _truncate_chart_query_results(
- data: Dict[str, Any], token_limit: int, advice: str
+ data: Dict[str, Any], max_bytes: int, advice: str
) -> list[str] | None:
"""Apply one response-wide row cap to every result of a multi-query chart."""
- from superset.utils import json as utils_json
-
query_results = data.get("query_results")
if not isinstance(query_results, list) or not query_results:
return None
@@ -1001,7 +891,7 @@ def _truncate_chart_query_results(
data["_truncation_notes"] = [
f"Result truncated: {original_count} of {original_count} rows returned "
f"across multiple queries "
- f"(limit ~{token_limit:,} tokens). {advice}"
+ f"(limit ~{max_bytes:,} bytes). {advice}"
]
lo, hi = 0, max(len(rows) for rows in originals)
@@ -1009,7 +899,7 @@ def _truncate_chart_query_results(
cap = (lo + hi + 1) // 2
for rows, original in zip(row_lists, originals, strict=False):
rows[:] = original[:cap]
- if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ if get_response_size_bytes(data) <= max_bytes:
lo = cap
else:
hi = cap - 1
@@ -1029,7 +919,7 @@ def _truncate_chart_query_results(
result["row_count"] = len(result["data"])
data["_truncation_notes"] = [
f"Result truncated: {kept_count} of {original_count} rows returned "
- f"across multiple queries (limit ~{token_limit:,} tokens). {advice}"
+ f"across multiple queries (limit ~{max_bytes:,} bytes). {advice}"
]
return data["_truncation_notes"]
@@ -1037,7 +927,7 @@ def _truncate_chart_query_results(
def _truncate_named_string_field(
data: Dict[str, Any],
field: str,
- token_limit: int,
+ max_bytes: int,
advice: str,
label: str | None = None,
suffix_marker: str = "",
@@ -1072,19 +962,19 @@ def _truncate_named_string_field(
# character count rather than a row count.
placeholder_note = (
f"{label} truncated: kept {original_len:,} of {original_len:,} "
- f"characters (limit ~{token_limit:,} tokens). {advice}"
+ f"characters (limit ~{max_bytes:,} bytes). {advice}"
)
data["_response_truncated"] = True
data["_truncation_notes"] = [placeholder_note]
kept_len = _bisect_string_length(
- data, field, value, token_limit, suffix=suffix_marker
+ data, field, value, max_bytes, suffix=suffix_marker
)
if kept_len < original_len:
notes = [
f"{label} truncated: kept {kept_len:,} of {original_len:,} "
- f"characters (limit ~{token_limit:,} tokens). {advice}"
+ f"characters (limit ~{max_bytes:,} bytes). {advice}"
]
data["_truncation_notes"] = notes
return notes
@@ -1096,10 +986,10 @@ def _truncate_named_string_field(
def truncate_query_result(
response: ToolResponse,
- token_limit: int,
+ max_bytes: int,
tool_name: str | None = None,
) -> tuple[ToolResponse, bool, list[str]]:
- """Truncate a data-query tool response to fit within the token limit.
+ """Truncate a data-query tool response to fit within the byte limit.
Unlike ``truncate_oversized_response`` (which targets info-tool dict
fields), this function targets the rows/data list directly: it performs
@@ -1116,15 +1006,13 @@ def truncate_query_result(
Args:
response: The tool response containing tabular row data.
- token_limit: Maximum estimated tokens allowed.
+ max_bytes: Maximum serialized response size in bytes.
tool_name: Name of the calling tool, used to tailor the truncation
advice (e.g. ``limit`` vs. ``row_limit`` vs. SQL ``LIMIT``).
Returns:
A tuple of (possibly-truncated response, was_truncated, list of notes).
"""
- from superset.utils import json as utils_json
-
advice = _ROW_LIMIT_ADVICE.get(tool_name or "", _DEFAULT_ROW_LIMIT_ADVICE)
if hasattr(response, "model_dump"):
@@ -1143,19 +1031,19 @@ def truncate_query_result(
if row_field is None:
# No recognised row field — fall back to generic field truncation.
- return truncate_oversized_response(response, token_limit)
+ return truncate_oversized_response(response, max_bytes)
- if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ if get_response_size_bytes(data) <= max_bytes:
return data, False, []
- notes = _truncate_chart_query_results(data, token_limit, advice)
+ notes = _truncate_chart_query_results(data, max_bytes, advice)
if notes is None:
- notes = _truncate_rows_field(data, row_field, token_limit, advice)
+ notes = _truncate_rows_field(data, row_field, max_bytes, advice)
if notes is None:
notes = _truncate_named_string_field(
data,
"csv_data",
- token_limit,
+ max_bytes,
advice,
label="CSV content",
suffix_marker=_STRING_FIELD_TRUNCATION_MARKERS.get("csv_data", ""),
@@ -1166,7 +1054,7 @@ def truncate_query_result(
def truncate_string_field_response(
response: ToolResponse,
- token_limit: int,
+ max_bytes: int,
field: str,
) -> tuple[ToolResponse, bool, list[str]]:
"""Truncate a response whose bulk is one string field with no size lever.
@@ -1179,14 +1067,12 @@ def truncate_string_field_response(
Args:
response: The tool response containing the oversized string field.
- token_limit: Maximum estimated tokens allowed.
+ max_bytes: Maximum serialized response size in bytes.
field: Name of the string field to truncate.
Returns:
A tuple of (possibly-truncated response, was_truncated, list of notes).
"""
- from superset.utils import json as utils_json
-
if hasattr(response, "model_dump"):
data = response.model_dump()
elif isinstance(response, dict):
@@ -1194,7 +1080,7 @@ def truncate_string_field_response(
else:
return response, False, []
- if estimate_token_count(utils_json.dumps(data)) <= token_limit:
+ if get_response_size_bytes(data) <= max_bytes:
return data, False, []
advice = (
@@ -1204,7 +1090,7 @@ def truncate_string_field_response(
notes = _truncate_named_string_field(
data,
field,
- token_limit,
+ max_bytes,
advice,
suffix_marker=_STRING_FIELD_TRUNCATION_MARKERS.get(field, ""),
)
@@ -1214,29 +1100,34 @@ def truncate_string_field_response(
def format_size_limit_error(
tool_name: str,
params: Dict[str, Any] | None,
- estimated_tokens: int,
- token_limit: int,
+ actual_bytes: int,
+ max_bytes: int,
response: ToolResponse | None = None,
) -> str:
"""
- Format a user-friendly error message when response exceeds token limit.
+ Format a user-friendly error message when response exceeds the byte limit.
Args:
tool_name: Name of the MCP tool
params: The tool parameters
- estimated_tokens: Estimated token count
- token_limit: Configured token limit
+ actual_bytes: Actual size of the response in bytes
+ max_bytes: Configured byte limit
response: Optional response for analysis
Returns:
Formatted error message with suggestions
"""
suggestions = generate_size_reduction_suggestions(
- tool_name, params, estimated_tokens, token_limit, response
+ tool_name, params, actual_bytes, max_bytes, response
)
+ size_text = (
+ "size could not be measured"
+ if actual_bytes == UNMEASURABLE_RESPONSE_BYTES
+ else f"{actual_bytes:,} bytes"
+ )
error_lines = [
- f"Response too large: ~{estimated_tokens:,} tokens (limit: {token_limit:,})",
+ f"Response too large: {size_text} (limit: {max_bytes:,})",
"",
"This response would overwhelm the LLM context window.",
"Please modify your query to reduce the response size:",
@@ -1247,9 +1138,7 @@ def format_size_limit_error(
error_lines.append(f"{i}. {suggestion}")
reduction_pct = (
- (estimated_tokens - token_limit) / estimated_tokens * 100
- if estimated_tokens
- else 0
+ (actual_bytes - max_bytes) / actual_bytes * 100 if actual_bytes else 0
)
error_lines.extend(
[
diff --git a/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py b/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py
index 27589a2e061..db4d799d4fe 100644
--- a/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py
+++ b/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py
@@ -36,13 +36,13 @@ from superset_core.semantic_layers.types import Dimension, Grains
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetSecurityException
from superset.mcp_service.app import mcp
-from superset.mcp_service.constants import DEFAULT_TOKEN_LIMIT
+from superset.mcp_service.constants import DEFAULT_MAX_RESPONSE_BYTES
from superset.mcp_service.middleware import (
create_response_size_guard_middleware,
ResponseSizeGuardMiddleware,
)
from superset.mcp_service.semantic_layer.schemas import ListMetricsRequest
-from superset.mcp_service.utils.token_utils import estimate_response_tokens
+from superset.mcp_service.utils.response_size_utils import get_response_size_bytes
from superset.semantic_layers.models import SemanticView
from superset.utils import json
@@ -157,18 +157,18 @@ def large_metric_catalog() -> Generator[MagicMock, None, None]:
@pytest.mark.asyncio
-async def test_list_metrics_default_page_token_bound(
+async def test_list_metrics_default_page_byte_bound(
mcp_server: FastMCP,
large_metric_catalog: MagicMock,
) -> None:
- """Default discovery leaves token headroom without embedded dimensions."""
+ """Default discovery leaves size headroom without embedded dimensions."""
async with Client(mcp_server) as client:
data: dict[str, Any] = json.loads(
(await client.call_tool("list_metrics", {})).content[0].text
)
assert data["success"] is True
assert data["total_count"] == 80
- assert estimate_response_tokens(data) <= DEFAULT_TOKEN_LIMIT // 2
+ assert get_response_size_bytes(data) <= DEFAULT_MAX_RESPONSE_BYTES // 2
assert data["page_size"] == 25
assert len(data["metrics"]) == 25
assert all(metric["compatible_dimensions"] == [] for metric in data["metrics"])
@@ -176,7 +176,7 @@ async def test_list_metrics_default_page_token_bound(
@pytest.mark.asyncio
-async def test_list_metrics_embedded_page_token_bound(
+async def test_list_metrics_embedded_page_byte_bound(
mcp_server: FastMCP,
large_metric_catalog: MagicMock,
) -> None:
@@ -190,7 +190,7 @@ async def test_list_metrics_embedded_page_token_bound(
"request": {
"view_id": 5,
"include_compatible_dimensions": True,
- "page_size": 8,
+ "page_size": 4,
}
},
)
@@ -199,13 +199,13 @@ async def test_list_metrics_embedded_page_token_bound(
.text
)
assert data["success"] is True
- assert len(data["metrics"]) == data["page_size"] == 8
+ assert len(data["metrics"]) == data["page_size"] == 4
assert data["total_count"] == 60
metric: dict[str, Any]
for metric in data["metrics"]:
assert len(metric["compatible_dimensions"]) == 40
assert len({dim["name"] for dim in metric["compatible_dimensions"]}) == 40
- assert estimate_response_tokens(data) < DEFAULT_TOKEN_LIMIT
+ assert get_response_size_bytes(data) < DEFAULT_MAX_RESPONSE_BYTES
@pytest.mark.asyncio
@@ -226,30 +226,30 @@ async def test_list_metrics_embedded_page_over_cap_rejected(
)
-@pytest.mark.parametrize("token_limit", [10000, 50000])
-def test_embedding_cap_is_independent_of_configured_token_limit(
- token_limit: int,
+@pytest.mark.parametrize("max_bytes", [10000, 50000])
+def test_embedding_cap_is_independent_of_configured_max_bytes(
+ max_bytes: int,
) -> None:
"""The guard honors overrides without changing the fixed embedding cap."""
app: MagicMock = MagicMock()
- app.config = {"MCP_RESPONSE_SIZE_CONFIG": {"token_limit": token_limit}}
+ app.config = {"MCP_RESPONSE_SIZE_CONFIG": {"max_bytes": max_bytes}}
with patch("superset.mcp_service.flask_singleton.get_flask_app", return_value=app):
guard: ResponseSizeGuardMiddleware | None = (
create_response_size_guard_middleware()
)
assert guard is not None
- assert guard.token_limit == token_limit
+ assert guard.max_bytes == max_bytes
assert (
- ListMetricsRequest(include_compatible_dimensions=True, page_size=8).page_size
- == 8
+ ListMetricsRequest(include_compatible_dimensions=True, page_size=4).page_size
+ == 4
)
error: pytest.ExceptionInfo[ValidationError]
with pytest.raises(ValidationError) as error:
- ListMetricsRequest(include_compatible_dimensions=True, page_size=9)
+ ListMetricsRequest(include_compatible_dimensions=True, page_size=5)
message: str = str(error.value)
- assert "MCP_RESPONSE_SIZE_CONFIG['token_limit']" in message
- assert "25k by default" in message
- assert "page_size <= 8" in message
+ assert "MCP_RESPONSE_SIZE_CONFIG['max_bytes']" in message
+ assert "50k by default" in message
+ assert "page_size <= 4" in message
assert "get_compatible_dimensions" in message
@@ -442,7 +442,7 @@ async def test_list_metrics_external_per_metric_compatible_dimensions(
"request": {
"view_id": 5,
"include_compatible_dimensions": True,
- "page_size": 8,
+ "page_size": 4,
}
},
)
@@ -639,7 +639,7 @@ async def test_builtin_embedded_metrics_reject_oversized_pages(
mock_ds: MagicMock = _make_dataset(42)
with _patched_dataset_lookup(mock_ds):
async with Client(mcp_server) as client:
- with pytest.raises(ToolError, match="page_size <= 8"):
+ with pytest.raises(ToolError, match="page_size <= 4"):
await client.call_tool(
"list_metrics",
{
@@ -653,7 +653,7 @@ async def test_builtin_embedded_metrics_reject_oversized_pages(
@pytest.mark.asyncio
-async def test_builtin_embedded_metrics_realistic_token_bound(
+async def test_builtin_embedded_metrics_realistic_byte_bound(
mcp_server: FastMCP,
) -> None:
"""Twenty metrics with thirty columns require bounded, lossless pages."""
@@ -665,7 +665,7 @@ async def test_builtin_embedded_metrics_realistic_token_bound(
metric: dict[str, Any]
with _patched_dataset_lookup(dataset):
async with Client(mcp_server) as client:
- with pytest.raises(ToolError, match="page_size <= 8"):
+ with pytest.raises(ToolError, match="page_size <= 4"):
await client.call_tool(
"list_metrics",
{
@@ -675,14 +675,14 @@ async def test_builtin_embedded_metrics_realistic_token_bound(
}
},
)
- for page in (1, 2, 3):
+ for page in (1, 2, 3, 4, 5):
result: Any = await client.call_tool(
"list_metrics",
{
"request": {
"dataset_id": 42,
"include_compatible_dimensions": True,
- "page_size": 8,
+ "page_size": 4,
"page": page,
}
},
@@ -690,10 +690,10 @@ async def test_builtin_embedded_metrics_realistic_token_bound(
data: dict[str, Any] = json.loads(result.content[0].text)
assert data["success"] is True
assert data["total_count"] == 20
- assert data["total_pages"] == 3
- assert data["page_size"] == 8
- assert len(data["metrics"]) == (4 if page == 3 else 8)
- assert estimate_response_tokens(data) < DEFAULT_TOKEN_LIMIT
+ assert data["total_pages"] == 5
+ assert data["page_size"] == 4
+ assert len(data["metrics"]) == 4
+ assert get_response_size_bytes(data) < DEFAULT_MAX_RESPONSE_BYTES
for metric in data["metrics"]:
assert len(metric["compatible_dimensions"]) == 30
names.append(metric["name"])
@@ -703,7 +703,7 @@ async def test_builtin_embedded_metrics_realistic_token_bound(
@pytest.mark.parametrize("scope", [{}, {"view_id": 42}])
def test_external_embedded_metrics_retain_page_ceiling(scope: dict[str, int]) -> None:
"""Unscoped and view-scoped requests can carry external dimensions."""
- with pytest.raises(ValidationError, match="page_size <= 8"):
+ with pytest.raises(ValidationError, match="page_size <= 4"):
ListMetricsRequest(**scope, include_compatible_dimensions=True, page_size=9)
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 b2c51b1a8e4..76f08979740 100644
--- a/tests/unit_tests/mcp_service/test_mcp_tool_registration.py
+++ b/tests/unit_tests/mcp_service/test_mcp_tool_registration.py
@@ -25,7 +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
+from superset.mcp_service.utils.response_size_utils import COMMITTED_WRITE_SPECS
# Patch target for the feature_flag_manager imported inside _apply_config_guards
_FFM_PATH = "superset.extensions.feature_flag_manager"
diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py
index b160d8d6e1e..ed7e882f5bd 100644
--- a/tests/unit_tests/mcp_service/test_middleware.py
+++ b/tests/unit_tests/mcp_service/test_middleware.py
@@ -47,7 +47,10 @@ from superset.mcp_service.middleware import (
StructuredContentStripperMiddleware,
ToolResultCompatibilityMiddleware,
)
-from superset.mcp_service.utils.token_utils import estimate_token_count
+from superset.mcp_service.utils.response_size_utils import (
+ get_response_size_bytes,
+ UNMEASURABLE_RESPONSE_BYTES,
+)
from superset.utils import json as utils_json
from superset.utils.log import DBEventLogger
@@ -58,21 +61,21 @@ class TestResponseSizeGuardMiddleware:
def test_init_default_values(self) -> None:
"""Should initialize with default values."""
middleware = ResponseSizeGuardMiddleware()
- assert middleware.token_limit == 25_000
+ assert middleware.max_bytes == 50_000
assert middleware.warn_threshold_pct == 80
- assert middleware.warn_threshold == 20000
+ assert middleware.warn_threshold == 40_000
assert middleware.excluded_tools == set()
assert middleware.max_list_items == 100
def test_init_custom_values(self) -> None:
"""Should initialize with custom values."""
middleware = ResponseSizeGuardMiddleware(
- token_limit=10000,
+ max_bytes=10000,
warn_threshold_pct=70,
excluded_tools=["health_check", "get_chart_preview"],
max_list_items=50,
)
- assert middleware.token_limit == 10000
+ assert middleware.max_bytes == 10000
assert middleware.warn_threshold_pct == 70
assert middleware.warn_threshold == 7000
assert middleware.excluded_tools == {"health_check", "get_chart_preview"}
@@ -95,8 +98,8 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_allows_small_response(self) -> None:
- """Should allow responses under token limit."""
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ """Should allow responses under the byte limit."""
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
# Create mock context
context = MagicMock()
@@ -118,8 +121,8 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_blocks_large_response(self) -> None:
- """Should block responses over token limit."""
- middleware = ResponseSizeGuardMiddleware(token_limit=100) # Very low limit
+ """Should block responses over the byte limit."""
+ middleware = ResponseSizeGuardMiddleware(max_bytes=100) # Very low limit
# Create mock context
context = MagicMock()
@@ -148,7 +151,7 @@ class TestResponseSizeGuardMiddleware:
async def test_skips_excluded_tools(self) -> None:
"""Should skip checking for excluded tools."""
middleware = ResponseSizeGuardMiddleware(
- token_limit=100, excluded_tools=["health_check"]
+ max_bytes=100, excluded_tools=["health_check"]
)
# Create mock context for excluded tool
@@ -168,14 +171,10 @@ class TestResponseSizeGuardMiddleware:
async def test_logs_warning_at_threshold(self) -> None:
"""Should log warning when approaching limit.
- Mocks the token estimator to return a specific value above the
- warn threshold but below the hard limit, decoupling the test
- from whichever tokenizer (tiktoken or char heuristic) happens
- to be loaded.
+ Mocks the size measurement to return a specific value above the
+ warn threshold but below the hard limit.
"""
- middleware = ResponseSizeGuardMiddleware(
- token_limit=1000, warn_threshold_pct=80
- )
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1000, warn_threshold_pct=80)
context = MagicMock()
context.message.name = "list_charts"
@@ -188,7 +187,7 @@ class TestResponseSizeGuardMiddleware:
patch("superset.mcp_service.middleware.get_user_id", return_value=1),
patch("superset.mcp_service.middleware.event_logger"),
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=850,
),
patch("superset.mcp_service.middleware.logger") as mock_logger,
@@ -203,7 +202,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_error_includes_suggestions(self) -> None:
"""Should include suggestions in error message."""
- middleware = ResponseSizeGuardMiddleware(token_limit=100)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=100)
context = MagicMock()
context.message.name = "list_charts"
@@ -225,10 +224,38 @@ class TestResponseSizeGuardMiddleware:
# Should suggest reducing page_size
assert "page_size" in error_message.lower() or "limit" in error_message.lower()
+ @pytest.mark.asyncio
+ async def test_unmeasurable_response_is_treated_as_oversized(self) -> None:
+ """A response whose size cannot be measured must not slip through.
+
+ The size helper never raises; it reports an unmeasurable response as
+ larger than any limit, so the guard takes the oversized path even
+ when ``max_bytes`` is configured above any fixed fallback value.
+ """
+ middleware = ResponseSizeGuardMiddleware(max_bytes=10_000_000)
+
+ context = MagicMock()
+ context.message.name = "list_charts"
+ context.message.arguments = {}
+ call_next = AsyncMock(return_value={"charts": []})
+
+ with (
+ patch("superset.mcp_service.middleware.get_user_id", return_value=1),
+ patch("superset.mcp_service.middleware.event_logger"),
+ patch(
+ "superset.mcp_service.middleware.get_response_size_bytes",
+ return_value=UNMEASURABLE_RESPONSE_BYTES,
+ ),
+ pytest.raises(ToolError) as exc_info,
+ ):
+ await middleware.on_call_tool(context, call_next)
+
+ assert "size could not be measured" in str(exc_info.value)
+
@pytest.mark.asyncio
async def test_logs_size_exceeded_event(self) -> None:
"""Should log to event logger when size exceeded."""
- middleware = ResponseSizeGuardMiddleware(token_limit=100)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=100)
context = MagicMock()
context.message.name = "list_charts"
@@ -252,7 +279,7 @@ class TestResponseSizeGuardMiddleware:
@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)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1000)
context = MagicMock()
context.message.name = "get_dataset_info"
@@ -278,10 +305,39 @@ class TestResponseSizeGuardMiddleware:
assert result["_response_truncated"] is True
assert "[truncated" in result["description"]
+ @pytest.mark.asyncio
+ async def test_truncates_info_tool_under_small_byte_budget(self) -> None:
+ """A budget below the fixed string clip must degrade, not block.
+
+ Clipping a string to a fixed 500 chars can never fit a 500-byte
+ budget, so the clip length has to follow the budget; otherwise the
+ info tool raises ToolError instead of returning a truncated response.
+ """
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
+
+ context = MagicMock()
+ context.message.name = "get_dataset_info"
+ context.message.arguments = {}
+ call_next = AsyncMock(
+ return_value={"id": 1, "table_name": "test", "description": "x" * 50000}
+ )
+
+ 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
+ assert "[truncated" in result["description"]
+ assert get_response_size_bytes(result) <= 500
+
@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)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "get_chart_info"
@@ -307,7 +363,7 @@ class TestResponseSizeGuardMiddleware:
@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)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=100)
context = MagicMock()
context.message.name = "list_charts" # Not an info tool
@@ -326,7 +382,7 @@ class TestResponseSizeGuardMiddleware:
@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)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1000)
context = MagicMock()
context.message.name = "get_dashboard_info"
@@ -359,7 +415,7 @@ class TestResponseSizeGuardMiddleware:
verifies the cap is now threaded through from the middleware
constructor rather than hardcoded.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=3000, max_list_items=50)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=6000, max_list_items=50)
context = MagicMock()
context.message.name = "get_dashboard_info"
@@ -389,7 +445,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_truncates_execute_sql_rows_instead_of_blocking(self) -> None:
"""execute_sql should truncate rows, not raise ToolError."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1500)
context = MagicMock()
context.message.name = "execute_sql"
@@ -419,7 +475,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_truncates_query_dataset_data_field(self) -> None:
"""query_dataset should truncate the 'data' list, not raise ToolError."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "query_dataset"
@@ -448,7 +504,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_truncates_get_chart_data_rows(self) -> None:
"""get_chart_data should truncate rows, not raise ToolError."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "get_chart_data"
@@ -477,8 +533,8 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_truncates_multi_query_chart_rows_across_whole_response(self) -> None:
- """All query results share the response's token budget."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ """All query results share the response's byte budget."""
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1500)
context = MagicMock()
context.message.name = "get_chart_data"
context.message.arguments = {}
@@ -512,8 +568,8 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_multi_query_truncation_result_fits_budget(self) -> None:
- """The final multi-query truncation note stays within the token budget."""
- middleware = ResponseSizeGuardMiddleware(token_limit=700)
+ """The final multi-query truncation note stays within the byte budget."""
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1500)
context = MagicMock()
context.message.name = "get_chart_data"
context.message.arguments = {}
@@ -537,13 +593,13 @@ class TestResponseSizeGuardMiddleware:
)
assert isinstance(result, dict)
- assert estimate_token_count(utils_json.dumps(result)) <= 700
+ assert get_response_size_bytes(result) <= 1500
assert " of 400 rows returned" in result["_truncation_notes"][0]
@pytest.mark.asyncio
async def test_data_query_truncation_updates_row_count(self) -> None:
"""row_count should reflect the truncated count, not the original."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "execute_sql"
@@ -569,7 +625,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_data_query_truncation_note_mentions_limit_clause(self) -> None:
"""Truncation note must tell the caller to add a LIMIT clause."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "execute_sql"
@@ -596,7 +652,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_data_query_truncation_logs_truncation_event(self) -> None:
"""Should log mcp_response_truncated (not size_exceeded) for query tools."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "execute_sql"
@@ -624,7 +680,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_truncates_get_chart_data_csv_export(self) -> None:
"""CSV exports (data=[], payload in csv_data) should be truncated too."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "get_chart_data"
@@ -659,10 +715,10 @@ class TestResponseSizeGuardMiddleware:
``_bisect_row_limit`` always keeps at least one row when the
original list is non-empty, even if that one row alone exceeds the
- token limit. The middleware must re-check the truncated size and
+ byte limit. The middleware must re-check the truncated size and
fall back to the hard error rather than treating this as success.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=50)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=50)
context = MagicMock()
context.message.name = "execute_sql"
@@ -686,7 +742,7 @@ class TestResponseSizeGuardMiddleware:
@pytest.mark.asyncio
async def test_data_query_under_limit_passes_through(self) -> None:
"""Small query results should pass through unchanged."""
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
context = MagicMock()
context.message.name = "execute_sql"
@@ -720,7 +776,7 @@ class TestResponseSizeGuardMiddleware:
payload. The guard must report success with a truncation marker so a
retrying agent doesn't replay an already-successful mutation.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -765,7 +821,7 @@ class TestResponseSizeGuardMiddleware:
degrade to a minimal success response rather than ever raising
ToolError -- unlike INFO_TOOLS, which fall through to a hard error
in this situation (see test_still_blocks_non_info_tools)."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -785,7 +841,7 @@ class TestResponseSizeGuardMiddleware:
# Keep every estimate over budget, including both _fits checks,
# to exercise the minimal fallback and its full shrink path.
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
@@ -803,9 +859,9 @@ class TestResponseSizeGuardMiddleware:
def test_minimal_response_already_fits_without_shrinking(self) -> None:
"""A minimal payload that fits must retain its fields without clipping."""
- from superset.mcp_service.utils.token_utils import COMMITTED_WRITE_SPECS
+ from superset.mcp_service.utils.response_size_utils import COMMITTED_WRITE_SPECS
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
minimal = {
"chart": {"id": 7, "description": "Keep this non-identity field"},
"success": True,
@@ -814,7 +870,7 @@ class TestResponseSizeGuardMiddleware:
}
original = utils_json.loads(utils_json.dumps(minimal))
- assert estimate_token_count(utils_json.dumps(minimal)) <= 500
+ assert get_response_size_bytes(minimal) <= 500
middleware._shrink_minimal_response(
minimal, COMMITTED_WRITE_SPECS["update_chart"]
)
@@ -826,12 +882,11 @@ class TestResponseSizeGuardMiddleware:
self, oversized_field: str
) -> None:
"""Real over-budget measurements must drive reduction, not mock errors."""
- from superset.mcp_service.utils.token_utils import (
+ from superset.mcp_service.utils.response_size_utils import (
COMMITTED_WRITE_SPECS,
- estimate_response_tokens,
)
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=2000)
minimal: dict[str, Any] = {
"chart": {"id": 7, "is_unsaved_state": True},
"success": False,
@@ -849,17 +904,17 @@ class TestResponseSizeGuardMiddleware:
else:
minimal["explore_url"] = "http://host/explore/?key=" + "k" * 40000
- assert estimate_response_tokens(minimal) > middleware.token_limit
+ assert get_response_size_bytes(minimal) > middleware.max_bytes
with patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
- wraps=estimate_response_tokens,
+ "superset.mcp_service.middleware.get_response_size_bytes",
+ wraps=get_response_size_bytes,
) as estimate:
middleware._shrink_minimal_response(
minimal, COMMITTED_WRITE_SPECS["update_chart"]
)
assert estimate.call_count == 2
- assert estimate_response_tokens(minimal) <= middleware.token_limit
+ assert get_response_size_bytes(minimal) <= middleware.max_bytes
assert minimal["chart"] == {
"id": 7,
"is_unsaved_state": True,
@@ -890,9 +945,9 @@ class TestResponseSizeGuardMiddleware:
self, changed_fields: list[Any], retained: bool
) -> None:
"""Keep useful patch confirmations without admitting unbounded lists."""
- from superset.mcp_service.utils.token_utils import COMMITTED_WRITE_SPECS
+ from superset.mcp_service.utils.response_size_utils import COMMITTED_WRITE_SPECS
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=2000)
spec = COMMITTED_WRITE_SPECS["update_dashboard"]
minimal = middleware._select_confirmation_fields(
{
@@ -909,7 +964,7 @@ class TestResponseSizeGuardMiddleware:
if retained:
assert minimal["changed_fields"] == changed_fields
assert minimal["dashboard"]["id"] == 7
- assert estimate_token_count(utils_json.dumps(minimal)) <= 500
+ assert get_response_size_bytes(minimal) <= 2000
@pytest.mark.asyncio
async def test_update_dashboard_committed_write_is_not_hard_blocked(
@@ -923,7 +978,7 @@ class TestResponseSizeGuardMiddleware:
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)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_dashboard"
@@ -969,7 +1024,7 @@ class TestResponseSizeGuardMiddleware:
dashboard response would protect nothing and clear the very field
that says which dashboard was written.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_dashboard"
@@ -1016,7 +1071,7 @@ class TestResponseSizeGuardMiddleware:
``success``; synthesizing them would hand the caller a payload its
own output schema does not describe.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_dashboard"
@@ -1035,7 +1090,7 @@ class TestResponseSizeGuardMiddleware:
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",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
@@ -1046,7 +1101,7 @@ class TestResponseSizeGuardMiddleware:
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 get_response_size_bytes(result) <= 2000
assert result["changed_fields"] == ["css"]
assert "re-read the dashboard" in result["_truncation_notes"][0]
@@ -1062,7 +1117,7 @@ class TestResponseSizeGuardMiddleware:
untruncated payload, so any one of them can keep the "minimal"
response far over budget.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -1090,7 +1145,7 @@ class TestResponseSizeGuardMiddleware:
# Keep every estimate over budget, including both _fits checks,
# to exercise the minimal fallback and its full shrink path.
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
@@ -1104,7 +1159,7 @@ class TestResponseSizeGuardMiddleware:
assert result["chart"]["uuid"] == "abc"
# ...but nothing unbounded rides along with it.
- assert estimate_token_count(utils_json.dumps(result)) <= 500
+ assert get_response_size_bytes(result) <= 2000
for value in (
result["chart"]["slice_name"],
result["error"],
@@ -1122,7 +1177,7 @@ class TestResponseSizeGuardMiddleware:
``validation_errors``. Treating ``error`` as a string would bound only
a shape the tools never emit and leave the real one to blow the limit.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -1148,14 +1203,14 @@ class TestResponseSizeGuardMiddleware:
patch("superset.mcp_service.middleware.get_user_id", return_value=1),
patch("superset.mcp_service.middleware.event_logger"),
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
result = await middleware.on_call_tool(context, call_next)
assert isinstance(result, dict)
- assert estimate_token_count(utils_json.dumps(result)) <= 500
+ assert get_response_size_bytes(result) <= 2000
# The error still identifies itself -- only the unbounded context goes.
assert result["error"]["error_type"] == "execution"
@@ -1175,7 +1230,7 @@ class TestResponseSizeGuardMiddleware:
chart to identifying fields must keep it, or the size guard turns a
preview into something indistinguishable from a committed update.
"""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -1198,7 +1253,7 @@ class TestResponseSizeGuardMiddleware:
patch("superset.mcp_service.middleware.get_user_id", return_value=1),
patch("superset.mcp_service.middleware.event_logger"),
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
@@ -1225,7 +1280,7 @@ class TestResponseSizeGuardMiddleware:
from fastmcp.tools.tool import ToolResult
from mcp.types import TextContent
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = tool_name
@@ -1238,7 +1293,7 @@ class TestResponseSizeGuardMiddleware:
patch("superset.mcp_service.middleware.get_user_id", return_value=1),
patch("superset.mcp_service.middleware.event_logger"),
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
patch.object(
@@ -1263,7 +1318,7 @@ class TestResponseSizeGuardMiddleware:
from fastmcp.tools.tool import ToolResult
from mcp.types import TextContent
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "update_chart"
@@ -1279,7 +1334,7 @@ class TestResponseSizeGuardMiddleware:
patch("superset.mcp_service.middleware.get_user_id", return_value=1),
patch("superset.mcp_service.middleware.event_logger"),
patch(
- "superset.mcp_service.middleware.estimate_response_tokens",
+ "superset.mcp_service.middleware.get_response_size_bytes",
return_value=600,
),
):
@@ -1295,7 +1350,7 @@ class TestResponseSizeGuardMiddleware:
"""get_chart_sql has no limit/row lever -- the oversized 'sql' field
itself must be bisected down instead of hard-blocking or emitting
the unactionable 'Reduction needed: ~0%' guidance."""
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "get_chart_sql"
@@ -1343,7 +1398,7 @@ class TestCreateResponseSizeGuardMiddleware:
"""Should create middleware when enabled in config."""
mock_config = {
"enabled": True,
- "token_limit": 30000,
+ "max_bytes": 30000,
"warn_threshold_pct": 75,
"excluded_tools": ["health_check"],
}
@@ -1359,7 +1414,7 @@ class TestCreateResponseSizeGuardMiddleware:
assert middleware is not None
assert isinstance(middleware, ResponseSizeGuardMiddleware)
- assert middleware.token_limit == 30000
+ assert middleware.max_bytes == 30000
assert middleware.warn_threshold_pct == 75
assert "health_check" in middleware.excluded_tools
@@ -1392,7 +1447,7 @@ class TestCreateResponseSizeGuardMiddleware:
middleware = create_response_size_guard_middleware()
assert middleware is not None
- assert middleware.token_limit == 25_000 # Default
+ assert middleware.max_bytes == 50_000 # Default
assert middleware.warn_threshold_pct == 80 # Default
def test_falls_back_to_default_when_max_list_items_is_none(self) -> None:
@@ -1629,7 +1684,7 @@ class TestToolResultWrapping:
from superset.utils import json
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1000)
context = MagicMock()
context.message.name = "get_dataset_info"
context.message.arguments = {}
@@ -1653,9 +1708,9 @@ class TestToolResultWrapping:
@pytest.mark.asyncio
async def test_small_tool_result_passes_through_unchanged(self) -> None:
- """Should return the original ToolResult when within the token limit."""
+ """Should return the original ToolResult when within the byte limit."""
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
context = MagicMock()
context.message.name = "get_chart_info"
context.message.arguments = {}
@@ -1675,7 +1730,7 @@ class TestToolResultWrapping:
@pytest.mark.asyncio
async def test_large_non_info_tool_result_is_blocked(self) -> None:
"""Should raise ToolError for a non-info ToolResult that exceeds the limit."""
- middleware = ResponseSizeGuardMiddleware(token_limit=100)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=100)
context = MagicMock()
context.message.name = "list_charts"
context.message.arguments = {}
@@ -1706,7 +1761,7 @@ class TestToolResultWrapping:
from superset.utils import json
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=500)
context = MagicMock()
context.message.name = "execute_sql"
context.message.arguments = {}
@@ -1738,7 +1793,7 @@ class TestToolResultWrapping:
from superset.utils import json
- middleware = ResponseSizeGuardMiddleware(token_limit=500)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=1000)
context = MagicMock()
context.message.name = "get_dashboard_info"
context.message.arguments = {}
@@ -1772,7 +1827,7 @@ class TestMiddlewareIntegration:
id: int
name: str
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
context = MagicMock()
context.message.name = "get_chart_info"
@@ -1792,7 +1847,7 @@ class TestMiddlewareIntegration:
@pytest.mark.asyncio
async def test_list_response(self) -> None:
"""Should handle list responses."""
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
context = MagicMock()
context.message.name = "list_charts"
@@ -1812,7 +1867,7 @@ class TestMiddlewareIntegration:
@pytest.mark.asyncio
async def test_string_response(self) -> None:
"""Should handle string responses."""
- middleware = ResponseSizeGuardMiddleware(token_limit=25000)
+ middleware = ResponseSizeGuardMiddleware(max_bytes=25000)
context = MagicMock()
context.message.name = "health_check"
diff --git a/tests/unit_tests/mcp_service/utils/test_token_utils.py b/tests/unit_tests/mcp_service/utils/test_response_size_utils.py
similarity index 87%
rename from tests/unit_tests/mcp_service/utils/test_token_utils.py
rename to tests/unit_tests/mcp_service/utils/test_response_size_utils.py
index fab8cfe8bd6..89ef613c05c 100644
--- a/tests/unit_tests/mcp_service/utils/test_token_utils.py
+++ b/tests/unit_tests/mcp_service/utils/test_response_size_utils.py
@@ -16,7 +16,7 @@
# under the License.
"""
-Unit tests for MCP service token utilities.
+Unit tests for MCP service response size utilities.
"""
from typing import Any, List
@@ -25,100 +25,35 @@ from unittest.mock import patch
import pytest
from pydantic import BaseModel
-from superset.mcp_service.utils import token_utils
-from superset.mcp_service.utils.token_utils import (
+from superset.mcp_service.utils.response_size_utils import (
_bisect_string_length,
_MAX_DICT_KEYS,
+ _MAX_STRING_CHARS,
+ _MIN_STRING_CHARS,
_replace_collections_with_summaries,
_STRING_FIELD_TRUNCATION_MARKERS,
_summarize_large_dicts,
_truncate_lists,
_truncate_strings,
_truncate_strings_recursive,
- CHARS_PER_TOKEN,
COMMITTED_WRITE_SPECS,
COMMITTED_WRITE_TOOLS,
- estimate_response_tokens,
- estimate_token_count,
extract_query_params,
format_size_limit_error,
generate_size_reduction_suggestions,
get_response_size_bytes,
INFO_TOOLS,
+ string_clip_chars,
STRING_FIELD_TRUNCATION_TOOLS,
truncate_oversized_response,
truncate_query_result,
truncate_string_field_response,
+ UNMEASURABLE_RESPONSE_BYTES,
)
-class TestEstimateTokenCount:
- """Test estimate_token_count function."""
-
- def test_estimate_string(self) -> None:
- """Should produce a positive non-zero estimate for a normal string.
-
- We don't assert on a specific number because the result depends on
- which tokenizer is loaded (tiktoken when available, char heuristic
- otherwise).
- """
- text = "Hello world"
- result = estimate_token_count(text)
- assert result > 0
-
- def test_estimate_bytes(self) -> None:
- """Bytes input should be decoded and produce the same count as the
- equivalent string."""
- text = "Hello world"
- assert estimate_token_count(text.encode("utf-8")) == estimate_token_count(text)
-
- def test_empty_string(self) -> None:
- """Should return 0 for empty string and empty bytes."""
- assert estimate_token_count("") == 0
- assert estimate_token_count(b"") == 0
-
- def test_json_like_content(self) -> None:
- """JSON content should produce a positive estimate."""
- json_str = '{"name": "test", "value": 123, "items": [1, 2, 3]}'
- assert estimate_token_count(json_str) > 0
-
- def test_long_text_roughly_scales_with_length(self) -> None:
- """A doubled string should produce roughly double the token count
- (within ±10%)."""
- small = "the quick brown fox jumps over the lazy dog. " * 20
- large = small * 2
- small_n = estimate_token_count(small)
- large_n = estimate_token_count(large)
- # Within 10% of 2x — both tokenizers (tiktoken and the char
- # fallback) preserve length monotonicity.
- assert 1.8 * small_n <= large_n <= 2.2 * small_n
-
- def test_fallback_uses_chars_per_token_when_tiktoken_unavailable(
- self,
- ) -> None:
- """When the tiktoken encoding is None (not installed), the
- function falls back to len/CHARS_PER_TOKEN math."""
- text = "x" * 100
- with patch.object(token_utils, "_ENCODING", None):
- result = estimate_token_count(text)
- assert result == int(100 / CHARS_PER_TOKEN)
-
- def test_fallback_when_tiktoken_encode_raises(self) -> None:
- """A misbehaving encoding should fall back to the char heuristic
- rather than raise — the size guard must never fail-open."""
-
- class BoomEncoding:
- def encode(self, text: str) -> list[int]:
- raise ValueError("simulated tiktoken failure")
-
- text = "abc" * 50
- with patch.object(token_utils, "_ENCODING", BoomEncoding()):
- result = estimate_token_count(text)
- assert result == int(len(text) / CHARS_PER_TOKEN)
-
-
-class TestEstimateResponseTokens:
- """Test estimate_response_tokens function."""
+class TestGetResponseSizeBytes:
+ """Test get_response_size_bytes function."""
class MockResponse(BaseModel):
"""Mock Pydantic response model."""
@@ -127,40 +62,6 @@ class TestEstimateResponseTokens:
value: int
items: List[Any]
- def test_estimate_pydantic_model(self) -> None:
- """Should estimate tokens for Pydantic model."""
- response = self.MockResponse(name="test", value=42, items=[1, 2, 3])
- result = estimate_response_tokens(response)
- assert result > 0
-
- def test_estimate_dict(self) -> None:
- """Should estimate tokens for dict."""
- response = {"name": "test", "value": 42}
- result = estimate_response_tokens(response)
- assert result > 0
-
- def test_estimate_list(self) -> None:
- """Should estimate tokens for list."""
- response = [{"name": "item1"}, {"name": "item2"}]
- result = estimate_response_tokens(response)
- assert result > 0
-
- def test_estimate_string(self) -> None:
- """Should estimate tokens for string response."""
- response = "Hello world"
- result = estimate_response_tokens(response)
- assert result > 0
-
- def test_estimate_large_response(self) -> None:
- """Should estimate tokens for large response."""
- response = {"items": [{"name": f"item{i}"} for i in range(1000)]}
- result = estimate_response_tokens(response)
- assert result > 1000 # Large response should have many tokens
-
-
-class TestGetResponseSizeBytes:
- """Test get_response_size_bytes function."""
-
def test_size_dict(self) -> None:
"""Should return size in bytes for dict."""
response = {"name": "test"}
@@ -179,6 +80,59 @@ class TestGetResponseSizeBytes:
result = get_response_size_bytes(response)
assert result == len(response)
+ def test_size_pydantic_model(self) -> None:
+ """Should return size in bytes for a Pydantic model."""
+ response = self.MockResponse(name="test", value=42, items=[1, 2, 3])
+ result = get_response_size_bytes(response)
+ assert result > 0
+
+ def test_size_list(self) -> None:
+ """Should return size in bytes for a list."""
+ response = [{"name": "item1"}, {"name": "item2"}]
+ result = get_response_size_bytes(response)
+ assert result > 0
+
+ def test_size_large_response(self) -> None:
+ """A larger response should measure a correspondingly larger size."""
+ small = {"items": [{"name": f"item{i}"} for i in range(10)]}
+ large = {"items": [{"name": f"item{i}"} for i in range(1000)]}
+ assert get_response_size_bytes(large) > get_response_size_bytes(small)
+
+ def test_unmeasurable_response_reads_as_oversized(self) -> None:
+ """A serialization failure must exceed any configurable limit.
+
+ A fixed fallback (say 1 MB) would count as "fits" for an operator who
+ set ``max_bytes`` at or above it, letting an unmeasured response
+ through the guard.
+ """
+ with patch("superset.utils.json.dumps", side_effect=MemoryError("boom")):
+ result = get_response_size_bytes({"data": "x"})
+ assert result == UNMEASURABLE_RESPONSE_BYTES
+ assert result > 1_000_000_000
+
+
+class TestStringClipChars:
+ """Test the budget-derived string clip length."""
+
+ def test_default_budget_keeps_full_clip_length(self) -> None:
+ """The 50 KB default budget must not change the historical clip."""
+ assert string_clip_chars(50_000) == _MAX_STRING_CHARS
+
+ def test_small_budget_scales_clip_length_down(self) -> None:
+ """A budget below ``4 * ceiling`` shrinks the clip proportionally."""
+ assert string_clip_chars(500) == 125
+ assert string_clip_chars(1000) == 250
+
+ def test_floor_keeps_strings_recognizable(self) -> None:
+ """Even a tiny budget leaves ``_MIN_STRING_CHARS`` of each string."""
+ assert string_clip_chars(50) == _MIN_STRING_CHARS
+ assert string_clip_chars(0) == _MIN_STRING_CHARS
+
+ def test_ceiling_is_respected(self) -> None:
+ """A caller-supplied ceiling (e.g. committed-write fields) is honored."""
+ assert string_clip_chars(100_000, ceiling=200) == 200
+ assert string_clip_chars(400, ceiling=200) == 100
+
class TestExtractQueryParams:
"""Test extract_query_params function."""
@@ -224,8 +178,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert any(
"page_size" in s.lower() or "limit" in s.lower() for s in suggestions
@@ -237,8 +191,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert any(
"limit" in s.lower() or "page_size" in s.lower() for s in suggestions
@@ -250,8 +204,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert any(
"select_columns" in s.lower() or "columns" in s.lower() for s in suggestions
@@ -263,8 +217,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert any("filter" in s.lower() for s in suggestions)
assert not any("editor" in s.lower() for s in suggestions)
@@ -275,8 +229,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="execute_sql",
params={"sql": "SELECT * FROM table"},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
combined = " ".join(suggestions)
# Should suggest SQL LIMIT clause
@@ -289,8 +243,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="execute_sql",
params={"sql": "SELECT * FROM table", "limit": 500},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
combined = " ".join(suggestions)
# Should still suggest SQL LIMIT
@@ -303,8 +257,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params={},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
# Should suggest excluding params or query_context
assert any(
@@ -316,8 +270,8 @@ class TestGenerateSizeReductionSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_dashboards",
params={},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert any("search" in s.lower() for s in suggestions)
@@ -325,13 +279,13 @@ class TestGenerateSizeReductionSuggestions:
class TestFormatSizeLimitError:
"""Test format_size_limit_error function."""
- def test_error_contains_token_counts(self) -> None:
- """Should include token counts in error message."""
+ def test_error_contains_byte_counts(self) -> None:
+ """Should include byte counts in error message."""
error = format_size_limit_error(
tool_name="list_charts",
params={},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert "50,000" in error
assert "25,000" in error
@@ -341,8 +295,8 @@ class TestFormatSizeLimitError:
error = format_size_limit_error(
tool_name="list_charts",
params={},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
assert "list_charts" in error
@@ -351,8 +305,8 @@ class TestFormatSizeLimitError:
error = format_size_limit_error(
tool_name="list_charts",
params={"page_size": 100},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
# Should have numbered suggestions
assert "1." in error
@@ -362,8 +316,8 @@ class TestFormatSizeLimitError:
error = format_size_limit_error(
tool_name="list_charts",
params={},
- estimated_tokens=50000,
- token_limit=25000,
+ actual_bytes=50000,
+ max_bytes=25000,
)
# 50% reduction needed
assert "50%" in error or "Reduction" in error
@@ -373,8 +327,8 @@ class TestFormatSizeLimitError:
error = format_size_limit_error(
tool_name="list_charts",
params={},
- estimated_tokens=100000,
- token_limit=10000,
+ actual_bytes=100000,
+ max_bytes=10000,
)
# Count numbered suggestions (1. through 5.)
suggestion_count = sum(1 for i in range(1, 10) if f"{i}." in error)
@@ -385,8 +339,8 @@ class TestFormatSizeLimitError:
error = format_size_limit_error(
tool_name="list_charts",
params={"page_size": 100},
- estimated_tokens=75000,
- token_limit=25000,
+ actual_bytes=75000,
+ max_bytes=25000,
)
# Should be multi-line and contain key information
lines = error.split("\n")
@@ -404,8 +358,8 @@ class TestCalculatedSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=50000, # 2x over limit
- token_limit=25000,
+ actual_bytes=50000, # 2x over limit
+ max_bytes=25000,
)
# Find the page_size suggestion
page_size_suggestion = next(
@@ -424,8 +378,8 @@ class TestCalculatedSuggestions:
suggestions = generate_size_reduction_suggestions(
tool_name="list_charts",
params=params,
- estimated_tokens=75000, # 3x over limit
- token_limit=25000,
+ actual_bytes=75000, # 3x over limit
+ max_bytes=25000,
)
# Should mention ~66% reduction needed (int truncation of 66.6%)
combined = " ".join(suggestions)
@@ -715,7 +669,12 @@ class TestTruncateOversizedResponse:
assert notes == []
def test_truncates_large_string_fields(self) -> None:
- """Should truncate long strings to fit."""
+ """Should truncate long strings to fit.
+
+ With a 500-byte budget, clipping a string to the fixed 500 chars
+ (plus its marker) can never fit, so the clip length has to follow
+ the budget instead of leaving the response over the limit.
+ """
response = {
"id": 1,
"description": "x" * 50000, # Very large description
@@ -725,6 +684,10 @@ class TestTruncateOversizedResponse:
assert isinstance(result, dict)
assert "[truncated" in result["description"]
assert any("description" in n for n in notes)
+ assert get_response_size_bytes(result) <= 500
+ assert result["description"].startswith("x" * _MIN_STRING_CHARS)
+ assert "[truncated from 50000 chars]" in result["description"]
+ assert notes == ["Field 'description' truncated from 50000 chars"]
def test_truncates_large_lists(self) -> None:
"""Should truncate lists when strings alone are not enough."""
@@ -795,7 +758,7 @@ class TestTruncateOversizedResponse:
result: Any
was_truncated: bool
notes: list[str]
- result, was_truncated, notes = truncate_oversized_response(response, 3000)
+ result, was_truncated, notes = truncate_oversized_response(response, 7000)
assert was_truncated is True
assert isinstance(result, dict)
assert len(result["charts"]) == 100
@@ -837,7 +800,7 @@ class TestTruncateOversizedResponse:
was_truncated: bool
notes: list[str]
result, was_truncated, notes = truncate_oversized_response(
- response, 200, max_list_items=5
+ response, 800, max_list_items=5
)
assert was_truncated is True
assert isinstance(result, dict)
@@ -939,7 +902,7 @@ class TestTruncateStringFieldResponse:
assert isinstance(result, dict)
assert 0 < len(result["sql"]) < len(response["sql"])
assert result["_response_truncated"] is True
- assert estimate_response_tokens(result) <= 500
+ assert get_response_size_bytes(result) <= 500
assert any("sql" in n for n in notes)
def test_truncated_sql_is_marked_unexecutable(self) -> None:
@@ -961,7 +924,7 @@ class TestTruncateStringFieldResponse:
assert result["sql"].endswith("DO NOT EXECUTE")
assert result["sql"].count("'") == 1
# The marker is inside the measured budget, not appended after it.
- assert estimate_response_tokens(result) <= 500
+ assert get_response_size_bytes(result) <= 500
def test_truncated_sql_is_rejected_whatever_the_cut_landed_in(self) -> None:
"""The marker must defeat every lexical state the bisect can end in.
@@ -1037,8 +1000,8 @@ class TestTruncateStringFieldResponse:
message = format_size_limit_error(
tool_name="get_chart_sql",
params={},
- estimated_tokens=20400,
- token_limit=20000,
+ actual_bytes=20400,
+ max_bytes=20000,
)
assert "no size-reduction parameter" in message
@@ -1074,7 +1037,7 @@ class TestTruncateQueryResult:
"""The final payload (rows + note metadata) must itself fit.
Regression test: the note is built from the kept row count, but
- that note text also consumes tokens. The bisection must reserve
+ that note text also consumes bytes. The bisection must reserve
room for it up front rather than measuring fit on bare rows and
appending the note afterward, which could push the final payload
back over the limit.
@@ -1083,7 +1046,7 @@ class TestTruncateQueryResult:
result, was_truncated, notes = truncate_query_result(response, 500)
assert was_truncated is True
assert isinstance(result, dict)
- assert estimate_response_tokens(result) <= 500
+ assert get_response_size_bytes(result) <= 500
assert result["row_count"] == len(result["rows"])
assert result["row_count"] < 200
@@ -1122,7 +1085,7 @@ class TestTruncateQueryResult:
assert isinstance(result, dict)
assert len(result["csv_data"]) < len(response["csv_data"])
assert result["csv_data"].endswith(marker)
- assert estimate_response_tokens(result) <= 500
+ assert get_response_size_bytes(result) <= 500
assert any("CSV" in n for n in notes)
def test_does_not_truncate_excel_binary_field(self) -> None: