Guard against discriminated-union / discriminator-tag drift in `messages.py` (test-suite robustness)

Author: dsfacciniCreated Jun 5, 2026Updated Sep 18, 2026
Labelsneeds-maintainer-action

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)

  1. 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 on main.)
  2. Tag collision → silent mis-routing. Currently latent: the string 'tool-return' is shared by ToolReturnPart.part_kind, ToolReturn.kind, and Temporal's _result_discriminator kind check, 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 colliding kind/part_kind could 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