11import functools
22import inspect
33import json
4+ import sys
45from collections .abc import Awaitable , Callable , Sequence
56from itertools import chain
67from types import GenericAlias
7- from typing import Annotated , Any , Union , cast , get_args , get_origin , get_type_hints
8+ from typing import Annotated , Any , Union , cast , get_args , get_origin
89
910import anyio
1011import anyio .to_thread
1112import pydantic_core
1213from mcp_types import CallToolResult , ContentBlock , InputRequiredResult , TextContent
13- from pydantic import BaseModel , ConfigDict , Field , PydanticUserError , WithJsonSchema , create_model
14+ from pydantic import BaseModel , ConfigDict , Field , PydanticUserError , TypeAdapter , WithJsonSchema , create_model
1415from pydantic .fields import FieldInfo
1516from pydantic .json_schema import GenerateJsonSchema , JsonSchemaWarningKind
16- from typing_extensions import is_typeddict
17+ from typing_extensions import NotRequired , TypedDict , get_type_hints , is_typeddict
1718from typing_inspection .introspection import (
1819 UNKNOWN ,
1920 AnnotationSource ,
@@ -85,9 +86,15 @@ def model_dump_one_level(self) -> dict[str, Any]:
8586class FuncMetadata (BaseModel ):
8687 arg_model : Annotated [type [ArgModelBase ], WithJsonSchema (None )]
8788 output_schema : dict [str , Any ] | None = None
88- output_model : Annotated [type [BaseModel ], WithJsonSchema (None )] | None = None
89+ output_model : Annotated [type [Any ], WithJsonSchema (None )] | None = None
8990 wrap_output : bool = False
9091
92+ @functools .cached_property
93+ def output_adapter (self ) -> TypeAdapter [Any ]:
94+ """Validates and serializes structured output against `output_model`."""
95+ assert self .output_model is not None , "Output model must be set if output schema is defined"
96+ return TypeAdapter (self .output_model )
97+
9198 def validate_arguments (self , arguments_to_validate : dict [str , Any ]) -> dict [str , Any ]:
9299 """Validate raw arguments into a one-level kwargs dict (no function call).
93100
@@ -144,8 +151,7 @@ def convert_result(self, result: Any) -> CallToolResult | InputRequiredResult:
144151 return result
145152 if isinstance (result , CallToolResult ):
146153 if self .output_schema is not None :
147- assert self .output_model is not None , "Output model must be set if output schema is defined"
148- self .output_model .model_validate (result .structured_content )
154+ self .output_adapter .validate_python (result .structured_content )
149155 return result
150156
151157 unstructured_content = _convert_to_content (result )
@@ -156,9 +162,12 @@ def convert_result(self, result: Any) -> CallToolResult | InputRequiredResult:
156162 if self .wrap_output :
157163 result = {"result" : result }
158164
159- assert self .output_model is not None , "Output model must be set if output schema is defined"
160- validated = self .output_model .model_validate (result )
161- structured_content = validated .model_dump (mode = "json" , by_alias = True )
165+ validated = self .output_adapter .validate_python (result )
166+ if isinstance (validated , BaseModel ):
167+ # Dump via the instance so a returned subclass keeps its own fields.
168+ structured_content = validated .model_dump (mode = "json" , by_alias = True )
169+ else :
170+ structured_content = self .output_adapter .dump_python (validated , mode = "json" , by_alias = True )
162171
163172 return CallToolResult (content = unstructured_content , structured_content = structured_content )
164173
@@ -238,7 +247,7 @@ def func_metadata(
238247 - BaseModel subclasses (used directly)
239248 - Primitive types (str, int, float, bool, bytes, None) - wrapped in a
240249 model with a 'result' field
241- - TypedDict - converted to a Pydantic model with same fields
250+ - TypedDict - used directly
242251 - Dataclasses and other annotated classes - converted to Pydantic models
243252 - Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field
244253 - Content blocks (TextContent, EmbeddedResource, ...), Image and Audio, bare or inside a
@@ -374,30 +383,41 @@ def func_metadata(
374383 # structured_output=True still forces one.
375384 return FuncMetadata (arg_model = arguments_model )
376385
377- output_model , output_schema , wrap_output = _try_create_model_and_schema (
378- original_annotation , return_type_expr , func .__name__
379- )
386+ output_model , wrap_output = _create_output_model (original_annotation , return_type_expr , func .__name__ )
387+
388+ if output_model is not None :
389+ meta = FuncMetadata (arg_model = arguments_model , output_model = output_model , wrap_output = wrap_output )
390+ try :
391+ # Building the validator here surfaces unsupported types at registration rather than on the
392+ # first call. StrictJsonSchema raises instead of emitting warnings.
393+ meta .output_schema = meta .output_adapter .json_schema (schema_generator = StrictJsonSchema )
394+ return meta
395+ except (
396+ PydanticUserError ,
397+ TypeError ,
398+ ValueError ,
399+ pydantic_core .SchemaError ,
400+ pydantic_core .ValidationError ,
401+ ) as e :
402+ # These are expected errors when a type can't be converted to a Pydantic schema
403+ # PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
404+ # subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
405+ # ValueError: When there are issues with the type definition (including our custom warnings)
406+ # SchemaError: When Pydantic can't build a schema
407+ # ValidationError: When validation fails
408+ logger .info (f"Cannot create schema for type { return_type_expr } in { func .__name__ } : { type (e ).__name__ } : { e } " )
380409
381- if output_model is None and structured_output is True :
410+ if structured_output is True :
382411 # Model creation failed or produced warnings - no structured output
383412 raise InvalidSignature (
384413 f"Function { func .__name__ } : return type { return_type_expr } is not serializable for structured output"
385414 )
386415
387- return FuncMetadata (
388- arg_model = arguments_model ,
389- output_schema = output_schema ,
390- output_model = output_model ,
391- wrap_output = wrap_output ,
392- )
416+ return FuncMetadata (arg_model = arguments_model )
393417
394418
395- def _try_create_model_and_schema (
396- original_annotation : Any ,
397- type_expr : Any ,
398- func_name : str ,
399- ) -> tuple [type [BaseModel ] | None , dict [str , Any ] | None , bool ]:
400- """Try to create a model and schema for the given annotation without warnings.
419+ def _create_output_model (original_annotation : Any , type_expr : Any , func_name : str ) -> tuple [type [Any ] | None , bool ]:
420+ """Pick the type structured output is validated against for the given return annotation.
401421
402422 Args:
403423 original_annotation: The original return annotation (may be wrapped in `Annotated`).
@@ -406,11 +426,11 @@ def _try_create_model_and_schema(
406426 func_name: The name of the function.
407427
408428 Returns:
409- tuple of (model or None, schema or None, wrap_output)
410- Model and schema are None if warnings occur or creation fails .
429+ tuple of (model or None, wrap_output)
430+ Model is None if the type cannot carry structured output .
411431 wrap_output is True if the result needs to be wrapped in {"result": ...}
412432 """
413- model = None
433+ model : type [ Any ] | None = None
414434 wrap_output = False
415435
416436 # First handle special case: None
@@ -446,9 +466,9 @@ def _try_create_model_and_schema(
446466 if issubclass (type_annotation , BaseModel ):
447467 model = type_annotation
448468
449- # Case 2: TypedDicts:
469+ # Case 2: TypedDicts (pydantic reads qualifiers, totality, docstring and `Annotated` metadata natively)
450470 elif is_typeddict (type_annotation ):
451- model = _create_model_from_typeddict (type_annotation )
471+ model = _pydantic_readable_typeddict (type_annotation )
452472
453473 # Case 3: Primitive types that need wrapping
454474 elif type_annotation in (str , int , float , bool , bytes , type (None )):
@@ -470,30 +490,7 @@ def _try_create_model_and_schema(
470490 model = _create_wrapped_model (func_name , original_annotation )
471491 wrap_output = True
472492
473- if model :
474- # If we successfully created a model, try to get its schema
475- # Use StrictJsonSchema to raise exceptions instead of warnings
476- try :
477- schema = model .model_json_schema (schema_generator = StrictJsonSchema )
478- except (
479- PydanticUserError ,
480- TypeError ,
481- ValueError ,
482- pydantic_core .SchemaError ,
483- pydantic_core .ValidationError ,
484- ) as e :
485- # These are expected errors when a type can't be converted to a Pydantic schema
486- # PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
487- # subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
488- # ValueError: When there are issues with the type definition (including our custom warnings)
489- # SchemaError: When Pydantic can't build a schema
490- # ValidationError: When validation fails
491- logger .info (f"Cannot create schema for type { type_expr } in { func_name } : { type (e ).__name__ } : { e } " )
492- return None , None , False
493-
494- return model , schema , wrap_output
495-
496- return None , None , False
493+ return model , wrap_output
497494
498495
499496_no_default = object ()
@@ -523,25 +520,29 @@ def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type
523520 return create_model (cls .__name__ , __config__ = ConfigDict (from_attributes = True ), ** model_fields )
524521
525522
526- def _create_model_from_typeddict (td_type : type [Any ]) -> type [BaseModel ]:
527- """Create a Pydantic model from a TypedDict.
523+ def _pydantic_readable_typeddict (td_type : type [Any ]) -> type [Any ]:
524+ """pydantic refuses `typing.TypedDict` below Python 3.12 (it needs `__orig_bases__`); rebuild those as an
525+ equivalent `typing_extensions.TypedDict` so tool authors don't have to know. Delete once 3.11 support goes."""
526+ if sys .version_info >= (3 , 12 ) or type (td_type ).__module__ != "typing" :
527+ return td_type
528+ return _as_typing_extensions_typeddict (td_type ) # pragma: lax no cover
528529
529- The created model will have the same name and fields as the TypedDict.
530- """
531- type_hints = get_type_hints (td_type )
532- required_keys = getattr (td_type , "__required_keys__" , set (type_hints .keys ()))
533530
534- model_fields : dict [str , Any ] = {}
535- for field_name , field_type in type_hints .items ():
536- if field_name not in required_keys :
537- # For optional TypedDict fields, set default=None
538- # This makes them not required in the Pydantic model
539- # The model should use exclude_unset=True when dumping to get TypedDict semantics
540- model_fields [field_name ] = (field_type , None )
541- else :
542- model_fields [field_name ] = field_type
543-
544- return create_model (td_type .__name__ , ** model_fields )
531+ def _as_typing_extensions_typeddict (td_type : type [Any ]) -> type [Any ]: # pragma: lax no cover
532+ items : dict [str , Any ] = {}
533+ for name , hint in get_type_hints (td_type , include_extras = True ).items ():
534+ key = inspect_annotation (hint , annotation_source = AnnotationSource .TYPED_DICT )
535+ item : Any = Annotated [(key .type , * key .metadata )] if key .metadata else key .type
536+ # pydantic's rule: an explicit qualifier wins over class totality. Needed because a stdlib TypedDict
537+ # this old computes `__required_keys__` without seeing `typing_extensions` qualifiers.
538+ required = (name in td_type .__required_keys__ or "required" in key .qualifiers ) and (
539+ "not_required" not in key .qualifiers
540+ )
541+ items [name ] = item if required else NotRequired [item ]
542+ # The functional form, spelled so type checkers don't try to evaluate it statically.
543+ rebuilt = cast ("Callable[[str, dict[str, Any]], type[Any]]" , TypedDict )(td_type .__name__ , items )
544+ rebuilt .__doc__ = td_type .__doc__
545+ return rebuilt
545546
546547
547548def _create_wrapped_model (func_name : str , annotation : Any ) -> type [BaseModel ]:
0 commit comments