Add mcp.warm(): opt-in prewarming for the deferred validators

The wire schemas load on the first message per protocol version and the
pydantic validators build on first use, which keeps imports fast but puts a
one-time bill on a process's first messages. A long-running host that
would rather pay that at startup than on its first request had only a DIY
recipe of model_rebuild() loops.

mcp_types.methods.warm(version=None, *, everything=False) is the supported
version: with no argument it builds the version-independent set (the
exported mcp_types models, the JSON-RPC envelopes and the routing union
adapters, ~50 ms); with a version it also imports that version's wire
package and builds the routing surface a connection at that version uses
(so its first messages then build nothing); everything=True covers every
known version. Model classes are always completed before the adapters
that reference them, so no schema is generated twice, and repeat calls
are no-ops. It returns a small WarmReport for logging and is re-exported
as mcp.warm. Nothing in the SDK calls it. The import-cost docs recipe now
uses it.

Author: maxisbey

docs_src/import_cost/tutorial001.py

@@ -1,18 +1,14 @@
-"""Prewarm the SDK's deferred startup work.
+"""Prewarm the SDK's deferred startup work in a host's startup hook.
 
 Imports stay fast because the heavy pieces load on first use: a protocol
 version's wire schemas load on the first message parsed for that version, and
 each pydantic model builds its validator the first time it is used. A host that
