Guard against discriminated-union / discriminator-tag drift in `messages.py` (test-suite robustness)
David's AICA here, filing on David's behalf.
Problem
messages.py serializes message history through several hand-maintained discriminated unions + callable Discriminators — ModelRequestPart, ModelResponsePart, MultiModalContent, ToolReturnContent, ModelMessage, and the delta/stream-event unions — plus Temporal's _result_discriminator, plus two Literal[...] part-kind lists (PartStartEvent.previous_part_kind, PartEndEvent.next_part_kind).
These carry implicit invariants that are not enforced anywhere and are invisible to line coverage: a union member is a type-declaration line, so dropping one doesn't reduce coverage or fail any test until that exact shape happens to be deserialized. Member-set drift therefore ships green.
[#5721] (fixed in [#5723]) is a concrete instance: the base ToolReturnPart was missing from the ModelResponsePart union, so a user-built ModelResponse carrying one failed to deserialize. The only existing guard of this kind is test_multi_modal_content_types_matches_union, which covers a single content union and no part union.
Two failure modes (worth naming separately)
- Missing member → loud
union_tag_invalid. The #5721 variant. Fails closed, but blocks valid persisted histories. (Note: this is a loud error, not silent corruption — verified by reproducing onmain.) - Tag collision → silent mis-routing. Currently latent: the string
'tool-return'is shared byToolReturnPart.part_kind,ToolReturn.kind, and Temporal's_result_discriminatorkindcheck, but those read different fields over disjoint member sets, so nothing cross-routes today. This is the dangerous mode and it is unguarded — a future part/wrapper introducing a collidingkind/part_kindcould silently mis-route.
Live exhibit (already drifted)
PartStartEvent.previous_part_kind / PartEndEvent.next_part_kind omit 4 kinds that are valid ModelResponsePart members: tool-search-call, builtin-tool-search-call, builtin-tool-search-return, capability-load-call. Either a latent streaming-event mis-grouping bug or intentional narrowing — it's undocumented, so it can't be told apart. Open question to resolve as part of this work.
Proposed guards (introspection-based, in tests/test_messages.py)
| # | Guard | Invariant | Catches |
|---|---|---|---|
| 1 | Round-trip-every-part | every concrete part survives ModelMessagesTypeAdapter dump→validate inside a ModelRequest and (where its tag is a member) a ModelResponse, asserting both type(...) and field equality |
#5721; metadata mangling; typed-subclass-shadows-base |
| 2 | Membership-matches-classes | every concrete *Part dataclass is a member of a part union or explicitly allow-listed with a reason (e.g. InstructionPart, which renders into ModelRequest.instructions: str and is never a serialized part) |
a new part the author forgot to wire into a union, caught at class-def altitude |
| 3 | Discriminator-tag ↔ member bijection | the tag set a discriminator can return (incl. _TYPED_PART_TAGS registry values) equals the member Tag set — no orphan tag, no unreachable member |
a registry entry pointing at a non-existent member tag |
| 4 | Literal ⊆ union |
previous_part_kind / next_part_kind args ⊆ ModelResponsePart tags |
a stale Literal naming a kind no longer in the union |
Guard 1 is the umbrella (it would have caught #5721 directly — confirmed it raises the exact union_tag_invalid on the pre-fix union). Guards 2–4 are cheap, precise localizers that name which invariant broke in one line when guard 1 goes red across many params. Guard 4 should ship as subset, not equality, pending the Literal-gap decision above (strict equality is red on the current tree).
Implementation note: use a non-image BinaryContent in the per-part factory so the BinaryImage.narrow_type AfterValidator doesn't re-type the content and trip field equality (same dodge test_tool_return_part_binary_content_round_trip already uses).
References
- #5721 / #5723 — motivating case (#5723 adds point regression tests, not the generalized guards)
- existing precedent:
test_multi_modal_content_types_matches_union
Source: pydantic/pydantic-ai