Resolve tool output-schema references within the schema document only (#3394)

Author: maxisbey

src/mcp/client/session.py

@@ -1126,7 +1126,8 @@ async def validate_tool_result(self, name: str, result: types.CallToolResult) ->
         """Revalidate a `CallToolResult` against the tool's declared output schema.
 
         Raises:
-            RuntimeError: Structured content is missing or does not conform to the schema.
+            RuntimeError: Structured content is missing or does not conform to the schema, or the
+                schema is invalid or has a `$ref` that does not resolve within the schema document.
         """
         if name not in self._tool_output_schemas:
             # refresh output schema cache
@@ -1140,17 +1141,22 @@ async def validate_tool_result(self, name: str, result: types.CallToolResult) ->
 
         if output_schema is not None:
             from jsonschema import exceptions as jsonschema_exceptions
+            from referencing.exceptions import Unresolvable
 
             if result.structured_content is None:
                 raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
             validator = self._output_schema_validator(name, output_schema)
             # `best_match` picks the same error the previous `jsonschema.validate()` call raised,
             # so the message a caller sees is unchanged. It is untyped upstream.
             errors = validator.iter_errors(result.structured_content)
-            error = cast(
-                "Exception | None",
-                jsonschema_exceptions.best_match(errors),  # pyright: ignore[reportUnknownMemberType]
-            )
+            try:
+                error = cast(
+                    "Exception | None",
+                    jsonschema_exceptions.best_match(errors),  # pyright: ignore[reportUnknownMemberType]
+                )
+            except Unresolvable as e:
+                # A `$ref` did not resolve within the schema document.
+                raise RuntimeError(f"Invalid schema for tool {name}: {e}") from e
             if error is not None:
                 raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") from error
 
@@ -1168,6 +1174,7 @@ def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) ->
         """
         from jsonschema import SchemaError
         from jsonschema.validators import validator_for
+        from referencing import Registry
 
         if (validator := self._tool_output_validators.get(name)) is not None:
             return validator
@@ -1177,9 +1184,8 @@ def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) ->
             validator_cls.check_schema(output_schema)
         except SchemaError as e:
             raise RuntimeError(f"Invalid schema for tool {name}: {e}")
-        # jsonschema ships no `py.typed`, so pyright reads typeshed's stub, which declares
-        # `registry` as required (concrete validators default it); cast to a schema-only ctor.
-        validator = cast("Callable[[dict[str, Any]], Validator]", validator_cls)(output_schema)
+        # An explicit empty registry: `$ref`s resolve within the schema document and the bundled metaschemas.
+        validator = validator_cls(output_schema, registry=Registry())
         self._tool_output_validators[name] = validator
         return validator
 

tests/client/test_output_schema_validation.py

@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 from typing import Any
 
 import pytest
@@ -10,6 +11,7 @@
     TextContent,
     Tool,
 )
+from referencing.exceptions import Unresolvable
 
 from mcp import Client
 from mcp.server import Server, ServerRequestContext
@@ -163,3 +165,26 @@ async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams)
         assert result.is_error is False
 
         assert "Tool mystery_tool not listed" in caplog.text
+
+
+# jsonschema's fallback retriever emits this DeprecationWarning; keep it a plain warning so the
+# assertions below decide the outcome rather than the suite's warnings-as-errors filter.
+@pytest.mark.filterwarnings("default:Automatically retrieving remote references:DeprecationWarning")
+@pytest.mark.anyio
+async def test_output_schema_ref_outside_the_document_is_rejected(tmp_path: Path):
+    """A `$ref` to a URI outside the output schema is not resolved, and a result whose validation
+    reaches one fails as an invalid schema (spec `$ref` resolution; applying it to `file:` URIs too
+    is SDK-defined)."""
+    target = tmp_path / "schema.json"
+    target.write_text("{}", encoding="utf-8")
+    server = _make_server(
+        tools=[Tool(name="probe", input_schema={"type": "object"}, output_schema={"$ref": target.as_uri()})],
+        structured_content={"v": 1},
+    )
+
+    async with Client(server) as client:
+        with pytest.raises(RuntimeError) as exc_info:
+            await client.call_tool("probe", {})
+    # SDK-authored prefix only; the tail is `referencing`'s text.
+    assert str(exc_info.value).startswith("Invalid schema for tool probe: ")
+    assert isinstance(exc_info.value.__cause__, Unresolvable)