fix(messages): add ToolReturnPart to ModelResponsePart union

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -1190,6 +1190,9 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:  # noqa
                     elif isinstance(part, _messages.CompactionPart):
                         if part.content:
                             compaction_text += part.content
+                    elif isinstance(part, _messages.ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history produce no streamed event here.
+                        pass
                     else:
                         assert_never(part)
 

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -32,6 +32,7 @@
     ToolCallPart,
     ToolCallPartDelta,
     ToolPartKind,
+    ToolReturnPart,
 )
 
 from ._utils import generate_tool_call_id as _generate_tool_call_id
@@ -526,6 +527,10 @@ def _resolve_provider_name(
         self, existing_part: ModelResponsePart | ToolCallPartDelta, provider_name: str | None
     ) -> str | None:
         """Return the provider name if it has not been set on previous parts."""
+        # `ToolReturnPart` is a valid `ModelResponsePart` member but is never tracked by the parts
+        # manager, and unlike the other members it carries no `provider_name`.
+        if isinstance(existing_part, ToolReturnPart):
+            return provider_name
         if existing_part.provider_name is None or provider_name != existing_part.provider_name:
             return provider_name
         return None

pydantic_ai_slim/pydantic_ai/messages.py

@@ -2066,6 +2066,7 @@ def _model_response_part_discriminator(v: Any) -> str | None:
     | Annotated[NativeToolCallPart, pydantic.Tag('builtin-tool-call')]
     | Annotated[NativeToolSearchReturnPart, pydantic.Tag('builtin-tool-search-return')]
     | Annotated[NativeToolReturnPart, pydantic.Tag('builtin-tool-return')]
+    | Annotated[ToolReturnPart, pydantic.Tag('tool-return')]
     | Annotated[ThinkingPart, pydantic.Tag('thinking')]
     | Annotated[CompactionPart, pydantic.Tag('compaction')]
     | Annotated[FilePart, pydantic.Tag('file')],
@@ -2285,7 +2286,7 @@ def new_event_body():
 
         return result
 
-    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:  # noqa: C901
         parts: list[_otel_messages.MessagePart] = []
         for part in self.parts:
             if isinstance(part, TextPart):
@@ -2333,6 +2334,10 @@ def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_me
                     return_part['result'] = serialize_any(part.content)
 
                 parts.append(return_part)
+            elif isinstance(part, ToolReturnPart):
+                # A user-defined tool return can appear here via user-constructed message history; map it
+                # like its request-side counterpart (no `builtin` flag) so it isn't dropped from telemetry.
+                parts.extend(part.otel_message_parts(settings))
             elif isinstance(part, CompactionPart):
                 # Compaction parts don't map to standard OTel message part types
                 pass
