diff --git a/src/uxarray_mcp/app.py b/src/uxarray_mcp/app.py index e28f9e0..040d46e 100644 --- a/src/uxarray_mcp/app.py +++ b/src/uxarray_mcp/app.py @@ -67,11 +67,39 @@ 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.""" + """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") + + 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/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/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/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 new file mode 100644 index 0000000..8fb4628 --- /dev/null +++ b/src/uxarray_mcp/typed_results.py @@ -0,0 +1,307 @@ +"""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. +#: +#: 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]] = { + "run_analysis": _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..f828efa --- /dev/null +++ b/tests/test_typed_results.py @@ -0,0 +1,306 @@ +"""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, +) + + +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(): + 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. + + 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("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("run_analysis")["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("run_analysis")["properties"]["preconditions"] + assert set(pre["properties"]["status"]["enum"]) == { + "satisfied", + "failed", + "not_evaluated", + } + + def test_interpretability_is_nullable(self): + status = output_schema_for("run_analysis")["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 "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") + 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 + + @needs_schema_adapter + 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("run_analysis") + 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 + + 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. + from uxarray_mcp import 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"} + + +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"]