Skip to content

Commit 08f4385

Browse files
committed
Defer building the monolith protocol models and routing adapters
Set defer_build=True on MCPModel, the JSON-RPC envelope models and every module-level TypeAdapter (the six union adapters in mcp_types, the message adapter, and the four in mcp.client.session), so importing mcp_types no longer pays for ~150 models' core schemas, validators and serializers: each is built once, on first use. The explicit model_rebuild() calls go away (they force an eager build), and the request-parameter models that carry InputResponses plus the ContentBlock members are defined before their users so their field metadata is complete without a first-use rebuild. Under defer_build pydantic derives a model's __signature__ only when it builds; a class-level lazy __signature__ (mcp_types._deferred, public pydantic API only) completes the one-time build via model_rebuild() on first inspect.signature() access, so introspection matches an eagerly-built model. MCPModel and the generated wire bases install it for all their subclasses; models deriving from BaseModel directly opt in with @deferred_model.
1 parent b8bc17d commit 08f4385

6 files changed

Lines changed: 327 additions & 110 deletions

File tree

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
"""Keep `inspect.signature()` accurate for `defer_build` models before their first use.
2+
3+
The SDK's model bases set `defer_build=True`, so pydantic builds a model's core
4+
schema, validator and serializer on first use instead of while the class
5+
statement runs; a module full of models then costs almost nothing to import in
6+
a process that never validates them.
7+
8+
The one observable pydantic ties to that build is the class's `__signature__`
9+
(the synthesized `__init__` signature that `inspect.signature(Model)` reports):
10+
under `defer_build` it exists only once the model has built, and until then a
11+
class reports the generic `(**data)`. `DeferredSignature` is a class-level
12+
`__signature__` that completes the (one-time) build via the public
13+
`model_rebuild()` on first signature access, so runtime introspection matches an
14+
eagerly-built model while nothing is built at import.
15+
"""
16+
17+
from __future__ import annotations
18+
19+
from inspect import Signature
20+
from typing import Any, TypeVar, cast
21+
22+
from pydantic import BaseModel
23+
24+
__all__ = ["DeferredSignature", "deferred_model", "install_deferred_signature", "new_deferred_signature"]
25+
26+
_ModelT = TypeVar("_ModelT", bound=BaseModel)
27+
28+
29+
class DeferredSignature:
30+
"""Class-level `__signature__` for a `defer_build=True` pydantic model.
31+
32+
First access completes the model build via `model_rebuild()` (public API,
33+
a no-op once complete); pydantic then binds its own signature on the class,
34+
replacing this attribute, and that is what gets returned. If the build cannot
35+
complete (an unresolvable forward reference), no signature results and
36+
`inspect` falls back to `__init__` exactly as for any unbuilt deferred model.
37+
"""
38+
39+
def __get__(self, obj: object | None, owner: type[BaseModel] | None = None) -> Any:
40+
if obj is not None or owner is None:
41+
# Instances take their signature from `__call__`; only the class has one.
42+
raise AttributeError("__signature__")
43+
# `raise_errors=False`: an annotation that cannot resolve yet leaves the
44+
# model unbuilt (and this attribute in place) instead of raising here,
45+
# so `inspect` falls back to `__init__` like it does for pydantic itself.
46+
owner.model_rebuild(raise_errors=False)
47+
bound = vars(owner).get("__signature__")
48+
if bound is None or bound is self:
49+
raise AttributeError("__signature__")
50+
return bound.__get__(None, owner) if hasattr(bound, "__get__") else bound
51+
52+
53+
def new_deferred_signature() -> Signature:
54+
"""A `DeferredSignature` typed as the `Signature` it stands in for.
55+
56+
Assign it in a model class body WITHOUT an annotation (`__signature__ =
57+
new_deferred_signature()`), matching how `BaseModel` itself binds
58+
`__signature__`, so it never appears in `get_type_hints(Model)`.
59+
"""
60+
return cast(Signature, DeferredSignature())
61+
62+
63+
def install_deferred_signature(cls: type[BaseModel]) -> None:
64+
"""Give a still-deferred model class its own lazily-completing `__signature__`.
65+
66+
Meant for a base's `__pydantic_init_subclass__` hook: every deferred subclass
67+
needs its OWN `__signature__` entry, otherwise it would inherit an
68+
already-built parent's signature. A class that pydantic already completed
69+
(e.g. one that turned `defer_build` off) keeps the real signature it has.
70+
"""
71+
if not cls.__pydantic_complete__:
72+
setattr(cls, "__signature__", new_deferred_signature())
73+
74+
75+
def deferred_model(cls: type[_ModelT]) -> type[_ModelT]:
76+
"""Class decorator for the root of a `defer_build` model hierarchy under `BaseModel`.
77+
78+
Installs the lazy `__signature__` on `cls` and, via a `__pydantic_init_subclass__`
79+
hook, on every model that later subclasses it (a subclass inherits the deferred
80+
build through the config). Models deriving from the SDK bases (`MCPModel`,
81+
`WireModel`, `WireRootModel`) get this from the base and do not need it.
82+
"""
83+
install_deferred_signature(cls)
84+
85+
def __pydantic_init_subclass__(sub: type[BaseModel], **kwargs: Any) -> None:
86+
install_deferred_signature(sub)
87+
# Run whatever hook `cls`'s own bases define (BaseModel's is a no-op).
88+
# (`super()` from a classmethod defined outside the class body needs
89+
# both arguments; the two-argument form is opaque to the type checker.)
90+
cast(Any, super(cls, sub)).__pydantic_init_subclass__(**kwargs)
91+
92+
setattr(cls, "__pydantic_init_subclass__", classmethod(__pydantic_init_subclass__))
93+
return cls

0 commit comments

Comments
 (0)