@@ -2710,7 +2715,16 @@ class PartStartEvent:
     """The newly started `ModelResponsePart`."""
 
     previous_part_kind: (
-        Literal['text', 'thinking', 'tool-call', 'builtin-tool-call', 'builtin-tool-return', 'compaction', 'file']
+        Literal[
+            'text',
+            'thinking',
+            'tool-call',
+            'builtin-tool-call',
+            'builtin-tool-return',
+            'tool-return',
+            'compaction',
+            'file',
+        ]
         | None
     ) = None
     """The kind of the previous part, if any.
@@ -2751,7 +2765,16 @@ class PartEndEvent:
     """The complete `ModelResponsePart`."""
 
     next_part_kind: (
-        Literal['text', 'thinking', 'tool-call', 'builtin-tool-call', 'builtin-tool-return', 'compaction', 'file']
+        Literal[
+            'text',
+            'thinking',
+            'tool-call',
+            'builtin-tool-call',
+            'builtin-tool-return',
+            'tool-return',
+            'compaction',
+            'file',
+        ]
         | None
     ) = None
     """The kind of the next part, if any.

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -1533,6 +1533,9 @@ async def _map_message(  # noqa: C901
                     elif isinstance(response_part, FilePart):  # pragma: no cover
                         # Files generated by models are not sent back to models that don't themselves generate files.
                         pass
+                    elif isinstance(response_part, ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                        pass
                     else:
                         assert_never(response_part)
                 if len(assistant_content_params) > 0:

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -1066,9 +1066,13 @@ async def _map_messages(  # noqa: C901
                     elif isinstance(item, CompactionPart | FilePart):
                         # Compaction and file parts are not sent back to models that don't support them.
                         pass  # pragma: no cover
-                    else:
-                        assert isinstance(item, ToolCallPart)
+                    elif isinstance(item, ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                        pass
+                    elif isinstance(item, ToolCallPart):
                         content.append(self._map_tool_call(item))
+                    else:
+                        assert_never(item)
                 if content:
                     bedrock_messages.append({'role': 'assistant', 'content': content})
             else:

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -253,7 +253,7 @@ def _map_messages(
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
                     elif isinstance(
-                        item, NativeToolCallPart | NativeToolReturnPart | FilePart | CompactionPart
+                        item, NativeToolCallPart | NativeToolReturnPart | FilePart | CompactionPart | ToolReturnPart
                     ):  # pragma: no cover
                         pass
                     else:

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -410,7 +410,7 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.RequestUsage:
                     response_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ToolCallPart | NativeToolCallPart):
                     response_tokens += 1 + _estimate_string_tokens(part.args_as_json_str())
-                elif isinstance(part, NativeToolReturnPart):
+                elif isinstance(part, NativeToolReturnPart | ToolReturnPart):
                     response_tokens += _estimate_string_tokens(part.model_response_str())
                 elif isinstance(part, FilePart):
                     response_tokens += _estimate_string_tokens([part.content])

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -683,6 +683,9 @@ def _content_model_response(m: ModelResponse) -> _GeminiContent:
         elif isinstance(item, CompactionPart):  # pragma: no cover
             # Compaction parts are not sent back to models that don't support compaction.
             pass
+        elif isinstance(item, ToolReturnPart):  # pragma: no cover
+            # User-defined tool returns in user-constructed message history are not replayed to the provider.
+            pass
         else:
             assert_never(item)
     return _GeminiContent(role='model', parts=parts)

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -1485,7 +1485,7 @@ def _map_code_execution_result(self, code_execution_result: CodeExecutionResult)
         assert self._code_execution_tool_call_id is not None
         return _map_code_execution_result(code_execution_result, self.provider_name, self._code_execution_tool_call_id)
 
-    def _handle_executable_code_streaming(self, executable_code: ExecutableCode) -> ModelResponsePart:
+    def _handle_executable_code_streaming(self, executable_code: ExecutableCode) -> NativeToolCallPart:
         """Handle executable code for streaming responses.
 
         Returns a NativeToolCallPart for file search or code execution.
@@ -1592,6 +1592,9 @@ def _content_model_response(
         elif isinstance(item, CompactionPart):  # pragma: no cover
             # Compaction parts are not sent back to models that don't support compaction.
             part = None
+        elif isinstance(item, ToolReturnPart):  # pragma: no cover
+            # User-defined tool returns in user-constructed message history are not replayed to the provider.
+            part = None
         else:
             assert_never(item)
 
@@ -1611,6 +1614,9 @@ def _decode_inline_thought_signature(
     Returns the raw signature bytes ready to embed in a `PartDict`, or `None` if no signature
     applies (either missing, or the response originated from a different provider).
     """
+    if isinstance(item, ToolReturnPart):  # pragma: no cover
+        # User-defined tool returns carry no provider signature.
+        return None
     if not item.provider_details:
         return None
     if m.provider_name not in accepted_provider_names and item.provider_name not in accepted_provider_names:

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -483,8 +483,8 @@ async def _map_messages(
                     elif isinstance(item, FilePart):  # pragma: no cover
                         # Files generated by models are not sent back to models that don't themselves generate files.
                         pass
-                    elif isinstance(item, CompactionPart):  # pragma: no cover
-                        # Compaction parts are not sent back to models that don't support compaction.
+                    elif isinstance(item, CompactionPart | ToolReturnPart):  # pragma: no cover
+                        # Compaction parts and user-constructed tool returns are not sent back to the provider.
                         pass
                     else:
                         assert_never(item)

pydantic_ai_slim/pydantic_ai/models/huggingface.py

@@ -400,8 +400,8 @@ async def _map_messages(
                     elif isinstance(item, FilePart):  # pragma: no cover
                         # Files generated by models are not sent back to models that don't themselves generate files.
                         pass
-                    elif isinstance(item, CompactionPart):  # pragma: no cover
-                        # Compaction parts are not sent back to models that don't support compaction.
+                    elif isinstance(item, CompactionPart | ToolReturnPart):  # pragma: no cover
+                        # Compaction parts and user-constructed tool returns are not sent back to the provider.
                         pass
                     else:
                         assert_never(item)

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -598,6 +598,9 @@ async def _map_messages(  # noqa: C901
                     elif isinstance(part, CompactionPart):  # pragma: no cover
                         # Compaction parts are not sent back to models that don't support compaction.
                         pass
+                    elif isinstance(part, ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                        pass
                     else:
                         assert_never(part)
                 if thinking_chunks:

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -1310,6 +1310,9 @@ def map_assistant_message(self, message: ModelResponse) -> chat.ChatCompletionAs
                 elif isinstance(item, CompactionPart):  # pragma: no cover
                     # Compaction parts are not sent back to the Chat Completions API.
                     pass
+                elif isinstance(item, ToolReturnPart):  # pragma: no cover
+                    # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                    pass
                 else:
                     assert_never(item)
             return self._into_message_param()
@@ -2769,6 +2772,9 @@ async def _map_messages(  # noqa: C901
                 file_search_item: responses.ResponseFileSearchToolCallParam | None = None
                 code_interpreter_item: responses.ResponseCodeInterpreterToolCallParam | None = None
                 for item in message.parts:
+                    if isinstance(item, ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                        continue
                     should_send_item_id = send_item_ids and (
                         item.provider_name == self.system
                         or (item.provider_name is None and message.provider_name == self.system)

pydantic_ai_slim/pydantic_ai/models/outlines.py

@@ -506,6 +506,9 @@ async def _format_prompt(  # noqa: C901
                     elif isinstance(part, CompactionPart):  # pragma: no cover
                         # Compaction parts are not sent back to models that don't support compaction.
                         pass
+                    elif isinstance(part, ToolReturnPart):  # pragma: no cover
+                        # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                        pass
                     else:
                         assert_never(part)
                 if len(text_parts) == 1 and len(image_parts) == 0:

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -365,6 +365,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             elif isinstance(part, CompactionPart):  # pragma: no cover
                 # NOTE: There's no way to reach this part of the code, since we don't generate CompactionPart on TestModel.
                 assert False, "This should be unreachable — we don't generate CompactionPart on TestModel."
+            elif isinstance(part, ToolReturnPart):  # pragma: no cover
+                # NOTE: There's no way to reach this part of the code, since we don't generate ToolReturnPart on TestModel.
+                assert False, "This should be unreachable — we don't generate ToolReturnPart on TestModel."
             else:
                 assert_never(part)
 

pydantic_ai_slim/pydantic_ai/models/xai.py

@@ -390,6 +390,9 @@ def _map_response_parts(self, parts: Sequence[ModelResponsePart]) -> list[chat_t
             elif isinstance(item, CompactionPart):  # pragma: no cover
                 # Compaction parts are not sent back to models that don't support compaction.
                 pass
+            elif isinstance(item, ToolReturnPart):  # pragma: no cover
+                # User-defined tool returns in user-constructed message history are not replayed to the provider.
+                pass
             else:
                 assert_never(item)
 

pydantic_ai_slim/pydantic_ai/ui/_event_stream.py

@@ -338,7 +338,7 @@ async def handle_event(self, event: NativeEvent) -> AsyncIterator[EventT]:  # no
             case _:
                 pass
 
-    async def handle_part_start(self, event: PartStartEvent) -> AsyncIterator[EventT]:
+    async def handle_part_start(self, event: PartStartEvent) -> AsyncIterator[EventT]:  # noqa: C901
         """Handle a `PartStartEvent`.
 
         This method dispatches to specific `handle_*` methods based on part type:
@@ -381,6 +381,9 @@ async def handle_part_start(self, event: PartStartEvent) -> AsyncIterator[EventT
             case CompactionPart():  # pragma: no cover
                 async for e in self.handle_compaction(part):
                     yield e
+            case ToolReturnPart():  # pragma: no cover
+                # User-defined tool returns in user-constructed message history have no UI start event.
+                pass
 
     async def handle_part_delta(self, event: PartDeltaEvent) -> AsyncIterator[EventT]:
         """Handle a PartDeltaEvent.
@@ -440,7 +443,7 @@ async def handle_part_end(self, event: PartEndEvent) -> AsyncIterator[EventT]:
             case NativeToolCallPart():
                 async for e in self.handle_builtin_tool_call_end(part):
                     yield e
-            case NativeToolReturnPart() | FilePart() | CompactionPart():  # pragma: no cover
+            case NativeToolReturnPart() | FilePart() | CompactionPart() | ToolReturnPart():  # pragma: no cover
                 # These don't have deltas, so they don't need to be ended.
                 pass
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -678,6 +678,8 @@ def flush() -> None:
                     )
             elif isinstance(part, CompactionPart):  # pragma: no cover
                 pass  # Compaction parts are not rendered in AG-UI
+            elif isinstance(part, ToolReturnPart):  # pragma: no cover
+                pass  # User-defined tool returns in user-constructed message history are not rendered in AG-UI
             else:
                 assert_never(part)
 

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py

@@ -682,6 +682,8 @@ def _dump_response_message(
                 ui_parts.extend(cls._dump_tool_call_part(part, tool_results, sdk_version))
             elif isinstance(part, CompactionPart):  # pragma: no cover
                 pass  # Compaction parts are not rendered in the UI
+            elif isinstance(part, ToolReturnPart):  # pragma: no cover
+                pass  # User-defined tool returns in user-constructed message history are not rendered in the UI
             else:
                 assert_never(part)
 

tests/models/test_model_function.py

@@ -21,7 +21,13 @@
     ToolReturnPart,
     UserPromptPart,
 )
-from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.models.function import (
+    AgentInfo,
+    DeltaToolCall,
+    DeltaToolCalls,
+    FunctionModel,
+    _estimate_usage,  # pyright: ignore[reportPrivateUsage]
+)
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.result import RunUsage
 from pydantic_ai.usage import RequestUsage
@@ -567,3 +573,25 @@ async def test_return_empty():
     with pytest.raises(ValueError, match='Stream function must return at least one item'):
         async with agent.run_stream(''):
             pass
+
+
+def test_estimate_usage_handles_tool_return_part_in_response():
+    """Regression for #5721: a `ToolReturnPart` stored on a `ModelResponse` must be handled.
+
+    Adding the base `ToolReturnPart` to the `ModelResponsePart` union means every consumer
+    that iterates `ModelResponse.parts` (here the usage estimator) must handle the new
+    variant rather than falling through to `assert_never`. The framework stores these parts
+    on responses for user-defined output tools, so they appear in real message history.
+    """
+    messages: list[ModelMessage] = [
+        ModelResponse(
+            parts=[
+                TextPart(content='hello'),
+                ToolReturnPart(tool_name='my_tool', content='tool result here', tool_call_id='call-1'),
+            ]
+        )
+    ]
+    estimated = _estimate_usage(messages)
+    # Text + the tool return content are both counted; the key assertion is that no
+    # `assert_never` is raised for the response-embedded `ToolReturnPart`.
+    assert estimated.output_tokens is not None and estimated.output_tokens > 0

tests/models/test_openai_responses.py

@@ -11525,7 +11525,9 @@ async def test_openai_responses_compact_stateful_mode_stream(allow_model_request
         e for e in all_events if isinstance(e, PartStartEvent) and isinstance(e.part, CompactionPart)
     ]
     assert compaction_start_events, 'expected PartStartEvent for CompactionPart during streaming'
-    assert compaction_start_events[0].part.provider_name == 'openai'
+    first_compaction_part = compaction_start_events[0].part
+    assert isinstance(first_compaction_part, CompactionPart)
+    assert first_compaction_part.provider_name == 'openai'
 
     # Verify final messages contain the CompactionPart with encrypted_content
     compaction_parts = [