aminghadersohi commented on code in PR #45081: URL: https://github.com/apache/superset/pull/45081#discussion_r4217619235
########## tests/unit_tests/mcp_service/test_native_tool_surface.py: ########## @@ -0,0 +1,662 @@ +# 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. + +"""Contract tests for the native named-tool surface and its compatibility mode. + +The MCP service can serve its catalog in two shapes, both built from the same +registered tool definitions and the same middleware stack: + +* native: ``tools/list`` advertises every permitted tool under its real name + and ``tools/call`` invokes it directly (``MCP_TOOL_SEARCH_CONFIG`` with + ``enabled=False``); +* compatibility (the default): ``tools/list`` advertises pinned tools plus the + synthetic ``search_tools``/``call_tool`` pair, and other tools are discovered + through search and invoked through the ``call_tool`` proxy. + +Every registered tool stays callable under its real name in both shapes. These +tests assemble each shape the way ``run_server`` does, from +``build_middleware_list`` and ``_apply_tool_search_transform``, and check that +both expose the same tools, schemas, results, errors, and authorization +decisions. + +``native_tool_inventory.json`` records the measured native catalog: each +tool's annotations, description length, schema sizes, and complete +``tools/list`` entry size with structured output disabled and enabled, plus +catalog totals. Sizes may shrink freely but must not grow beyond a small +headroom. To accept an intentional change, regenerate the report with:: + + SUPERSET_MCP_UPDATE_TOOL_INVENTORY=1 pytest \\ + tests/unit_tests/mcp_service/test_native_tool_surface.py +""" + +import math +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import mcp.types as mt +import pytest +from fastmcp import Client, FastMCP +from fastmcp.client.client import CallToolResult +from fastmcp.tools.tool import Tool +from fastmcp.utilities.json_schema import dereference_refs +from jsonschema import Draft202012Validator + +from superset.mcp_service.app import ALLOWED_UNPROTECTED, mcp +from superset.mcp_service.chart.schemas import CHART_TYPE_VALUES +from superset.mcp_service.mcp_config import MCP_TOOL_SEARCH_CONFIG +from superset.mcp_service.server import ( + _apply_tool_search_transform, + build_middleware_list, +) +from superset.utils import json + +INVENTORY_PATH = Path(__file__).with_name("native_tool_inventory.json") +UPDATE_INVENTORY_ENV = "SUPERSET_MCP_UPDATE_TOOL_INVENTORY" + +SEARCH_TOOL = MCP_TOOL_SEARCH_CONFIG["search_tool_name"] +CALL_TOOL = MCP_TOOL_SEARCH_CONFIG["call_tool_name"] +PINNED_TOOLS = set(MCP_TOOL_SEARCH_CONFIG["always_visible"]) + +# Output modes of MCP_STRUCTURED_OUTPUT_ENABLED, keyed as in the inventory. +OUTPUT_MODES = {"text_only": False, "structured": True} + +# MCP gateways cap a single list page at 100 KB; see test_tool_inventory.py. +GATEWAY_PAGE_BYTE_LIMIT = 100_000 + +# The registry-derived chart_type enum grows with every registered chart type, +# so it is excluded from size budgets as in test_tool_inventory.py. +CHART_TYPE_ENUM = json.dumps( + CHART_TYPE_VALUES, ensure_ascii=False, separators=(",", ":") +) + +ANNOTATION_KEYS = ("readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint") + + +def compact_bytes(value: Any) -> int: + """Measure compact UTF-8 JSON, excluding the chart_type enum.""" + text = json.dumps(value, ensure_ascii=False, separators=(",", ":")) + enum_bytes = len(CHART_TYPE_ENUM.encode("utf-8")) * text.count(CHART_TYPE_ENUM) + return len(text.encode("utf-8")) - enum_bytes + + +def size_budget(recorded: int) -> int: + """Allow incidental growth above a recorded size: 2%, at least 200 bytes.""" + return recorded + max(200, math.ceil(recorded * 0.02)) + + +def wire_entry(tool: mt.Tool) -> dict[str, Any]: + """Return a tool definition exactly as serialized in ``tools/list``.""" + return tool.model_dump(by_alias=True, mode="json", exclude_none=True) + + +async def canonical_tools() -> list[Tool]: + """Return the registered tool definitions, without request middleware.""" + return list(await mcp.list_tools(run_middleware=False)) + + +async def build_server( + *, structured_output_enabled: bool, compatibility: bool +) -> FastMCP: + """Assemble a server from the registered tools as ``run_server`` does.""" + server = FastMCP( + "superset-mcp-surface-test", + middleware=build_middleware_list( + structured_output_enabled=structured_output_enabled + ), + ) + for tool in await canonical_tools(): + server.add_tool(tool) + if compatibility: + _apply_tool_search_transform(server, dict(MCP_TOOL_SEARCH_CONFIG)) + return server + + +async def list_native(structured_output_enabled: bool) -> mt.ListToolsResult: + """List the native catalog over the MCP protocol.""" + server = await build_server( + structured_output_enabled=structured_output_enabled, compatibility=False + ) + async with Client(server) as client: + return await client.list_tools_mcp() + + +def unresolved_refs(schema: dict[str, Any]) -> set[str]: + """Return local ``$ref`` targets that do not resolve inside ``schema``.""" + missing: set[str] = set() + + def resolves(ref: str) -> bool: + node: Any = schema + for part in ref.removeprefix("#/").split("/"): + part = part.replace("~1", "/").replace("~0", "~") + if not isinstance(node, dict) or part not in node: + return False + node = node[part] + return True + + def walk(node: Any) -> None: + if isinstance(node, dict): + ref = node.get("$ref") + if isinstance(ref, str) and (not ref.startswith("#") or not resolves(ref)): + missing.add(ref) + for value in node.values(): + walk(value) + elif isinstance(node, list): + for value in node: + walk(value) + + walk(schema) + return missing + + +async def measure_inventory() -> dict[str, Any]: + """Measure the native catalog in both structured-output modes.""" + listings = { + mode: await list_native(enabled) for mode, enabled in OUTPUT_MODES.items() + } + structured = {tool.name: tool for tool in listings["structured"].tools} + tools: dict[str, Any] = {} + for tool in listings["text_only"].tools: + annotations = tool.annotations.model_dump() if tool.annotations else {} + output_schema = structured[tool.name].outputSchema + tools[tool.name] = { + "annotations": {key: annotations.get(key) for key in ANNOTATION_KEYS}, + "description_chars": len(tool.description or ""), + "input_schema_bytes": compact_bytes(tool.inputSchema), + "output_schema_bytes": compact_bytes(output_schema) if output_schema else 0, + "entry_bytes": { + "text_only": compact_bytes(wire_entry(tool)), + "structured": compact_bytes(wire_entry(structured[tool.name])), + }, + } + totals: dict[str, Any] = {"tool_count": len(tools), "catalog_bytes": {}} + totals["largest_entry"] = {} + for mode, listing in listings.items(): + entries = {tool.name: wire_entry(tool) for tool in listing.tools} + totals["catalog_bytes"][mode] = compact_bytes(list(entries.values())) + name = max(entries, key=lambda key: compact_bytes(entries[key])) + totals["largest_entry"][mode] = { + "name": name, + "bytes": compact_bytes(entries[name]), + } + return {"tools": dict(sorted(tools.items())), "totals": totals} + + [email protected] +async def test_native_inventory_report_matches_registered_tools() -> None: + """The recorded report covers every tool and bounds every measured size.""" + measured = await measure_inventory() + if os.environ.get(UPDATE_INVENTORY_ENV): + INVENTORY_PATH.write_text(json.dumps(measured, indent=2, sort_keys=True) + "\n") + recorded = json.loads(INVENTORY_PATH.read_text()) + + assert set(measured["tools"]) == set(recorded["tools"]) + assert measured["totals"]["tool_count"] == recorded["totals"]["tool_count"] + + grown: dict[str, tuple[int, int]] = {} + for name, entry in measured["tools"].items(): + expected = recorded["tools"][name] + assert entry["annotations"] == expected["annotations"], name + for key in ("description_chars", "input_schema_bytes", "output_schema_bytes"): + if entry[key] > size_budget(expected[key]): + grown[f"{name}.{key}"] = (entry[key], expected[key]) + for mode in OUTPUT_MODES: + actual = entry["entry_bytes"][mode] + if actual > size_budget(expected["entry_bytes"][mode]): + grown[f"{name}.entry_bytes.{mode}"] = ( + actual, + expected["entry_bytes"][mode], + ) + for mode in OUTPUT_MODES: + actual = measured["totals"]["catalog_bytes"][mode] + expected_total = recorded["totals"]["catalog_bytes"][mode] + if actual > size_budget(expected_total): + grown[f"catalog_bytes.{mode}"] = (actual, expected_total) + assert not grown, ( + f"Native tool definitions grew beyond the recorded inventory; review the " + f"growth and rerun with {UPDATE_INVENTORY_ENV}=1: {grown}" + ) + + [email protected] [email protected]("mode", OUTPUT_MODES) +async def test_largest_native_entry_fits_one_gateway_page(mode: str) -> None: + """Any single tool definition fits a gateway list page on its own.""" + listing = await list_native(OUTPUT_MODES[mode]) + sizes = {tool.name: compact_bytes(wire_entry(tool)) for tool in listing.tools} + largest = max(sizes, key=sizes.__getitem__) + assert sizes[largest] <= GATEWAY_PAGE_BYTE_LIMIT, (largest, sizes[largest]) Review Comment: Agreed on the facts: the native listing is one unpaginated page of 218 KB text-only / 524 KB structured (recorded in the inventory totals), so the per-entry check alone doesn't guarantee gateway-safe discovery. Native mode is opt-in, and deployments behind a 100 KB gateway are expected to stay on the default compatibility mode. Fixed in e388bda075d893b09494ada7c2308bc26f08fd1c: the per-entry test's docstring now says it bounds entries only and that the native catalog spans several gateway pages, and a new test_compatibility_listing_fits_one_gateway_page asserts the default tools/list is a single page under 100 KB in both structured-output modes. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
