perf(client): cache compiled jsonschema validator per tool output schema Fixes #3133

src/mcp/client/session.py

@@ -6,7 +6,7 @@
 from functools import reduce
 from operator import or_
 from types import TracebackType
-from typing import Annotated, Any, Final, Literal, Protocol, cast, overload
+from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, cast, overload
 
 import anyio
 import anyio.abc
@@ -57,6 +57,9 @@
 from mcp.shared.subscriptions import SUBSCRIPTION_ID_META_KEY, event_from_wire
 from mcp.shared.transport_context import TransportContext
 
+if TYPE_CHECKING:
+    from jsonschema.protocols import Validator
+
 DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
 DISCOVER_TIMEOUT_SECONDS = 10.0
 _NOTIFICATION_QUEUE_SIZE: Final = 256
@@ -373,6 +376,7 @@ def __init__(
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
+        self._tool_output_validators: dict[str, Validator] = {}
         self._x_mcp_header_maps: dict[str, dict[tuple[str, ...], str]] = {}
         self._initialize_result: types.InitializeResult | None = None
         self._discover_result: types.DiscoverResult | None = None
@@ -1057,16 +1061,42 @@ async def validate_tool_result(self, name: str, result: types.CallToolResult) ->
             logger.warning(f"Tool {name} not listed by server, cannot validate any structured content")
 
         if output_schema is not None:
-            from jsonschema import SchemaError, ValidationError, validate
+            # jsonschema's stub omits best_match's return type, so its import (and any
+            # unannotated use of its result) is reported as partially-unknown under strict
+            # pyright; cast() the result below rather than suppressing checks on this file.
+            from jsonschema.exceptions import (
+                SchemaError,
+                best_match,  # pyright: ignore[reportUnknownVariableType]
+            )
+            from jsonschema.exceptions import ValidationError as JSONSchemaValidationError
+            from jsonschema.validators import validator_for
 
             if result.structured_content is None:
                 raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
-            try:
-                validate(result.structured_content, output_schema)
-            except ValidationError as e:
-                raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}")
-            except SchemaError as e:  # pragma: no cover
-                raise RuntimeError(f"Invalid schema for tool {name}: {e}")  # pragma: no cover
+
+            validator = self._tool_output_validators.get(name)
+            if validator is None:
+                # Build and cache the validator once per (tool, schema) instead of on every
+                # call: jsonschema.validate() re-derives the validator class, re-checks the
+                # schema, and re-crawls its $refs each time it's invoked, which is wasted
+                # work once we already know the schema is valid and unchanged.
+                validator_cls = validator_for(output_schema)
+                try:
+                    validator_cls.check_schema(output_schema)
+                except SchemaError as e:  # pragma: no cover
+                    raise RuntimeError(f"Invalid schema for tool {name}: {e}")  # pragma: no cover
+                # The Validator protocol's stub requires `registry`, but every concrete
+                # validator (as returned by validator_for) defaults it, same as plain
+                # jsonschema.validate() relies on.
+                validator = validator_cls(output_schema)  # pyright: ignore[reportCallIssue]
+                self._tool_output_validators[name] = validator
+
+            # Mirror jsonschema.validate()'s error selection (best_match over all errors)
+            # rather than Validator.validate()'s "raise the first one", so error messages
+            # are unchanged from before this caching was introduced.
+            error = cast(JSONSchemaValidationError | None, best_match(validator.iter_errors(result.structured_content)))
+            if error is not None:
+                raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}")
 
     async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult:
         """Send a prompts/list request.
@@ -1197,6 +1227,10 @@ def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool)
 
         # Cache tool output schemas for future validation; cursor pages only ever add.
         for tool in result.tools:
+            if tool.name not in self._tool_output_schemas or self._tool_output_schemas[tool.name] != tool.output_schema:
+                # Schema is new or changed since we last saw this tool: drop any compiled
+                # validator built from the old schema so validate_tool_result recompiles it.
+                self._tool_output_validators.pop(tool.name, None)
             self._tool_output_schemas[tool.name] = tool.output_schema
 
         if complete:
@@ -1205,6 +1239,7 @@ def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool)
             names = {tool.name for tool in result.tools}
             self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names}
             self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names}
+            self._tool_output_validators = {k: v for k, v in self._tool_output_validators.items() if k in names}
 
         return result
 

tests/client/test_session_promotions.py

@@ -64,3 +64,49 @@ async def test_validate_tool_result_raises_on_schema_mismatch() -> None:
         # Stable SDK prefix only: the message tail is jsonschema text that shifts with the dependency.
         with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
             await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": "no"}))
+
+
+@pytest.mark.anyio
+async def test_validate_tool_result_reuses_cached_validator() -> None:
+    """The compiled jsonschema validator is built once per tool and reused, not rebuilt on every call."""
+    server = _make_server({"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]})
+    async with Client(server) as client:
+        await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": 1}))
+        first = client.session._tool_output_validators["t"]  # pyright: ignore[reportPrivateUsage]
+
+        await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": 2}))
+        second = client.session._tool_output_validators["t"]  # pyright: ignore[reportPrivateUsage]
+
+        assert first is second
+
+
+@pytest.mark.anyio
+async def test_validate_tool_result_rebuilds_validator_when_schema_changes() -> None:
+    """A fresh `list_tools()` that reports a different output schema drops the stale cached validator."""
+    schema_holder: dict[str, dict[str, object]] = {
+        "t": {"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]},
+    }
+
+    async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
+        return ListToolsResult(
+            tools=[Tool(name="t", input_schema={"type": "object"}, output_schema=schema_holder["t"])]
+        )
+
+    server = Server("test-server", on_list_tools=on_list_tools)
+    async with Client(server) as client:
+        await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": 1}))
+        stale_validator = client.session._tool_output_validators["t"]  # pyright: ignore[reportPrivateUsage]
+
+        # The tool's output schema now requires a string instead of an integer.
+        schema_holder["t"] = {"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]}
+        await client.session.list_tools()
+        assert "t" not in client.session._tool_output_validators  # pyright: ignore[reportPrivateUsage]
+
+        # The old value fails the new schema; a value matching the new schema passes,
+        # and rebuilds (rather than reuses) the cached validator.
+        with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
+            await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": 1}))
+        await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": "ok"}))
+
+        fresh_validator = client.session._tool_output_validators["t"]  # pyright: ignore[reportPrivateUsage]
+        assert fresh_validator is not stale_validator