Serialize the deferred first build of the SDK's pydantic models

Released pydantic (through 2.13) does not lock model_rebuild(), so the
first use of a defer_build model from several threads at once - sync tools
running on worker threads, or one client session per thread at process
start - could raise (AttributeError: __pydantic_core_schema__ and friends)
or install a stale validator instead of just building twice. A narrow
version of the same hazard already existed for the handful of models
that were incomplete at import; deferring every model made it reachable
from every entry point.

One process-wide RLock now serializes that one-time build. The
@deferred_model decorator becomes the single mechanism for a deferred
model root: it installs the lazy __signature__, the subclass hook, a
locked model_rebuild (public pydantic API; the extra wrapper frame is
accounted for via _parent_namespace_depth + 1) and a model_json_schema
that completes the class under the same lock before pydantic reads it
(pydantic re-reads the mock schema around a concurrent build).
MCPModel, WireModel and WireRootModel are decorated rather than
hand-writing the hook, and the remaining private deferred roots get the
decorator too, so no deferred model in the SDK builds outside the lock.
The decorator refuses a class that lacks defer_build or that defines its
own __pydantic_init_subclass__ (which it would otherwise replace).
Steady state is untouched: an already-built model never takes the lock.

A subprocess ratchet test races 8 threads over the monolith models, a
wire package, the JSON-RPC envelopes, the SDK's own models, the
module-level adapters and first schema/signature access, and asserts
zero exceptions (it fails without the lock).

Author: maxisbey

docs/migration.md

@@ -2724,11 +2724,13 @@ Two incidental things did change:
 
 The SDK's models and `TypeAdapter`s defer their validator build to first construction or
 validation instead of building at import; `inspect.signature(Model)`, `model_fields` and
-`model_json_schema()` report what they always did. A model you subclass from the SDK
-(`Tool`, `Resource`, `AuthSettings`, ...) inherits the deferral, so an annotation that cannot
-resolve in your subclass now surfaces at first use rather than at class creation. Set
-`model_config = ConfigDict(defer_build=False)` in your subclass, or call `Model.model_rebuild()`
-after defining it, if you want the build (and its errors) up front.
+`model_json_schema()` report what they always did, and the one-time build is serialized across
+threads (fixing a rare failure earlier v2 releases could hit when several threads first-used the
+same type at once). A model you subclass from the SDK (`Tool`, `Resource`, `AuthSettings`, ...)
+inherits the deferral, so an annotation that cannot resolve in your subclass now surfaces at
+first use rather than at class creation, and `Model.__pydantic_complete__` reads `False` until
+that first use. Set `model_config = ConfigDict(defer_build=False)` in your subclass, or call
+`Model.model_rebuild()` after defining it, if you want the build (and its errors) up front.
 
 ## Testing utilities
 

docs/whats-new.md

