fix(mcp): wire up compact schema serialization for search_tools results (#39229)

(cherry picked from commit e17cf3c)

Author: aminghadersohi

superset/mcp_service/mcp_config.py

@@ -248,9 +248,19 @@
 # - "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.
+#
 # Rollback:
 # ---------
-# Set enabled=False in superset_config.py for instant 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.
 # =============================================================================
 MCP_TOOL_SEARCH_CONFIG: Dict[str, Any] = {
     "enabled": True,  # Enabled by default — reduces initial context by ~70%
@@ -262,6 +272,8 @@
     ],
     "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 and simplify $ref in search results
+    "max_description_length": 300,  # Truncate tool descriptions (0 = no truncation)
 }
 
 

superset/mcp_service/server.py

@@ -175,6 +175,84 @@ 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 _compact_schema(obj: Any) -> 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 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) for item in obj]
+    if not isinstance(obj, dict):
+        return obj
+
+    # Direct $ref → generic object type
+    if "$ref" in obj:
+        replacement: dict[str, Any] = {"type": "object"}
+        if desc := obj.get("description"):
+            replacement["description"] = desc
+        return replacement
+
+    result: dict[str, Any] = {}
+    for key, value in obj.items():
+        if key == "$defs":
+            continue
+        result[key] = _compact_schema(value)
+
+    return _simplify_optional_union(result)
+
+
+def _truncate_description(text: str, max_length: int) -> str:
+    """Truncate a tool description for search results.
+
+    Cuts at the last sentence boundary before *max_length*, or at
+    *max_length* with an ellipsis if no sentence boundary is found.
+    """
+    if not text or len(text) <= max_length:
+        return text
+    # Try to cut at the last sentence boundary
+    truncated = text[:max_length]
+    last_period = truncated.rfind(". ")
+    if last_period > max_length // 2:
+        return truncated[: last_period + 1]
+    return truncated.rstrip() + "..."
+
+
 def _serialize_tools_without_output_schema(
     tools: Sequence[Any],
 ) -> list[dict[str, Any]]:
@@ -194,6 +272,46 @@ def _serialize_tools_without_output_schema(
     return results
 
 
+def _create_search_result_serializer(
+    config: dict[str, Any],
+) -> Any:
+    """Build a search-result serializer from the tool-search config.
+
+    When ``compact_schemas`` is enabled (default), the serializer applies
+    additional compaction on top of the base serialization:
+
+    * ``$defs`` sections and ``$ref`` pointers are collapsed
+      (see :func:`_compact_schema`).
+    * Tool descriptions are truncated to ``max_description_length`` chars.
+
+    This reduces per-search-call token cost by ~40-60 % while keeping
+    enough detail for the LLM to identify the right tool and construct
+    a basic invocation.
+    """
+    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)
+
+    if not compact and 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 "description" in data:
+                data["description"] = _truncate_description(
+                    data["description"], max_desc
+                )
+        return results
+
+    return _serializer
+
+
 def _fix_call_tool_arguments(tool: Any) -> Any:
     """Fix anyOf schema in call_tool ``arguments`` for MCP bridge compatibility.
 
@@ -270,7 +388,7 @@ def _apply_tool_search_transform(mcp_instance: Any, config: dict[str, Any]) -> N
         "always_visible": config.get("always_visible", []),
         "search_tool_name": config.get("search_tool_name", "search_tools"),
         "call_tool_name": config.get("call_tool_name", "call_tool"),
-        "search_result_serializer": _serialize_tools_without_output_schema,
+        "search_result_serializer": _create_search_result_serializer(config),
     }
 
     def _make_normalizing_call_tool(transform: Any) -> Tool: