fix(ag-ui): mixed tool handling

Leverage message snapshots to ensure agents have access to the entire
context of a conversation, include the details of non AG-UI tool
requests in the message history.

This is needed when the agent processes a non-AG-UI tool request before
an AG-UI tool request, which requires multiple UI to agent interactions
to complete.

Author: stevenh

docs/ag-ui.md

@@ -194,12 +194,6 @@ convenience method, it accepts the same arguments as the
 AG-UI tools are seamlessly provided to the PydanticAI agent, enabling rich
 use experiences with frontend user interfaces.
 
-!!! warning "Requests requiring PydanticAI and AG-UI tools"
-    Agent requests which require a PydanticAI tool use followed by an
-    AG-UI tool will currently fail to process correctly, as the response from
-    AG-UI will not include details about the PydanticAI tool request / response.
-    This will be fixed in a future release
-
 #### Events
 
 The adapter provides the ability for PydanticAI tools to send

pydantic_ai_ag_ui/pydantic_ai_ag_ui/adapter.py

@@ -14,7 +14,9 @@
     BaseEvent,
     DeveloperMessage,
     EventType,
+    FunctionCall,
     Message,
+    MessagesSnapshotEvent,
     RunAgentInput,
     RunErrorEvent,
     RunFinishedEvent,
@@ -24,6 +26,7 @@
     TextMessageEndEvent,
     TextMessageStartEvent,
     Tool as ToolAGUI,
+    ToolCall,
     ToolCallArgsEvent,
     ToolCallEndEvent,
     ToolCallStartEvent,
@@ -34,6 +37,7 @@
 
 from pydantic_ai import Agent, ModelRequestNode, models
 from pydantic_ai._output import OutputType
