modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cache compiled output-schema validators on ClientSession (#3134)

Browse source 
Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
Jeremiah Lowin 2026-07-26 06:29:10 -04:00 • committed by GitHub
commit f599cdfcf9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 134 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

78
src/mcp/client/session.py
View file

@ -1,12 +1,13 @@
from __future__ import annotations
import json
import logging
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import reduce
from operator import or_
from types import TracebackType
from typing import Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, overload
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, overload
import anyio
import anyio.abc
@ -56,6 +57,11 @@ from mcp.shared.message import ClientMessageMetadata, SessionMessage
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
@ -70,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.
@ -378,6 +395,9 @@ class ClientSession:
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, 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
@ -1062,16 +1082,49 @@ class ClientSession:
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}") from error
def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> Validator:
"""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 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
failed compile is never cached.
"""
from jsonschema import SchemaError
from jsonschema.validators import validator_for
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}")
# 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] = validator
return validator
async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult:
"""Send a prompts/list request.
@ -1200,8 +1253,14 @@ class ClientSession:
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:
@ -1210,6 +1269,7 @@ class ClientSession:
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

4
tests/client/test_client.py
View file

@ -657,12 +657,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:

61
tests/client/test_session_promotions.py
View file

@ -64,3 +64,64 @@ 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.
# 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)
@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_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."""
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)