This is an automated email from the ASF dual-hosted git repository.
aminghadersohi 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 1e08d61204f fix(mcp): preserve deduplicated chart tool input schemas
(#44575)
1e08d61204f is described below
commit 1e08d61204fff26a94f054736dc2041bfea49d8b
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Fri Sep 25 07:44:43 2026 -0400
fix(mcp): preserve deduplicated chart tool input schemas (#44575)
Co-authored-by: Claude Haiku 4.5 <[email protected]>
---
docs/admin_docs/configuration/mcp-server.mdx | 8 +-
superset/mcp_service/mcp_config.py | 21 +-
superset/mcp_service/server.py | 131 +-------
.../mcp_service/test_chart_tool_inventory.py | 189 +++++++++++
.../unit_tests/mcp_service/test_tool_inventory.py | 157 ++++++++++
.../mcp_service/test_tool_search_transform.py | 347 +++------------------
6 files changed, 417 insertions(+), 436 deletions(-)
diff --git a/docs/admin_docs/configuration/mcp-server.mdx
b/docs/admin_docs/configuration/mcp-server.mdx
index a0a3e3db147..2dea57a0d6c 100644
--- a/docs/admin_docs/configuration/mcp-server.mdx
+++ b/docs/admin_docs/configuration/mcp-server.mdx
@@ -762,8 +762,8 @@ MCP_TOOL_SEARCH_CONFIG = {
],
"search_tool_name": "search_tools",
"call_tool_name": "call_tool",
- "include_schemas": False, # False=summary mode (name + parameters_hint)
- "compact_schemas": True, # Strip $defs (only applies when
include_schemas=True)
+ "include_schemas": True, # False=summary mode (name + parameters_hint)
+ "compact_schemas": True, # Default description limit when
include_schemas=True
"max_description_length": 300,
}
```
@@ -774,8 +774,8 @@ MCP_TOOL_SEARCH_CONFIG = {
| `strategy` | `"bm25"` | Search ranking algorithm. `"bm25"`
supports natural language; `"regex"` supports pattern matching
|
| `max_results` | `5` | Maximum tools returned per search
query
|
| `always_visible` | See above | Tools that always appear in
`list_tools`, regardless of search
|
-| `include_schemas` | `False` | When `False` (default, "summary
mode"), search results omit `inputSchema` entirely and include a lightweight
`parameters_hint` listing top-level parameter names. Set to `True` to include
the full `inputSchema` in search results. Full schemas are always used when a
tool is actually invoked via `call_tool`. |
-| `compact_schemas` | `True` | Strip `$defs` / `$ref` and replace
with `{"type": "object"}` in search results to reduce token cost. Only takes
effect when `include_schemas=True` — ignored in summary mode.
|
+| `include_schemas` | `True` | Include the full `inputSchema` in
search results by default. When `False` ("summary mode"), search results omit
`inputSchema` entirely and include a lightweight `parameters_hint` listing
top-level parameter names. Full schemas are always used when a tool is actually
invoked via `call_tool`. |
+| `compact_schemas` | `True` | Legacy setting selecting the default
description limit (300 when `True`, 0 when `False`) if `max_description_length`
is omitted. Input schemas preserve `$defs`, `$ref`, nullable unions, and
constraints in either mode; titles are omitted. Clients should resolve
references within each tool's `inputSchema`.
|
| `max_description_length` | `300` | Truncate tool descriptions in search
results (0 = no truncation). Applies in both summary and full-schema modes.
|
:::tip
diff --git a/superset/mcp_service/mcp_config.py
b/superset/mcp_service/mcp_config.py
index 502ffad1506..34583723de9 100644
--- a/superset/mcp_service/mcp_config.py
+++ b/superset/mcp_service/mcp_config.py
@@ -463,18 +463,17 @@ MCP_RESPONSE_SIZE_CONFIG: dict[str, Any] = {
# - "bm25": Natural language search using BM25 ranking (recommended)
# - "regex": Pattern-based search using regular expressions
#
-# Schema Compaction:
-# ------------------
-# When compact_schemas=True, search results strip $defs sections and replace
-# $ref pointers with {"type": "object"}, and truncate tool descriptions.
-# This reduces per-search token cost by ~40-60%. Full schemas remain
-# available when the tool is actually invoked via call_tool.
+# Schema Serialization:
+# ---------------------
+# Input schemas preserve $defs, $ref, nullable unions, and validation
constraints;
+# titles and output schemas are omitted. Clients should resolve references
within
+# each tool's inputSchema. Inlining references duplicates shared chart models.
+# The legacy compact_schemas setting only selects the default description limit
+# (300 when True, 0 when False) if max_description_length is omitted.
#
# Rollback:
# ---------
# - Set enabled=False to disable tool search entirely (full catalog exposed).
-# - Set compact_schemas=False to disable schema compaction only (full $defs
-# and descriptions in search results, tool search still active).
# - Set max_description_length=0 to disable description truncation only.
#
# Summary Mode (include_schemas):
@@ -489,8 +488,8 @@ MCP_RESPONSE_SIZE_CONFIG: dict[str, Any] = {
# so LLMs can see structured/discriminated-union configs (e.g. chart
# generation) without a second round trip. Set include_schemas=False to
# switch to summary mode if search_tools response size becomes a problem
-# again; compact_schemas is ignored when include_schemas=False (no schema to
-# compact); max_description_length still applies in summary mode.
+# again; compact_schemas is ignored when include_schemas=False.
+# max_description_length still applies in summary mode.
# =============================================================================
MCP_TOOL_SEARCH_CONFIG: dict[str, Any] = {
"enabled": True, # Enabled by default — reduces initial context by ~70%
@@ -502,7 +501,7 @@ MCP_TOOL_SEARCH_CONFIG: dict[str, Any] = {
],
"search_tool_name": "search_tools", # Name of the search tool
"call_tool_name": "call_tool", # Name of the call proxy tool
- "compact_schemas": True, # Strip $defs/$ref (requires
include_schemas=True)
+ "compact_schemas": True, # Legacy default description limit for full
schemas
"max_description_length": 300, # Truncate tool descriptions (0 = no
truncation)
"include_schemas": True, # full inputSchema in search results
}
diff --git a/superset/mcp_service/server.py b/superset/mcp_service/server.py
index 7727b33f047..45765560a54 100644
--- a/superset/mcp_service/server.py
+++ b/superset/mcp_service/server.py
@@ -299,111 +299,6 @@ def _strip_titles(obj: Any, in_properties_map: bool =
False) -> Any:
return obj
-def _simplify_optional_union(result: dict[str, Any]) -> dict[str, Any]:
- """Collapse ``anyOf``/``oneOf`` with exactly one non-null variant.
-
- Pydantic encodes ``Optional[X]`` as ``{"anyOf": [<X>, {"type": "null"}]}``.
- This replaces the union with the non-null variant while preserving any
- ``description`` or ``default`` from the parent node.
- """
- for union_key in ("anyOf", "oneOf"):
- variants = result.get(union_key)
- if not isinstance(variants, list) or len(variants) != 2:
- continue
- non_null = [v for v in variants if v.get("type") != "null"]
- if len(non_null) != 1:
- continue
- simplified = dict(non_null[0])
- for keep in ("description", "default"):
- if keep in result and keep not in simplified:
- simplified[keep] = result[keep]
- result.pop(union_key)
- result.pop("description", None)
- result.pop("default", None)
- result.update(simplified)
- return result
-
-
-def _resolve_ref(
- obj: dict[str, Any],
- defs: dict[str, Any],
- resolving: frozenset[str],
-) -> Any:
- """Resolve a ``$ref`` pointer by inlining its definition from *defs*.
-
- Falls back to ``{"type": "object"}`` when the definition is missing
- or would cause a circular reference.
- """
- ref_path: str = obj["$ref"]
- ref_name = ref_path.rsplit("/", 1)[-1] if "/" in ref_path else ""
- definition = defs.get(ref_name) if ref_name else None
-
- if definition is not None and ref_name not in resolving:
- inlined = _compact_schema(
- definition,
- _defs=defs,
- _resolving=resolving | {ref_name},
- )
- if isinstance(inlined, dict):
- if desc := obj.get("description"):
- inlined.setdefault("description", desc)
- return inlined
-
- replacement: dict[str, Any] = {"type": "object"}
- if desc := obj.get("description"):
- replacement["description"] = desc
- return replacement
-
-
-def _compact_schema(
- obj: Any,
- *,
- _defs: dict[str, Any] | None = None,
- _resolving: frozenset[str] | None = None,
-) -> Any:
- """Collapse ``$defs`` and ``$ref`` pointers in a JSON Schema.
-
- Search results only need enough schema detail for the LLM to identify
- which tool to call and construct a basic invocation. Full schemas
- (with all nested model definitions) are still available when the tool
- is actually invoked via ``call_tool``.
-
- Transformations applied:
-
- * ``$defs`` sections are removed entirely.
- * ``{"$ref": "..."}`` is resolved by inlining the referenced
- definition from ``$defs``. If the definition cannot be found
- (or would cause a circular reference), the ref is replaced with
- ``{"type": "object"}``.
- * ``anyOf``/``oneOf`` lists containing only a ``$ref`` and
- ``{"type": "null"}`` (Pydantic's Optional encoding) are collapsed
- to the simplified non-null variant.
- """
- if isinstance(obj, list):
- return [
- _compact_schema(item, _defs=_defs, _resolving=_resolving) for item
in obj
- ]
- if not isinstance(obj, dict):
- return obj
-
- # On the first (top-level) call, extract $defs for later resolution.
- if _defs is None:
- _defs = obj.get("$defs", {})
- if _resolving is None:
- _resolving = frozenset()
-
- if "$ref" in obj:
- return _resolve_ref(obj, _defs, _resolving)
-
- result: dict[str, Any] = {}
- for key, value in obj.items():
- if key == "$defs":
- continue
- result[key] = _compact_schema(value, _defs=_defs,
_resolving=_resolving)
-
- return _simplify_optional_union(result)
-
-
def _truncate_description(text: str, max_length: int) -> str:
"""Truncate a tool description for search results.
@@ -528,21 +423,20 @@ def _create_search_result_serializer(
) -> Any:
"""Build a search-result serializer from the tool-search config.
- When ``include_schemas`` is False (default), delegates to
+ When ``include_schemas`` is False, delegates to
:func:`_build_summary_serializer`, which strips ``inputSchema``
entirely and adds a ``parameters_hint`` field with comma-separated
top-level parameter names. This reduces per-search token cost by
~80% vs compact mode while still conveying what parameters a tool
accepts.
- When ``include_schemas`` is True, the full ``compact_schemas``/
- ``max_description_length`` pipeline applies (existing behavior):
-
- * ``$defs`` sections and ``$ref`` pointers are collapsed when
- ``compact_schemas`` is True (see :func:`_compact_schema`).
- * Tool descriptions are truncated to ``max_description_length`` chars.
+ When ``include_schemas`` is True, input schemas retain their definitions,
+ references, and validation constraints. Inlining references duplicates
shared
+ chart models and can make a single tool exceed client result limits.
- Full schemas remain available when the tool is invoked via ``call_tool``.
+ Titles and output schemas are stripped by the base serializer. The legacy
+ ``compact_schemas`` setting only selects the default description limit;
+ ``max_description_length`` explicitly controls description truncation.
"""
include_schemas = config.get("include_schemas", False)
@@ -550,23 +444,18 @@ def _create_search_result_serializer(
max_desc = config.get("max_description_length", 300)
return _build_summary_serializer(max_desc)
- # include_schemas=True: apply full compact_schemas/max_description_length
pipeline
compact = config.get("compact_schemas", True)
# Description truncation defaults to 300 when compact_schemas is on,
# but is disabled when compact_schemas is off (unless explicitly set).
- max_desc_default = 300 if compact else 0
- max_desc = config.get("max_description_length", max_desc_default)
+ max_desc = config.get("max_description_length", 300 if compact else 0)
- if not compact and not max_desc:
+ if not max_desc:
return _serialize_tools_without_output_schema
def _serializer(tools: Sequence[Any]) -> list[dict[str, Any]]:
results = _serialize_tools_without_output_schema(tools)
for data in results:
- if compact:
- if input_schema := data.get("inputSchema"):
- data["inputSchema"] = _compact_schema(input_schema)
- if max_desc and (desc := data.get("description")):
+ if desc := data.get("description"):
data["description"] = _truncate_description(desc, max_desc)
return results
diff --git a/tests/unit_tests/mcp_service/test_chart_tool_inventory.py
b/tests/unit_tests/mcp_service/test_chart_tool_inventory.py
new file mode 100644
index 00000000000..7db3fd81582
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_chart_tool_inventory.py
@@ -0,0 +1,189 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Size and schema-first invocation regressions for chart tool inventory
entries."""
+
+from collections.abc import Callable, Iterator
+from copy import deepcopy
+from typing import Any
+
+import pytest
+from fastmcp.tools import FunctionTool
+from jsonschema import Draft202012Validator
+from pydantic import BaseModel
+
+from superset.mcp_service.app import mcp
+from superset.mcp_service.chart.schemas import (
+ GenerateChartRequest,
+ GenerateExploreLinkRequest,
+ UpdateChartRequest,
+)
+from superset.mcp_service.mcp_config import MCP_TOOL_SEARCH_CONFIG
+from superset.mcp_service.server import (
+ _create_search_result_serializer,
+ _strip_titles,
+)
+from superset.utils import json
+
+# Compact JSON, measured as UTF-8 bytes, including tool metadata.
+# Apache f8f293d2 -> reference-preserving inventory:
+# generate_chart: 97013 B -> 49531 B
+# update_chart: 104163 B -> 53395 B
+# generate_explore_link: 96581 B -> 49230 B
+# update_chart exceeds the proposed 50 kB target; 55 kB retains the entire
+# contract and stays below the 100 kB delivery cap. Byte budgets catch
reference
+# inlining without a tokenizer vocabulary download in the unit-test path.
+TOOL_BUDGETS = [
+ ("generate_chart", 50_000),
+ ("update_chart", 55_000),
+ ("generate_explore_link", 50_000),
+]
+
+
+def _references(node: Any) -> Iterator[str]:
+ """Collect schema references, including discriminator mapping targets."""
+ if isinstance(node, dict):
+ if "$ref" in node:
+ yield node["$ref"]
+ if "discriminator" in node:
+ yield from node["discriminator"].get("mapping", {}).values()
+ for value in node.values():
+ yield from _references(value)
+ elif isinstance(node, list):
+ for value in node:
+ yield from _references(value)
+
+
+def _resolve_pointer(schema: dict[str, Any], ref: str) -> Any:
+ """Resolve a local JSON Pointer against this tool's input schema."""
+ assert ref.startswith("#/")
+ target: Any = schema
+ for part in ref[2:].split("/"):
+ target = target[part.replace("~1", "/").replace("~0", "~")]
+ return target
+
+
[email protected]
[email protected](("name", "byte_budget"), TOOL_BUDGETS)
+async def test_chart_tool_inventory_size(name: str, byte_budget: int) -> None:
+ """Measure each real registered inventory entry, not a hand-built
schema."""
+ tool = await mcp.get_tool(name)
+ serializer = _create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)
+ entry = serializer([tool])[0]
+ assert "inputSchema" in entry # Summary mode must not mask a size
regression.
+ text = json.dumps(entry, ensure_ascii=False, separators=(",", ":"))
+ byte_count = len(text.encode("utf-8"))
+
+ assert byte_count <= byte_budget, (name, byte_count)
+
+
[email protected]
[email protected]("name", [row[0] for row in TOOL_BUDGETS])
+async def test_chart_tool_inventory_preserves_complete_schema(name: str) ->
None:
+ """All constraints and reference targets survive the production
serializer."""
+ tool = await mcp.get_tool(name)
+ original = deepcopy(tool.parameters)
+ entry = _create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)([tool])[0]
+ schema = entry["inputSchema"]
+
+ assert schema == _strip_titles(original)
+ assert tool.parameters == original
+ Draft202012Validator.check_schema(schema)
+ refs = list(_references(schema))
+ assert refs
+ for ref in refs:
+ assert _resolve_pointer(schema, ref) == _strip_titles(
+ _resolve_pointer(original, ref)
+ )
+
+ validator = Draft202012Validator(schema)
+ request: dict[str, Any] = {
+ "identifier" if name == "update_chart" else "dataset_id": 1,
+ "config": {"chart_type": "table", "columns": [{"name": "region"}]},
+ }
+ validator.validate({"request": request})
+ request["config"]["columns"] = [{"name": ""}]
+ assert not validator.is_valid({"request": request})
+
+
+def _generate_fixture(request: GenerateChartRequest) -> str:
+ """Accept a typed request without creating charts or querying a
database."""
+ return request.config.chart_type
+
+
+def _update_fixture(request: UpdateChartRequest) -> str:
+ """Accept a typed update without changing a saved chart."""
+ assert request.config is not None
+ return request.config.chart_type
+
+
+def _explore_fixture(request: GenerateExploreLinkRequest) -> str:
+ """Accept a typed request without writing to the permalink cache."""
+ assert request.config is not None
+ return request.config.chart_type
+
+
[email protected]
[email protected](
+ ("name", "fixture", "model"),
+ [
+ ("generate_chart", _generate_fixture, GenerateChartRequest),
+ ("update_chart", _update_fixture, UpdateChartRequest),
+ ("generate_explore_link", _explore_fixture,
GenerateExploreLinkRequest),
+ ],
+)
+async def test_chart_tool_schema_first_invocation(
+ name: str, fixture: Callable[..., str], model: type[BaseModel]
+) -> None:
+ """Follow published refs to build an invocation for a controlled tool
body."""
+ tool = await mcp.get_tool(name)
+ schema =
_create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)([tool])[0][
+ "inputSchema"
+ ]
+ request_schema = _resolve_pointer(schema,
schema["properties"]["request"]["$ref"])
+ assert request_schema == _strip_titles(schema["$defs"][model.__name__])
+ # Resolve the advertised table variant and column model rather than
inlining
+ # or substituting a generic object when a client encounters a reference.
+ table_ref = next(
+ ref
+ for ref in _references(request_schema["properties"]["config"])
+ if _resolve_pointer(schema, ref)
+ .get("properties", {})
+ .get("chart_type", {})
+ .get("const")
+ == "table"
+ )
+ table_schema = _resolve_pointer(schema, table_ref)
+ column_schema = _resolve_pointer(
+ schema, table_schema["properties"]["columns"]["items"]["$ref"]
+ )
+ assert "columns" in table_schema["required"]
+ assert "name" in column_schema["properties"]
+ request: dict[str, Any] = {
+ key: 1 for key in request_schema["required"] if key != "config"
+ }
+ request["config"] = {
+ "chart_type": table_schema["properties"]["chart_type"]["const"],
+ "columns": [{"name": "region"}],
+ }
+ arguments = {"request": request}
+ Draft202012Validator(schema).validate(arguments)
+ # Use the same request model through FastMCP's real argument validation,
+ # but replace persistence/query execution with the explicitly typed
fixture.
+ controlled_tool = FunctionTool.from_function(fixture)
+ result = await controlled_tool.run(arguments)
+ assert result.content[0].text == "table"
diff --git a/tests/unit_tests/mcp_service/test_tool_inventory.py
b/tests/unit_tests/mcp_service/test_tool_inventory.py
new file mode 100644
index 00000000000..03a451eaf9a
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_tool_inventory.py
@@ -0,0 +1,157 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Per-tool size and schema fidelity budgets for the entire registered
inventory."""
+
+from copy import deepcopy
+
+import pytest
+from jsonschema import Draft202012Validator
+
+from superset.mcp_service.app import mcp
+from superset.mcp_service.mcp_config import MCP_TOOL_SEARCH_CONFIG
+from superset.mcp_service.server import _create_search_result_serializer,
_strip_titles
+from superset.utils import json
+
+# Compact JSON including tool metadata, measured as UTF-8 bytes.
+# Small-tool budgets are fixed snapshots: ceil(measured_bytes / 100) * 100 +
100,
+# leaving 100-199 bytes for incidental description edits. Do not recompute
limits
+# at test time: they must catch schema growth. Large chart tools retain their
+# explicit delivery budgets, well below the sizes produced by reference
inlining.
+TOOL_BUDGETS = {
+ "add_chart_to_existing_dashboard": 1_500,
+ "apply_dashboard_filters": 2_900,
+ "create_dataset": 1_800,
+ "create_theme": 1_100,
+ "create_virtual_dataset": 3_700,
+ "delete_chart": 1_100,
+ "delete_dashboard": 1_100,
+ "duplicate_dashboard": 1_900,
+ "execute_sql": 2_100,
+ "find_users": 1_500,
+ "generate_bug_report": 2_600,
+ "generate_chart": 50_000,
+ "generate_dashboard": 3_400,
+ "generate_explore_link": 50_000,
+ "get_annotation_layer_info": 1_000,
+ "get_chart_data": 2_900,
+ "get_chart_info": 3_600,
+ "get_chart_preview": 3_400,
+ "get_chart_sql": 2_100,
+ "get_chart_type_schema": 900,
+ "get_compatible_dimensions": 1_500,
+ "get_compatible_metrics": 1_500,
+ "get_dashboard_data": 2_400,
+ "get_dashboard_datasets": 1_100,
+ "get_dashboard_info": 3_100,
+ "get_dashboard_layout": 1_600,
+ "get_database_info": 1_400,
+ "get_dataset_info": 2_400,
+ "get_instance_info": 900,
+ "get_layer_annotation_info": 1_100,
+ "get_query_info": 1_100,
+ "get_report_info": 1_400,
+ "get_rls_filter_info": 1_000,
+ "get_role_info": 900,
+ "get_saved_query_info": 1_200,
+ "get_schema": 1_100,
+ "get_table": 3_900,
+ "get_tag_info": 1_000,
+ "get_task_info": 1_100,
+ "get_theme_info": 1_000,
+ "get_user_info": 1_000,
+ "health_check": 700,
+ "list_annotation_layers": 2_700,
+ # Include the deleted_state edit/restore audience and under-enumeration
+ # caveats from #44128: 5,149 and 4,626 bytes, plus the headroom above.
+ # Keep the complete-schema parity test below alongside these size limits.
+ "list_charts": 5_300,
+ "list_dashboards": 4_800,
+ "list_databases": 3_500,
+ "list_datasets": 4_600,
+ "list_layer_annotations": 2_900,
+ "list_metrics": 1_900,
+ "list_queries": 3_000,
+ "list_reports": 3_900,
+ "list_rls_filters": 2_600,
+ "list_roles": 2_700,
+ "list_saved_queries": 3_000,
+ "list_tags": 3_100,
+ "list_tasks": 2_700,
+ "list_themes": 3_000,
+ "list_users": 2_900,
+ "manage_dashboard_certification": 1_900,
+ "manage_dashboard_owners": 2_200,
+ "manage_dashboard_roles": 1_900,
+ "manage_native_filters": 6_700,
+ "open_sql_lab_with_context": 1_800,
+ "query_dataset": 3_700,
+ "remove_chart_from_dashboard": 1_300,
+ "restore_chart": 1_100,
+ "restore_dashboard": 1_000,
+ "save_sql_query": 1_600,
+ "update_chart": 55_000,
+ "update_chart_preview": 55_000,
+ "update_dashboard": 4_100,
+ "update_dataset_metric": 3_100,
+}
+
+
[email protected]
+async def test_inventory_budgets_cover_every_registered_tool() -> None:
+ """New or renamed tools must get explicit budgets instead of escaping
checks."""
+ tools = await mcp.list_tools(run_middleware=False)
+ assert {tool.name for tool in tools} == set(TOOL_BUDGETS)
+
+
[email protected]
[email protected]("name", TOOL_BUDGETS)
+async def test_tool_inventory_size(name: str) -> None:
+ """Limit every real serialized entry, including non-chart tools and
metadata."""
+ tool = await mcp.get_tool(name)
+ assert tool is not None
+ entry = _create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)([tool])[0]
+ assert "inputSchema" in entry # Summary mode must not hide schema growth.
+ text = json.dumps(entry, ensure_ascii=False, separators=(",", ":"))
+ byte_count = len(text.encode("utf-8"))
+ byte_budget = TOOL_BUDGETS[name]
+ assert byte_count <= byte_budget, (name, byte_count, byte_budget)
+
+
[email protected]
[email protected]("name", TOOL_BUDGETS)
+async def test_tool_inventory_preserves_complete_schema(name: str) -> None:
+ """Size optimizations must not discard nullable unions, refs or
constraints."""
+ tool = await mcp.get_tool(name)
+ assert tool is not None
+ original = deepcopy(tool.parameters)
+ entry = _create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)([tool])[0]
+ schema = entry["inputSchema"]
+ assert schema == _strip_titles(original)
+ assert tool.parameters == original
+ Draft202012Validator.check_schema(schema)
+
+
[email protected]
+async def test_inventory_budget_allows_incidental_description_edit() -> None:
+ """A small wording change must not exhaust the chart preview byte
budget."""
+ tool = await mcp.get_tool("get_chart_preview")
+ assert tool is not None
+ entry = _create_search_result_serializer(MCP_TOOL_SEARCH_CONFIG)([tool])[0]
+ entry["description"] += " A chart preview."
+ text = json.dumps(entry, ensure_ascii=False, separators=(",", ":"))
+ assert len(text.encode("utf-8")) <= TOOL_BUDGETS["get_chart_preview"]
diff --git a/tests/unit_tests/mcp_service/test_tool_search_transform.py
b/tests/unit_tests/mcp_service/test_tool_search_transform.py
index 2e7e70eaa23..bc51229894f 100644
--- a/tests/unit_tests/mcp_service/test_tool_search_transform.py
+++ b/tests/unit_tests/mcp_service/test_tool_search_transform.py
@@ -21,6 +21,7 @@ import logging
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
+import pytest
from fastmcp.server.transforms.search import BM25SearchTransform,
RegexSearchTransform
from flask import Flask, g
@@ -29,7 +30,6 @@ from superset.mcp_service.mcp_config import
MCP_TOOL_SEARCH_CONFIG
from superset.mcp_service.privacy import requires_data_model_metadata_access
from superset.mcp_service.server import (
_apply_tool_search_transform,
- _compact_schema,
_create_search_result_serializer,
_extract_parameter_names,
_filter_tools_by_current_user_permission,
@@ -312,317 +312,66 @@ def test_normalize_ignores_keys_not_in_schema():
assert isinstance(result["unknown_key"], dict)
-# -- _compact_schema tests --
+# -- search schema fidelity tests --
-def test_compact_schema_removes_defs():
- """$defs section is stripped and refs are inlined from definitions."""
+def test_search_schema_preserves_references_and_constraints() -> None:
+ """Shared and recursive definitions remain resolvable without losing
constraints."""
schema = {
"type": "object",
"properties": {
- "filters": {"items": {"$ref": "#/$defs/MyFilter"}, "type":
"array"},
+ "first": {"$ref": "#/$defs/Node", "description": "First",
"maxItems": 2},
+ "second": {"$ref": "#/$defs/Node"},
},
+ "required": ["first"],
+ "additionalProperties": False,
"$defs": {
- "MyFilter": {
- "type": "object",
- "properties": {"col": {"type": "string"}},
- }
- },
- }
-
- result = _compact_schema(schema)
-
- assert "$defs" not in result
- # $ref is resolved by inlining the definition from $defs
- assert result["properties"]["filters"]["items"] == {
- "type": "object",
- "properties": {"col": {"type": "string"}},
- }
-
-
-def test_compact_schema_replaces_ref_with_object():
- """Direct $ref is replaced with {"type": "object"}."""
- schema = {"$ref": "#/$defs/SomeModel"}
-
- result = _compact_schema(schema)
-
- assert result == {"type": "object"}
-
-
-def test_compact_schema_preserves_ref_description():
- """$ref replacement preserves sibling description if present."""
- schema = {"$ref": "#/$defs/SomeModel", "description": "A model"}
-
- result = _compact_schema(schema)
-
- assert result == {"type": "object", "description": "A model"}
-
-
-def test_compact_schema_simplifies_optional_ref():
- """anyOf with $ref and null is collapsed to the non-null variant."""
- schema = {
- "anyOf": [
- {"$ref": "#/$defs/SomeModel"},
- {"type": "null"},
- ],
- "description": "Optional model",
- }
-
- result = _compact_schema(schema)
-
- assert "anyOf" not in result
- assert result["type"] == "object"
- assert result["description"] == "Optional model"
-
-
-def test_compact_schema_simplifies_optional_primitive():
- """anyOf with primitive type and null is collapsed."""
- schema = {
- "anyOf": [
- {"type": "integer"},
- {"type": "null"},
- ],
- "default": None,
- }
-
- result = _compact_schema(schema)
-
- assert "anyOf" not in result
- assert result["type"] == "integer"
- assert result["default"] is None
-
-
-def test_compact_schema_preserves_multi_variant_anyof():
- """anyOf with >2 variants or no null is left unchanged."""
- schema = {
- "anyOf": [
- {"type": "integer"},
- {"type": "string"},
- ]
- }
-
- result = _compact_schema(schema)
-
- assert "anyOf" in result
- assert len(result["anyOf"]) == 2
-
-
-def test_compact_schema_nested_in_items():
- """$ref nested inside items/properties falls back to object when no
$defs."""
- schema = {
- "type": "object",
- "properties": {
- "filters": {
- "type": "array",
- "items": {"$ref": "#/$defs/Filter"},
- },
- "name": {"type": "string"},
- },
- }
-
- result = _compact_schema(schema)
-
- # No $defs in schema → falls back to {"type": "object"}
- assert result["properties"]["filters"]["items"] == {"type": "object"}
- assert result["properties"]["name"] == {"type": "string"}
-
-
-def test_compact_schema_passthrough_simple():
- """Simple schema without $defs/$ref passes through unchanged."""
- schema = {
- "type": "object",
- "properties": {
- "page": {"type": "integer", "default": 1},
- "search": {"type": "string"},
- },
- }
-
- result = _compact_schema(schema)
-
- assert result == schema
-
-
-def test_compact_schema_handles_non_dict():
- """Non-dict inputs pass through unchanged."""
- assert _compact_schema("hello") == "hello"
- assert _compact_schema(42) == 42
- assert _compact_schema(None) is None
- assert _compact_schema([1, 2]) == [1, 2]
-
-
-def test_compact_schema_inlines_columnref_like_model() -> None:
- """$ref to a model with fields is inlined, preserving field structure."""
- schema = {
- "type": "object",
- "properties": {
- "metrics": {
+ "Node": {
"type": "array",
- "items": {"$ref": "#/$defs/ColumnRef"},
- },
- },
- "$defs": {
- "ColumnRef": {
- "type": "object",
- "properties": {
- "name": {"type": "string", "description": "Column name"},
- "aggregate": {"type": "string", "description": "SQL
aggregate"},
- "saved_metric": {"type": "boolean", "default": False},
+ "minItems": 1,
+ "items": {
+ "anyOf": [{"$ref": "#/$defs/Node"}, {"type": "integer"}],
},
- "required": ["name"],
- }
- },
- }
-
- result = _compact_schema(schema)
-
- assert "$defs" not in result
- inlined = result["properties"]["metrics"]["items"]
- assert inlined["type"] == "object"
- assert "name" in inlined["properties"]
- assert "aggregate" in inlined["properties"]
- assert "saved_metric" in inlined["properties"]
- assert inlined["required"] == ["name"]
-
-
-def test_compact_schema_inlines_nested_refs() -> None:
- """Nested $ref chains are resolved transitively."""
- schema = {
- "type": "object",
- "properties": {
- "config": {"$ref": "#/$defs/ChartConfig"},
- },
- "$defs": {
- "ChartConfig": {
- "type": "object",
- "properties": {
- "metrics": {
- "type": "array",
- "items": {"$ref": "#/$defs/ColumnRef"},
- },
- },
- },
- "ColumnRef": {
- "type": "object",
- "properties": {
- "name": {"type": "string"},
- },
- "required": ["name"],
},
},
}
+ serializer = _create_search_result_serializer({"include_schemas": True})
+ result = serializer([_make_mock_tool("tree", "A tree.", schema)])
- result = _compact_schema(schema)
+ assert result[0]["inputSchema"] == schema
- assert "$defs" not in result
- config = result["properties"]["config"]
- assert config["type"] == "object"
- col_ref = config["properties"]["metrics"]["items"]
- assert col_ref["type"] == "object"
- assert "name" in col_ref["properties"]
-
-def test_compact_schema_circular_ref_fallback() -> None:
- """Circular $ref falls back to {"type": "object"} to avoid infinite
recursion."""
+def test_search_schema_preserves_nullable_unions() -> None:
+ """Nullable unions, siblings, defaults, and discriminator mappings are
guidance."""
schema = {
"type": "object",
"properties": {
- "node": {"$ref": "#/$defs/TreeNode"},
- },
- "$defs": {
- "TreeNode": {
- "type": "object",
- "properties": {
- "children": {
- "type": "array",
- "items": {"$ref": "#/$defs/TreeNode"},
- },
+ "config": {
+ "oneOf": [{"$ref": "#/$defs/Config"}, {"type": "null"}],
+ "description": "Optional config",
+ "default": None,
+ "discriminator": {
+ "propertyName": "kind",
+ "mapping": {"table": "#/$defs/Config"},
},
- }
- },
- }
-
- result = _compact_schema(schema)
-
- assert "$defs" not in result
- node = result["properties"]["node"]
- assert node["type"] == "object"
- # The self-reference falls back to {"type": "object"}
- assert node["properties"]["children"]["items"] == {"type": "object"}
-
-
-def test_compact_schema_inline_with_sibling_description() -> None:
- """Inlined $ref preserves sibling description via setdefault."""
- schema = {
- "type": "object",
- "properties": {
- "col": {
- "$ref": "#/$defs/ColumnRef",
- "description": "The column to use",
},
- },
- "$defs": {
- "ColumnRef": {
- "type": "object",
- "properties": {"name": {"type": "string"}},
- }
- },
- }
-
- result = _compact_schema(schema)
-
- col = result["properties"]["col"]
- assert col["type"] == "object"
- assert "name" in col["properties"]
- assert col["description"] == "The column to use"
-
-
-def test_compact_schema_inline_optional_ref() -> None:
- """anyOf with $ref and null inlines the definition, not just {"type":
"object"}."""
- schema = {
- "type": "object",
- "properties": {
- "col": {
- "anyOf": [
- {"$ref": "#/$defs/ColumnRef"},
- {"type": "null"},
- ],
- "description": "Optional column",
+ "count": {
+ "anyOf": [{"type": "integer", "minimum": 1}, {"type": "null"}],
+ "default": None,
},
},
"$defs": {
- "ColumnRef": {
+ "Config": {
"type": "object",
- "properties": {"name": {"type": "string"}},
- }
- },
- }
-
- result = _compact_schema(schema)
-
- col = result["properties"]["col"]
- assert "anyOf" not in col
- assert col["type"] == "object"
- assert "name" in col["properties"]
- assert col["description"] == "Optional column"
-
-
-def test_compact_schema_inlines_empty_def() -> None:
- """Empty $defs entry ({}) is inlined as-is, not downgraded to {"type":
"object"}."""
- schema = {
- "type": "object",
- "properties": {
- "anything": {"$ref": "#/$defs/Anything"},
- },
- "$defs": {
- "Anything": {},
+ "properties": {"kind": {"const": "table"}},
+ "required": ["kind"],
+ },
},
}
+ serializer = _create_search_result_serializer({"include_schemas": True})
+ result = serializer([_make_mock_tool("nullable", "Nullable.", schema)])
- result = _compact_schema(schema)
-
- assert "$defs" not in result
- # Empty dict is a valid schema — should be inlined as {}, not {"type":
"object"}
- assert result["properties"]["anything"] == {}
+ assert result[0]["inputSchema"] == schema
# -- _truncate_description tests --
@@ -680,7 +429,7 @@ def _make_mock_tool(name, description, input_schema):
def test_create_serializer_compacts_schemas():
- """Compact serializer strips $defs and replaces $ref."""
+ """Search serialization keeps shared definitions and their references."""
tool = _make_mock_tool(
"list_charts",
"List charts with filtering.",
@@ -708,12 +457,8 @@ def test_create_serializer_compacts_schemas():
assert len(result) == 1
schema = result[0]["inputSchema"]
- assert "$defs" not in schema
- # $ref is inlined from $defs, preserving field structure
- assert schema["properties"]["filters"]["items"] == {
- "type": "object",
- "properties": {"col": {"type": "string"}},
- }
+ assert schema["$defs"]["ChartFilter"]["properties"] == {"col": {"type":
"string"}}
+ assert schema["properties"]["filters"]["items"] == {"$ref":
"#/$defs/ChartFilter"}
def test_create_serializer_truncates_descriptions():
@@ -731,8 +476,9 @@ def test_create_serializer_truncates_descriptions():
assert len(result[0]["description"]) <= 53 # 50 + potential "..."
-def test_create_serializer_disabled():
- """When compact_schemas=False and max_description_length=0, no
compaction."""
[email protected]("compact", [False, True])
+def test_create_serializer_disabled(compact: bool) -> None:
+ """An explicit zero description limit disables truncation in either
mode."""
tool = _make_mock_tool(
"test_tool",
"A long description " * 20,
@@ -743,7 +489,11 @@ def test_create_serializer_disabled():
)
serializer = _create_search_result_serializer(
- {"include_schemas": True, "compact_schemas": False,
"max_description_length": 0}
+ {
+ "include_schemas": True,
+ "compact_schemas": compact,
+ "max_description_length": 0,
+ }
)
result = serializer([tool])
@@ -1207,7 +957,7 @@ def
test_create_serializer_include_schemas_true_restores_full_schema():
def test_create_serializer_include_schemas_true_with_compact():
- """include_schemas=True + compact_schemas=True still compacts the
schema."""
+ """The legacy compact setting must not expand or weaken input schemas."""
schema = {
"type": "object",
"properties": {
@@ -1223,10 +973,7 @@ def
test_create_serializer_include_schemas_true_with_compact():
result = serializer([tool])
assert "inputSchema" in result[0]
- assert "$defs" not in result[0]["inputSchema"]
- assert result[0]["inputSchema"]["properties"]["filters"]["items"] == {
- "type": "object"
- }
+ assert result[0]["inputSchema"] == schema
# -- search_tools optional query tests --
