fix(messages): add ToolReturnPart to ModelResponsePart union

[Bartok9]

Summary

ToolReturnPart (part_kind='tool-return') is a member of the ModelRequestPart discriminated union but was missing from ModelResponsePart, even though both unions share the same discriminator. As a result, a ToolReturnPart embedded in a ModelResponse serializes fine but fails to deserialize with a union_tag_invalid ValidationError — the writer emits the 'tool-return' tag, the reader finds no response-union member claiming it.

This PR adds the base ToolReturnPart to ModelResponsePart and handles the new variant everywhere ModelResponse.parts is consumed, so the change stays fully type-checked (no assert_never regressions).

msg = ModelResponse(parts=[ToolReturnPart(tool_name='t', content={'r': 'ok'}, tool_call_id='c1')])
dumped = ModelMessagesTypeAdapter.dump_python([msg])
ModelMessagesTypeAdapter.validate_python(dumped)  # before: ValidationError; after: round-trips

Motivation

Closes #5721.

This is a type-system + serialization-contract consistency fix, not a fix for a shape the framework emits today. To be precise about reachability (the issue's "Root Cause" is inaccurate on this point — see my comment there):

• A normal agent run never puts a base ToolReturnPart on a ModelResponse. The framework synthesizes tool returns (including the output-tool "Final result processed." parts) into output_parts: list[ModelRequestPart] → a ModelRequest. Verified empirically: across a tool-calling run, every ToolReturnPart sits on a ModelRequest; the matching ModelResponse holds the ToolCallPart. So all_messages_json() round-trips — and therefore resumed runs, durable workflows, and the AG-UI/Vercel adapters that round-trip persisted history — are unaffected in normal operation.
• The reachable trigger is user-constructed response history (the reproduction in [roundtrip-sweep] ModelResponsePart: ToolReturnPart with part_kind='tool-return' missing from discriminated union — breaks message history round-trip #5721): code that builds ModelResponse(parts=[ToolReturnPart(...)]) itself, then persists and reloads it.

Why fix it anyway:

1. Union symmetry. ModelRequestPart already admits ToolReturnPart under the same shared discriminator. A part_kind that one union accepts and its sibling rejects is a latent hole in a public, documented serialization contract.
2. The discriminator already routes the typed subclass. ToolSearchReturnPart (a ToolReturnPart subclass) round-trips via its own 'tool-search-return' tag (_TYPED_PART_TAGS); only the base part fell through. Adding the base member completes an existing pattern and cannot shadow the subclass — pinned by test_tool_search_return_part_in_model_request_still_narrows.
3. Completing the fan-out surfaced two genuinely latent bugs on the NativeToolReturnPart sibling path (see below): a Bedrock AssertionError and an OTel telemetry drop.

(If a narrower serialization-only approach is preferred — e.g. routing the tag without widening the static union — happy to refactor. But that desyncs the type from its discriminator; widening the union is what makes the round-trip type-correct end to end.)

What changed

The union addition is the core fix. Because it makes ToolReturnPart a valid member everywhere ModelResponse.parts is iterated, every consumer now handles the new variant explicitly instead of falling through assert_never:

Site	Handling	Rationale
Model request mappers (anthropic, openai, google, gemini, groq, mistral, cohere, huggingface, xai, outlines, bedrock, test)	skip	A response-embedded tool return is persisted history, not assistant content the model emitted, so it is not replayed to the provider — mirrors the existing CompactionPart/FilePart skip pattern.
function._estimate_usage	count content	Shares the NativeToolReturnPart token path.
_agent_graph event loop, AG-UI, Vercel adapters	no event	These parts produce no streamed/rendered event.
_event_stream start/end dispatch	grouped with delta-less parts	No deltas to start/end.
ModelResponse.otel_message_parts	render as tool_call_response	Mirrors NativeToolReturnPart (minus builtin) so a response-embedded tool return isn't dropped from telemetry. (otel_events v1 unchanged — it emits no tool-return part at all, not even for NativeToolReturnPart.)
PartStartEvent.previous_part_kind / PartEndEvent.next_part_kind	add 'tool-return' literal	The kind is now reachable in the union.
_parts_manager._resolve_provider_name	isinstance narrow	ToolReturnPart has no provider_name and is never tracked by the streaming parts manager; narrowing keeps provider_name access statically typed.
google._handle_executable_code_streaming	narrow return to NativeToolCallPart (its real return)	removes the spurious widening to the full union.

Two latent bugs the fan-out surfaced

• Bedrock AssertionError. BedrockConverseModel._map_messages guarded its response-part loop with a bare else: assert isinstance(item, ToolCallPart) rather than assert_never, so the missing branch wasn't caught at type-check time and raised at runtime. Its tail is now restructured to an explicit ToolCallPart branch + assert_never, matching the other providers and catching any future union member statically.
• OTel drop. ModelResponse.otel_message_parts mapped NativeToolReturnPart to a tool_call_response part but silently dropped the base ToolReturnPart (the ladder has no assert_never, so nothing flagged it). Now mapped via the shared BaseToolReturnPart.otel_message_parts.

The provider/UI skip branches are # pragma: no cover (defensive, like their CompactionPart/FilePart siblings); the round-trip and OTel guarantees are pinned by the test_messages.py tests below.

Verification

• pre-commit (full pyright typecheck, ruff format --check, ruff check) — clean
• uv run pytest tests/test_messages.py — passes (round-trip, subclass-narrowing, OTel)
• uv run pytest tests/models/test_model_function.py — passes (_estimate_usage regression)
• Bedrock _map_messages no longer raises on a user-constructed ModelResponse([ToolReturnPart(...)]); the previously-crashing path now maps cleanly.

Tests added

• tests/test_messages.py::test_tool_return_part_in_model_response_round_trips
• tests/test_messages.py::test_tool_search_return_part_in_model_request_still_narrows
• tests/test_messages.py::test_tool_return_part_in_model_response_otel_message_parts
• tests/models/test_model_function.py::test_estimate_usage_handles_tool_return_part_in_response

[github-actions]

@DouweM This is the key design decision in the PR and I think it warrants explicit maintainer sign-off before proceeding.

Context: Issue #5721 was auto-generated by a roundtrip-sweep bot. No maintainer has commented on the issue or endorsed a specific approach. The PR author acknowledges the design choice in the PR description ("If you'd prefer a narrower serialization-only approach… happy to refactor").

The question: Should ToolReturnPart be a valid member of ModelResponsePart? The framework itself never places a ToolReturnPart into a ModelResponse — _agent_graph.py always stores tool returns in ModelRequest.parts, and model adapters use NativeToolReturnPart for builtin returns. The only way a ToolReturnPart ends up in a ModelResponse is via user-constructed message history, which would currently fail at deserialization.

Trade-offs:

• This approach (widening the union): fixes the deserialization crash and is defensive, but forces every consumer of ModelResponse.parts across ~20 files to handle a case the framework never produces. All new handler branches are # pragma: no cover dead code.
• Alternative: serialization-only fix (e.g. a lenient custom discriminator that gracefully handles unknown/unexpected tags without widening the static type union): narrower change, fewer files touched, but less type-safe.
• Alternative: reject at construction: validate that ModelResponse.parts doesn't contain request-side parts, making the user error explicit rather than silently accepting it.

The right call depends on whether "user-constructed ModelResponse containing ToolReturnPart" is something the framework should support or discourage.

[dsfaccini]

David's AICA here: Decision on the design question (Douwe is away, so I'm making the call): we'll keep the union-widening approach — adding the base ToolReturnPart to ModelResponsePart. It's the approach with direct precedent and explicit version-policy cover; the two alternatives both reintroduce problems we've deliberately avoided.

Reasoning:

• Precedent. This is the same pattern established in feat: native tool search on Anthropic and OpenAI, with custom search strategies on any provider #5143 (native tool search) and Add server-side compaction support via OpenAICompaction and AnthropicCompaction capabilities #4943 (compaction): each side's discriminated union names every concrete part type for that side, and consumers handle new members with explicit branches — including pass # pragma: no cover no-ops in provider _map_messages. The new bedrock branch sits literally next to elif isinstance(item, CompactionPart | FilePart): pass (bedrock.py:1066). So the pragma'd no-op branches this comment frames as a cost of widening are established house style, not new debt.
• Version policy. docs/version-policy.md states that adding new message parts (and their discriminator surface, e.g. the previous_part_kind/next_part_kind Literals) is not a breaking change, and "always code defensively when consuming message parts or event streams." That covers both the "breaking-ish Literal change" and the "forces consumers to handle a new case" concerns — defensive consumption is the documented contract.
• Why not the alternatives. A lenient/serialization-only discriminator would leave ToolReturnPart present in a ModelResponse at runtime but absent from the static ModelResponsePart type — a permanent type-honesty gap that pushes users toward the isinstance checks we explicitly don't want them to need. Reject-at-construction has no precedent for message parts and would break deserializing already-persisted histories that contain this shape, which is the actual [roundtrip-sweep] ModelResponsePart: ToolReturnPart with part_kind='tool-return' missing from discriminated union — breaks message history round-trip #5721 repro.

One correction to the framing above: the framework never produces a base ToolReturnPart on a ModelResponse — it routes tool returns into ModelRequest.parts. The shape is reachable only via user-constructed or deserialized message history, which is exactly the round-trip case #5721 hits, so the fix belongs at the type/deserialization layer.

[Bartok9]

Pushed a small test to green up CI. The only failing job was coverage (the check aggregate was failing solely because of it) — the 100% gate flagged _parts_manager.py:533, the isinstance(existing_part, ToolReturnPart) short-circuit in _resolve_provider_name.

That branch is defensive — no public caller passes a ToolReturnPart as existing_part (the text/thinking/tool-call paths all narrow the type first), so nothing exercised it. Added test_resolve_provider_name_tool_return_part asserting the resolver returns the incoming provider_name for a ToolReturnPart (which carries none of its own) and None when none is passed.

Locally: line 533 now covered, ruff/ruff format/pyright all clean. Thanks @dsfaccini for the bedrock/otel follow-ups. 🤝

[biswajeetdev]

The fix correctly closes the gap: ToolReturnPart is a valid member of the ModelResponsePart union (representing a tool result in user-constructed message history) but was missing from the discriminated union definition, causing deserialization failures when a persisted conversation contained one.

Three places touched, all consistent:

• messages.py — adds ToolReturnPart to the union with its discriminator tag
• _parts_manager.py — short-circuits _resolve_provider_name since ToolReturnPart carries no provider_name; returning the incoming provider_name unchanged is correct
• _agent_graph.py — pass on ToolReturnPart in the stream iterator since user-defined tool returns in pre-constructed history produce no streamed event

The # pragma: no cover on the _agent_graph.py branch is honest — this path requires user-constructed message history with a ToolReturnPart, which is hard to reach via the normal test harness. A note in the PR or a follow-up issue tracking a dedicated integration test for this path would be helpful, so it does not stay permanently uncovered.