add schema checking on client side via jsonschema

Author: davemssavage

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "sse-starlette>=1.6.1",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1; sys_platform != 'emscripten'",
+    "jsonschema==4.23.0",
 ]
 
 [project.optional-dependencies]

src/mcp/client/session.py

@@ -1,8 +1,10 @@
+from collections.abc import Awaitable, Callable
 from datetime import timedelta
-from typing import Any, Protocol
+from typing import Any, Protocol, TypeAlias
 
 import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from jsonschema import ValidationError, validate
 from pydantic import AnyUrl, TypeAdapter
 
 import mcp.types as types
@@ -11,6 +13,14 @@
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
+
+class ToolOutputValidator:
+    async def validate(
+        self, request: types.CallToolRequest, result: types.CallToolResult
+    ) -> bool:
+        raise RuntimeError("Not implemented")
+
+
 DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
 
 
@@ -77,6 +87,25 @@ async def _default_logging_callback(
     pass
 
 
+ToolOutputValidatorProvider: TypeAlias = Callable[
+    ...,
+    Awaitable[ToolOutputValidator],
+]
+
+
+# this bag of spanners is required in order to
+# enable the client session to be parsed to the validator
+async def _python_circularity_hell(arg: Any) -> ToolOutputValidator:
+    # in any sane version of the universe this should never happen
+    # of course in any sane programming language class circularity
+    # dependencies shouldn't be this hard to manage
+    raise RuntimeError(
+        "Help I'm stuck in python circularity hell, please send biscuits"
+    )
+
+
+_default_tool_output_validator: ToolOutputValidatorProvider = _python_circularity_hell
+
 ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
     types.ClientResult | types.ErrorData
 )
@@ -101,6 +130,7 @@ def __init__(
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
+        tool_output_validator_provider: ToolOutputValidatorProvider | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -114,6 +144,7 @@ def __init__(
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
+        self._tool_output_validator_provider = tool_output_validator_provider
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability()
@@ -154,6 +185,11 @@ async def initialize(self) -> types.InitializeResult:
             )
         )
 
+        tool_output_validator_provider = (
+            self._tool_output_validator_provider or _default_tool_output_validator
+        )
+        self._tool_output_validator = await tool_output_validator_provider(self)
+
         return result
 
     async def send_ping(self) -> types.EmptyResult:
@@ -271,24 +307,33 @@ async def call_tool(
         arguments: dict[str, Any] | None = None,
         read_timeout_seconds: timedelta | None = None,
         progress_callback: ProgressFnT | None = None,
+        validate_result: bool = True,
     ) -> types.CallToolResult:
         """Send a tools/call request with optional progress callback support."""
 
-        return await self.send_request(
-            types.ClientRequest(
-                types.CallToolRequest(
-                    method="tools/call",
-                    params=types.CallToolRequestParams(
-                        name=name,
-                        arguments=arguments,
-                    ),
-                )
+        request = types.CallToolRequest(
+            method="tools/call",
+            params=types.CallToolRequestParams(
+                name=name,
+                arguments=arguments,
             ),
+        )
+
+        result = await self.send_request(
+            types.ClientRequest(request),
             types.CallToolResult,
             request_read_timeout_seconds=read_timeout_seconds,
             progress_callback=progress_callback,
         )
 
+        if validate_result:
+            valid = await self._tool_output_validator.validate(request, result)
+
+            if not valid:
+                raise RuntimeError("Server responded with invalid result: " f"{result}")
+        # not validating or is valid
+        return result
+
     async def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult:
         """Send a prompts/list request."""
         return await self.send_request(
@@ -404,3 +449,67 @@ async def _received_notification(
                 await self._logging_callback(params)
             case _:
                 pass
+
+
+class SimpleCachingToolOutputValidator(ToolOutputValidator):
+    _schema_cache: dict[str, dict[str, Any] | bool]
+
+    def __init__(self, session: ClientSession):
+        self._session = session
+        self._schema_cache = {}
+        self._refresh_cache = True
+
+    async def validate(
+        self, request: types.CallToolRequest, result: types.CallToolResult
+    ) -> bool:
+        if result.isError:
+            # allow errors to be propagated
+            return True
+        else:
+            if self._refresh_cache:
+                await self._refresh_schema_cache()
+
+            schema = self._schema_cache.get(request.params.name)
+
+            if schema is None:
+                raise RuntimeError(f"Unknown tool {request.params.name}")
+            elif schema is False:
+                # no schema
+                # TODO add logging
+                return result.structuredContent is None
+            else:
+                try:
+                    # TODO opportunity to build jsonschema.protocol.Validator
+                    # and reuse rather than build every time
+                    validate(result.structuredContent, schema)
+                    return True
+                except ValidationError:
+                    # TODO log this
+                    return False
+
+    async def _refresh_schema_cache(self):
+        cursor = None
+        first = True
+        while first or cursor is not None:
+            first = False
+            tools_result = await self._session.list_tools(cursor)
+            for tool in tools_result.tools:
+                # store a flag to be able to later distinguish between
+                # no schema for tool and unknown tool which can't be verified
+                schema_or_flag = (
+                    False if tool.outputSchema is None else tool.outputSchema
+                )
+                self._schema_cache[tool.name] = schema_or_flag
+            cursor = tools_result.nextCursor
+            continue
+
+        self._refresh_cache = False
+
+
+async def _escape_from_circular_python_hell(
+    session: ClientSession,
+) -> ToolOutputValidator:
+    return SimpleCachingToolOutputValidator(session)
+
+
+_default_tool_output_validator = _escape_from_circular_python_hell

src/mcp/server/lowlevel/server.py

@@ -487,6 +487,7 @@ def handle_error(e: Exception):
                 )
 
             if schema_func is None:
+
                 async def handler(req: types.CallToolRequest):
                     try:
                         result = await func(
@@ -496,6 +497,7 @@ async def handler(req: types.CallToolRequest):
                     except Exception as e:
                         return handle_error(e)
             else:
+
                 async def handler(req: types.CallToolRequest):
                     try:
                         result = await func(

tests/issues/test_88_random_error.py

@@ -15,6 +15,7 @@
     EmbeddedResource,
     ImageContent,
     TextContent,
+    Tool,
 )
 
 
@@ -55,6 +56,10 @@ async def slow_tool(
             return [TextContent(type="text", text=f"fast {request_count}")]
         return [TextContent(type="text", text=f"unknown {request_count}")]
 
+    @server.list_tools()
+    async def list_tools() -> list[Tool]:
+        return [Tool(name="fast", inputSchema={}), Tool(name="slow", inputSchema={})]
+
     async def server_handler(
         read_stream,
         write_stream,