From 7ffaae894a1071f37997dbfb64cc5bffada6e37e Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:06:34 -0400 Subject: [PATCH 1/2] Cache compiled output-schema validators on ClientSession `validate_tool_result` called one-shot `jsonschema.validate()` on every tool result, which rebuilt the validator class, re-ran `check_schema`, and re-crawled `$ref`s each time. Compile the validator once per schema instead, keyed on the identity of the cached schema object so a relisted tool can never reuse a validator built from its previous schema. Refs #3133. --- src/mcp/client/session.py | 59 +++++++++++++++++++++---- tests/client/test_session_promotions.py | 45 +++++++++++++++++++ 2 files changed, 96 insertions(+), 8 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 097ade1c91..171dfce15a 100644 --- a/src/mcp/client/session.py +++ b/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 @@ -56,6 +56,11 @@ from mcp.shared.subscriptions import SUBSCRIPTION_ID_META_KEY, event_from_wire from mcp.shared.transport_context import TransportContext +if TYPE_CHECKING: + # `jsonschema` is imported lazily inside `validate_tool_result`: pulling it (and its + # `attrs`/`referencing` tree) in at module scope costs every client that never validates. + 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 @@ -357,6 +362,9 @@ 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] = {} + # Compiled output-schema validators, each paired with the schema object it was built + # from so a re-listed tool can't be served a validator for its previous schema. + self._tool_output_validators: dict[str, tuple[dict[str, Any], 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 @@ -1032,16 +1040,51 @@ 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 + from jsonschema import exceptions as jsonschema_exceptions 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._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] + ) + if error is not None: + raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") + + def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> Validator: + """Compiled validator for `output_schema`, built once and reused across calls. + + Compiling is ~60x the cost of validating, so a one-shot `jsonschema.validate()` per + result dominates `call_tool`. The cache holds the schema object it compiled alongside + the validator and only reuses it on an identity match: `_absorb_tool_listing` assigns a + freshly parsed schema object on every listing, so a server that changes a tool's schema + forces a rebuild rather than reusing the stale validator. Holding the reference also + keeps the allocator from recycling that identity behind our back. + + Raises: + RuntimeError: The schema is not a valid JSON Schema. Raised on every call, since a + failed compile is never cached. + """ + from jsonschema import SchemaError + from jsonschema.validators import validator_for + + cached = self._tool_output_validators.get(name) + if cached is not None and cached[0] is output_schema: + return cached[1] + + 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}") + # The `Validator` protocol declares `registry` without a default; concrete classes have one. + validator = cast("Callable[[dict[str, Any]], Validator]", validator_cls)(output_schema) + self._tool_output_validators[name] = (output_schema, validator) + return validator async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult: """Send a prompts/list request. diff --git a/tests/client/test_session_promotions.py b/tests/client/test_session_promotions.py index e4a62732b9..b6b82ba7ce 100644 --- a/tests/client/test_session_promotions.py +++ b/tests/client/test_session_promotions.py @@ -64,3 +64,48 @@ 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_raises_on_an_unusable_output_schema() -> None: + """A schema that isn't valid JSON Schema is reported as such, on every call.""" + server = _make_server({"type": "not-a-json-schema-type"}) + async with Client(server) as client: + result = CallToolResult(content=[], structured_content={"x": 1}) + for _ in range(2): + # Compiling is never cached on failure, so the second call raises like the first. + with pytest.raises(RuntimeError, match="Invalid schema for tool t"): + await client.session.validate_tool_result("t", result) + + +@pytest.mark.anyio +async def test_validate_tool_result_compiles_the_output_schema_once_per_tool() -> None: + """Regression guard: compiling dominates validating, so the validator must outlive one call.""" + server = _make_server({"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]}) + async with Client(server) as client: + result = CallToolResult(content=[], structured_content={"x": 1}) + await client.session.validate_tool_result("t", result) + compiled = client.session._tool_output_validators["t"] + await client.session.validate_tool_result("t", result) + assert client.session._tool_output_validators["t"] is compiled + + +@pytest.mark.anyio +async def test_validate_tool_result_recompiles_when_the_server_changes_the_schema() -> None: + """A relisted tool must not be validated against the schema it used to declare.""" + 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))]) + + server = Server("test-server", on_list_tools=on_list_tools) + async with Client(server) as client: + integer_result = CallToolResult(content=[], structured_content={"x": 1}) + await client.session.validate_tool_result("t", integer_result) + + 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) From 2e2fcec922e81f627dbc99394e32e59b110e00ce Mon Sep 17 00:00:00 2001 From: Max Isbey <224885523+maxisbey@users.noreply.github.com> Date: Sun, 26 Jul 2026 09:49:08 +0000 Subject: [PATCH 2/2] Key validator reuse to schema value and let listing absorption own it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _absorb_tool_listing now owns the compiled-validator lifecycle: an unchanged output schema (a re-listing, or the response cache re-absorbing a served hit) keeps its validator, a changed schema evicts it, and a complete listing prunes validators for tools the server no longer lists, alongside the other per-tool state. Reusing on schema value rather than object identity keeps validators across relistings — a re-parsed or deep-copied listing yields a fresh but equal schema object, which the identity check treated as a change and rebuilt for. Schemas are compared as canonical JSON so `const: true` and `const: 1` stay distinct. Also chains the structured-content RuntimeError to the underlying ValidationError so its structured fields remain reachable via __cause__. --- src/mcp/client/session.py | 49 +++++++++++++++++-------- tests/client/test_client.py | 4 ++ tests/client/test_session_promotions.py | 16 ++++++++ 3 files changed, 53 insertions(+), 16 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 171dfce15a..b49d42dbdf 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import logging from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass @@ -75,6 +76,17 @@ def _clamp_inbound_ttl(raw: dict[str, Any]) -> None: raw["ttlMs"] = 0 +def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool: + """JSON equality for two output schemas. + + Python `==` is not JSON equality: it conflates `True`/`1` and `False`/`0`, which JSON + Schema keeps distinct (`const: true` vs `const: 1`). Canonical serialization compares as + JSON does; where it is stricter (`1` vs `1.0`), erring toward "changed" only costs a + recompile, never a stale validator. + """ + return json.dumps(a, sort_keys=True) == json.dumps(b, sort_keys=True) + + def _preconnect_stamp(data: dict[str, Any], opts: CallOptions) -> None: # initialize/discover forbid cancellation; other pre-handshake requests (lowlevel # ClientSession callers may skip the handshake entirely) keep the courtesy cancel. @@ -362,9 +374,9 @@ 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] = {} - # Compiled output-schema validators, each paired with the schema object it was built - # from so a re-listed tool can't be served a validator for its previous schema. - self._tool_output_validators: dict[str, tuple[dict[str, Any], Validator]] = {} + # Compiled output-schema validators, derived from `_tool_output_schemas` and owned by + # `_absorb_tool_listing`, which evicts a tool's entry whenever its schema changes. + 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 @@ -1053,17 +1065,15 @@ async def validate_tool_result(self, name: str, result: types.CallToolResult) -> jsonschema_exceptions.best_match(errors), # pyright: ignore[reportUnknownMemberType] ) if error is not None: - raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") + raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") from error def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> Validator: - """Compiled validator for `output_schema`, built once and reused across calls. + """Compiled validator for the tool's cached output schema, built once per schema value. Compiling is ~60x the cost of validating, so a one-shot `jsonschema.validate()` per - result dominates `call_tool`. The cache holds the schema object it compiled alongside - the validator and only reuses it on an identity match: `_absorb_tool_listing` assigns a - freshly parsed schema object on every listing, so a server that changes a tool's schema - forces a rebuild rather than reusing the stale validator. Holding the reference also - keeps the allocator from recycling that identity behind our back. + result dominates `call_tool`; the compiled validator is cached instead. It stays valid + because `_absorb_tool_listing` evicts a tool's validator whenever it absorbs a different + schema for that tool, so a cached entry always matches `output_schema`. Raises: RuntimeError: The schema is not a valid JSON Schema. Raised on every call, since a @@ -1072,18 +1082,18 @@ def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> from jsonschema import SchemaError from jsonschema.validators import validator_for - cached = self._tool_output_validators.get(name) - if cached is not None and cached[0] is output_schema: - return cached[1] + 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}") - # The `Validator` protocol declares `registry` without a default; concrete classes have one. + # 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) - self._tool_output_validators[name] = (output_schema, validator) + self._tool_output_validators[name] = validator return validator async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult: @@ -1213,8 +1223,14 @@ def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) kept.append(tool) result.tools = kept - # Cache tool output schemas for future validation; cursor pages only ever add. + # Cache tool output schemas for future validation; cursor pages only ever add. A + # changed schema evicts its compiled validator; an unchanged one (a re-listing, or the + # response cache re-absorbing a served hit) keeps it. Only validated tools pay the check. for tool in result.tools: + if tool.name in self._tool_output_validators and not _same_schema( + self._tool_output_schemas.get(tool.name), tool.output_schema + ): + del self._tool_output_validators[tool.name] self._tool_output_schemas[tool.name] = tool.output_schema if complete: @@ -1223,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 diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 6c78503b97..b3f9f4db6e 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -639,12 +639,16 @@ async def test_a_complete_listing_prunes_per_tool_state_for_tools_it_no_longer_c with anyio.fail_after(5): async with Client(server) as client: await client.session.list_tools() + # Compile the retired tool's output-schema validator so its eviction is observable. + await client.session.validate_tool_result("retired", CallToolResult(content=[], structured_content={})) assert set(client.session._x_mcp_header_maps) == {"retired", "survivor"} assert set(client.session._tool_output_schemas) == {"retired", "survivor"} + assert set(client.session._tool_output_validators) == {"retired"} await client.session.list_tools() assert set(client.session._x_mcp_header_maps) == {"survivor"} assert set(client.session._tool_output_schemas) == {"survivor"} + assert client.session._tool_output_validators == {} async def test_a_complete_listing_prunes_output_schemas_on_a_legacy_session_too() -> None: diff --git a/tests/client/test_session_promotions.py b/tests/client/test_session_promotions.py index b6b82ba7ce..6d6b6bc8dc 100644 --- a/tests/client/test_session_promotions.py +++ b/tests/client/test_session_promotions.py @@ -74,6 +74,7 @@ async def test_validate_tool_result_raises_on_an_unusable_output_schema() -> Non result = CallToolResult(content=[], structured_content={"x": 1}) for _ in range(2): # Compiling is never cached on failure, so the second call raises like the first. + # Stable SDK prefix only: the message tail is jsonschema text that shifts with the dependency. with pytest.raises(RuntimeError, match="Invalid schema for tool t"): await client.session.validate_tool_result("t", result) @@ -90,6 +91,21 @@ async def test_validate_tool_result_compiles_the_output_schema_once_per_tool() - assert client.session._tool_output_validators["t"] is compiled +@pytest.mark.anyio +async def test_validate_tool_result_keeps_the_validator_across_a_relisting_of_the_same_schema() -> None: + """SDK-defined: an unchanged schema on a re-listing (or a re-absorbed cache hit) keeps its + compiled validator, so relisting between calls doesn't reintroduce the per-call compile.""" + server = _make_server({"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]}) + async with Client(server) as client: + result = CallToolResult(content=[], structured_content={"x": 1}) + await client.session.validate_tool_result("t", result) + compiled = client.session._tool_output_validators["t"] + + await client.session.list_tools() # a second listing carrying an equal schema + await client.session.validate_tool_result("t", result) + assert client.session._tool_output_validators["t"] is compiled + + @pytest.mark.anyio async def test_validate_tool_result_recompiles_when_the_server_changes_the_schema() -> None: """A relisted tool must not be validated against the schema it used to declare."""