Fix initial tool call args not being streamed with AG-UI (pydantic#2303)

Author: DouweM

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -380,21 +380,26 @@ async def _handle_model_request_event(
                 yield TextMessageStartEvent(
                     message_id=message_id,
                 )
-                stream_ctx.part_end = TextMessageEndEvent(
-                    message_id=message_id,
-                )
                 if part.content:  # pragma: no branch
                     yield TextMessageContentEvent(
                         message_id=message_id,
                         delta=part.content,
                     )
+                stream_ctx.part_end = TextMessageEndEvent(
+                    message_id=message_id,
+                )
             elif isinstance(part, ToolCallPart):  # pragma: no branch
                 message_id = stream_ctx.message_id or stream_ctx.new_message_id()
                 yield ToolCallStartEvent(
                     tool_call_id=part.tool_call_id,
                     tool_call_name=part.tool_name,
                     parent_message_id=message_id,
                 )
+                if part.args:
+                    yield ToolCallArgsEvent(
+                        tool_call_id=part.tool_call_id,
+                        delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
+                    )
                 stream_ctx.part_end = ToolCallEndEvent(
                     tool_call_id=part.tool_call_id,
                 )
@@ -403,6 +408,8 @@ async def _handle_model_request_event(
                 yield ThinkingTextMessageStartEvent(
                     type=EventType.THINKING_TEXT_MESSAGE_START,
                 )
+                # Always send the content even if it's empty, as it may be
+                # used to indicate the start of thinking.
                 yield ThinkingTextMessageContentEvent(
                     type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
                     delta=part.content,

tests/test_ag_ui.py

@@ -303,8 +303,8 @@ async def stream_function(
     ) -> AsyncIterator[DeltaToolCalls | str]:
         if len(messages) == 1:
             # First call - make a tool call
-            yield {0: DeltaToolCall(name='get_weather')}
-            yield {0: DeltaToolCall(json_args='{"location": "Paris"}')}
+            yield {0: DeltaToolCall(name='get_weather', json_args='{"location": ')}
+            yield {0: DeltaToolCall(json_args='"Paris"}')}
         else:
             # Second call - return text result
             yield '{"get_weather": "Tool result"}'
@@ -369,8 +369,9 @@ async def stream_function(
             {
                 'type': 'TOOL_CALL_ARGS',
                 'toolCallId': tool_call_id,
-                'delta': '{"location": "Paris"}',
+                'delta': '{"location": ',
             },
+            {'type': 'TOOL_CALL_ARGS', 'toolCallId': tool_call_id, 'delta': '"Paris"}'},
             {'type': 'TOOL_CALL_END', 'toolCallId': tool_call_id},
             {
                 'type': 'RUN_FINISHED',