diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 644e938..78887b4 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -401,3 +401,11 @@ def get_output( if output_type is None: return self.text return output_type.model_validate_json(self.text) + + +# ``MessageBundle`` forward-references ``Message``, which is defined later in +# this module, so its schema (and the adapter built from it) is incomplete at +# class-creation time. Rebuild both once ``Message`` exists so importers never +# see a half-built model. +MessageBundle.model_rebuild() +_SPECIAL_TOOL_RESULT_ADAPTER.rebuild(force=True) diff --git a/tests/types/test_messages.py b/tests/types/test_messages.py index 08dee49..b8847cc 100644 --- a/tests/types/test_messages.py +++ b/tests/types/test_messages.py @@ -2,6 +2,8 @@ from __future__ import annotations +import subprocess +import sys from typing import Any import pytest @@ -121,6 +123,70 @@ def test_from_bytes_unknown_raises() -> None: # --------------------------------------------------------------------------- +def test_models_fully_defined_at_import() -> None: + """No model may escape ``import ai`` with an unbuilt schema. + + ``MessageBundle`` forward-references ``Message``, which is defined later + in the module, so its schema is deferred at class creation and must be + completed by the module-end ``model_rebuild()``. Without that it leaks + out of the import incomplete, and whether a downstream consumer works + then depends on *when* and *where* the lazy rebuild fires: embedding the + class in another model usually heals it, but e.g. a durable-workflow + module re-import broke with "``SessionState`` is not fully defined; you + should define ``Message``" in the seal backend. + + This must run in a fresh interpreter: pydantic heals an incomplete model + on its first direct validation, so any earlier test that builds a + ``MessageBundle`` (e.g. the ai_sdk adapter tests, which sort first) + silently masks the regression in-process. + """ + code = ( + "import inspect\n" + "import pydantic\n" + "import ai\n" + "import ai.types.messages as m\n" + "bad = [\n" + " name\n" + " for name, obj in vars(m).items()\n" + " if inspect.isclass(obj)\n" + " and issubclass(obj, pydantic.BaseModel)\n" + " and obj.__module__ == m.__name__\n" + " and not obj.__pydantic_complete__\n" + "]\n" + "assert not bad, f'models escaped import with unbuilt schemas: {bad}'\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) + + +def test_message_bundle_embeddable_without_sys_modules_entry() -> None: + """Embedding ``MessageBundle`` must not require a ``sys.modules`` lookup. + + If ``MessageBundle`` escapes import incomplete, a consumer model + embedding it makes pydantic re-evaluate the deferred ``"Message"`` + annotation via ``sys.modules[MessageBundle.__module__]``. Environments + that rebuild ``sys.modules`` break then: the vercel workflow sandbox + clears ``sys.modules`` and re-registers only what sandboxed code + explicitly imports — a bare ``import ai`` does *not* re-register the + ``ai.types.messages`` key — so the lookup misses, the consumer fails + with "name 'Message' is not defined", and every model embedding it is + poisoned (this took down the seal backend's ``SessionState``). A + complete-at-import ``MessageBundle`` never needs the lookup. + + Runs in a fresh interpreter for the same masking reason as above. + """ + code = ( + "import sys\n" + "import pydantic\n" + "import ai\n" + "from ai.types import messages\n" + "del sys.modules['ai.types.messages']\n" + "class Holder(pydantic.BaseModel):\n" + " x: messages.MessageBundle | None = None\n" + "Holder(x=messages.MessageBundle(messages=()))\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) + + def test_tool_result_content_output_with_file_part_round_trip() -> None: """FilePart inside ContentOutput survives JSON round-trip.""" fp = messages.FilePart(data=b"fake-image-data", media_type="image/png")