-would rather pay all of that once, before serving traffic, can trigger it here.
+would rather pay that once, before serving traffic, calls `mcp.warm()`.
 """
 
-import mcp_types
-from mcp_types.methods import CLIENT_REQUESTS
+import mcp
 
-# Load the wire-schema package for each protocol version you will speak.
-# Reading any row of a surface map imports that version's package.
-_ = CLIENT_REQUESTS[("initialize", "2025-11-25")]
-
-# Build the validators for the models on your hot path up front.
-for model in (mcp_types.CallToolRequest, mcp_types.CallToolResult, mcp_types.Tool):
-    model.model_rebuild()
+# Build the deferred validators now: the version-independent set, plus the
+# routing surface of each protocol version this host will serve.
+report = mcp.warm("2025-11-25")
+print(f"prewarmed {report.models} models and {report.adapters} adapters in {report.elapsed_ms} ms")

src/mcp-types/mcp_types/methods.py

@@ -15,15 +15,17 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterator, Mapping
+import time
+from collections.abc import Iterable, Iterator, Mapping
 from functools import cache
 from importlib import import_module
 from types import MappingProxyType, ModuleType, UnionType
-from typing import TYPE_CHECKING, Any, Final, Literal, TypeGuard, TypeVar, cast, get_args
+from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeGuard, TypeVar, cast, get_args
 
 from pydantic import BaseModel, TypeAdapter
 
 import mcp_types as types
+from mcp_types import jsonrpc
 from mcp_types.version import KNOWN_PROTOCOL_VERSIONS
 
 if TYPE_CHECKING:
@@ -52,6 +54,7 @@
     "SERVER_RESULTS",
     "SPEC_CLIENT_METHODS",
     "SPEC_CLIENT_NOTIFICATION_METHODS",
+    "WarmReport",
     "is_input_required",
     "parse_client_notification",
     "parse_client_request",
@@ -64,6 +67,7 @@
     "validate_client_request",
     "validate_client_result",
     "validate_server_result",
+    "warm",
 ]
 
 
@@ -849,3 +853,130 @@ def parse_client_result(
     validate_client_result(method, version, data, surface=surface)
     result: types.Result = _adapter(_monolith_row(monolith, method)).validate_python(data, by_name=False)
     return result
+
+
+# --- Opt-in prewarm ---
+
+_ALL_SURFACES: Final = (
+    CLIENT_REQUESTS,
+    CLIENT_NOTIFICATIONS,
+    CLIENT_RESULTS,
+    SERVER_REQUESTS,
+    SERVER_NOTIFICATIONS,
+    SERVER_RESULTS,
+)
+"""Every surface map, for `warm(version)` (rows of one version import one wire package)."""
+
+_ALL_MONOLITHS: Final = (MONOLITH_REQUESTS, MONOLITH_NOTIFICATIONS, MONOLITH_RESULTS)
+"""Every monolith map, for `warm()`."""
+
+
+class WarmReport(NamedTuple):
+    """What one `warm()` call built (already-built work is skipped and not counted)."""
+
+    models: int
+    """Deferred model classes whose validators this call built."""
+
+    adapters: int
+    """Deferred TypeAdapters (union / routing adapters) this call built."""
+
+    elapsed_ms: float
+    """Wall-clock milliseconds the call took."""
+
+
+def _model_classes(target: object) -> Iterator[type[BaseModel]]:
+    """The model classes in a surface/monolith row: the class itself, or a union's arms."""
+    if isinstance(target, type):
+        if issubclass(target, BaseModel):
+            yield target
+        return
+    for arg in get_args(target):
+        yield from _model_classes(arg)
+
+
+def _warm_models(rows: Iterable[object]) -> int:
+    """Complete every not-yet-built model class among `rows`; return how many were built."""
+    built = 0
+    for row in rows:
+        for cls in _model_classes(row):
+            if not cls.__pydantic_complete__:
+                cls.model_rebuild()
+                built += 1
+    return built
+
+
+def _warm_adapter(adapter: TypeAdapter[Any]) -> int:
+    """Build a deferred TypeAdapter's validator if it is not built yet; return 0 or 1."""
+    return 0 if adapter.rebuild() is None else 1
+
+
+def warm(version: str | None = None, *, everything: bool = False) -> WarmReport:
+    """Build the SDK's deferred validators up front, before the first message needs them.
+
+    Importing the SDK stays fast because pydantic validators build on first use
+    and each protocol version's wire schemas load with the first message parsed
+    at that version. A long-running host that would rather pay that once, at
+    startup, calls this from a startup hook; nothing in the SDK calls it, and
+    steady-state behaviour is identical either way.
+
+    * `warm()` builds the version-independent set: the monolith `mcp_types`
+      models a session routes through, the JSON-RPC envelopes and the
+      module-level union adapters.
+    * `warm(version)` additionally imports that protocol version's wire
+      package and builds its routing surface (the wire models and the cached
+      adapters a connection at that version routes through).
+    * `warm(everything=True)` does that for every known protocol version.
+
+    Model classes are completed before the adapters that reference them are
+    built, so no schema is generated twice; anything already built is skipped,
+    so the call is idempotent and repeat calls are cheap.
+
+    Args:
+        version: The protocol version whose routing surface to build too.
+        everything: Warm every known protocol version (benchmarks, tests).
+
+    Returns:
+        Counts of what this call built and the milliseconds it took.
+
+    Raises:
+        ValueError: `version` is not a known protocol version.
+    """
+    started = time.perf_counter()
+    versions: list[str] = []
+    if everything:
+        versions = sorted(KNOWN_PROTOCOL_VERSIONS)
+    elif version is not None:
+        _check_known_version(version)
+        versions = [version]
+
+    # Model classes first (the exported monolith models, the JSON-RPC
+    # envelopes, then the wire packages and their routing rows), so the
+    # adapters below never generate a referenced model's schema inline.
+    models = _warm_models(getattr(types, name) for name in types.__all__)
+    models += _warm_models(row for monolith in _ALL_MONOLITHS for row in monolith.values())
+    models += _warm_models(get_args(jsonrpc.JSONRPCMessage))
+    rows: list[Any] = []
+    for warmed_version in versions:
+        package = _wire_package(warmed_version)  # imports the package once
+        models += _warm_models(vars(package).values())
+        for surface in _ALL_SURFACES:
+            rows.extend(surface[key] for key in surface if key[1] == warmed_version)
+    models += _warm_models(rows)
+
+    adapters = _warm_adapter(jsonrpc.jsonrpc_message_adapter)
+    for adapter in (
+        types.client_request_adapter,
+        types.client_notification_adapter,
+        types.client_result_adapter,
+        types.server_request_adapter,
+        types.server_notification_adapter,
+        types.server_result_adapter,
+    ):
+        adapters += _warm_adapter(adapter)
+    for row in rows:
+        # `_adapter` is the cached routing-adapter builder the parse/serialize
+        # functions use; warming it here (and completing the adapter's own
+        # core schema from the now-complete row model) leaves nothing for the
+        # first message at that version to build.
+        adapters += _warm_adapter(_adapter(row))
+    return WarmReport(models=models, adapters=adapters, elapsed_ms=round((time.perf_counter() - started) * 1000, 1))

