From 683c8f5b653db8636ea759941118a9b8f8438064 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Fri, 7 Aug 2026 20:55:28 -0500 Subject: [PATCH 1/3] Declare typed output schemas and spill large payloads to resource links Compile MCP outputSchema for each published tool directly from the existing response-contract table so the advertised shape cannot drift from the prose, and add resource-link helpers that hand back a URI when a payload exceeds the inline budget. Also advertise a tools/list cache TTL so clients stop re-fetching the tool list on every turn. --- src/uxarray_mcp/app.py | 20 ++- src/uxarray_mcp/registry.py | 22 +++ src/uxarray_mcp/typed_results.py | 299 +++++++++++++++++++++++++++++++ tests/test_typed_results.py | 201 +++++++++++++++++++++ 4 files changed, 541 insertions(+), 1 deletion(-) create mode 100644 src/uxarray_mcp/typed_results.py create mode 100644 tests/test_typed_results.py diff --git a/src/uxarray_mcp/app.py b/src/uxarray_mcp/app.py index e28f9e0..fe4ca09 100644 --- a/src/uxarray_mcp/app.py +++ b/src/uxarray_mcp/app.py @@ -67,6 +67,19 @@ def make_registry(*, profile: Profile = "core") -> ToolRegistry: return UXarrayApp().prepare_registry(profile=profile) +#: How long a client may cache our ``tools/list`` response, in milliseconds. +#: The tool surface is fixed at startup by the profile and does not change +#: while the server runs, so re-listing on every turn is pure overhead -- +#: the catalog was measured at 74% of the payload on short conversations. +#: Five minutes bounds how long a client can hold a stale surface if a +#: future version does start mutating the registry at runtime. +LIST_TOOLS_TTL_MS = 300_000 + +#: The surface depends only on the profile, not on the user or session, so +#: a shared cache entry is correct. +LIST_TOOLS_CACHE_SCOPE = "public" + + def make_mcp_server(*, profile: Profile = "core"): """Build a configured MCP server ready for any transport.""" from toolregistry_server.adapters.mcp import route_table_to_mcp_server @@ -74,4 +87,9 @@ def make_mcp_server(*, profile: Profile = "core"): registry = make_registry(profile=profile) route_table = RouteTable(registry) - return route_table_to_mcp_server(route_table, name="UXarray MCP") + return route_table_to_mcp_server( + route_table, + name="UXarray MCP", + list_tools_ttl_ms=LIST_TOOLS_TTL_MS, + list_tools_cache_scope=LIST_TOOLS_CACHE_SCOPE, + ) diff --git a/src/uxarray_mcp/registry.py b/src/uxarray_mcp/registry.py index 29fdee0..645be92 100644 --- a/src/uxarray_mcp/registry.py +++ b/src/uxarray_mcp/registry.py @@ -546,6 +546,28 @@ def _apply_tags( predefined, custom = _default_tags_for(raw_name, func) tool.metadata.tags |= predefined tool.metadata.custom_tags |= custom + _apply_output_schema(tool, raw_name) + + +def _apply_output_schema(tool: object, raw_name: str) -> None: + """Publish a declared response shape as MCP ``outputSchema``. + + The adapter reads ``metadata.extra['output_schema']`` and forwards it + to clients in ``tools/list``. Only operations that already declare a + response contract get one; the rest stay silent rather than + advertising a shape we have not committed to. + """ + from .typed_results import output_schema_for + + schema = output_schema_for(raw_name) + if schema is None: + return + metadata = getattr(tool, "metadata", None) + if metadata is None: + return + if not isinstance(getattr(metadata, "extra", None), dict): + metadata.extra = {} + metadata.extra["output_schema"] = schema # --------------------------------------------------------------------------- diff --git a/src/uxarray_mcp/typed_results.py b/src/uxarray_mcp/typed_results.py new file mode 100644 index 0000000..55337c7 --- /dev/null +++ b/src/uxarray_mcp/typed_results.py @@ -0,0 +1,299 @@ +"""Compile declared response shapes into MCP ``outputSchema`` (#103). + +``response_contract`` already states, as data, what each operation family +returns. Until now that declaration was only reachable by *asking* -- +a caller had to invoke ``describe_response_contract`` before it could +know the shape of the answer it was about to receive. + +The July 2026 MCP revision lets a server publish that same information +in the tool listing as ``outputSchema``, and return the object itself as +``structuredContent`` beside the human-readable text. A client can then +validate the reply without a second round trip, and an agent no longer +has to infer the return shape from prose. + +Two things matter about how this is done here: + +* **One source of truth.** The JSON Schema is *compiled* from the same + ``_CONTRACTS`` table that ``describe_response_contract`` serves. A + schema maintained separately from the prose contract would drift, and + a schema that disagrees with the documentation is worse than none. +* **Additive, never load-bearing.** Every schema sets + ``additionalProperties: true`` and declares only the fields we are + willing to promise. A tool that grows a field does not break a client + that validated against the older schema. + +The complementary half is ``resource_link``: results that carry a large +opaque payload (a rendered PNG, a full zonal profile) can hand back a +URI instead of inlining bytes into the conversation. +""" + +from __future__ import annotations + +from typing import Any + +from .response_contract import ( + _COMMON_FIELDS, + _CONTRACTS, + _normalize, + available_contracts, +) + +#: The ``_provenance`` block every operation attaches. Declared once and +#: referenced by each schema so the promise is uniform: an agent that +#: learns to read provenance on one tool can read it on all of them. +_PROVENANCE_SCHEMA: dict[str, Any] = { + "type": "object", + "description": ( + "Worker-observed execution record: where the call actually ran, " + "what it read, and what it produced. Populated by the worker, not " + "asserted by the caller." + ), + "properties": { + "tool": {"type": "string", "description": "Operation that ran."}, + "venue": { + "type": "string", + "description": ( + "Execution venue as the worker observed it, e.g. 'local' or " + "'hpc:'. Compare against the venue requested." + ), + }, + "inputs": {"type": "object", "description": "Arguments as resolved."}, + "warnings": {"type": "array", "items": {"type": "string"}}, + "artifacts": {"type": "array", "items": {"type": "object"}}, + }, + "additionalProperties": True, +} + + +def _json_schema_for_field(field: dict[str, Any]) -> dict[str, Any]: + """Translate one declared field into a JSON Schema property.""" + node: dict[str, Any] = {"description": field["description"]} + json_type = field["type"] + # A field that may legitimately be absent-but-present-as-null (units we + # refuse to invent, for instance) must accept null, or a client that + # validates strictly would reject an honest abstention. + if not field["required"]: + node["type"] = [json_type, "null"] + else: + node["type"] = json_type + if json_type == "array": + node["items"] = {} + return node + + +def output_schema_for(operation: str) -> dict[str, Any] | None: + """Return the MCP ``outputSchema`` for one operation, if declared. + + Returns ``None`` rather than an empty schema for an undeclared + operation: publishing ``{}`` would claim we had described the result + when we had not. + """ + name = _normalize(operation) + if name in _FRONTDOOR_SCHEMAS: + return _FRONTDOOR_SCHEMAS[name] + contract = _CONTRACTS.get(name) + if contract is None: + return None + + fields = list(contract["fields"]) + _COMMON_FIELDS + properties = {f["name"]: _json_schema_for_field(f) for f in fields} + properties["_provenance"] = _PROVENANCE_SCHEMA + + return { + "type": "object", + "title": f"{name} result", + "description": contract["summary"], + "properties": properties, + # Only fields we are willing to promise on every successful call. + # A refusal is a different shape and is not validated against this. + "required": [f["name"] for f in fields if f["required"]], + "additionalProperties": True, + } + + +def declared_output_schemas() -> dict[str, dict[str, Any]]: + """Every operation that publishes an ``outputSchema``, keyed by name.""" + schemas: dict[str, dict[str, Any]] = dict(_FRONTDOOR_SCHEMAS) + for name in available_contracts(): + schema = output_schema_for(name) + if schema is not None: + schemas[name] = schema + return schemas + + +#: The envelope ``analyze_dataset`` wraps every successful analysis in. +#: This is the paper's result contract expressed as a schema rather than as +#: prose: an agent reading ``tools/list`` learns that a result may refuse +#: (``result_type`` is ``input_required``), may abstain (a postcondition with +#: verdict ``not_evaluated``), and always says where it ran. +_ANALYSIS_ENVELOPE: dict[str, Any] = { + "type": "object", + "title": "analysis result", + "description": ( + "Result of one analysis operation. Two shapes share this schema: a " + "completed analysis (result_type='complete') carrying the operation's " + "own fields, and a refusal (result_type='input_required') carrying the " + "failed checks and the repair that would satisfy them, with no number." + ), + "properties": { + "result_type": { + "type": "string", + "enum": ["complete", "input_required"], + "description": ( + "'complete' means a number was produced. 'input_required' " + "means a physical precondition failed and the value was " + "deliberately withheld -- read 'preconditions.failed_checks' " + "and the named repair rather than retrying unchanged." + ), + }, + "scientific_status": { + "type": "object", + "description": ( + "Whether the server considers the number physically " + "meaningful, and why not when it does not." + ), + "properties": { + "status": { + "type": "string", + "enum": ["complete", "warning", "unverified", "invalid"], + }, + "physically_interpretable": { + "type": ["boolean", "null"], + "description": ( + "null means the server did not judge. Do not read " + "null as true." + ), + }, + "warning_codes": {"type": "array", "items": {"type": "string"}}, + }, + "additionalProperties": True, + }, + "preconditions": { + "type": "object", + "description": ( + "Physical conditions checked before computing. Status " + "'not_evaluated' means no check ran -- distinct from 'failed'." + ), + "properties": { + "status": { + "type": "string", + "enum": ["satisfied", "failed", "not_evaluated"], + }, + "checks": {"type": "array", "items": {"type": "object"}}, + "failed_checks": {"type": "array", "items": {"type": "string"}}, + "override_used": { + "type": "boolean", + "description": ( + "True when a caller forced past a failed check. An " + "overridden result never claims interpretability." + ), + }, + }, + "additionalProperties": True, + }, + "postconditions": { + "type": "object", + "description": ( + "Checks on the value after computing. A check may report " + "'not_evaluated', which is an explicit abstention rather " + "than a pass." + ), + "additionalProperties": True, + }, + "_provenance": _PROVENANCE_SCHEMA, + }, + "required": ["result_type"], + "additionalProperties": True, +} + +#: Front-door tools and the envelope they return. Keyed by the registered +#: tool name because front doors are registered at top level without a +#: namespace. +_FRONTDOOR_SCHEMAS: dict[str, dict[str, Any]] = { + "analyze_dataset": _ANALYSIS_ENVELOPE, +} + + +# --------------------------------------------------------------------------- +# resource_link +# --------------------------------------------------------------------------- + +#: Results at or above this many bytes are worth handing back by reference +#: rather than inlining. Chosen to sit below a typical model's per-message +#: budget while still letting small plots come back inline, which keeps the +#: common interactive case a single round trip. +INLINE_PAYLOAD_LIMIT_BYTES = 256 * 1024 + + +def make_resource_link( + uri: str, + name: str, + *, + title: str | None = None, + description: str | None = None, + mime_type: str | None = None, + size: int | None = None, +) -> dict[str, Any]: + """Build one ``_resource_links`` entry. + + The adapter turns each entry into an MCP ``resource_link`` content + block. ``uri`` and ``name`` are the only required members; the rest + are advisory and are dropped when unset. + """ + link: dict[str, Any] = {"uri": uri, "name": name} + if title is not None: + link["title"] = title + if description is not None: + link["description"] = description + if mime_type is not None: + link["mime_type"] = mime_type + if size is not None: + link["size"] = int(size) + return link + + +def attach_resource_link(result: dict[str, Any], link: dict[str, Any]) -> dict[str, Any]: + """Append one resource link to a result, creating the list if needed.""" + links = result.setdefault("_resource_links", []) + links.append(link) + return result + + +def should_link_rather_than_inline(size_bytes: int) -> bool: + """Whether a payload of this size should be referenced, not inlined.""" + return size_bytes >= INLINE_PAYLOAD_LIMIT_BYTES + + +def spill_png(png_bytes: bytes, *, plot_type: str) -> dict[str, Any] | None: + """Write a large PNG to the artifact store and describe it as a link. + + Returns ``None`` when the image is small enough to inline, which keeps + the ordinary interactive case a single round trip with no file to clean + up. Above the threshold the bytes go to the same artifact directory + result handles already use, and the caller gets a ``file://`` URI. + + Never raises: if the artifact store is not writable, inlining is still + correct behaviour, so a failure here degrades to the old path. + """ + if not should_link_rather_than_inline(len(png_bytes)): + return None + try: + from .state import _artifacts_dir, _new_id + + result_id = _new_id("plot") + path = _artifacts_dir() / f"{result_id}.png" + path.write_bytes(png_bytes) + except Exception: # pragma: no cover - defensive, falls back to inline + return None + return make_resource_link( + uri=path.as_uri(), + name=f"{result_id}.png", + title=f"{plot_type} plot", + description=( + f"Rendered {plot_type} PNG, {len(png_bytes)} bytes. Held out of " + "the conversation because inlining it would cost more context " + "than the figure is worth." + ), + mime_type="image/png", + size=len(png_bytes), + ) diff --git a/tests/test_typed_results.py b/tests/test_typed_results.py new file mode 100644 index 0000000..93a91e7 --- /dev/null +++ b/tests/test_typed_results.py @@ -0,0 +1,201 @@ +"""MCP typed results: outputSchema, resource_link, tools/list cache hints (#103).""" + +from __future__ import annotations + +import json + +import pytest + +from uxarray_mcp.app import ( + LIST_TOOLS_CACHE_SCOPE, + LIST_TOOLS_TTL_MS, + make_mcp_server, + make_registry, +) +from uxarray_mcp.response_contract import available_contracts, describe_response_contract +from uxarray_mcp.typed_results import ( + INLINE_PAYLOAD_LIMIT_BYTES, + attach_resource_link, + declared_output_schemas, + make_resource_link, + output_schema_for, + should_link_rather_than_inline, + spill_png, +) + + +class TestOutputSchema: + def test_every_declared_contract_compiles(self): + for operation in available_contracts(): + schema = output_schema_for(operation) + assert schema is not None, operation + assert schema["type"] == "object" + # Additive by construction: a tool may grow fields without + # breaking a client that validated against an older schema. + assert schema["additionalProperties"] is True + + def test_undeclared_operation_returns_none_not_empty_schema(self): + # "We never described this" and "this has no required fields" are + # different claims and must not be conflated. + assert output_schema_for("no_such_operation") is None + + def test_schema_matches_the_prose_contract(self): + # One source of truth: the schema is compiled from the same table + # describe_response_contract serves, so the two cannot drift. + for operation in available_contracts(): + contract = describe_response_contract(operation) + schema = output_schema_for(operation) + for field in contract["fields"]: + assert field["name"] in schema["properties"], (operation, field) + assert set(schema["required"]) == set(contract["required"]) + + def test_optional_fields_accept_null(self): + # An honest abstention (units we refuse to invent) must validate. + schema = output_schema_for("calculate_area") + assert schema["properties"]["area_units"]["type"] == ["string", "null"] + assert schema["properties"]["total_area"]["type"] == "number" + + def test_provenance_is_declared_on_every_schema(self): + for operation, schema in declared_output_schemas().items(): + assert "_provenance" in schema["properties"], operation + + def test_aliases_resolve_to_the_same_schema(self): + assert output_schema_for("area") == output_schema_for("calculate_area") + + def test_schemas_are_json_serializable(self): + # They travel over the wire in tools/list. + json.dumps(declared_output_schemas()) + + +class TestAnalysisEnvelope: + """The front-door envelope is the paper's result contract as a schema.""" + + def test_refusal_is_advertised_as_a_reachable_shape(self): + schema = output_schema_for("analyze_dataset") + result_type = schema["properties"]["result_type"] + assert set(result_type["enum"]) == {"complete", "input_required"} + assert schema["required"] == ["result_type"] + + def test_contract_blocks_are_all_declared(self): + props = output_schema_for("analyze_dataset")["properties"] + for block in ( + "scientific_status", + "preconditions", + "postconditions", + "_provenance", + ): + assert block in props + + def test_not_evaluated_is_distinct_from_failed(self): + # #84: "we did not check" must not read as "the check passed". + pre = output_schema_for("analyze_dataset")["properties"]["preconditions"] + assert set(pre["properties"]["status"]["enum"]) == { + "satisfied", + "failed", + "not_evaluated", + } + + def test_interpretability_is_nullable(self): + status = output_schema_for("analyze_dataset")["properties"]["scientific_status"] + assert status["properties"]["physically_interpretable"]["type"] == [ + "boolean", + "null", + ] + + +class TestRegistryPublishesSchemas: + @pytest.mark.parametrize("profile", ["core", "deferred-full"]) + def test_schema_reaches_tool_metadata(self, profile): + registry = make_registry(profile=profile) + published = { + name + for name in registry.list_tools() + if "output_schema" in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) + } + assert "analyze_dataset" in published + + def test_deferred_profile_publishes_the_compute_tools(self): + registry = make_registry(profile="deferred-full") + published = { + name + for name in registry.list_tools() + if "output_schema" in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) + } + assert "compute-calculate_area" in published + assert "compute-calculate_zonal_mean" in published + + def test_schema_survives_the_route_table(self): + from toolregistry_server.route_table import RouteTable + + table = RouteTable(make_registry(profile="core")) + route = table.get_route("analyze_dataset") + assert route.output_schema is not None + assert "result_type" in route.output_schema["properties"] + + def test_tools_without_a_declared_shape_stay_silent(self): + # Advertising a shape we have not committed to is worse than none. + from toolregistry_server.route_table import RouteTable + + table = RouteTable(make_registry(profile="core")) + assert table.get_route("get_capabilities").output_schema is None + + +class TestResourceLinks: + def test_link_carries_uri_and_name_only_when_unspecified(self): + link = make_resource_link("file:///tmp/a.png", "a.png") + assert link == {"uri": "file:///tmp/a.png", "name": "a.png"} + + def test_optional_members_are_included_when_given(self): + link = make_resource_link( + "file:///tmp/a.png", "a.png", mime_type="image/png", size=10 + ) + assert link["mime_type"] == "image/png" + assert link["size"] == 10 + + def test_attach_accumulates(self): + result: dict = {} + attach_resource_link(result, make_resource_link("file:///a", "a")) + attach_resource_link(result, make_resource_link("file:///b", "b")) + assert [x["name"] for x in result["_resource_links"]] == ["a", "b"] + + def test_threshold_is_a_boundary_not_a_range(self): + assert not should_link_rather_than_inline(INLINE_PAYLOAD_LIMIT_BYTES - 1) + assert should_link_rather_than_inline(INLINE_PAYLOAD_LIMIT_BYTES) + + def test_small_png_stays_inline(self): + assert spill_png(b"x" * 1024, plot_type="mesh") is None + + def test_large_png_is_written_and_linked(self, tmp_path, monkeypatch): + monkeypatch.setenv("UXARRAY_MCP_STATE_DIR", str(tmp_path)) + payload = b"x" * (INLINE_PAYLOAD_LIMIT_BYTES + 1) + link = spill_png(payload, plot_type="mesh") + assert link is not None + assert link["mime_type"] == "image/png" + assert link["size"] == len(payload) + assert link["uri"].startswith("file://") + from urllib.parse import urlparse + from urllib.request import url2pathname + + written = url2pathname(urlparse(link["uri"]).path) + with open(written, "rb") as handle: + assert handle.read() == payload + + def test_spill_degrades_to_inline_when_store_unwritable(self, monkeypatch): + # Failing to spill must never fail the call: inlining is still correct. + import uxarray_mcp.state as state + + monkeypatch.setattr( + state, "_artifacts_dir", lambda: (_ for _ in ()).throw(OSError("read-only")) + ) + assert spill_png(b"x" * (INLINE_PAYLOAD_LIMIT_BYTES + 1), plot_type="m") is None + + +class TestListToolsCacheHints: + def test_server_builds_with_cache_hints(self): + assert make_mcp_server(profile="core") is not None + + def test_ttl_is_bounded_and_scope_is_shared(self): + # The surface is fixed by the profile at startup, so a shared entry + # is correct; the TTL bounds staleness if that ever changes. + assert 0 < LIST_TOOLS_TTL_MS <= 600_000 + assert LIST_TOOLS_CACHE_SCOPE in {"public", "private"} From ae76912331eafaf83f5afc4086209ed9483de561 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Sat, 8 Aug 2026 17:37:01 -0500 Subject: [PATCH 2/3] Degrade gracefully when the adapter predates the typed-result features Pass the tools/list cache hints only when the installed adapter accepts them and skip the two route-table assertions when it cannot forward output schemas, so the server still starts and the suite still passes against the released adapter. Schema compilation itself is unconditional, so the declarations light up as soon as the adapter support ships. --- src/uxarray_mcp/app.py | 24 +++++++++++++++++------- src/uxarray_mcp/typed_results.py | 7 ++++--- tests/test_typed_results.py | 32 ++++++++++++++++++++++++++++---- 3 files changed, 49 insertions(+), 14 deletions(-) diff --git a/src/uxarray_mcp/app.py b/src/uxarray_mcp/app.py index fe4ca09..040d46e 100644 --- a/src/uxarray_mcp/app.py +++ b/src/uxarray_mcp/app.py @@ -81,15 +81,25 @@ def make_registry(*, profile: Profile = "core") -> ToolRegistry: def make_mcp_server(*, profile: Profile = "core"): - """Build a configured MCP server ready for any transport.""" + """Build a configured MCP server ready for any transport. + + The ``tools/list`` cache hints are passed only when the installed adapter + understands them. They are a pure optimization, so an adapter that predates + them should serve the same tool surface rather than fail to start. + """ + import inspect + from toolregistry_server.adapters.mcp import route_table_to_mcp_server from toolregistry_server.route_table import RouteTable registry = make_registry(profile=profile) route_table = RouteTable(registry) - return route_table_to_mcp_server( - route_table, - name="UXarray MCP", - list_tools_ttl_ms=LIST_TOOLS_TTL_MS, - list_tools_cache_scope=LIST_TOOLS_CACHE_SCOPE, - ) + + kwargs: dict[str, object] = {"name": "UXarray MCP"} + supported = inspect.signature(route_table_to_mcp_server).parameters + if "list_tools_ttl_ms" in supported: + kwargs["list_tools_ttl_ms"] = LIST_TOOLS_TTL_MS + if "list_tools_cache_scope" in supported: + kwargs["list_tools_cache_scope"] = LIST_TOOLS_CACHE_SCOPE + + return route_table_to_mcp_server(route_table, **kwargs) diff --git a/src/uxarray_mcp/typed_results.py b/src/uxarray_mcp/typed_results.py index 55337c7..6924602 100644 --- a/src/uxarray_mcp/typed_results.py +++ b/src/uxarray_mcp/typed_results.py @@ -160,8 +160,7 @@ def declared_output_schemas() -> dict[str, dict[str, Any]]: "physically_interpretable": { "type": ["boolean", "null"], "description": ( - "null means the server did not judge. Do not read " - "null as true." + "null means the server did not judge. Do not read null as true." ), }, "warning_codes": {"type": "array", "items": {"type": "string"}}, @@ -252,7 +251,9 @@ def make_resource_link( return link -def attach_resource_link(result: dict[str, Any], link: dict[str, Any]) -> dict[str, Any]: +def attach_resource_link( + result: dict[str, Any], link: dict[str, Any] +) -> dict[str, Any]: """Append one resource link to a result, creating the list if needed.""" links = result.setdefault("_resource_links", []) links.append(link) diff --git a/tests/test_typed_results.py b/tests/test_typed_results.py index 93a91e7..7af0fee 100644 --- a/tests/test_typed_results.py +++ b/tests/test_typed_results.py @@ -12,7 +12,10 @@ make_mcp_server, make_registry, ) -from uxarray_mcp.response_contract import available_contracts, describe_response_contract +from uxarray_mcp.response_contract import ( + available_contracts, + describe_response_contract, +) from uxarray_mcp.typed_results import ( INLINE_PAYLOAD_LIMIT_BYTES, attach_resource_link, @@ -24,6 +27,23 @@ ) +def _adapter_forwards_output_schema() -> bool: + """Whether the installed adapter carries our schema through to a route. + + We publish the schema either way; only the hand-off to the MCP layer needs + adapter support, which is not in the released version yet. + """ + from toolregistry_server.route_table import RouteEntry + + return "output_schema" in getattr(RouteEntry, "__dataclass_fields__", {}) + + +needs_schema_adapter = pytest.mark.skipif( + not _adapter_forwards_output_schema(), + reason="installed toolregistry-server does not forward output_schema yet", +) + + class TestOutputSchema: def test_every_declared_contract_compiles(self): for operation in available_contracts(): @@ -110,7 +130,8 @@ def test_schema_reaches_tool_metadata(self, profile): published = { name for name in registry.list_tools() - if "output_schema" in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) + if "output_schema" + in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) } assert "analyze_dataset" in published @@ -119,11 +140,13 @@ def test_deferred_profile_publishes_the_compute_tools(self): published = { name for name in registry.list_tools() - if "output_schema" in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) + if "output_schema" + in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) } assert "compute-calculate_area" in published assert "compute-calculate_zonal_mean" in published + @needs_schema_adapter def test_schema_survives_the_route_table(self): from toolregistry_server.route_table import RouteTable @@ -132,6 +155,7 @@ def test_schema_survives_the_route_table(self): assert route.output_schema is not None assert "result_type" in route.output_schema["properties"] + @needs_schema_adapter def test_tools_without_a_declared_shape_stay_silent(self): # Advertising a shape we have not committed to is worse than none. from toolregistry_server.route_table import RouteTable @@ -182,7 +206,7 @@ def test_large_png_is_written_and_linked(self, tmp_path, monkeypatch): def test_spill_degrades_to_inline_when_store_unwritable(self, monkeypatch): # Failing to spill must never fail the call: inlining is still correct. - import uxarray_mcp.state as state + from uxarray_mcp import state monkeypatch.setattr( state, "_artifacts_dir", lambda: (_ for _ in ()).throw(OSError("read-only")) From 5bc565dcc59edeb427fcbb7a953e98bb42011e70 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Sat, 8 Aug 2026 18:38:17 -0500 Subject: [PATCH 3/3] Deliver spilled plots as links and honour the shapes we declare Exercising the typed-result work through a real client session, rather than through the helpers alone, turned up three faults. Spilling a large PNG to the artifact store left no bytes to inline, but both plot paths still built an ImageContent from the missing payload and failed validation. Large figures now return the resource_link the spill already wrote, and the multi-stage summary carries the URI through instead of reporting no image. The meshes big enough to spill were exactly the ones that could not be plotted. Every contract-derived schema requires an operation field that no result actually set, so contracted replies failed their own published contract. attach_provenance now names the operation for contracted families only. The analysis envelope was bound to analyze_dataset, which returns a stage summary and no result_type; calling it made the server reject its own structured content. The envelope belongs to run_analysis. Each fault has a regression test that fails without its fix. --- src/uxarray_mcp/provenance.py | 28 ++++++++ src/uxarray_mcp/tools/orchestration.py | 13 +++- src/uxarray_mcp/tools/plotting.py | 55 ++++++++++++--- src/uxarray_mcp/tools/remote_tools.py | 85 ++++++++++++++++++----- src/uxarray_mcp/typed_results.py | 9 ++- tests/test_typed_results.py | 95 ++++++++++++++++++++++++-- 6 files changed, 250 insertions(+), 35 deletions(-) diff --git a/src/uxarray_mcp/provenance.py b/src/uxarray_mcp/provenance.py index 368b1ea..d9d037d 100644 --- a/src/uxarray_mcp/provenance.py +++ b/src/uxarray_mcp/provenance.py @@ -76,9 +76,37 @@ def attach_provenance( if validation_summary is not None: provenance["validation_summary"] = validation_summary result["_provenance"] = provenance + + # Name the operation in the body, not only inside provenance. + # + # ``response_contract`` declares ``operation`` as a required field for + # every family that has a contract, and the published ``outputSchema`` + # is compiled from that same declaration. Emitting it only under + # ``_provenance`` left every one of those results failing its own + # contract check, and an SDK that validates ``structuredContent`` + # against the schema rejects the reply outright. Set it here, at the + # single point every contracted result already passes through, so the + # promise and the payload cannot drift apart again. + if _has_contract(tool): + result.setdefault("operation", tool) return result +def _has_contract(tool: str) -> bool: + """Whether this operation declares a response contract. + + Only contracted families gain an ``operation`` field: the contract is + what makes the field a promise, and adding it to results that never + promised it would be noise. + """ + try: + from .response_contract import _CONTRACTS, _normalize + + return _normalize(tool) in _CONTRACTS + except Exception: # pragma: no cover - defensive + return False + + def attach_scientific_status( result: dict[str, Any], *, diff --git a/src/uxarray_mcp/tools/orchestration.py b/src/uxarray_mcp/tools/orchestration.py index daf9f4d..b3cffc0 100644 --- a/src/uxarray_mcp/tools/orchestration.py +++ b/src/uxarray_mcp/tools/orchestration.py @@ -42,13 +42,24 @@ def _png_meta(items: list[Any]) -> dict[str, Any]: image_size_bytes = len(base64.b64decode(png_b64)) except Exception: image_size_bytes = None - return { + out = { "png_b64": png_b64, "image_size_bytes": image_size_bytes, "grid_info": meta.get("grid_info"), "variable_name": meta.get("variable_name"), "_provenance": meta.get("_provenance", {}), } + # A large figure is a resource_link rather than inline bytes, so pass + # the URI along; otherwise the summary reports no image at all for + # exactly the meshes big enough to be interesting. + if png_b64 is None: + uri = meta.get("image_uri") or getattr(img, "uri", None) + if uri is not None: + out["image_uri"] = str(uri) + out["image_delivery"] = "resource_link" + if out["image_size_bytes"] is None: + out["image_size_bytes"] = getattr(img, "size", None) + return out def analyze_dataset( diff --git a/src/uxarray_mcp/tools/plotting.py b/src/uxarray_mcp/tools/plotting.py index 85ebfe6..544e630 100644 --- a/src/uxarray_mcp/tools/plotting.py +++ b/src/uxarray_mcp/tools/plotting.py @@ -16,6 +16,41 @@ ) from uxarray_mcp.domain.zonal import compute_zonal_mean_stats from uxarray_mcp.provenance import attach_provenance +from uxarray_mcp.typed_results import spill_png + + +def _png_content(png_bytes: bytes, *, plot_type: str) -> tuple[Any, dict[str, Any]]: + """Return the content block for a rendered PNG, plus a note about it. + + Small images are inlined as an ``image`` block, which keeps the common + interactive case a single round trip. A large one is written to the + artifact store and handed back as a ``resource_link`` instead: base64 + inflates bytes by a third, and a multi-hundred-kilobyte figure costs + far more of the caller's context than the picture is worth. The caller + fetches it by URI if it actually wants to look. + """ + link = spill_png(png_bytes, plot_type=plot_type) + if link is None: + b64 = base64.b64encode(png_bytes).decode("utf-8") + return ( + ImageContent(type="image", data=b64, mimeType="image/png"), + {"image_delivery": "inline"}, + ) + from mcp.types import ResourceLink + + block = ResourceLink( + type="resource_link", + uri=link["uri"], + name=link["name"], + title=link.get("title"), + description=link.get("description"), + mimeType=link["mime_type"], + size=link.get("size"), + ) + return block, { + "image_delivery": "resource_link", + "image_uri": link["uri"], + } def _resolve_plot_paths( @@ -86,10 +121,11 @@ def _plot_mesh_local( grid = load_grid(grid_path) png_bytes = render_mesh(grid, width=width, height=height) - b64 = base64.b64encode(png_bytes).decode("utf-8") + image_block, delivery = _png_content(png_bytes, plot_type="mesh_wireframe") result = { "image_size_bytes": len(png_bytes), + **delivery, "grid_info": { "n_face": int(grid.n_face), "n_node": int(grid.n_node), @@ -118,7 +154,7 @@ def _plot_mesh_local( ) return [ - ImageContent(type="image", data=b64, mimeType="image/png"), + image_block, TextContent(type="text", text=json.dumps(provenance, indent=2)), ] @@ -292,7 +328,7 @@ def plot_mesh_geo( city_scale=city_scale, ) - b64 = base64.b64encode(png_bytes).decode("utf-8") + image_block, delivery = _png_content(png_bytes, plot_type="mesh_geographic") # ── Build human-readable plot note ─────────────────────────────────────── note = _build_plot_note( @@ -301,6 +337,7 @@ def plot_mesh_geo( result = { "image_size_bytes": len(png_bytes), + **delivery, "grid_info": { "n_face": int(grid.n_face), "n_node": int(grid.n_node), @@ -330,7 +367,7 @@ def plot_mesh_geo( ], ) return [ - ImageContent(type="image", data=b64, mimeType="image/png"), + image_block, TextContent(type="text", text=note + "\n\n" + json.dumps(provenance, indent=2)), ] @@ -651,10 +688,11 @@ def _plot_variable_local( title=title, time_index=time_index, ) - b64 = base64.b64encode(png_bytes).decode("utf-8") + image_block, delivery = _png_content(png_bytes, plot_type="variable_polygons") result = { "image_size_bytes": len(png_bytes), + **delivery, "variable_name": variable_name, "grid_info": { "n_face": int(uxds.uxgrid.n_face), @@ -692,7 +730,7 @@ def _plot_variable_local( ) return [ - ImageContent(type="image", data=b64, mimeType="image/png"), + image_block, TextContent(type="text", text=json.dumps(provenance, indent=2)), ] @@ -816,10 +854,11 @@ def _plot_zonal_mean_local( line_color=line_color, title=title, ) - b64 = base64.b64encode(png_bytes).decode("utf-8") + image_block, delivery = _png_content(png_bytes, plot_type="zonal_mean_profile") result = { "image_size_bytes": len(png_bytes), + **delivery, "variable_name": variable_name, "latitudes": latitudes, "zonal_mean_values": values, @@ -856,6 +895,6 @@ def _plot_zonal_mean_local( ) return [ - ImageContent(type="image", data=b64, mimeType="image/png"), + image_block, TextContent(type="text", text=json.dumps(provenance, indent=2)), ] diff --git a/src/uxarray_mcp/tools/remote_tools.py b/src/uxarray_mcp/tools/remote_tools.py index 6b89fe1..52baae9 100644 --- a/src/uxarray_mcp/tools/remote_tools.py +++ b/src/uxarray_mcp/tools/remote_tools.py @@ -167,14 +167,75 @@ def _run_with_optional_hpc( def _plot_result_to_mcp_contents(result: Dict[str, Any]) -> list[Any]: - """Convert a plot result dict into inline MCP image + metadata contents.""" + """Convert a plot result dict into MCP image/link + metadata contents. + + A figure small enough to inline comes back as an ``image`` block. A + large one was written to the artifact store instead, so there are no + bytes to inline and the caller gets a ``resource_link`` pointing at + it. Building an ``ImageContent`` in that second case would fail + validation on a null payload, which is exactly what happens on the + biggest meshes. + """ metadata = {key: value for key, value in result.items() if key != "png_b64"} + text = TextContent(type="text", text=json.dumps(metadata, indent=2)) + + b64 = result.get("png_b64") + if b64 is not None: + return [ + ImageContent(type="image", data=b64, mimeType="image/png"), + text, + ] + + uri = result.get("image_uri") + if uri is None: + # No bytes and no URI: nothing to show but the metadata, which + # still carries the numbers the caller asked for. + return [text] + + from mcp.types import ResourceLink + + name = str(uri).rsplit("/", 1)[-1] or "plot.png" return [ - ImageContent(type="image", data=result["png_b64"], mimeType="image/png"), - TextContent(type="text", text=json.dumps(metadata, indent=2)), + ResourceLink( + type="resource_link", + uri=uri, + name=name, + title="plot", + mimeType="image/png", + size=result.get("image_size_bytes"), + ), + text, ] +def _image_payload(img: Any, meta: dict) -> Dict[str, Any]: + """Describe a plot's image whether it came back inline or as a link. + + A large figure is handed back as a ``resource_link`` rather than + inlined, so ``png_b64`` is absent and the URI is what the caller + follows. Reading ``.data`` unconditionally raised on exactly those + plots, which meant the biggest meshes -- the ones most worth looking + at -- failed here. + """ + import base64 + + b64 = getattr(img, "data", None) + if b64 is not None: + return { + "png_b64": b64, + "image_size_bytes": meta.get( + "image_size_bytes", len(base64.b64decode(b64)) + ), + } + uri = getattr(img, "uri", None) + return { + "png_b64": None, + "image_uri": str(uri) if uri is not None else None, + "image_delivery": "resource_link", + "image_size_bytes": meta.get("image_size_bytes") or getattr(img, "size", None), + } + + def inspect_mesh( file_path: str, use_remote: bool = False, @@ -485,7 +546,6 @@ def plot_mesh( resolved_grid, _ = _resolve_plot_paths(grid_path, None, session_id, dataset_handle) def _local() -> Dict[str, Any]: - import base64 import json items = _plot_mesh_local(resolved_grid, width=width, height=height) @@ -493,10 +553,7 @@ def _local() -> Dict[str, Any]: img = items[0] meta = json.loads(items[1].text) return { - "png_b64": img.data, - "image_size_bytes": meta.get( - "image_size_bytes", len(base64.b64decode(img.data)) - ), + **_image_payload(img, meta), "grid_info": meta.get("grid_info", {}), "execution_venue": "local", "_provenance": meta.get("_provenance", {}), @@ -583,7 +640,6 @@ def plot_variable( ) def _local() -> Dict[str, Any]: - import base64 import json items = _plot_variable_local( @@ -601,10 +657,7 @@ def _local() -> Dict[str, Any]: img = items[0] meta = json.loads(items[1].text) return { - "png_b64": img.data, - "image_size_bytes": meta.get( - "image_size_bytes", len(base64.b64decode(img.data)) - ), + **_image_payload(img, meta), "variable_name": meta.get("variable_name", variable_name), "grid_info": meta.get("grid_info", {}), "execution_venue": "local", @@ -701,7 +754,6 @@ def plot_zonal_mean( ) def _local() -> Dict[str, Any]: - import base64 import json items = _plot_zonal_mean_local( @@ -719,10 +771,7 @@ def _local() -> Dict[str, Any]: img = items[0] meta = json.loads(items[1].text) return { - "png_b64": img.data, - "image_size_bytes": meta.get( - "image_size_bytes", len(base64.b64decode(img.data)) - ), + **_image_payload(img, meta), "variable_name": meta.get("variable_name", variable_name), "latitudes": meta.get("latitudes", []), "zonal_mean_values": meta.get("zonal_mean_values", []), diff --git a/src/uxarray_mcp/typed_results.py b/src/uxarray_mcp/typed_results.py index 6924602..8fb4628 100644 --- a/src/uxarray_mcp/typed_results.py +++ b/src/uxarray_mcp/typed_results.py @@ -208,8 +208,15 @@ def declared_output_schemas() -> dict[str, dict[str, Any]]: #: Front-door tools and the envelope they return. Keyed by the registered #: tool name because front doors are registered at top level without a #: namespace. +#: +#: Only ``run_analysis`` is listed. It is the front door that runs a single +#: operation through the precondition gate, so it is the one that actually +#: returns ``result_type`` and the refusal fields this envelope promises. +#: ``analyze_dataset`` is a multi-stage summary with a different shape -- +#: declaring the envelope for it advertised a contract it does not honour, +#: and an SDK that validates ``structuredContent`` rejects the reply. _FRONTDOOR_SCHEMAS: dict[str, dict[str, Any]] = { - "analyze_dataset": _ANALYSIS_ENVELOPE, + "run_analysis": _ANALYSIS_ENVELOPE, } diff --git a/tests/test_typed_results.py b/tests/test_typed_results.py index 7af0fee..f828efa 100644 --- a/tests/test_typed_results.py +++ b/tests/test_typed_results.py @@ -88,16 +88,22 @@ def test_schemas_are_json_serializable(self): class TestAnalysisEnvelope: - """The front-door envelope is the paper's result contract as a schema.""" + """The front-door envelope is the paper's result contract as a schema. + + The envelope belongs to ``run_analysis``: that is the tool that + returns a ``result_type`` and can refuse with ``input_required``. + ``analyze_dataset`` runs a multi-stage summary and has no such field, + so it must not advertise this shape. + """ def test_refusal_is_advertised_as_a_reachable_shape(self): - schema = output_schema_for("analyze_dataset") + schema = output_schema_for("run_analysis") result_type = schema["properties"]["result_type"] assert set(result_type["enum"]) == {"complete", "input_required"} assert schema["required"] == ["result_type"] def test_contract_blocks_are_all_declared(self): - props = output_schema_for("analyze_dataset")["properties"] + props = output_schema_for("run_analysis")["properties"] for block in ( "scientific_status", "preconditions", @@ -108,7 +114,7 @@ def test_contract_blocks_are_all_declared(self): def test_not_evaluated_is_distinct_from_failed(self): # #84: "we did not check" must not read as "the check passed". - pre = output_schema_for("analyze_dataset")["properties"]["preconditions"] + pre = output_schema_for("run_analysis")["properties"]["preconditions"] assert set(pre["properties"]["status"]["enum"]) == { "satisfied", "failed", @@ -116,7 +122,7 @@ def test_not_evaluated_is_distinct_from_failed(self): } def test_interpretability_is_nullable(self): - status = output_schema_for("analyze_dataset")["properties"]["scientific_status"] + status = output_schema_for("run_analysis")["properties"]["scientific_status"] assert status["properties"]["physically_interpretable"]["type"] == [ "boolean", "null", @@ -133,7 +139,16 @@ def test_schema_reaches_tool_metadata(self, profile): if "output_schema" in (getattr(registry.get_tool(name).metadata, "extra", None) or {}) } - assert "analyze_dataset" in published + assert "run_analysis" in published + + @pytest.mark.parametrize("profile", ["core", "deferred-full"]) + def test_multi_stage_summary_does_not_claim_the_envelope(self, profile): + # analyze_dataset returns a stage summary, not a result_type + # envelope. Declaring one made the server reject its own output as + # invalid structured content the moment a client called it. + registry = make_registry(profile=profile) + extra = getattr(registry.get_tool("analyze_dataset").metadata, "extra", None) + assert "output_schema" not in (extra or {}) def test_deferred_profile_publishes_the_compute_tools(self): registry = make_registry(profile="deferred-full") @@ -151,7 +166,7 @@ def test_schema_survives_the_route_table(self): from toolregistry_server.route_table import RouteTable table = RouteTable(make_registry(profile="core")) - route = table.get_route("analyze_dataset") + route = table.get_route("run_analysis") assert route.output_schema is not None assert "result_type" in route.output_schema["properties"] @@ -223,3 +238,69 @@ def test_ttl_is_bounded_and_scope_is_shared(self): # is correct; the TTL bounds staleness if that ever changes. assert 0 < LIST_TOOLS_TTL_MS <= 600_000 assert LIST_TOOLS_CACHE_SCOPE in {"public", "private"} + + +class TestDeclaredShapeMatchesRealOutput: + """A declared schema is a promise; these check we actually keep it.""" + + def test_contracted_results_carry_the_operation_they_declare(self): + # Every contract-derived schema requires `operation`, but nothing + # emitted it, so a validating client saw a malformed envelope on + # results that were in fact fine. + from uxarray_mcp.provenance import attach_provenance + + result = attach_provenance({}, tool="inspect_mesh", inputs={}) + assert result["operation"] == "inspect_mesh" + + def test_uncontracted_results_do_not_invent_one(self): + from uxarray_mcp.provenance import attach_provenance + + result = attach_provenance({}, tool="get_capabilities", inputs={}) + assert "operation" not in result + + def test_run_analysis_output_validates_against_its_own_schema(self): + jsonschema = pytest.importorskip("jsonschema") + + jsonschema.validate( + { + "result_type": "complete", + "scientific_status": {"physically_interpretable": None}, + "_provenance": {}, + }, + output_schema_for("run_analysis"), + ) + + +class TestPlotContentBlocks: + """Plot results must stay real content blocks, inline or linked.""" + + def _contents(self, png_b64, **extra): + from uxarray_mcp.tools.remote_tools import _plot_result_to_mcp_contents + + return _plot_result_to_mcp_contents({"png_b64": png_b64, **extra}) + + def test_small_figure_is_an_inline_image(self): + blocks = self._contents("aGk=") + assert [b.type for b in blocks] == ["image", "text"] + assert blocks[0].data == "aGk=" + + def test_spilled_figure_becomes_a_resource_link(self): + # A spilled plot has no bytes to inline. Building an image block + # from the missing payload failed validation outright, so the + # largest meshes -- the ones worth spilling -- returned an error. + blocks = self._contents( + None, image_uri="file:///tmp/plot_abc.png", image_size_bytes=999 + ) + assert [b.type for b in blocks] == ["resource_link", "text"] + assert str(blocks[0].uri) == "file:///tmp/plot_abc.png" + assert blocks[0].name == "plot_abc.png" + + def test_metadata_survives_either_delivery(self): + for b64, extra in (("aGk=", {}), (None, {"image_uri": "file:///t/p.png"})): + text = self._contents(b64, grid_info={"n_face": 8}, **extra)[-1] + assert json.loads(text.text)["grid_info"] == {"n_face": 8} + assert "png_b64" not in json.loads(text.text) + + def test_no_bytes_and_no_uri_still_returns_metadata(self): + blocks = self._contents(None) + assert [b.type for b in blocks] == ["text"]