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(