@@ -126,7 +126,7 @@ On those types, every Python attribute is now snake_case: `result.is_error`, `to
 
 ### Imports pay for what you use
 
-`import mcp` is now nearly free, and every entry point loads only its own stack: a stdio server never loads the web framework, client code never loads the server, and the protocol's wire schemas load per negotiated version. The work did not vanish, it moved to first use (the first URL `Client`, the first HTTP app, the first message per protocol version), each a one-time cost. **[Imports & startup time](advanced/import-cost.md)** lists what loads when, and how to prewarm all of it at startup if that is where you want it.
+`import mcp` is now nearly free, and every entry point loads only its own stack: a stdio server never loads the web framework, client code never loads the server, and the protocol's wire schemas load per negotiated version. The work did not vanish, it moved to first use (the first URL `Client`, the first HTTP app, the first message per protocol version), each a one-time cost. That first use is also thread-safe now: the SDK serializes each model's one-time build, which fixes an occasional failure v2 could hit when several threads first-used the same protocol type at once (for example, one client session per thread at process start). **[Imports & startup time](advanced/import-cost.md)** lists what loads when, and how to prewarm all of it at startup if that is where you want it.
 
 ### Behavior that changes without an import error
 

src/mcp-types/mcp_types/_deferred.py

@@ -1,30 +1,61 @@
-"""Keep `inspect.signature()` accurate for `defer_build` models before their first use.
+"""Support for the SDK's `defer_build=True` pydantic models.
 
 The SDK's model bases set `defer_build=True`, so pydantic builds a model's core
 schema, validator and serializer on first use instead of while the class
 statement runs; a module full of models then costs almost nothing to import in
-a process that never validates them.
-
-The one observable pydantic ties to that build is the class's `__signature__`
-(the synthesized `__init__` signature that `inspect.signature(Model)` reports):
-under `defer_build` it exists only once the model has built, and until then a
-class reports the generic `(**data)`. `DeferredSignature` is a class-level
-`__signature__` that completes the (one-time) build via the public
-`model_rebuild()` on first signature access, so runtime introspection matches an
-eagerly-built model while nothing is built at import.
+a process that never validates them. Two things need help until that first
+build happens:
+
+* `inspect.signature(Model)`. Under `defer_build` the synthesized `__init__`
+  signature (`__signature__`) exists only once the model has built; until then
+  a class would report the generic `(**data)`. `DeferredSignature` is a
+  class-level `__signature__` that completes the (one-time) build via the
+  public `model_rebuild()` on first signature access, so runtime introspection
+  matches an eagerly-built model while nothing is built at import.
+* The first build under concurrency. Released pydantic (<= 2.13) does not
+  lock `model_rebuild()`, so several threads first-using one model at once
+  (sync tools on worker threads, one session per thread) can raise
+  `AttributeError: __pydantic_core_schema__` or install a stale validator
+  (pydantic#13419; fixed unreleased in pydantic#13438). One process-wide
+  reentrant lock around the build closes that window: a late thread waits,
+  then finds the model complete and its own rebuild is a no-op.
+
+`deferred_model` installs both on the root class of a deferred model hierarchy
+(the type-layer bases `MCPModel` / `WireModel` / `WireRootModel`, the JSON-RPC
+envelopes, and the `mcp` package's own model roots); subclasses inherit them.
+
+This is a private module of `mcp-types`, but the `mcp` package (which
+exact-pins its `mcp-types` sibling) imports `deferred_model` for its own model
+roots too: keep the signatures stable across both packages.
 """
 
 from __future__ import annotations
 
+import threading
 from inspect import Signature
-from typing import Any, TypeVar, cast
+from typing import Any, Final, TypeVar, cast
 
 from pydantic import BaseModel
 
-__all__ = ["DeferredSignature", "deferred_model", "install_deferred_signature", "new_deferred_signature"]
+__all__ = [
+    "REBUILD_LOCK",
+    "DeferredSignature",
+    "deferred_model",
+    "install_deferred_signature",
+    "new_deferred_signature",
+]
 
 _ModelT = TypeVar("_ModelT", bound=BaseModel)
 
+REBUILD_LOCK: Final = threading.RLock()
+"""Serializes the first (deferred) build of every SDK model across threads.
+
+Reentrant on purpose: completing one model rebuilds the models it references
+on the same thread, re-entering `model_rebuild()`. Held only while a model
+actually builds; a completed model's `model_rebuild()` returns before the
+lock, and no per-message path takes it.
+"""
+
 
 class DeferredSignature:
     """Class-level `__signature__` for a `defer_build=True` pydantic model.
@@ -53,41 +84,80 @@ def __get__(self, obj: object | None, owner: type[BaseModel] | None = None) -> A
 def new_deferred_signature() -> Signature:
     """A `DeferredSignature` typed as the `Signature` it stands in for.
 
-    Assign it in a model class body WITHOUT an annotation (`__signature__ =
-    new_deferred_signature()`), matching how `BaseModel` itself binds
-    `__signature__`, so it never appears in `get_type_hints(Model)`.
+    Bound in a class namespace WITHOUT an annotation (matching how `BaseModel`
+    itself binds `__signature__`), so it never appears in `get_type_hints(Model)`.
     """
     return cast(Signature, DeferredSignature())
 
 
 def install_deferred_signature(cls: type[BaseModel]) -> None:
     """Give a still-deferred model class its own lazily-completing `__signature__`.
 
-    Meant for a base's `__pydantic_init_subclass__` hook: every deferred subclass
-    needs its OWN `__signature__` entry, otherwise it would inherit an
-    already-built parent's signature. A class that pydantic already completed
-    (e.g. one that turned `defer_build` off) keeps the real signature it has.
+    Every deferred subclass needs its OWN `__signature__` entry, otherwise it
+    would inherit an already-built parent's signature. A class that pydantic
+    already completed (e.g. one that turned `defer_build` off) keeps the real
+    signature it has.
     """
     if not cls.__pydantic_complete__:
         setattr(cls, "__signature__", new_deferred_signature())
 
 
 def deferred_model(cls: type[_ModelT]) -> type[_ModelT]:
-    """Class decorator for the root of a `defer_build` model hierarchy under `BaseModel`.
-
-    Installs the lazy `__signature__` on `cls` and, via a `__pydantic_init_subclass__`
-    hook, on every model that later subclasses it (a subclass inherits the deferred
-    build through the config). Models deriving from the SDK bases (`MCPModel`,
-    `WireModel`, `WireRootModel`) get this from the base and do not need it.
+    """Class decorator for the root of a `defer_build=True` model hierarchy.
+
+    Installs, on `cls` and (through inheritance) on every model that later
+    subclasses it - a subclass inherits the deferred build through the config:
+
+    * the lazy `__signature__`, via a `__pydantic_init_subclass__` hook that
+      stamps a fresh one on each still-deferred subclass;
+    * a `model_rebuild` classmethod that takes `REBUILD_LOCK`, so the first
+      build of any model in the hierarchy is serialized across threads (public
+      pydantic API; `_parent_namespace_depth + 1` accounts for this extra
+      frame);
+    * a `model_json_schema` classmethod that completes the (deferred) model
+      under that lock before pydantic generates the schema, so concurrent first
+      schema calls do not race on the not-yet-built class.
+
+    Raises:
+        TypeError: `cls` does not set `defer_build=True` in its `model_config`,
+            or defines its own `__pydantic_init_subclass__` (which the hook
+            installed here would silently replace).
     """
+    if not cls.model_config.get("defer_build"):
+        raise TypeError(f"@deferred_model expects {cls.__name__} to set model_config['defer_build'] = True")
+    if "__pydantic_init_subclass__" in vars(cls):
+        raise TypeError(f"@deferred_model would replace {cls.__name__}'s own __pydantic_init_subclass__")
+
     install_deferred_signature(cls)
 
     def __pydantic_init_subclass__(sub: type[BaseModel], **kwargs: Any) -> None:
         install_deferred_signature(sub)
-        # Run whatever hook `cls`'s own bases define (BaseModel's is a no-op).
-        # (`super()` from a classmethod defined outside the class body needs
-        # both arguments; the two-argument form is opaque to the type checker.)
+        # Chain to whatever hook `cls`'s bases define (BaseModel's is a no-op).
+        # (`super()` from a classmethod defined outside the class body needs the
+        # two-argument form, which is opaque to the type checker.)
         cast(Any, super(cls, sub)).__pydantic_init_subclass__(**kwargs)
 
+    def model_rebuild(sub: type[BaseModel], *args: Any, _parent_namespace_depth: int = 2, **kwargs: Any) -> Any:
+        # An already-built model is the common case: answer it like pydantic
+        # does, without touching the lock.
+        if sub.__pydantic_complete__ and not (args or kwargs.get("force")):
+            return None
+        # One process-wide lock serializes the deferred first build across
+        # threads; a late thread finds the model complete and no-ops.
+        with REBUILD_LOCK:
+            return cast(Any, super(cls, sub)).model_rebuild(
+                *args, _parent_namespace_depth=_parent_namespace_depth + 1, **kwargs
+            )
+
+    def model_json_schema(sub: type[BaseModel], *args: Any, **kwargs: Any) -> Any:
+        # Complete the class before pydantic reads its (mock) core schema:
+        # pydantic re-reads that attribute around the build, which is not safe
+        # against a concurrent first build; a complete class is only read.
+        if not sub.__pydantic_complete__:
+            sub.model_rebuild(raise_errors=False)
+        return cast(Any, super(cls, sub)).model_json_schema(*args, **kwargs)
+
     setattr(cls, "__pydantic_init_subclass__", classmethod(__pydantic_init_subclass__))
+    setattr(cls, "model_rebuild", classmethod(model_rebuild))
+    setattr(cls, "model_json_schema", classmethod(model_json_schema))
     return cls

src/mcp-types/mcp_types/_types.py

@@ -21,7 +21,7 @@
 from pydantic.alias_generators import to_camel
 from typing_extensions import NotRequired, Self, TypedDict
 
-from mcp_types._deferred import install_deferred_signature, new_deferred_signature
+from mcp_types._deferred import deferred_model
 from mcp_types.jsonrpc import RequestId
 
 DEFAULT_NEGOTIATED_VERSION: Final[str] = "2025-03-26"
@@ -43,24 +43,17 @@
 """Theme an icon is designed for. Wire values of `Icon.theme` (2025-11-25+)."""
 
 
+@deferred_model
 class MCPModel(BaseModel):
     """Base class for all MCP protocol types."""
 
     # defer_build: the core schema / validator / serializer for each model is
     # built on first use instead of at class creation, so importing this module
-    # does not pay for ~150 models most processes never touch.
+    # does not pay for ~150 models most processes never touch. `deferred_model`
+    # keeps `inspect.signature()` accurate before that first use and makes the
+    # first build thread-safe (see `mcp_types._deferred`).
     model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, defer_build=True)
 
-    # Complete the deferred build on the first `inspect.signature()` access, so a
-    # never-used model still reports its real field signature. (Bound without an
-    # annotation, like `BaseModel.__signature__`, so it is not a type hint.)
-    __signature__ = new_deferred_signature()
-
-    @classmethod
-    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
-        install_deferred_signature(cls)
-        super().__pydantic_init_subclass__(**kwargs)
-
 
 Meta: TypeAlias = dict[str, Any]
 

src/mcp-types/mcp_types/_wire_base.py

@@ -4,33 +4,29 @@
 and serializer construction at class creation and builds each model once, on
 its first validation. A generated package defines ~175 models but a connection
 touches only the handful its negotiated version and methods use, so importing a
-package is class-body execution only and the schema build is paid lazily,
-per model, exactly once. `inspect.signature()` on a not-yet-built model is kept
-accurate by the lazy `__signature__` from `mcp_types._deferred`.
+package is class-body execution only and the schema build is paid lazily, per
+model, exactly once. `@deferred_model` keeps `inspect.signature()` on a
+not-yet-built model accurate and serializes that first build across threads
+(see `mcp_types._deferred`).
 """
 
-from typing import Any, Generic, TypeVar
+from typing import Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict, RootModel
 
-from mcp_types._deferred import install_deferred_signature, new_deferred_signature
+from mcp_types._deferred import deferred_model
 
 _RootT = TypeVar("_RootT")
 
 
+@deferred_model
 class WireModel(BaseModel):
     """Base for generated wire models: enables `populate_by_name`; subclasses set `extra` themselves."""
 
     model_config = ConfigDict(populate_by_name=True, defer_build=True)
 
-    __signature__ = new_deferred_signature()
-
-    @classmethod
-    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
-        install_deferred_signature(cls)
-        super().__pydantic_init_subclass__(**kwargs)
-
 
+@deferred_model
 class WireRootModel(RootModel[_RootT], Generic[_RootT]):
     """Base for generated named-alias root models (`X(WireRootModel[T])`).
 
@@ -41,10 +37,3 @@ class WireRootModel(RootModel[_RootT], Generic[_RootT]):
     """
 
     model_config = ConfigDict(defer_build=True)
-
-    __signature__ = new_deferred_signature()
-
-    @classmethod
-    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
-        install_deferred_signature(cls)
-        super().__pydantic_init_subclass__(**kwargs)

src/mcp/client/session_group.py

@@ -111,6 +111,7 @@ class ClientSessionGroup:
         ```
     """
 
+    @deferred_model
     class _ComponentNames(BaseModel):
         """Used for reverse index to find components."""
 

src/mcp/server/mcpserver/resolve.py

@@ -66,6 +66,7 @@
     Tool,
     ToolChoice,
 )
+from mcp_types._deferred import deferred_model as _deferred_model
 from mcp_types.version import is_version_at_least
 from pydantic import BaseModel, ConfigDict, ValidationError
 from typing_extensions import TypeVar
@@ -727,6 +728,7 @@ def _result_type(
 # import (import-time cost); the build is transparent at first construction/validation.
 
 
+@_deferred_model
 class _StateEntry(BaseModel):
     """One resolver's recorded outcome inside `request_state`."""
 
@@ -749,6 +751,7 @@ def _request_digest(request: InputRequest) -> str:
     return base64.urlsafe_b64encode(digest).decode().rstrip("=")
 
 
+@_deferred_model
 class _State(BaseModel):
     """The decoded `request_state`: resolver progress from earlier rounds."""