From f1141841983dbd218802e830b39bbcc5db9a0e4e Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Fri, 11 Sep 2026 22:41:50 -0500 Subject: [PATCH] Stop serving 2,400 tokens of schema nobody asked for Nothing this server runs reaches a tool through Tool.get_schema(), which is where upstream cleans parameter schemas. RouteTable._tool_to_route reads tool.parameters directly, so both the MCP and REST surfaces receive the uncleaned dict, and two pieces of dead weight rode along on every request. toolcall_reason: Tool.model_post_init writes the key into parameters on every construction, unconditionally. The switch meant to govern it lives in get_schema() and defaults to off, which is the registry built here. The feature has been disabled since this registry was written and the schema shipped regardless, on all 33 core tools, for an argument execution discards. Pydantic title: every generated property carried a title restating its own name in title case, which no client reads. Both are trimmed at registry build, in _trim_wire_schema. The core surface falls from 11,361 to 8,974 cl100k_base tokens, 21% off the fixed cost of every request. The toolcall_reason pop becomes a no-op if upstream closes the gap; neither pass forks anything. The title trim recurses into property values only. Upstream's blanket key filter does not, which is why plot_dataset's real title argument is missing from registry.get_schemas() -- a test pins the distinction so the fix cannot regress into the bug it works around. The budget that should have caught all of this was measuring the same unserved surface: TestToolSpecBudget and scripts/measure_payload.py both read get_schemas(), understating the catalog by 4,263 bytes and passing. Both now measure RouteTable, with budgets ratcheted onto the trimmed figures. Separately, list_datasets reported an archive root as empty. A non-recursive scan of a directory whose datasets sit one level down returned total_files: 0 and no groups, which reads as a missing dataset rather than a wrong flag. That layout is GDEX and most model archives. The zero-file case now counts the subdirectories and names the first ten. Applied to the worker copy too, which is the one that meets GDEX. --- CHANGELOG.md | 32 +++++++++ scripts/measure_payload.py | 18 +++++- src/uxarray_mcp/registry.py | 70 +++++++++++++++++++- src/uxarray_mcp/tools/catalog.py | 63 +++++++++++++++--- tests/test_catalog.py | 63 ++++++++++++++++++ tests/test_payload_budget.py | 108 +++++++++++++++++++++++++++++-- 6 files changed, 335 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a6854ab..aa41d2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -156,6 +156,38 @@ built against; see `docs/release.md`. Versions through `0.3.1` were SemVer. issue, and a human merge is what reaches PyPI. ### Fixed +- Every tool schema we serve carried ~2,400 tokens of dead weight, on every + request. Two causes, one root: nothing we run reaches tools through + `Tool.get_schema()`, which is where upstream does its cleaning. + `RouteTable._tool_to_route` reads `tool.parameters` raw, so both surfaces get + the uncleaned dict. First, `toolcall_reason`: `Tool.model_post_init` writes + it into `parameters` on every construction, unconditionally, while the switch + meant to govern it lives in `get_schema()` and defaults to off. Thought- + augmented tool calling has been disabled here since the registry was written + and the schema shipped anyway, on all 33 core tools, for an argument + execution then discards. Second, a Pydantic `title` on every generated + property, restating the property's own name in title case. Both are now + trimmed at registry build. The core surface went from 11,361 to 8,974 + `cl100k_base` tokens, 21% off the fixed cost of every request. +- `plot_dataset`'s `title` argument was invisible to any client reading + `registry.get_schemas()`. Upstream strips Pydantic `title` annotations with a + blanket key filter (`Tool._EXTRA_STRIP_KEYS`) that descends into `properties` + and deletes the parameter *named* `title` along with them. Our trim recurses + into property values only, so the annotation goes and the argument stays; a + test pins the distinction. The served MCP surface was never affected, because + it never called the function doing the stripping. +- The tool-specification budget measured a surface nobody receives. + `TestToolSpecBudget` and `scripts/measure_payload.py` both read + `registry.get_schemas()`, understating the served catalog by 4,263 bytes and + passing throughout. Both now measure `RouteTable`, and the budgets are + ratcheted onto the trimmed figures: 42000 to 36000 bytes overall, 6000 to + 5000 for `run_analysis`, 4200 to 3500 for `get_capabilities`. +- `list_datasets` reported an archive root as empty. A non-recursive scan of a + directory whose datasets all sit one level down returned `total_files: 0` and + no groups, which reads as a missing dataset rather than a wrong flag -- and + that layout is GDEX, and most model archives. The zero-file case now counts + the subdirectories and names the first ten, so the caller can fix it in one + move. Applied to the worker copy as well, which is the one that meets GDEX. - The worker-payload drift guard read source text, so it failed on formatting and could pass on real drift. `test_every_dispatch_offers_the_same_input_kinds` counted the substrings `startswith("healpix:")` and `".shp", ".geojson"` and diff --git a/scripts/measure_payload.py b/scripts/measure_payload.py index 8516ff3..92e2247 100644 --- a/scripts/measure_payload.py +++ b/scripts/measure_payload.py @@ -130,12 +130,24 @@ def measure(grid_file: str, data_file: str) -> dict[str, dict[str, Any]]: def measure_catalog() -> dict[str, Any]: - """How much context the tool catalog itself occupies before any call.""" + """How much context the tool catalog itself occupies before any call. + + Measured through ``RouteTable``, which is what every surface we run + actually serves. ``registry.get_schemas()`` returns ``get_schema()``'s + cleaned output instead, which no client receives, and reporting it + understated the catalog by 4,263 bytes. + """ + from toolregistry_server.route_table import RouteTable + from uxarray_mcp.app import make_registry schemas = { - schema.get("function", schema)["name"]: schema - for schema in make_registry().get_schemas() + route.tool_name: { + "name": route.tool_name, + "description": route.description, + "parameters": route.parameters_schema, + } + for route in RouteTable(make_registry()).list_routes() } by_name = {name: _size(schema) for name, schema in schemas.items()} total = sum(by_name.values()) diff --git a/src/uxarray_mcp/registry.py b/src/uxarray_mcp/registry.py index ecafedf..dc58916 100644 --- a/src/uxarray_mcp/registry.py +++ b/src/uxarray_mcp/registry.py @@ -618,6 +618,71 @@ def _make_results_wire_safe(tool: Any) -> None: callable_.fn = _wire_safe(inner) +def _trim_wire_schema(tool: Any) -> None: + """Drop schema nobody asked for from what we actually serve. + + Two separate pieces of dead weight ride on ``tool.parameters``, and + both reach clients because the MCP and REST surfaces read that dict + directly -- ``RouteTable._tool_to_route`` at + ``toolregistry_server/route_table.py:186`` -- rather than going + through ``Tool.get_schema()``, which is where upstream does its + cleaning. Measured together at ~2,400 tokens per request on the + 33-tool core surface, or 21% of it. + + ``toolcall_reason``: ``Tool.model_post_init`` writes this key into + ``parameters`` on every construction (``toolregistry/tool.py:253``), + unconditionally. The switch meant to govern it lives in + ``get_schema()`` (``tool.py:436``) and defaults to off, which is the + registry we build -- so the feature is disabled and the schema ships + anyway. Execution already discards the argument + (``tool_registry.py:361``), so nothing depends on clients sending it. + + Pydantic ``title``: every generated property carries a ``"title"`` + restating its own name in title case. It is display metadata no + client needs. Upstream agrees it is noise and strips it in + ``get_schema()``, but with a blanket key filter + (``Tool._EXTRA_STRIP_KEYS``) that descends into ``properties`` and + deletes the *parameter named* ``title`` along with it -- which is why + ``plot_dataset``'s real ``title`` argument is missing from + ``get_schemas()``. Recurse into property values only, so the + annotation goes and the parameter stays. + + Both passes are idempotent, and the first becomes a no-op if upstream + closes the gap. + """ + params = getattr(tool, "parameters", None) + if not isinstance(params, dict): + return + + props = params.get("properties") + if isinstance(props, dict): + props.pop("toolcall_reason", None) + required = params.get("required") + if isinstance(required, list) and "toolcall_reason" in required: + params["required"] = [r for r in required if r != "toolcall_reason"] + + _drop_title_annotations(params) + + +def _drop_title_annotations(schema: Any) -> None: + """Remove Pydantic ``title`` metadata in place, keeping parameter names. + + Only the values under ``properties`` are recursed into. The keys of + that mapping are parameter names and are never inspected, which is + the whole difference between this and a blanket key strip. + """ + if not isinstance(schema, dict): + return + schema.pop("title", None) + props = schema.get("properties") + if isinstance(props, dict): + for spec in props.values(): + _drop_title_annotations(spec) + items = schema.get("items") + if isinstance(items, dict): + _drop_title_annotations(items) + + def _apply_tags( registry: ToolRegistry, registered_name: str, @@ -810,12 +875,13 @@ def build_registry( # ``enable_tool_discovery`` registers ``discover_tools`` itself, so it # never passes through the loops above. Sweep the whole surface rather # than name that one tool: anything the library registers on its own - # belongs behind the same boundary, and ``_make_results_wire_safe`` is - # idempotent, so re-running it over already-wrapped tools costs nothing. + # belongs behind the same boundary, and both sweeps are idempotent, so + # re-running them over already-treated tools costs nothing. for name in registry.list_tools(): tool = registry.get_tool(name) if tool is not None: _make_results_wire_safe(tool) + _trim_wire_schema(tool) _verify_coverage(registered, profile) return registry diff --git a/src/uxarray_mcp/tools/catalog.py b/src/uxarray_mcp/tools/catalog.py index 729b551..b1d03f3 100644 --- a/src/uxarray_mcp/tools/catalog.py +++ b/src/uxarray_mcp/tools/catalog.py @@ -59,6 +59,36 @@ def _classify(name: str) -> str: return "unknown" +def _empty_scan_hint(root: Path, directory: str, recursive: bool) -> str: + """Explain a zero-file scan, naming subdirectories when they exist. + + A non-recursive scan of an archive root is the common way to get zero + hits, and ``total_files: 0`` with an empty ``groups`` reads as "nothing + here" when the truth is "everything is one level down". GDEX is laid + out exactly that way, so the failure looks like a missing dataset + rather than a wrong flag. Name the subdirectories and the caller can + fix it in one move instead of guessing. + """ + if not recursive: + try: + subdirs = sorted(p.name for p in root.iterdir() if p.is_dir()) + except OSError: + subdirs = [] + if subdirs: + shown = ", ".join(subdirs[:10]) + more = f" (+{len(subdirs) - 10} more)" if len(subdirs) > 10 else "" + return ( + f"No mesh or data files directly in {directory}, but it holds " + f"{len(subdirs)} subdirector" + f"{'y' if len(subdirs) == 1 else 'ies'}: {shown}{more}. " + "Re-run with recursive=True, or point at one subdirectory." + ) + return ( + f"No mesh or data files found in {directory}. " + "Try recursive=True or check the path." + ) + + def list_datasets( directory: str, recursive: bool = False, @@ -171,10 +201,7 @@ def list_datasets( # Build recommendations recommendations: List[str] = [] if not all_files: - recommendations.append( - f"No mesh or data files found in {directory}. " - "Try recursive=True or check the path." - ) + recommendations.append(_empty_scan_hint(root, directory, recursive)) else: if grid_found and data_found: recommendations.append( @@ -325,10 +352,30 @@ def _classify(name: str) -> str: recommendations = [] if not all_files: - recommendations.append( - f"No mesh or data files found in {directory}. " - "Try recursive=True or check the path." - ) + # Same reasoning as ``_empty_scan_hint``, inlined because this + # function is shipped to the worker by value and cannot reach + # module scope. This is the copy that matters most: GDEX is an + # archive of subdirectories and lives on the remote filesystem. + subdirs = [] + if not recursive: + try: + subdirs = sorted(p.name for p in root.iterdir() if p.is_dir()) + except OSError: + subdirs = [] + if subdirs: + shown = ", ".join(subdirs[:10]) + more = f" (+{len(subdirs) - 10} more)" if len(subdirs) > 10 else "" + recommendations.append( + f"No mesh or data files directly in {directory}, but it holds " + f"{len(subdirs)} subdirector" + f"{'y' if len(subdirs) == 1 else 'ies'}: {shown}{more}. " + "Re-run with recursive=True, or point at one subdirectory." + ) + else: + recommendations.append( + f"No mesh or data files found in {directory}. " + "Try recursive=True or check the path." + ) else: if grid_found and data_found: recommendations.append( diff --git a/tests/test_catalog.py b/tests/test_catalog.py index ae91fff..800ec8f 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -217,3 +217,66 @@ def test_size_mb_present(self, tmp_path): result = list_datasets(str(tmp_path)) entry = result["groups"][0]["files"][0] assert "size_mb" in entry + + +class TestListDatasetsArchiveRoot: + """A directory whose datasets are all one level down. + + This is how GDEX and most model archives are laid out, and the + non-recursive scan of such a root is the ordinary way to get zero hits. + Reporting only ``total_files: 0`` made a wrong flag look like a missing + dataset, so the recommendation names what is actually there. + """ + + def test_subdirectories_are_named_when_no_files_match(self, tmp_path): + _make_files(tmp_path, ["d651007/data.nc", "d651014/data.nc"]) + result = list_datasets(str(tmp_path), recursive=False) + combined = " ".join(result["recommendations"]) + assert result["total_files"] == 0 + assert "d651007" in combined + assert "d651014" in combined + assert "recursive=True" in combined + + def test_the_subdirectory_count_is_reported(self, tmp_path): + _make_files(tmp_path, [f"member_{i:02d}/data.nc" for i in range(14)]) + result = list_datasets(str(tmp_path), recursive=False) + combined = " ".join(result["recommendations"]) + assert "14 subdirectories" in combined + # Only the first ten are named; the rest are counted, not listed. + assert "(+4 more)" in combined + + def test_a_single_subdirectory_reads_as_singular(self, tmp_path): + _make_files(tmp_path, ["only/data.nc"]) + result = list_datasets(str(tmp_path), recursive=False) + combined = " ".join(result["recommendations"]) + assert "1 subdirectory:" in combined + + def test_a_genuinely_empty_directory_still_says_so(self, tmp_path): + result = list_datasets(str(tmp_path), recursive=False) + combined = " ".join(result["recommendations"]) + assert "No mesh or data files found" in combined + assert "subdirector" not in combined + + def test_a_recursive_scan_that_finds_nothing_does_not_blame_subdirs(self, tmp_path): + """recursive=True already looked inside, so the hint would be wrong.""" + _make_files(tmp_path, ["sub/notes.txt"]) + result = list_datasets(str(tmp_path), recursive=True) + combined = " ".join(result["recommendations"]) + assert result["total_files"] == 0 + assert "subdirector" not in combined + + def test_the_remote_scan_gives_the_same_hint(self, tmp_path): + """The worker copy is the one that meets GDEX; it must not drift. + + ``_remote_catalog_fn`` is shipped to the endpoint by value and + cannot call module scope, so the logic is duplicated there. Test + the duplicate rather than trust it. + """ + from uxarray_mcp.tools.catalog import _remote_catalog_fn + + _make_files(tmp_path, ["d651007/data.nc", "d651014/data.nc"]) + result = _remote_catalog_fn(str(tmp_path), False, 200) + combined = " ".join(result["recommendations"]) + assert result["total_files"] == 0 + assert "d651007" in combined + assert "recursive=True" in combined diff --git a/tests/test_payload_budget.py b/tests/test_payload_budget.py index e24ffb9..eb8d2ca 100644 --- a/tests/test_payload_budget.py +++ b/tests/test_payload_budget.py @@ -13,6 +13,7 @@ from __future__ import annotations import json +import pathlib import warnings import numpy as np @@ -79,11 +80,17 @@ SIGNAL_FRACTION_FLOOR = 0.30 #: Upper bound on the serialized core tool specification, sent every request. -TOOL_SPEC_BYTE_BUDGET = 42000 +#: Measured on the served surface (``RouteTable``), which is 4,263 bytes +#: larger than the ``get_schemas()`` output this used to read. Ratcheted from +#: 42000 after #146 trimmed ``toolcall_reason`` and the Pydantic ``title`` +#: annotations: 40748 served bytes before the trim, 35475 after. +TOOL_SPEC_BYTE_BUDGET = 36000 -#: Upper bound for the two largest individual tool schemas (#89). -RUN_ANALYSIS_SCHEMA_BUDGET = 6000 -GET_CAPABILITIES_SCHEMA_BUDGET = 4200 +#: Upper bound for the two largest individual tool schemas (#89). Ratcheted +#: from 6000/4200 on the same measurement; ``run_analysis`` alone was 5994 +#: served bytes before the trim, six under a budget it was never tested on. +RUN_ANALYSIS_SCHEMA_BUDGET = 5000 +GET_CAPABILITIES_SCHEMA_BUDGET = 3500 def _measure(result: dict) -> tuple[int, int, float]: @@ -148,13 +155,28 @@ def test_discovery_keys_never_appear_in_results(self, analysis_results, operatio class TestToolSpecBudget: + """Budgets on the specification clients receive, not the one upstream cleans. + + This fixture read ``registry.get_schemas()`` until #146. That is + ``Tool.get_schema()``'s output -- flattened, gated, and served to + nobody: every surface we run reaches tools through ``RouteTable``, + which reads ``tool.parameters`` raw. The budget was therefore measuring + a schema 4,263 bytes smaller than the wire, and passing. + """ + @pytest.fixture(scope="class") def schemas(self): + from toolregistry_server.route_table import RouteTable + from uxarray_mcp.app import make_registry return { - schema.get("function", schema)["name"]: schema - for schema in make_registry().get_schemas() + route.tool_name: { + "name": route.tool_name, + "description": route.description, + "parameters": route.parameters_schema, + } + for route in RouteTable(make_registry()).list_routes() } def test_core_tool_specification_stays_within_budget(self, schemas): @@ -180,6 +202,80 @@ def test_largest_schemas_stay_within_budget(self, schemas, name, budget): ) +class TestServedSchemaIsTrimmed: + """Guard the surface clients actually receive. + + Everything above measures ``registry.get_schemas()``, which runs + upstream's cleaning and so has always looked tidy. Nothing we serve + goes through it: ``RouteTable._tool_to_route`` reads ``tool.parameters`` + directly, so the MCP and REST surfaces get the raw dict, uncleaned. + The gap hid ~2,400 tokens per request -- ``toolcall_reason`` on all 33 + tools for a disabled feature, plus a Pydantic ``title`` on every + property. ``_trim_wire_schema`` closes it; these tests keep it closed. + """ + + @pytest.fixture(scope="class") + def routes(self): + from toolregistry_server.route_table import RouteTable + + from uxarray_mcp.app import make_registry + + return RouteTable(make_registry()).list_routes() + + def test_no_served_tool_advertises_toolcall_reason(self, routes): + offenders = sorted( + r.tool_name + for r in routes + if "toolcall_reason" in (r.parameters_schema.get("properties") or {}) + ) + assert not offenders, ( + f"{len(offenders)} served tools advertise toolcall_reason: " + f"{offenders}. Thought-augmented tool calling is off, so this is " + "schema nobody asked for, re-sent on every request." + ) + + def test_no_served_property_carries_a_title_annotation(self, routes): + offenders = sorted( + f"{r.tool_name}.{name}" + for r in routes + for name, spec in (r.parameters_schema.get("properties") or {}).items() + if isinstance(spec, dict) and "title" in spec + ) + assert not offenders, ( + f"{len(offenders)} served properties carry a Pydantic title " + f"annotation: {offenders[:5]}. It restates the property name in " + "title case and no client reads it." + ) + + def test_a_parameter_named_title_survives_the_trim(self, routes): + """The bug the blanket strip causes, pinned so we never repeat it. + + Upstream's ``get_schema()`` filters the key ``title`` at every + depth, which deletes ``plot_dataset``'s real ``title`` argument + along with the annotations -- confirm with + ``make_registry().get_schemas()``, where it is missing. Our trim + recurses into property *values* only, so the argument stays. + """ + plot = next(r for r in routes if r.tool_name == "plot_dataset") + props = plot.parameters_schema.get("properties") or {} + assert "title" in props, ( + "plot_dataset lost its `title` argument to the trim. A caller " + "cannot label a figure it cannot see." + ) + assert props["title"].get("type") == "string" + + +def test_measurement_script_reports_the_served_surface(): + """``scripts/measure_payload.py`` must not measure a surface we do not send.""" + source = ( + pathlib.Path(__file__).resolve().parents[1] / "scripts" / "measure_payload.py" + ).read_text() + assert "RouteTable" in source, ( + "measure_payload.py still reads get_schemas(); that is upstream's " + "cleaned schema, not the one RouteTable puts on the wire." + ) + + def test_measurement_helper_counts_signal_only(): """Guard the measurement itself, so a budget cannot pass by miscounting.""" total, signal, fraction = _measure(