Fix new_messages() and capture_run_messages() when history processors are used (#2921)

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -187,9 +187,8 @@ async def run(  # noqa: C901
                 messages = ctx_messages.messages
                 ctx_messages.used = True
 
-        message_history = _clean_message_history(ctx.state.message_history)
-        # Add message history to the `capture_run_messages` list, which will be empty at this point
-        messages.extend(message_history)
+        # Replace the `capture_run_messages` list with the message history
+        messages[:] = _clean_message_history(ctx.state.message_history)
         # Use the `capture_run_messages` list as the message history so that new messages are added to it
         ctx.state.message_history = messages
         ctx.deps.new_message_index = len(messages)
@@ -456,7 +455,18 @@ async def _prepare_request(
         # This will raise errors for any tool name conflicts
         ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
 
-        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
+        original_history = ctx.state.message_history[:]
+        message_history = await _process_message_history(original_history, ctx.deps.history_processors, run_context)
+        # Never merge the new `ModelRequest` with the one preceding it, to keep `new_messages()` from accidentally including part of the existing message history
+        message_history = [*_clean_message_history(message_history[:-1]), message_history[-1]]
+        # `ctx.state.message_history` is the same list used by `capture_run_messages`, so we should replace its contents, not the reference
+        ctx.state.message_history[:] = message_history
+        # Update the new message index to ensure `result.new_messages()` returns the correct messages
+        ctx.deps.new_message_index -= len(original_history) - len(message_history)
+
+        # Do one more cleaning pass to merge possible consecutive trailing `ModelRequest`s into one, with tool call parts before user parts,
+        # but don't store it in the message history on state.
+        # See `tests/test_tools.py::test_parallel_tool_return_with_deferred` for an example where this is necessary
         message_history = _clean_message_history(message_history)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
@@ -1087,12 +1097,11 @@ def build_agent_graph(
 
 
 async def _process_message_history(
-    state: GraphAgentState,
+    messages: list[_messages.ModelMessage],
     processors: Sequence[HistoryProcessor[DepsT]],
     run_context: RunContext[DepsT],
 ) -> list[_messages.ModelMessage]:
     """Process message history through a sequence of processors."""
-    messages = state.message_history
     for processor in processors:
         takes_ctx = is_takes_ctx(processor)
 
@@ -1110,8 +1119,12 @@ async def _process_message_history(
                 sync_processor = cast(_HistoryProcessorSync, processor)
                 messages = await run_in_executor(sync_processor, messages)
 
-    # Replaces the message history in the state with the processed messages
-    state.message_history = messages
+    if len(messages) == 0:
+        raise exceptions.UserError('Processed history cannot be empty.')
+
+    if not isinstance(messages[-1], _messages.ModelRequest):
+        raise exceptions.UserError('Processed history must end with a `ModelRequest`.')
+
     return messages
 
 

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -614,12 +614,10 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
             instrumentation_settings = None
             tracer = NoOpTracer()
 
-        graph_deps = _agent_graph.GraphAgentDeps[
-            AgentDepsT, RunOutputDataT
-        ](
+        graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunOutputDataT](
             user_deps=deps,
             prompt=user_prompt,
-            new_message_index=0,  # This will be set in `UserPromptNode` based on the length of the cleaned message history
+            new_message_index=len(message_history) if message_history else 0,
             model=model_used,
             model_settings=model_settings,
             usage_limits=usage_limits,

pydantic_ai_slim/pydantic_ai/result.py

@@ -322,9 +322,7 @@ def all_messages_json(self, *, output_tool_return_content: str | None = None) ->
             self.all_messages(output_tool_return_content=output_tool_return_content)
         )
 
-    def new_messages(
-        self, *, output_tool_return_content: str | None = None
-    ) -> list[_messages.ModelMessage]:  # pragma: no cover
+    def new_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
         """Return new messages associated with this run.
 
         Messages from older runs are excluded.