+from pydantic_ai._parts_manager import ModelResponsePartsManager
 from pydantic_ai.agent import RunOutputDataT
 from pydantic_ai.mcp import ToolResult
 from pydantic_ai.messages import (
@@ -43,6 +47,7 @@
     ModelRequest,
     ModelRequestPart,
     ModelResponse,
+    ModelResponsePart,
     PartDeltaEvent,
     PartStartEvent,
     SystemPromptPart,
@@ -108,13 +113,6 @@ class Adapter(Generic[AgentDepsT, OutputDataT]):
     an adapter for running agents with Server-Sent Event (SSE) streaming
     responses using the AG-UI protocol.
 
-    # Warning
-
-    Agent requests which require a PydanticAI tool use followed by an AG-UI
-    tool will currently fail to process correctly, as the response from AG-UI
-    will not include details about the PydanticAI tool request / response.
-    This will be fixed in a future release
-
     Examples:
     This is an example of base usage with FastAPI.
     .. code-block:: python
@@ -239,14 +237,9 @@ async def run(
             if isinstance(deps, StateHandler):
                 deps.set_state(run_input.state)
 
-            prompt: str | None = None
-            if isinstance(run_input.messages[-1], UserMessage):
-                prompt = run_input.messages[-1].content
-                run_input.messages.pop()
-
             run: AgentRun[AgentDepsT, Any]
             async with self.agent.iter(
-                user_prompt=prompt,
+                user_prompt=None,
                 output_type=output_type,
                 message_history=_convert_history(run_input.messages),
                 model=model,
@@ -257,10 +250,21 @@ async def run(
                 infer_name=infer_name,
                 additional_tools=run_tools,
             ) as run:
-                async for event in self._agent_stream(tool_names, run):
+                parts_manager: ModelResponsePartsManager = ModelResponsePartsManager()
+                async for event in self._agent_stream(tool_names, run, parts_manager):
                     if event is None:
                         # Tool call signals early return, so we stop processing.
                         self.logger.debug('tool call early return')
+
+                        # TODO(steve): Remove this workaround, it's only needed as AG-UI doesn't
+                        # currently have a way to add server side tool calls to the message history
+                        # via events. To workaround this we create a full snapshot of the messages
+                        # and send that.
+                        snapshot: MessagesSnapshotEvent | None = self._message_snapshot(
+                            run, run_input.messages, parts_manager
+                        )
+                        if snapshot is not None:
+                            yield encoder.encode(snapshot)
                         break
 
                     yield encoder.encode(event)
@@ -285,6 +289,102 @@ async def run(
 
         self.logger.info('done thread_id=%s run_id=%s', run_input.thread_id, run_input.run_id)
 
+    def _message_snapshot(
+        self, run: AgentRun[AgentDepsT, Any], messages: list[Message], parts_manager: ModelResponsePartsManager
+    ) -> MessagesSnapshotEvent | None:
+        """Create a message snapshot to replicate the current state of the run.
+
+        This method collects all messages from the run's state and the parts
+        manager, converting them into AG-UI messages.
+
+        Args:
+            run: The agent run instance.
+            messages: The initial messages from the run input.
+            parts_manager: The parts manager containing the response parts.
+
+        Returns:
+            A full snapshot of the messages so far in the run if local tool
+            calls were made, otherwise `None`.
+        """
+        new_messages: list[ModelMessage] = run.ctx.state.message_history[len(messages) :]
+        if not any(
+            isinstance(request_part, ToolReturnPart)
+            for msg in new_messages
+            if isinstance(msg, ModelRequest)
+            for request_part in msg.parts
+        ):
+            # No tool calls were made, so we don't need a snapshot.
+            return None
+
+        # Tool calls were made, so we need to create a snapshot.
+        for msg in new_messages:
+            match msg:
+                case ModelRequest():
+                    for request_part in msg.parts:
+                        if isinstance(request_part, ToolReturnPart):
+                            messages.append(
+                                ToolMessage(
+                                    id='result-' + request_part.tool_call_id,
+                                    role=Role.TOOL,
+                                    content=request_part.content,
+                                    tool_call_id=request_part.tool_call_id,
+                                )
+                            )
+                case ModelResponse():
+                    self._convert_response_parts(msg.parts, messages)
+
+        self._convert_response_parts(parts_manager.get_parts(), messages)
+
+        return MessagesSnapshotEvent(
+            type=EventType.MESSAGES_SNAPSHOT,
+            messages=messages,
+        )
+
+    def _convert_response_parts(self, parts: list[ModelResponsePart], messages: list[Message]) -> None:
+        """Convert model response parts to AG-UI messages.
+
+        Args:
+            parts: The list of model response parts to convert.
+            messages: The list of messages to append the converted parts to.
+        """
+        response_part: ModelResponsePart
+        for response_part in parts:
+            match response_part:
+                case TextPart():  # pragma: no cover
+                    # This is not expected, but we handle it gracefully.
+                    messages.append(
+                        AssistantMessage(
+                            id=uuid.uuid4().hex,
+                            role=Role.ASSISTANT,
+                            content=response_part.content,
+                        )
+                    )
+                case ToolCallPart():
+                    args: str = (
+                        json.dumps(response_part.args)
+                        if isinstance(response_part.args, dict)
+                        else response_part.args or '{}'
+                    )
+                    messages.append(
+                        AssistantMessage(
+                            id=uuid.uuid4().hex,
+                            role=Role.ASSISTANT,
+                            tool_calls=[
+                                ToolCall(
+                                    id=response_part.tool_call_id,
+                                    type='function',
+                                    function=FunctionCall(
+                                        name=response_part.tool_name,
+                                        arguments=args,
+                                    ),
+                                )
+                            ],
+                        ),
+                    )
+                case ThinkingPart():  # pragma: no cover
+                    # No AG-UI equivalent for thinking parts, so we skip them.
+                    pass
+
     async def _tool_events(self, parts: list[ModelRequestPart]) -> AsyncGenerator[BaseEvent | None, None]:
         """Check for tool call results that are AG-UI events.
 
@@ -371,12 +471,14 @@ async def _agent_stream(
         self,
         tool_names: dict[str, str],
         run: AgentRun[AgentDepsT, Any],
+        parts_manager: ModelResponsePartsManager,
     ) -> AsyncGenerator[BaseEvent | None, None]:
         """Run the agent streaming responses using AG-UI protocol events.
 
         Args:
             tool_names: A mapping of tool names to their AG-UI names.
             run: The agent run to process.
+            parts_manager: The parts manager to handle tool call parts.
 
         Yields:
             AG-UI Server-Sent Events (SSE).
@@ -399,7 +501,7 @@ async def _agent_stream(
             async with node.stream(run.ctx) as request_stream:
                 agent_event: AgentStreamEvent
                 async for agent_event in request_stream:
-                    async for msg in self._handle_agent_event(tool_names, stream_ctx, agent_event):
+                    async for msg in self._handle_agent_event(tool_names, stream_ctx, agent_event, parts_manager):
                         yield msg
 
                 for part_end in stream_ctx.part_ends:
@@ -410,6 +512,7 @@ async def _handle_agent_event(
         tool_names: dict[str, str],
         stream_ctx: _RequestStreamContext,
         agent_event: AgentStreamEvent,
+        parts_manager: ModelResponsePartsManager,
     ) -> AsyncGenerator[BaseEvent | None, None]:
         """Handle an agent event and yield AG-UI protocol events.
 
@@ -418,6 +521,7 @@ async def _handle_agent_event(
             tool_names: A mapping of tool names to their AG-UI names.
             stream_ctx: The request stream context to manage state.
             agent_event: The agent event to process.
+            parts_manager: The parts manager to handle tool call parts.
 
         Yields:
             AG-UI Server-Sent Events (SSE) based on the agent event.
@@ -454,9 +558,16 @@ async def _handle_agent_event(
                     case ToolCallPart():  # pragma: no branch
                         tool_name: str | None = tool_names.get(agent_event.part.tool_name)
                         if not tool_name:
+                            # Local tool calls are not sent as events to the UI.
                             stream_ctx.local_tool_calls.add(agent_event.part.tool_call_id)
                             return
 
+                        parts_manager.handle_tool_call_part(
+                            vendor_part_id=None,
+                            tool_name=agent_event.part.tool_name,
+                            args=agent_event.part.args,
+                            tool_call_id=agent_event.part.tool_call_id,
+                        )
                         stream_ctx.last_tool_call_id = agent_event.part.tool_call_id
                         yield ToolCallStartEvent(
                             type=EventType.TOOL_CALL_START,
@@ -483,9 +594,15 @@ async def _handle_agent_event(
                         )
                     case ToolCallPartDelta():  # pragma: no branch
                         if agent_event.delta.tool_call_id in stream_ctx.local_tool_calls:
-                            # Local tool calls are not sent to the UI.
+                            # Local tool calls are not sent as events to the UI.
                             return
 
+                        parts_manager.handle_tool_call_delta(
+                            vendor_part_id=None,
+                            tool_name=None,
+                            args=agent_event.delta.args_delta,
+                            tool_call_id=agent_event.delta.tool_call_id,
+                        )
                         yield ToolCallArgsEvent(
                             type=EventType.TOOL_CALL_ARGS,
                             tool_call_id=agent_event.delta.tool_call_id

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -45,6 +45,33 @@ class _WrappedToolOutput:
     value: Any | None
 
 
+@dataclass
+class TestToolCallPart:
+    """Represents a tool call in the test model."""
+
+    call_tools: list[str] | Literal['all'] = 'all'
+    deltas: bool = False
+
+
+@dataclass
+class TestTextPart:
+    """Represents a text part in the test model."""
+
+    text: str
+
+
+TestPart = TestTextPart | TestToolCallPart
+"""A part of the test model response."""
+
+
+@dataclass
+class TestNode:
+    """A node in the test model."""
+
+    parts: list[TestPart]
+    id: str = field(default_factory=_utils.generate_tool_call_id)
+
+
 @dataclass
 class TestModel(Model):
     """A model specifically for testing purposes.
@@ -65,6 +92,8 @@ class TestModel(Model):
     """List of tools to call. If `'all'`, all tools will be called."""
     tool_call_deltas: set[str] = field(default_factory=set)
     """A set of tool call names which should result in tool call part deltas."""
+    custom_response_nodes: list[TestNode] | None = None
+    """A list of nodes which defines a custom model response."""
     custom_output_text: str | None = None
     """If set, this text is returned as the final output."""
     custom_output_args: Any | None = None
@@ -155,23 +184,71 @@ def _get_output(self, model_request_parameters: ModelRequestParameters) -> _Wrap
         else:
             return _WrappedTextOutput(None)
 
+    def _node_response(
+        self,
+        messages: list[ModelMessage],
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse | None:
+        """Returns a ModelResponse based on configured nodes.
+
+        Args:
+            messages: The messages sent to the model.
+            model_request_parameters: The parameters for the model request.
+
+        Returns:
+            The response from the model, or `None` if no nodes configured or
+            all nodes have already been processed.
+        """
+        if not self.custom_response_nodes:
+            # No nodes configured, follow the default behaviour.
+            return None
+
+        # Pick up where we left off by counting the number of ModelResponse messages in the stream.
+        # This allows us to stream the response in chunks, simulating a real model response.
+        node: TestNode
+        count: int = sum(isinstance(m, ModelResponse) for m in messages)
+        if count < len(self.custom_response_nodes):
+            node: TestNode = self.custom_response_nodes[count]
+            assert node.parts, 'Node parts should not be empty.'
+
+            parts: list[ModelResponsePart] = []
+            part: TestPart
+            for part in node.parts:
+                if isinstance(part, TestTextPart):
+                    assert model_request_parameters.allow_text_output, (
+                        'Plain response not allowed, but `part` is a `TestText`.'
+                    )
+                    parts.append(TextPart(part.text))
+                elif isinstance(part, TestToolCallPart):
+                    tool_calls = self._get_tool_calls(model_request_parameters)
+                    if part.call_tools == 'all':
+                        parts.extend(ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls)
+                    else:
+                        parts.extend(
+                            ToolCallPart(name, self.gen_tool_args(args))
+                            for name, args in tool_calls
+                            if name in part.call_tools
+                        )
+            return ModelResponse(vendor_id=node.id, parts=parts, model_name=self._model_name)
+
     def _request(
         self,
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
-        tool_calls = self._get_tool_calls(model_request_parameters)
-        output_wrapper = self._get_output(model_request_parameters)
-        output_tools = model_request_parameters.output_tools
+        if (response := self._node_response(messages, model_request_parameters)) is not None:
+            return response
 
-        # if there are tools, the first thing we want to do is call all of them
+        tool_calls = self._get_tool_calls(model_request_parameters)
         if tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
             return ModelResponse(
                 parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls],
                 model_name=self._model_name,
             )
 
+        output_wrapper = self._get_output(model_request_parameters)
+        output_tools = model_request_parameters.output_tools
         if messages:  # pragma: no branch
             last_message = messages[-1]
             assert isinstance(last_message, ModelRequest), 'Expected last message to be a `ModelRequest`.'