feat: Add structured output support for tool functions (#993)

src/mcp/client/session.py

@@ -1,8 +1,10 @@
+import logging
 from datetime import timedelta
 from typing import Any, Protocol
 
 import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from jsonschema import SchemaError, ValidationError, validate
 from pydantic import AnyUrl, TypeAdapter
 
 import mcp.types as types
@@ -13,6 +15,9 @@ from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("client")
+
 
 class SamplingFnT(Protocol):
     async def __call__(
@@ -128,6 +133,7 @@ class ClientSession(
         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_schemas: dict[str, dict[str, Any] | None] = {}
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -285,7 +291,7 @@ class ClientSession(
     ) -> types.CallToolResult:
         """Send a tools/call request with optional progress callback support."""
 
-        return await self.send_request(
+        result = await self.send_request(
             types.ClientRequest(
                 types.CallToolRequest(
                     method="tools/call",
@@ -300,6 +306,33 @@ class ClientSession(
             progress_callback=progress_callback,
         )
 
+        if not result.isError:
+            await self._validate_tool_result(name, result)
+
+        return result
+
+    async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
+        """Validate the structured content of a tool result against its output schema."""
+        if name not in self._tool_output_schemas:
+            # refresh output schema cache
+            await self.list_tools()
+
+        output_schema = None
+        if name in self._tool_output_schemas:
+            output_schema = self._tool_output_schemas.get(name)
+        else:
+            logger.warning(f"Tool {name} not listed by server, cannot validate any structured content")
+
+        if output_schema is not None:
+            if result.structuredContent is None:
+                raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
+            try:
+                validate(result.structuredContent, output_schema)
+            except ValidationError as e:
+                raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}")
+            except SchemaError as e:
+                raise RuntimeError(f"Invalid schema for tool {name}: {e}")
+
     async def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult:
         """Send a prompts/list request."""
         return await self.send_request(
@@ -351,7 +384,7 @@ class ClientSession(
 
     async def list_tools(self, cursor: str | None = None) -> types.ListToolsResult:
         """Send a tools/list request."""
-        return await self.send_request(
+        result = await self.send_request(
             types.ClientRequest(
                 types.ListToolsRequest(
                     method="tools/list",
@@ -361,6 +394,13 @@ class ClientSession(
             types.ListToolsResult,
         )
 
+        # Cache tool output schemas for future validation
+        # Note: don't clear the cache, as we may be using a cursor
+        for tool in result.tools:
+            self._tool_output_schemas[tool.name] = tool.outputSchema
+
+        return result
+
     async def send_roots_list_changed(self) -> None:
         """Send a roots/list_changed notification."""
         await self.send_notification(