Generate a deferred model's JSON schema under the rebuild lock too

Completing a deferred model under the process-wide rebuild lock and then
letting pydantic generate its JSON schema outside it left one interleaving:
schema generation for a model reads the core schema of every model it
references, and for the mutually recursive wire JSON models (JSONObject /
JSONArray / JSONValue) a sibling could still be mid-build on another
thread, surfacing as a rare KeyError inside pydantic's schema generation
on some pydantic versions when many threads generated their first schema
of that family at once. No SDK-internal path did this (registration is
single-threaded), but user code may.

The generation now runs under the same lock as the build. Schema
generation is a cold path, so it just serializes with the deferred first
builds; steady state is unchanged. The concurrency ratchet also covers a
concurrent first schema over the recursive JSON family and the 2026 wire
package.

Author: maxisbey

src/mcp-types/mcp_types/_deferred.py

@@ -115,8 +115,9 @@ def deferred_model(cls: type[_ModelT]) -> type[_ModelT]:
       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.
+      and generates the schema under that lock, so concurrent first schema
+      calls race neither on the not-yet-built class nor on a referenced
+      (e.g. mutually recursive) model that another thread is completing.
 
     Raises:
         TypeError: `cls` does not set `defer_build=True` in its `model_config`,
@@ -150,12 +151,17 @@ def model_rebuild(sub: type[BaseModel], *args: Any, _parent_namespace_depth: int
             )
 
     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)
+        # Complete the class AND generate under the lock: pydantic re-reads a
+        # class's (mock) core schema around a concurrent first build, and the
+        # generation for one model reads the core schema of every model it
+        # references (a mutually recursive family such as the wire JSON models
+        # references its siblings) while another thread may still be
+        # completing one of those. Schema generation is a cold path, so it
+        # simply serializes with the deferred first builds.
+        with REBUILD_LOCK:
+            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))

tests/types/test_deferred.py

@@ -126,6 +126,7 @@ class Later(MCPModel):  # defined only afterwards: too late for `Pending`
 import pydantic
 import mcp_types._types as _types
 import mcp_types._v2025_11_25 as v2025
+import mcp_types._v2026_07_28 as v2026
 import mcp_types.jsonrpc as jsonrpc
 from mcp_types import methods
 import mcp.client.session as session          # deferred module-level TypeAdapters
@@ -140,7 +141,8 @@ def own_models(mod):
     return [o for o in vars(mod).values() if isinstance(o, type) and issubclass(o, pydantic.BaseModel)
             and o.__module__ == mod.__name__ and not o.__pydantic_complete__]
 
-MODULES = (_types, v2025, jsonrpc, mcpserver, prompts_base, tools_base, resource_types, auth_settings, shared_auth)
+MODULES = (_types, v2025, v2026, jsonrpc, mcpserver, prompts_base, tools_base, resource_types, auth_settings,
+           shared_auth)
 CLASSES = [cls for mod in MODULES for cls in own_models(mod)]
 ADAPTERS = [v for m in (_types, jsonrpc, session, prompts_base) for v in vars(m).values()
             if isinstance(v, pydantic.TypeAdapter)]
@@ -149,6 +151,11 @@ def own_models(mod):
 errors, lock = [], threading.Lock()
 
 def first_use_everything():
+    # The mutually recursive 2026 JSON models first: a FIRST json-schema generation
+    # reads the core schema of each sibling it references while another thread may
+    # still be completing that sibling.
+    for cls in (v2026.JSONObject, v2026.JSONArray, v2026.JSONValue):
+        cls.model_json_schema()
     for cls in CLASSES:
         try:
             cls.model_validate({"__probe__": 1})
@@ -185,9 +192,10 @@ def worker():
 
 
 def test_concurrent_first_use_of_deferred_models_never_raises() -> None:
-    """Eight threads first-using a broad set of never-built models (monolith, wire package,
-    JSON-RPC envelopes, the SDK's own models, the module-level adapters, first schema and
-    signature access) build every class exactly once and raise nothing.
+    """Eight threads first-using a broad set of never-built models (monolith, both wire
+    packages incl. the mutually recursive JSON models' first schema, JSON-RPC envelopes, the
+    SDK's own models, the module-level adapters, first schema and signature access) build
+    every class exactly once and raise nothing.
 
     The models must be unbuilt when the race starts, so the probe runs in fresh interpreters.
     """