Defer building the generated wire models and emit them in dependency order

The generated _v2025_11_25 and _v2026_07_28 packages spent ~100-115 ms
each at import building pydantic core schemas, validators and serializers
for ~175 models per version, of which a connection touches only a handful.

Both generated bases now carry defer_build=True: WireModel, and a new
generic WireRootModel[T] the RootModel aliases derive from (deferring only
the object models made things slower, as the eager union roll-ups then
regenerated every deferred model's schema inline). The generator drops
codegen's trailing model_rebuild() epilogue and emits classes in dependency
order - failing if any class names a class defined later - so every
non-recursive generated class still exposes complete field metadata at
import. Building each model's validator now happens once, on first use.

Author: maxisbey

pyproject.toml

@@ -221,6 +221,8 @@ max-complexity = 24 # Default is 10
 "src/mcp/types/*.py" = ["F403"]
 # Generated by scripts/gen_surface_types.py: raw datamodel-codegen output (TID251 lifts the repo-wide RootModel ban for these generated validators).
 "src/mcp-types/mcp_types/_v*/__init__.py" = ["D212", "E501", "I001", "TID251", "UP007", "UP037"]
+# The deferred RootModel base the generated validators derive from.
+"src/mcp-types/mcp_types/_wire_base.py" = ["TID251"]
 "tests/server/mcpserver/test_func_metadata.py" = ["E501"]
 "tests/shared/test_progress_notifications.py" = ["PLW0603"]
 

scripts/gen_surface_types.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import argparse
+import ast
 import difflib
 import hashlib
 import json
@@ -196,6 +197,139 @@ def run_codegen(schema_path: Path, output_path: Path) -> None:
         raise SystemExit(f"datamodel-codegen failed:\n{result.stderr}")
 
 
+def use_deferred_bases(source: str) -> str:
+    """Route generated `RootModel[T]` classes through the deferred `WireRootModel[T]` base.
+
+    Object models already derive from `WireModel` via codegen's `--base-class`;
+    RootModel classes get no such hook, so rewrite them here. Both bases carry
+    `defer_build=True` (see `mcp_types/_wire_base.py`), which is only a win if
+    EVERY class in the package is deferred: an eagerly-built union rollup would
+    regenerate the schemas of every deferred model it references.
+    """
+    source = source.replace("RootModel[", "WireRootModel[")
+    source = source.replace(
+        "from mcp_types._wire_base import WireModel", "from mcp_types._wire_base import WireModel, WireRootModel"
+    )
+    source = re.sub(r"^from pydantic import (.*), RootModel$", r"from pydantic import \1", source, flags=re.MULTILINE)
+    assert "RootModel[" not in source.replace("WireRootModel[", ""), "unrewritten RootModel base"
+    return source
+
+
+def strip_rebuild_epilogue(source: str) -> str:
+    """Drop codegen's trailing `X.model_rebuild()` calls.
+
+    codegen emits them to force-resolve forward references eagerly; with the
+    deferred bases each model is built once on first use, when its whole module
+    namespace already exists, so the calls would only re-introduce the eager
+    build cost the deferral removes.
+    """
+    return re.sub(r"^\w+\.model_rebuild\(\)\n", "", source, flags=re.MULTILINE)
+
+
+def define_before_use(source: str) -> str:
+    """Reorder the generated classes so every referenced class is defined first.
+
+    codegen's ordering leaves whole subtrees referring to not-yet-defined
+    classes (it relied on the eager `model_rebuild()` epilogue to patch them
+    up). A deferred model whose annotations name an undefined class collects
+    incomplete `model_fields` (an `Annotated[...]` alias trapped in an
+    unevaluated string) until its first-use build; defining dependencies first
+    keeps field metadata complete at import for every class. Only genuine
+    recursion cannot be linearized (the mutually-recursive JSONValue trio in
+    2026-07-28); such a cycle is kept together in codegen's own order, its
+    string self-references resolving on first use.
+    """
+    tree = ast.parse(source)
+    lines = source.splitlines(keepends=True)
+    classes = [node for node in tree.body if isinstance(node, ast.ClassDef)]
+    if not classes:
+        return source
+    names = {node.name for node in classes}
+    index = {node.name: i for i, node in enumerate(classes)}
+
+    def referenced(node: ast.ClassDef, *, quoted: bool) -> set[str]:
+        refs: set[str] = set()
+        for sub in ast.walk(node):
+            if isinstance(sub, ast.Name) and sub.id in names:
+                refs.add(sub.id)
+            elif quoted and isinstance(sub, ast.Constant) and isinstance(sub.value, str) and sub.value in names:
+                refs.add(sub.value)  # quoted forward reference
+        refs.discard(node.name)
+        return refs
+
+    deps = {node.name: referenced(node, quoted=True) for node in classes}
+
+    # Tarjan's SCC: a cycle (real recursion) is placed as one unit.
+    counter: list[int] = [0]
+    stack: list[str] = []
+    on_stack: set[str] = set()
+    ids: dict[str, int] = {}
+    low: dict[str, int] = {}
+    comps: list[list[str]] = []
+
+    def strongconnect(v: str) -> None:
+        ids[v] = low[v] = counter[0]
+        counter[0] += 1
+        stack.append(v)
+        on_stack.add(v)
+        for w in sorted(deps[v], key=index.__getitem__):
+            if w not in ids:
+                strongconnect(w)
+                low[v] = min(low[v], low[w])
+            elif w in on_stack:
+                low[v] = min(low[v], ids[w])
+        if low[v] == ids[v]:
+            comp: list[str] = []
+            while True:
+                w = stack.pop()
+                on_stack.discard(w)
+                comp.append(w)
+                if w == v:
+                    break
+            comps.append(sorted(comp, key=index.__getitem__))
+
+    for name in sorted(names, key=index.__getitem__):
+        if name not in ids:
+            strongconnect(name)
+
+    # Emit each SCC once, dependencies first, stable on codegen's original order.
+    comp_of = {name: i for i, comp in enumerate(comps) for name in comp}
+    emitted: list[str] = []
+    done: set[int] = set()
+
+    def emit(ci: int) -> None:
+        if ci in done:
+            return
+        done.add(ci)
+        for name in comps[ci]:
+            for dep in sorted(deps[name], key=index.__getitem__):
+                if comp_of[dep] != ci:
+                    emit(comp_of[dep])
+        emitted.extend(comps[ci])
+
+    for name in sorted(names, key=index.__getitem__):
+        emit(comp_of[name])
+    assert set(emitted) == names
+    # Invariant: every class named by a bare identifier is defined earlier;
+    # only quoted references inside a genuine recursive cycle may still point
+    # forward. A regression here silently reintroduces incomplete
+    # `model_fields` at import, so fail generation instead.
+    position = {name: i for i, name in enumerate(emitted)}
+    unresolved = {
+        (node.name, ref)
+        for node in classes
+        for ref in referenced(node, quoted=False)
+        if position[ref] > position[node.name]
+    }
+    assert not unresolved, f"classes still name a later-defined class: {sorted(unresolved)}"
+
+    node_of = {node.name: node for node in classes}
+    blocks = ["".join(lines[node_of[name].lineno - 1 : node_of[name].end_lineno]) for name in emitted]
+    head = "".join(lines[: classes[0].lineno - 1])
+    tail = "".join(lines[max(node.end_lineno or 0 for node in classes) :])
+    return head + "\n\n\n".join(block.strip("\n") + "\n" for block in blocks) + tail
+
+
 def allow_open_class_extras(source: str, open_classes: frozenset[str]) -> str:
     """Restore `extra="allow"` on `open_classes` only.
 
@@ -243,12 +377,11 @@ def build(entry: dict[str, str]) -> str:
     # strict mkdocs link validation.
     source = source.replace("](/", "](https://modelcontextprotocol.io/")
     source = allow_open_class_extras(source, OPEN_CLASSES[version])
+    source = use_deferred_bases(source)
+    source = strip_rebuild_epilogue(source)
+    source = define_before_use(source)
     if epilogue := EPILOGUES.get(version, ""):
-        # Insert before the trailing model_rebuild() block: pyright's evaluation
-        # order for the recursive RootModel block is sensitive to placement.
-        match = re.search(r"^\w+\.model_rebuild\(\)$", source, flags=re.MULTILINE)
-        cut = match.start() if match else len(source)
-        source = f"{source[:cut]}{epilogue}\n\n{source[cut:]}"
+        source = f"{source.rstrip()}\n\n\n{epilogue}"
     source = HEADER.format(version=version, sha=entry["sha256"]) + source
 
     staging = TYPES_DIR / f"_staging_{version}.py"

src/mcp-types/mcp_types/_v2025_11_25/__init__.py

@@ -8,8 +8,8 @@
 
 from typing import Annotated, Any, Literal
 
-from mcp_types._wire_base import WireModel
-from pydantic import ConfigDict, Field, RootModel
+from mcp_types._wire_base import WireModel, WireRootModel
+from pydantic import ConfigDict, Field
 
 
 class BaseMetadata(WireModel):
@@ -285,7 +285,7 @@ class CompleteResult(WireModel):
     completion: Completion
 
 
-class Cursor(RootModel[str]):
+class Cursor(WireRootModel[str]):
     root: str
     """
     An opaque token used to represent a cursor for pagination.
@@ -557,7 +557,7 @@ class LegacyTitledEnumSchema(WireModel):
 
 
 class LoggingLevel(
-    RootModel[
+    WireRootModel[
         Literal[
             "alert",
             "critical",
@@ -723,7 +723,7 @@ class PaginatedResult(WireModel):
     """
 
 
-class ProgressToken(RootModel[str | int]):
+class ProgressToken(WireRootModel[str | int]):
     root: str | int
     """
     A progress token, used to associate progress notifications with the original request.
@@ -853,7 +853,7 @@ class Request(WireModel):
     params: dict[str, Any] | None = None
 
 
-class RequestId(RootModel[str | int]):
+class RequestId(WireRootModel[str | int]):
     root: str | int
     """
     A uniquely identifying ID for a request in JSON-RPC.
@@ -970,7 +970,7 @@ class Result(WireModel):
     """
 
 
-class Role(RootModel[Literal["assistant", "user"]]):
+class Role(WireRootModel[Literal["assistant", "user"]]):
     root: Literal["assistant", "user"]
     """
     The sender or recipient of messages and data in a conversation.
@@ -1216,7 +1216,7 @@ class TaskMetadata(WireModel):
     """
 
 
-class TaskStatus(RootModel[Literal["cancelled", "completed", "failed", "input_required", "working"]]):
+class TaskStatus(WireRootModel[Literal["cancelled", "completed", "failed", "input_required", "working"]]):
     root: Literal["cancelled", "completed", "failed", "input_required", "working"]
     """
     The status of a task.
@@ -1867,12 +1867,12 @@ class EmbeddedResource(WireModel):
     type: Literal["resource"]
 
 
-class EmptyResult(RootModel[Result]):
+class EmptyResult(WireRootModel[Result]):
     root: Result
 
 
 class EnumSchema(
-    RootModel[
+    WireRootModel[
         UntitledSingleSelectEnumSchema
         | TitledSingleSelectEnumSchema
         | UntitledMultiSelectEnumSchema
@@ -2115,7 +2115,7 @@ class LoggingMessageNotification(WireModel):
     params: LoggingMessageNotificationParams
 
 
-class MultiSelectEnumSchema(RootModel[UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema]):
+class MultiSelectEnumSchema(WireRootModel[UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema]):
     root: UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
 
 
@@ -2153,7 +2153,7 @@ class PingRequest(WireModel):
 
 
 class PrimitiveSchemaDefinition(
-    RootModel[
+    WireRootModel[
         StringSchema
         | NumberSchema
         | BooleanSchema
@@ -2499,7 +2499,7 @@ class SetLevelRequest(WireModel):
     params: SetLevelRequestParams
 
 
-class SingleSelectEnumSchema(RootModel[UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema]):
+class SingleSelectEnumSchema(WireRootModel[UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema]):
     root: UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
 
 
@@ -2793,7 +2793,7 @@ class CompleteRequest(WireModel):
     params: CompleteRequestParams
 
 
-class ContentBlock(RootModel[TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource]):
+class ContentBlock(WireRootModel[TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource]):
     root: TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
 
 
@@ -2812,7 +2812,7 @@ class CreateTaskResult(WireModel):
     task: Task
 
 
-class ElicitRequestParams(RootModel[ElicitRequestURLParams | ElicitRequestFormParams]):
+class ElicitRequestParams(WireRootModel[ElicitRequestURLParams | ElicitRequestFormParams]):
     root: ElicitRequestURLParams | ElicitRequestFormParams
     """
     The parameters for a request to elicit additional information from the user via the client.
@@ -2857,14 +2857,16 @@ class InitializeRequest(WireModel):
     params: InitializeRequestParams
 
 
-class JSONRPCMessage(RootModel[JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse]):
+class JSONRPCMessage(
+    WireRootModel[JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse]
+):
     root: JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse
     """
     Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent.
     """
 
 
-class JSONRPCResponse(RootModel[JSONRPCResultResponse | JSONRPCErrorResponse]):
+class JSONRPCResponse(WireRootModel[JSONRPCResultResponse | JSONRPCErrorResponse]):
     root: JSONRPCResultResponse | JSONRPCErrorResponse
     """
     A response to a request, containing either the result or error.
@@ -3174,7 +3176,7 @@ class CallToolResult(WireModel):
 
 
 class ClientNotification(
-    RootModel[
+    WireRootModel[
         CancelledNotification
         | InitializedNotification
         | ProgressNotification
@@ -3192,7 +3194,7 @@ class ClientNotification(
 
 
 class ClientRequest(
-    RootModel[
+    WireRootModel[
         InitializeRequest
         | PingRequest
         | ListResourcesRequest
@@ -3267,13 +3269,13 @@ class GetPromptResult(WireModel):
 
 
 class SamplingMessageContentBlock(
-    RootModel[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
+    WireRootModel[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
 ):
     root: TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
 
 
 class ServerNotification(
-    RootModel[
+    WireRootModel[
         CancelledNotification
         | ProgressNotification
         | ResourceListChangedNotification
@@ -3299,7 +3301,7 @@ class ServerNotification(
 
 
 class ServerResult(
-    RootModel[
+    WireRootModel[
         Result
         | InitializeResult
         | ListResourcesResult
@@ -3399,7 +3401,7 @@ class SamplingMessage(WireModel):
 
 
 class ClientResult(
-    RootModel[
+    WireRootModel[
         Result
         | GetTaskResult
         | GetTaskPayloadResult
@@ -3503,7 +3505,7 @@ class CreateMessageRequest(WireModel):
 
 
 class ServerRequest(
-    RootModel[
+    WireRootModel[
         PingRequest
         | GetTaskRequest
         | GetTaskPayloadRequest