src/mcp/__init__.py

@@ -67,6 +67,7 @@
         UnsubscribeRequest,
     )
     from mcp_types import Role as SamplingRole
+    from mcp_types.methods import warm as warm
 
     # Bind the `mcp.types` submodule on the package, as v1's `from .types import
     # ...` did, so `import mcp` followed by `mcp.types.Tool` keeps working. The
@@ -159,6 +160,7 @@
     "InvalidUriTemplate",
     "stdio_client",
     "stdio_server",
+    "warm",
 ]
 
 # ---------------------------------------------------------------------------
@@ -253,6 +255,7 @@
     "UrlElicitationRequiredError": ("mcp.shared.exceptions", "UrlElicitationRequiredError"),
     "InvalidUriTemplate": ("mcp.shared.uri_template", "InvalidUriTemplate"),
     "UriTemplate": ("mcp.shared.uri_template", "UriTemplate"),
+    "warm": ("mcp_types.methods", "warm"),
 }
 
 # Submodules a bare `import mcp` used to bind on the package as a side effect

tests/docs_src/test_import_cost.py

@@ -2,9 +2,12 @@
 
 import json
 import os
+import pathlib
 import subprocess
 import sys
 
+REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
+
 # The recipe's effect is only observable in a fresh interpreter: this process has long
 # since loaded the wire packages and built the models.
 PROBE = """
@@ -25,18 +28,19 @@
 
 
 def test_prewarm_recipe_loads_the_wire_package_and_builds_the_models() -> None:
-    """tutorial001: before the recipe nothing is loaded/built; after it, the negotiated
-    version's wire package is imported and the named models are built."""
+    """tutorial001: before the recipe nothing is loaded/built; after `mcp.warm(version)` the
+    named version's wire package is imported and the models are built."""
     result = subprocess.run(
         [sys.executable, "-c", PROBE],
         capture_output=True,
         text=True,
         check=False,
         timeout=60,
+        cwd=REPO_ROOT,  # the recipe module is imported by its docs_src path
         # pydantic plugins (e.g. logfire) building models is the environment's cost, not ours
         env={**os.environ, "PYDANTIC_DISABLE_PLUGINS": "__all__"},
     )
     assert result.returncode == 0, result.stderr
-    observed = json.loads(result.stdout)
+    observed = json.loads(result.stdout.strip().splitlines()[-1])
     assert observed["before"] == {"wire_2025_loaded": False, "tool_built": False}
     assert observed["after"] == {"wire_2025_loaded": True, "tool_built": True}

tests/types/test_methods.py

@@ -1,6 +1,8 @@
 """Tests for the wire-method maps and two-step parse functions in `mcp_types.methods`."""
 
 import importlib.util
+import json
+import os
 import subprocess
 import sys
 from collections.abc import Mapping
@@ -12,6 +14,7 @@
 import mcp_types._v2026_07_28 as v2026
 import pydantic
 import pytest
+from inline_snapshot import snapshot
 from mcp_types import methods
 from mcp_types.version import KNOWN_PROTOCOL_VERSIONS
 from pydantic import BaseModel
@@ -594,7 +597,7 @@ def test_surface_rows_for_a_version_with_no_wire_package_fail_at_map_constructio
     """SDK-defined: an unregistered protocol version fails loudly at import, not as a version-gate KeyError."""
     with pytest.raises(ValueError) as excinfo:
         methods._LazyWireMap({("ping", "2099-01-01"): "PingRequest"})
-    assert "2099-01-01" in str(excinfo.value)
+    assert str(excinfo.value) == snapshot("no wire package registered for protocol version(s) ['2099-01-01']")
 
 
 def test_cacheable_methods_mirror_the_cacheable_method_literal():
