From dc6635ea7d9316f25f39b8c75bdd5964e265b3e0 Mon Sep 17 00:00:00 2001 From: James Yang Date: Tue, 11 Aug 2026 13:21:04 -0400 Subject: [PATCH] Add ClientSession.register_tool_schema for dynamic tool discovery. Lets clients register output schemas for tools that never appear in list_tools(), so call_tool can validate structuredContent without touching private caches. Mirrored on Client; complete listings still prune unlisted registrations. Fixes #3145 --- docs/client/index.md | 2 + src/mcp/client/client.py | 8 ++ src/mcp/client/session.py | 33 ++++- tests/client/test_register_tool_schema.py | 156 ++++++++++++++++++++++ 4 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 tests/client/test_register_tool_schema.py diff --git a/docs/client/index.md b/docs/client/index.md index 1e1df3c01b..72c0719398 100644 --- a/docs/client/index.md +++ b/docs/client/index.md @@ -106,6 +106,8 @@ That is why `main` narrows with `isinstance(block, TextContent)` before touching When both are present they say the same thing twice on purpose: `content` is for a model, `structured_content` is for code. Where the structured half comes from, and how to control it, is the **[Structured Output](../servers/structured-output.md)** page. +The client validates `structured_content` against schemas learned from `list_tools()`. If a server exposes tools through a search or catalog API instead of listing them all, call `register_tool_schema(name, output_schema)` (on `Client` or `ClientSession`) before `call_tool` so those results are validated the same way. A later complete `list_tools()` that omits that name drops the registration — re-register if you still need it. + ### `is_error`: whether the tool failed A tool that raises does **not** raise in your client. It comes back as an ordinary result with `is_error=True`. diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index ed7c40f123..2aafb40ae8 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -943,6 +943,14 @@ async def list_tools( ), ) + def register_tool_schema(self, name: str, output_schema: dict[str, Any] | None = None) -> None: + """Register a tool's output schema for result validation. + + Delegates to `ClientSession.register_tool_schema`. Use when tools are discovered + dynamically and will not appear in `list_tools()` responses. + """ + self.session.register_tool_schema(name, output_schema) + @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) async def send_roots_list_changed(self) -> None: """Send a notification that the roots list has changed.""" diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 895339ca18..97cdf7aac9 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -405,8 +405,8 @@ def __init__( self._log_level: types.LoggingLevel | None = log_level 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. + # Compiled output-schema validators, derived from `_tool_output_schemas`. Evicted by + # `_absorb_tool_listing` and `register_tool_schema` whenever a tool's 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 @@ -1077,6 +1077,30 @@ def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dic return {} return mcp_param_headers(header_map, arguments) + def register_tool_schema(self, name: str, output_schema: dict[str, Any] | None = None) -> None: + """Register a tool's output schema for result validation. + + Use this when tools are discovered dynamically (for example via a catalog search API) + and will not appear in `list_tools()` responses. Writes into the same cache that + `list_tools()` / `_absorb_tool_listing` populate, and evicts any compiled validator + when the registered schema differs from the previously cached one. + + A later complete (uncursored, single-page) `list_tools()` that omits `name` drops the + registration, the same prune path used for listing-absorbed schemas. Re-register after + such a listing if the tool is still in use. If the listing includes `name`, the listed + `outputSchema` replaces this registration. + + Args: + name: Tool name as passed to `call_tool`. + output_schema: JSON Schema for `structuredContent`, or `None` when the tool has no + output schema (suppresses the "not listed" warning without validating). + """ + if name in self._tool_output_validators and not _same_schema( + self._tool_output_schemas.get(name), output_schema + ): + del self._tool_output_validators[name] + self._tool_output_schemas[name] = output_schema + async def validate_tool_result(self, name: str, result: types.CallToolResult) -> None: """Revalidate a `CallToolResult` against the tool's declared output schema. @@ -1114,8 +1138,9 @@ def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> 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`. + because `_absorb_tool_listing` and `register_tool_schema` evict a tool's validator + whenever they store 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 diff --git a/tests/client/test_register_tool_schema.py b/tests/client/test_register_tool_schema.py new file mode 100644 index 0000000000..d79801729f --- /dev/null +++ b/tests/client/test_register_tool_schema.py @@ -0,0 +1,156 @@ +"""`ClientSession.register_tool_schema` for tools absent from `list_tools`.""" + +import logging + +import pytest +from mcp_types import ( + CallToolRequestParams, + CallToolResult, + ListToolsResult, + PaginatedRequestParams, + Tool, +) + +from mcp.client.client import Client +from mcp.server import Server, ServerRequestContext + +_SCORE_SCHEMA: dict[str, object] = { + "type": "object", + "properties": {"score": {"type": "integer"}}, + "required": ["score"], +} +_SCORE_AS_STRING_SCHEMA: dict[str, object] = { + "type": "object", + "properties": {"score": {"type": "string"}}, + "required": ["score"], +} + + +def _dynamic_tool_server(*, structured_content: dict[str, object]) -> Server: + """`list_tools` advertises only a search meta-tool; `analyze` is callable but unlisted.""" + + async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult: + return ListToolsResult(tools=[Tool(name="search", input_schema={"type": "object"})]) + + async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + assert params.name == "analyze" + return CallToolResult(content=[], structured_content=structured_content) + + return Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool) + + +@pytest.mark.anyio +async def test_register_tool_schema_lets_call_tool_validate_an_unlisted_tool() -> None: + """SDK-defined: a schema registered for a tool absent from list_tools is used by call_tool.""" + server = _dynamic_tool_server(structured_content={"score": 1}) + async with Client(server) as client: + client.register_tool_schema("analyze", _SCORE_SCHEMA) + result = await client.call_tool("analyze", {}) + assert result.structured_content == {"score": 1} + + +@pytest.mark.anyio +async def test_register_tool_schema_makes_call_tool_reject_nonconforming_structured_content() -> None: + """Without registration, an unlisted tool skips validation; with it, mismatches raise.""" + server = _dynamic_tool_server(structured_content={"score": "no"}) + async with Client(server) as client: + # Unregistered: validation is skipped (tool never appears in list_tools). + skipped = await client.call_tool("analyze", {}) + assert skipped.structured_content == {"score": "no"} + + client.register_tool_schema("analyze", _SCORE_SCHEMA) + # 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 analyze"): + await client.call_tool("analyze", {}) + + +@pytest.mark.anyio +async def test_register_tool_schema_with_none_suppresses_the_unlisted_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + """SDK-defined: registering None marks the tool known without validating structuredContent.""" + server = _dynamic_tool_server(structured_content={"anything": True}) + async with Client(server) as client: + client.register_tool_schema("analyze", None) + with caplog.at_level(logging.WARNING, logger="client"): + result = await client.call_tool("analyze", {}) + assert result.structured_content == {"anything": True} + assert "not listed by server" not in caplog.text + + +@pytest.mark.anyio +async def test_register_tool_schema_evicts_the_compiled_validator_when_the_schema_changes() -> None: + """SDK-defined: a changed registration must not reuse a validator compiled for the old schema.""" + server = _dynamic_tool_server(structured_content={"score": 1}) + async with Client(server) as client: + client.register_tool_schema("analyze", _SCORE_SCHEMA) + await client.session.validate_tool_result( + "analyze", CallToolResult(content=[], structured_content={"score": 1}) + ) + compiled = client.session._tool_output_validators["analyze"] + + client.register_tool_schema("analyze", _SCORE_AS_STRING_SCHEMA) + assert "analyze" not in client.session._tool_output_validators + + with pytest.raises(RuntimeError, match="Invalid structured content returned by tool analyze"): + await client.session.validate_tool_result( + "analyze", CallToolResult(content=[], structured_content={"score": 1}) + ) + assert client.session._tool_output_validators["analyze"] is not compiled + + +@pytest.mark.anyio +async def test_register_tool_schema_keeps_the_validator_when_the_schema_is_unchanged() -> None: + """SDK-defined: re-registering an equal schema keeps the compiled validator.""" + server = _dynamic_tool_server(structured_content={"score": 1}) + async with Client(server) as client: + client.register_tool_schema("analyze", _SCORE_SCHEMA) + result = CallToolResult(content=[], structured_content={"score": 1}) + await client.session.validate_tool_result("analyze", result) + compiled = client.session._tool_output_validators["analyze"] + + client.register_tool_schema("analyze", dict(_SCORE_SCHEMA)) + await client.session.validate_tool_result("analyze", result) + assert client.session._tool_output_validators["analyze"] is compiled + + +@pytest.mark.anyio +async def test_a_complete_list_tools_prunes_a_manually_registered_schema() -> None: + """SDK-defined: a complete listing is still the full tool universe for prune — a registered + tool omitted from that listing is dropped, same as listing-absorbed schemas.""" + server = _dynamic_tool_server(structured_content={"score": 1}) + async with Client(server) as client: + client.register_tool_schema("analyze", _SCORE_SCHEMA) + assert "analyze" in client.session._tool_output_schemas + + await client.session.list_tools() + assert "analyze" not in client.session._tool_output_schemas + assert set(client.session._tool_output_schemas) == {"search"} + + +@pytest.mark.anyio +async def test_list_tools_that_includes_a_registered_name_replaces_the_registered_schema() -> None: + """SDK-defined: when the same name later appears in a listing, the listed schema wins.""" + + async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult: + return ListToolsResult( + tools=[ + Tool(name="analyze", input_schema={"type": "object"}, output_schema=_SCORE_AS_STRING_SCHEMA), + ] + ) + + async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + raise NotImplementedError + + server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool) + async with Client(server) as client: + client.register_tool_schema("analyze", _SCORE_SCHEMA) + await client.session.validate_tool_result( + "analyze", CallToolResult(content=[], structured_content={"score": 1}) + ) + + await client.session.list_tools() + with pytest.raises(RuntimeError, match="Invalid structured content returned by tool analyze"): + await client.session.validate_tool_result( + "analyze", CallToolResult(content=[], structured_content={"score": 1}) + )