Hack around a pydantic bug preventing GCing modules

src/ai/models/core/model.py

@@ -1,17 +1,29 @@
 """Model metadata types."""
 
 import os
-from typing import Any, Self
+import weakref
+from typing import TYPE_CHECKING, Any, Self, cast
 
 import pydantic
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 from ... import _modelsdev
 from ...errors import ConfigurationError
 from ...providers import base
 
 _DEFAULT_MODEL_ENV = "AI_SDK_DEFAULT_MODEL"
 
 
+def _exclude_if_none(v: Any) -> bool:
+    # pydantic-core's SchemaSerializer doesn't GC-traverse ``exclude_if``
+    # callables, so a closure stored there directly forms an uncollectable
+    # cycle back through this module; wrap it in weakref.proxy at the call
+    # site to keep the serializer's reference weak.
+    return v is None
+
+
 class Model(pydantic.BaseModel):
     """Lightweight reference to a model on a specific provider.
 
@@ -23,7 +35,10 @@ class Model(pydantic.BaseModel):
     id: str
     provider: base.Provider[Any]
     protocol: base.ProviderProtocol[Any] | None = pydantic.Field(
-        default=None, exclude_if=lambda v: v is None
+        default=None,
+        exclude_if=cast(
+            "Callable[[Any], bool]", weakref.proxy(_exclude_if_none)
+        ),
     )
 
     def __repr__(self) -> str:

src/ai/providers/base.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import weakref
 from collections.abc import Mapping
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Self, cast
@@ -17,7 +18,7 @@
 from ..errors import UnsupportedProviderError
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Sequence
+    from collections.abc import AsyncGenerator, Callable, Sequence
 
     import modelsdotdev
 
@@ -30,6 +31,14 @@
 ClientT = TypeVar("ClientT", default=Any)
 
 
+def _exclude_if_none(v: Any) -> bool:
+    # pydantic-core's SchemaSerializer doesn't GC-traverse ``exclude_if``
+    # callables, so a closure stored there directly forms an uncollectable
+    # cycle back through this module; wrap it in weakref.proxy at the call
+    # site to keep the serializer's reference weak.
+    return v is None
+
+
 def _generic_origin(cls: type[Any]) -> type[Any]:
     # if this is a generic class, return the original class
     return (
@@ -169,10 +178,16 @@ class Provider(pydantic.BaseModel, Generic[ClientT]):
     name: str  # models.dev identity
     default_base_url: str
     protocol_override: ProviderProtocol[Any] | None = pydantic.Field(
-        default=None, exclude_if=lambda v: v is None
+        default=None,
+        exclude_if=cast(
+            "Callable[[Any], bool]", weakref.proxy(_exclude_if_none)
+        ),
     )
     api_key_value: str | None = pydantic.Field(
-        default=None, exclude_if=lambda v: v is None
+        default=None,
+        exclude_if=cast(
+            "Callable[[Any], bool]", weakref.proxy(_exclude_if_none)
+        ),
     )
     api_key_env: str | None = None
     base_url_env: str | None = None

src/ai/types/messages.py

@@ -3,6 +3,7 @@
 import contextvars
 import functools
 import random
+import weakref
 from collections.abc import AsyncIterator, Callable, Iterator, Sequence
 from typing import Annotated, Any, Literal, Protocol, Self, cast, overload
 
@@ -217,6 +218,23 @@ class _ModelInputUnset:
 
 _MODEL_INPUT_UNSET: Any = _ModelInputUnset()
 
+
+def _exclude_if_model_input_unset(v: Any) -> bool:
+    return isinstance(v, _ModelInputUnset)
+
+
+def _exclude_if_none(v: Any) -> bool:
+    # pydantic-core's SchemaSerializer doesn't GC-traverse ``exclude_if``
+    # callables, so a closure stored there directly forms an uncollectable
+    # cycle back through this module; wrap it in weakref.proxy at the call
+    # site to keep the serializer's reference weak.
+    return v is None
+
+
+def _exclude_if_falsy(v: Any) -> bool:
+    return not v
+
+
 # Coarse tag for the shape of ``ToolResultPart.result``.
 # ``"special"`` means a :class:`SpecialToolResult`; ``"error"`` flags
 # an error result; ``"json"`` (the default) is any plain value.
@@ -264,7 +282,10 @@ class ToolResultPart(pydantic.BaseModel):
     # again, though.
     model_input: Any = pydantic.Field(
         default_factory=lambda: _MODEL_INPUT_UNSET,
-        exclude_if=lambda v: isinstance(v, _ModelInputUnset),
+        exclude_if=cast(
+            "Callable[[Any], bool]",
+            weakref.proxy(_exclude_if_model_input_unset),
+        ),
         repr=False,
     )
 
@@ -362,7 +383,9 @@ class ToolCallPart(pydantic.BaseModel):
     # running the tool body again.
     cached_result: ToolResultPart | None = pydantic.Field(
         default=None,
-        exclude_if=lambda v: v is None,
+        exclude_if=cast(
+            "Callable[[Any], bool]", weakref.proxy(_exclude_if_none)
+        ),
         repr=False,
     )
 
@@ -463,7 +486,9 @@ class Message(pydantic.BaseModel):
     # producing a duplicate turn.
     replay: bool = pydantic.Field(
         default=False,
-        exclude_if=lambda v: not v,
+        exclude_if=cast(
+            "Callable[[Any], bool]", weakref.proxy(_exclude_if_falsy)
+        ),
         repr=False,
     )
 

tests/types/test_exclude_if_gc.py

@@ -0,0 +1,96 @@
+"""Guard against the pydantic-core ``exclude_if`` GC leak.
+
+pydantic-core's ``SchemaSerializer`` stores ``Field(exclude_if=...)``
+callables but omits them from ``tp_traverse``.  A callable whose
+``__globals__`` can reach its own model class therefore forms a
+reference cycle the garbage collector can never prove unreachable, and
+the module's entire object graph leaks.  Passing ``weakref.proxy`` of a
+named module-level function keeps the serializer's (untraversed)
+reference weak, so no strong cycle forms.
+"""
+
+from __future__ import annotations
+
+import gc
+import importlib
+import pathlib
+import sys
+import weakref
+
+import ai.types.messages as messages_mod
+from ai.types.messages import Message, ToolCallPart, ToolResultPart
+
+_AI_SRC = pathlib.Path(messages_mod.__file__).parents[2]
+
+
+def test_exclude_if_behavior_through_proxy() -> None:
+    """The weakref.proxy indirection must not change serialization."""
+    part = ToolResultPart(tool_call_id="tc", tool_name="t", result=1)
+    dumped = part.model_dump()
+    # model_input is unset -> excluded by _exclude_if_model_input_unset
+    assert "model_input" not in dumped
+
+    part = ToolResultPart(
+        tool_call_id="tc", tool_name="t", result=1, model_input="visible"
+    )
+    assert part.model_dump()["model_input"] == "visible"
+
+    call = ToolCallPart(tool_call_id="tc", tool_name="t", tool_args="{}")
+    # cached_result is None -> excluded by _exclude_if_none
+    assert "cached_result" not in call.model_dump()
+
+    msg = Message(role="user", parts=[])
+    # replay is False -> excluded by _exclude_if_falsy
+    dumped = msg.model_dump()
+    assert "replay" not in dumped
+    assert Message(role="user", parts=[], replay=True).model_dump()["replay"]
+
+
+def test_exclude_if_models_are_collectable(tmp_path: pathlib.Path) -> None:
+    """A module-level model using the proxy pattern must be collectable.
+
+    With a bare lambda this cycle is uncollectable:
+    lambda -> __globals__ -> class -> __pydantic_serializer__ -> lambda.
+    """
+    sys.path.insert(0, str(tmp_path))
+    (tmp_path / "excl_gc_probe.py").write_text(
+        "import weakref\n"
+        "import pydantic\n"
+        "\n"
+        "def _exclude_if_none(v: object) -> bool:\n"
+        "    return v is None\n"
+        "\n"
+        "class Probe(pydantic.BaseModel):\n"
+        "    x: int | None = pydantic.Field(\n"
+        "        default=None, exclude_if=weakref.proxy(_exclude_if_none)\n"
+        "    )\n"
+    )
+    try:
+        mod = importlib.import_module("excl_gc_probe")
+        assert mod.Probe(x=None).model_dump() == {}
+        assert mod.Probe(x=3).model_dump() == {"x": 3}
+        ref = weakref.ref(mod.Probe)
+        del sys.modules["excl_gc_probe"], mod
+        gc.collect()
+        gc.collect()
+        assert ref() is None, "model class leaked despite weakref.proxy"
+    finally:
+        sys.path.remove(str(tmp_path))
+        sys.modules.pop("excl_gc_probe", None)
+
+
+def test_no_exclude_if_lambdas_in_source() -> None:
+    """Fail if the leaking ``exclude_if=lambda`` pattern reappears."""
+    offenders = [
+        f"{path.relative_to(_AI_SRC)}:{lineno}"
+        for path in sorted(_AI_SRC.rglob("*.py"))
+        for lineno, line in enumerate(
+            path.read_text(encoding="utf-8").splitlines(), start=1
+        )
+        if "exclude_if=lambda" in line
+    ]
+    assert not offenders, (
+        "exclude_if must not be passed a lambda (pydantic-core does not "
+        "GC-traverse it; use weakref.proxy of a named function): "
+        f"{offenders}"
+    )