@@ -1026,3 +1029,94 @@ def test_importing_the_module_builds_no_adapters_and_identical_rows_share_one():
     # Identical row values at another version: no new adapters.
     fresh.parse_server_result("ping", "2024-11-05", {})
     assert fresh._adapter.cache_info().currsize == 2
+
+
+# --- warm(): the opt-in prewarm helper ---
+
+_WARM_PROBE = """
+import json, sys
+import pydantic
+import mcp_types
+from mcp_types import methods
+
+report = methods.warm({args}) if {args!r} != "None" else None
+loaded_wire = sorted(m for m in sys.modules if m.startswith("mcp_types._v"))
+
+# From here on nothing may build: count model completions and adapter builds.
+built = {{"models": 0, "adapters": 0}}
+import pydantic._internal._model_construction as construction
+_complete = construction.complete_model_class
+def counting_complete(*a, **k):
+    built["models"] += 1
+    return _complete(*a, **k)
+construction.complete_model_class = counting_complete
+_init = pydantic.TypeAdapter._init_core_attrs
+def counting_init(self, *a, **k):
+    result = _init(self, *a, **k)
+    if result:
+        built["adapters"] += 1
+    return result
+pydantic.TypeAdapter._init_core_attrs = counting_init
+
+methods.parse_client_request(
+    "tools/call", "2026-07-28",
+    {{"name": "t", "arguments": {{}},
+      "_meta": {{"io.modelcontextprotocol/protocolVersion": "2026-07-28",
+                 "io.modelcontextprotocol/clientCapabilities": {{}}}}}},
+)
+methods.serialize_server_result(
+    "tools/list", "2026-07-28",
+    {{"tools": [], "cacheScope": "private", "ttlMs": 0, "resultType": "complete"}},
+)
+mcp_types.jsonrpc_message_adapter.validate_python({{"jsonrpc": "2.0", "id": 1, "method": "ping"}})
+print(json.dumps({{"report": report._asdict() if report else None, "wire": loaded_wire, "built": built}}))
+"""
+
+
+def _run_warm_probe(args: str) -> dict[str, Any]:
+    """Warm (with `args`) in a fresh interpreter, then report what the first messages built.
+
+    The probe counts builds through pydantic internals purely as a test observer.
+    """
+    result = subprocess.run(
+        [sys.executable, "-c", _WARM_PROBE.format(args=args)],
+        capture_output=True,
+        text=True,
+        check=False,
+        timeout=120,
+        # a pydantic plugin (e.g. logfire) building models is the environment's cost, not ours
+        env={**os.environ, "PYDANTIC_DISABLE_PLUGINS": "__all__"},
+    )
+    assert result.returncode == 0, result.stderr
+    return json.loads(result.stdout.strip().splitlines()[-1])
+
+
+def test_first_messages_build_validators_lazily_without_a_warm_up():
+    """SDK-defined: with no `warm()`, nothing is loaded or built up front; the first messages
+    at a version import that version's wire package and build what they route through."""
+    observed = _run_warm_probe("None")
+    assert observed["report"] is None
+    assert observed["wire"] == []  # no wire package loaded before the first message
+    assert observed["built"]["models"] > 0
+    assert observed["built"]["adapters"] > 0
+
+
+def test_warming_a_version_leaves_nothing_for_its_first_messages_to_build():
+    """SDK-defined: `warm(version)` imports that version's wire package and builds its models
+    and routing adapters up front, so the first messages at that version build nothing."""
+    observed = _run_warm_probe('"2026-07-28"')
+    assert observed["report"]["models"] > 0
+    assert observed["report"]["adapters"] > 0
+    assert observed["wire"] == ["mcp_types._v2026_07_28"]
+    assert observed["built"] == {"models": 0, "adapters": 0}
+
+
+def test_warm_is_idempotent_and_gates_unknown_versions():
+    """SDK-defined: repeating a `warm()` builds nothing more; an unknown version is a ValueError."""
+    methods.warm(everything=True)
+    again = methods.warm(everything=True)
+    assert (again.models, again.adapters) == (0, 0)
+    assert again.elapsed_ms >= 0
+    with pytest.raises(ValueError) as excinfo:
+        methods.warm("2099-01-01")
+    assert "2099-01-01" in str(excinfo.value)