Max Isbey
33471c77fd
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.
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
"""`docs/advanced/import-cost.md`: the prewarm recipe does what the page says."""
|
|
|
|
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 = """
|
|
import json, sys
|
|
import mcp_types
|
|
|
|
before = {
|
|
"wire_2025_loaded": "mcp_types._v2025_11_25" in sys.modules,
|
|
"tool_built": mcp_types.Tool.__pydantic_complete__,
|
|
}
|
|
import docs_src.import_cost.tutorial001 # the recipe under test
|
|
after = {
|
|
"wire_2025_loaded": "mcp_types._v2025_11_25" in sys.modules,
|
|
"tool_built": mcp_types.Tool.__pydantic_complete__,
|
|
}
|
|
print(json.dumps({"before": before, "after": after}))
|
|
"""
|
|
|
|
|
|
def test_prewarm_recipe_loads_the_wire_package_and_builds_the_models() -> None:
|
|
"""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.strip().splitlines()[-1])
|
|
assert observed["before"] == {"wire_2025_loaded": False, "tool_built": False}
|
|
assert observed["after"] == {"wire_2025_loaded": True, "tool_built": True}
|