Hand TypedDict tool results to pydantic natively (#3331)

tests/server/mcpserver/test_func_metadata.py

@@ -5,17 +5,21 @@
 # pyright: reportUnknownLambdaType=false
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Annotated, Any, Final, NamedTuple, TypedDict
+from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, TypedDict
 
 import annotated_types
 import pytest
 from dirty_equals import IsPartialDict
 from mcp_types import CallToolResult, ContentBlock, EmbeddedResource, InputRequiredResult, TextContent
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
+from typing_extensions import NotRequired, ReadOnly, Required
 
 from mcp.server.mcpserver import Audio, Image
 from mcp.server.mcpserver.exceptions import InvalidSignature
-from mcp.server.mcpserver.utilities.func_metadata import func_metadata
+from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata, func_metadata
+
+if TYPE_CHECKING:
+    from decimal import Decimal
 
 
 class SomeInputModelA(BaseModel):
@@ -773,22 +777,28 @@ def test_structured_output_dataclass():
 def test_structured_output_typeddict():
     """Test structured output with TypedDict return types"""
 
+    # stdlib TypedDict with a qualifier: exercises the typing_extensions rebuild below Python 3.12
     class PersonTypedDictOptional(TypedDict, total=False):
-        name: str
+        name: Required[str]
         age: int
 
-    def func_returning_typeddict_optional() -> PersonTypedDictOptional:  # pragma: no cover
+    def func_returning_typeddict_optional() -> PersonTypedDictOptional:
         return {"name": "Dave"}  # Only returning one field to test partial dict
 
     meta = func_metadata(func_returning_typeddict_optional)
     assert meta.output_schema == {
         "type": "object",
         "properties": {
-            "name": {"title": "Name", "type": "string", "default": None},
-            "age": {"title": "Age", "type": "integer", "default": None},
+            "name": {"title": "Name", "type": "string"},
+            "age": {"title": "Age", "type": "integer"},
         },
+        "required": ["name"],
         "title": "PersonTypedDictOptional",
     }
+    # An optional key the tool leaves out is absent, not null, so it validates against the schema above
+    result = meta.convert_result(func_returning_typeddict_optional())
+    assert isinstance(result, CallToolResult)
+    assert result.structured_content == {"name": "Dave"}
 
     # Test with total=True (all required)
     class PersonTypedDictRequired(TypedDict):
@@ -812,6 +822,78 @@ def test_structured_output_typeddict():
     }
 
 
+def test_structured_output_typeddict_qualifiers_and_metadata():
+    """PEP 655/705 qualifiers register on every supported Python and decide `required`; the docstring and
+    `Annotated` field metadata reach the schema like they do for a BaseModel, and an alias names the wire key.
+    A stdlib TypedDict, so below 3.12 all of this goes through the typing_extensions rebuild."""
+
+    class Forecast(TypedDict, total=False):
+        """Tomorrow's weather."""
+
+        city: Required[Annotated[str, Field(alias="cityName", description="City name")]]
+        high: ReadOnly[Required[float]]
+        low: float
+        summary: NotRequired[Annotated[str, Field(max_length=80)]]
+
+    def forecast() -> Forecast:
+        return {"city": "Berlin", "high": 21.5}
+
+    with pytest.warns(UserWarning, match="ReadOnly"):  # pydantic notes it won't enforce ReadOnly
+        meta = func_metadata(forecast)
+    result = meta.convert_result(forecast())
+    assert isinstance(result, CallToolResult)
+    assert result.structured_content == {"cityName": "Berlin", "high": 21.5}
+    assert meta.output_schema == {
+        "type": "object",
+        "title": "Forecast",
+        "description": "Tomorrow's weather.",
+        "properties": {
+            "cityName": {"title": "Cityname", "type": "string", "description": "City name"},
+            "high": {"title": "High", "type": "number"},
+            "low": {"title": "Low", "type": "number"},
+            "summary": {"title": "Summary", "type": "string", "maxLength": 80},
+        },
+        "required": ["cityName", "high"],
+    }
+
+
+def test_func_metadata_built_by_hand_keeps_a_given_output_schema():
+    """A hand-built FuncMetadata publishes the schema it was given but still validates results against output_model."""
+    meta = FuncMetadata(arg_model=ArgModelBase, output_model=SomeInputModelB.InnerModel, output_schema={"x": 1})
+    assert meta.output_schema == {"x": 1}
+    with pytest.raises(ValidationError):
+        meta.convert_result({"x": "not an int"})
+
+
+def test_func_metadata_output_fields_are_read_live():
+    """Code in the wild switches structured output off or on by assigning the fields after registration."""
+
+    def total() -> int:
+        return 3
+
+    meta = func_metadata(total)
+    meta.output_schema = None
+    off = meta.convert_result(total())
+    meta.output_schema, meta.output_model, meta.wrap_output = {"type": "object"}, SomeInputModelB.InnerModel, False
+    on = meta.convert_result({"x": 3})
+    assert isinstance(off, CallToolResult) and isinstance(on, CallToolResult)
+    assert (off.structured_content, on.structured_content) == (None, {"x": 3})
+
+
+def test_structured_output_typeddict_with_unresolvable_annotation_is_unstructured():
+    """An annotation only importable under TYPE_CHECKING degrades the same way on every supported Python."""
+
+    class Report(TypedDict):
+        total: "Decimal"
+
+    def report() -> Report:  # pragma: no cover
+        raise NotImplementedError
+
+    assert func_metadata(report).output_schema is None
+    with pytest.raises(InvalidSignature):
+        func_metadata(report, structured_output=True)
+
+
 def test_structured_output_ordinary_class():
     """Test structured output with ordinary annotated classes"""
 

tests/server/mcpserver/test_server.py

@@ -44,6 +44,7 @@ from mcp_types import (
 from pydantic import BaseModel
 from starlette.applications import Starlette
 from starlette.routing import Mount, Route
+from typing_extensions import NotRequired, TypedDict
 
 from mcp.client import Client
 from mcp.server.context import ServerRequestContext
@@ -715,6 +716,32 @@ class TestServerTools:
             assert "Unknown tool" in content.text
 
 
+async def test_typeddict_tool_omitting_optional_keys_passes_client_validation():
+    """The client validates structured content against the tool's output schema, so a `NotRequired`
+    key the tool leaves out must be absent from `structured_content` rather than null."""
+
+    class Person(TypedDict):
+        name: str
+        age: NotRequired[int]
+
+    mcp = MCPServer()
+
+    @mcp.tool()
+    def get_person() -> Person:
+        return {"name": "Dave"}
+
+    async with Client(mcp) as client:
+        (tool,) = (await client.list_tools()).tools
+        assert tool.output_schema == {
+            "type": "object",
+            "title": "Person",
+            "properties": {"name": {"title": "Name", "type": "string"}, "age": {"title": "Age", "type": "integer"}},
+            "required": ["name"],
+        }
+        result = await client.call_tool("get_person", {})
+        assert result.structured_content == {"name": "Dave"}
+
+
 class TestServerResources:
     async def test_init_with_resources(self):
         def get_text() -> str: