Recover streamed nested agent output before cancellation fallback (#2714)

Co-authored-by: Kazuhiro Sera <seratch@openai.com>

Author: elainegan-openai

src/agents/agent.py

@@ -789,25 +789,37 @@ async def dispatch_stream_events() -> None:
                                 break
 
                     dispatch_task = asyncio.create_task(dispatch_stream_events())
+                    stream_iteration_cancelled = False
 
                     try:
                         from .stream_events import AgentUpdatedStreamEvent
 
                         current_agent = run_result_streaming.current_agent
-                        async for event in run_result_streaming.stream_events():
-                            if isinstance(event, AgentUpdatedStreamEvent):
-                                current_agent = event.new_agent
-
-                            payload: AgentToolStreamEvent = {
-                                "event": event,
-                                "agent": current_agent,
-                                "tool_call": context.tool_call,
-                            }
-                            await event_queue.put(payload)
+                        try:
+                            async for event in run_result_streaming.stream_events():
+                                if isinstance(event, AgentUpdatedStreamEvent):
+                                    current_agent = event.new_agent
+
+                                payload: AgentToolStreamEvent = {
+                                    "event": event,
+                                    "agent": current_agent,
+                                    "tool_call": context.tool_call,
+                                }
+                                await event_queue.put(payload)
+                        except asyncio.CancelledError:
+                            stream_iteration_cancelled = True
+                            raise
                     finally:
-                        await event_queue.put(None)
-                        await event_queue.join()
-                        await dispatch_task
+                        if stream_iteration_cancelled:
+                            dispatch_task.cancel()
+                            try:
+                                await dispatch_task
+                            except asyncio.CancelledError:
+                                pass
+                        else:
+                            await event_queue.put(None)
+                            await event_queue.join()
+                            await dispatch_task
                     run_result = run_result_streaming
                 else:
                     run_result = await Runner.run(

src/agents/items.py

@@ -675,6 +675,19 @@ def extract_last_text(cls, message: TResponseOutputItem) -> str | None:
 
         return None
 
+    @classmethod
+    def extract_text(cls, message: TResponseOutputItem) -> str | None:
+        """Extracts all text content from a message, if any. Ignores refusals."""
+        if not isinstance(message, ResponseOutputMessage):
+            return None
+
+        text = ""
+        for content_item in message.content:
+            if isinstance(content_item, ResponseOutputText):
+                text += content_item.text
+
+        return text or None
+
     @classmethod
     def input_to_new_input_list(
         cls, input: str | list[TResponseInputItem]

src/agents/run_internal/turn_resolution.py

@@ -615,7 +615,7 @@ async def execute_tools_and_side_effects(
 
     message_items = [item for item in new_step_items if isinstance(item, MessageOutputItem)]
     potential_final_output_text = (
-        ItemHelpers.extract_last_text(message_items[-1].raw_item) if message_items else None
+        ItemHelpers.extract_text(message_items[-1].raw_item) if message_items else None
     )
 
     if not processed_response.has_tools_or_approvals_to_run():