Keep validator cache updates on the event loop

Author: Kludex

src/mcp/client/session.py

@@ -1140,6 +1140,8 @@ async def validate_tool_result(self, name: str, result: types.CallToolResult) ->
             if validator is None:
                 # First compilation lazily reads jsonschema's bundled schemas.
                 validator = await anyio.to_thread.run_sync(self._output_schema_validator, name, output_schema)
+                if _same_schema(self._tool_output_schemas.get(name), output_schema):
+                    self._tool_output_validators[name] = validator
 
             from jsonschema import exceptions as jsonschema_exceptions
             from referencing.exceptions import Unresolvable
@@ -1174,18 +1176,13 @@ def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) ->
         from jsonschema.validators import validator_for
         from referencing import Registry
 
-        if (validator := self._tool_output_validators.get(name)) is not None:
-            return validator
-
         validator_cls = validator_for(output_schema)
         try:
             validator_cls.check_schema(output_schema)
         except SchemaError as e:
             raise RuntimeError(f"Invalid schema for tool {name}: {e}")
         # 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
+        return validator_cls(output_schema, registry=Registry())
 
     async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult:
         """Send a prompts/list request.

tests/client/test_session_promotions.py

@@ -1,7 +1,11 @@
 """`dispatch_input_request` and `validate_tool_result` are public `ClientSession` API."""
 
+import anyio
+import anyio.from_thread
 import mcp_types as types
 import pytest
+from jsonschema.protocols import Validator
+from jsonschema.validators import validator_for
 from mcp_types import (
     CallToolResult,
     ErrorData,
@@ -125,3 +129,46 @@ async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParam
         await client.session.list_tools()
         with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
             await client.session.validate_tool_result("t", integer_result)
+
+
+@pytest.mark.anyio
+async def test_schema_change_during_compilation_does_not_cache_the_old_validator(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """SDK-defined: a concurrent tool relisting cannot leave a stale compiled validator cached."""
+    schemas = [
+        {"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]},
+        {"type": "object", "properties": {"x": {"type": "string"}}, "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=schemas.pop(0))])
+
+    compilation_started = anyio.Event()
+    continue_compilation = anyio.Event()
+
+    def paused_validator_for(schema: dict[str, object], default: type[Validator] | None = None) -> type[Validator]:
+        validator = validator_for(schema) if default is None else validator_for(schema, default=default)
+        if not compilation_started.is_set():
+            anyio.from_thread.run_sync(compilation_started.set)
+            anyio.from_thread.run(continue_compilation.wait)
+        return validator
+
+    monkeypatch.setattr("jsonschema.validators.validator_for", paused_validator_for)
+    server = Server("test-server", on_list_tools=on_list_tools)
+    async with Client(server) as client:
+        await client.session.list_tools()
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as task_group:
+                task_group.start_soon(
+                    client.session.validate_tool_result,
+                    "t",
+                    CallToolResult(content=[], structured_content={"x": 1}),
+                )
+                await compilation_started.wait()
+                try:
+                    await client.session.list_tools()
+                finally:
+                    continue_compilation.set()
+
+        await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": "yes"}))