aminghadersohi commented on code in PR #44575: URL: https://github.com/apache/superset/pull/44575#discussion_r4088422285
########## tests/unit_tests/mcp_service/test_tool_inventory.py: ########## @@ -0,0 +1,147 @@ +# 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 +import tiktoken +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 independently as UTF-8 bytes +# and tiktoken 0.14.0 / cl100k_base tokens (BPE estimates, not Claude counts). +# Small tools round up to 100 bytes / 25 tokens; large chart tools retain their +# explicit delivery budgets. Keep both limits: bytes are not a token estimate. +TOOL_BUDGETS = { + "add_chart_to_existing_dashboard": (1_400, 325), + "apply_dashboard_filters": (2_800, 625), + "create_dataset": (1_700, 375), + "create_theme": (1_000, 225), + "create_virtual_dataset": (3_600, 800), + "delete_chart": (1_000, 225), + "delete_dashboard": (1_000, 225), + "duplicate_dashboard": (1_800, 375), + "execute_sql": (2_000, 450), + "find_users": (1_400, 325), + "generate_bug_report": (2_500, 575), + "generate_chart": (50_000, 20_000), + "generate_dashboard": (3_300, 725), + "generate_explore_link": (50_000, 20_000), + "get_annotation_layer_info": (900, 200), + "get_chart_data": (2_800, 625), + "get_chart_info": (3_500, 800), + "get_chart_preview": (3_300, 750), Review Comment: Fixed in b5a91a3f4ae94a0c3fb994bb2a5e12f2ba30cd15. Regenerated the fixed small-tool byte budgets as ceil(measured_bytes / 100) * 100 + 100, so they have 100–199 bytes of headroom rather than just 1–100. All 72 tools have at least 100 bytes of measured slack; the large chart caps remain 50/55 kB. Your exact wording change measures 3,302 bytes against a 3,400-byte budget and passes. Added a regression test for an incidental description edit. Token budgets are removed as discussed in the dependency thread. Validation: pytest tests/unit_tests/mcp_service: 5,205 passed, 4 skipped (518.56s). The skips are three Vega/Node prerequisites and one unavailable optional sqlalchemy_bigquery import; the latter accounts for the extra skip versus your baseline. Focused inventory/search tests: 209 passed. Pre-commit and explicit mypy hooks passed. ########## pyproject.toml: ########## @@ -327,6 +327,7 @@ development = [ "sqloxide", "statsd", "syntaqlite>=0.9.0,<0.10.0", + "tiktoken==0.14.0", # reproducible MCP inventory token budgets Review Comment: Fixed in b5a91a3f4ae94a0c3fb994bb2a5e12f2ba30cd15. Removed tiktoken imports/assertions from both inventory test files and its development dependency, including the lockfile entry and tokenizer-only regex dependency. I chose removal rather than skip/cache logic because the byte checks independently reject the regression: running the f8f293d2 serializer against the registered tools produced 97,013 / 104,163 / 96,581 bytes for generate_chart / update_chart / generate_explore_link, above the unchanged 50,000 / 55,000 / 50,000 caps. update_chart_preview is also caught (102,352 > 55,000). Both inventory files passed with socket connections blocked and tiktoken imports deliberately made unavailable: 155 passed. This keeps the checks active offline rather than conditionally skipping them; no vocabulary download or approximate token substitution is needed. ########## docs/admin_docs/configuration/mcp-server.mdx: ########## @@ -775,7 +775,7 @@ MCP_TOOL_SEARCH_CONFIG = { | `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. | +| `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`. | Review Comment: Fixed in b5a91a3f4ae94a0c3fb994bb2a5e12f2ba30cd15. Updated mcp_config.py's schema serialization block to describe preservation of $defs, $ref, nullable unions and constraints, and local reference resolution. Removed the obsolete compact_schemas=False rollback bullet and replaced the inline “Strip $defs/$ref” comment. The legacy setting is documented as selecting only the default description limit when max_description_length is omitted, consistent with the MDX documentation. ########## docs/admin_docs/configuration/mcp-server.mdx: ########## @@ -775,7 +775,7 @@ MCP_TOOL_SEARCH_CONFIG = { | `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`. | Review Comment: Fixed in b5a91a3f4ae94a0c3fb994bb2a5e12f2ba30cd15. The include_schemas table row now says True and explains that full inputSchema is the shipped default, with False opting into summary mode. Also removed the misleading “False (default)” wording from the serializer docstring. The empty-config fallback remains unchanged. ########## superset/mcp_service/server.py: ########## @@ -535,22 +430,20 @@ def _create_search_result_serializer( ~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) if not include_schemas: 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) Review Comment: Fixed in b5a91a3f4ae94a0c3fb994bb2a5e12f2ba30cd15. Applied the simplified max_desc default and early return for any falsy max_desc; the wrapper only runs when truncation is enabled. This preserves serialized output: I compared the 90fe26612 serializer with this version across 36 configurations (include_schemas omitted/False/True, compact_schemas omitted/False/True, max_description_length omitted/0/50/300) for all 72 registered tools and obtained identical results. Parameterized the explicit-zero test over both compact_schemas values. Focused inventory/search tests: 209 passed; full MCP suite: 5,205 passed, 4 skipped. -- 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]
