Recover streamed agent-tool text on cancellation

Co-authored-by: Codex <noreply@openai.com>

Author: 2 people

src/agents/agent.py

@@ -29,6 +29,7 @@
 from .exceptions import ModelBehaviorError, UserError
 from .guardrail import InputGuardrail, OutputGuardrail
 from .handoffs import Handoff
+from .items import ItemHelpers
 from .logger import logger
 from .mcp import MCPUtil
 from .model_settings import ModelSettings
@@ -118,6 +119,36 @@ def _validate_codex_tool_name_collisions(tools: list[Tool]) -> None:
         )
 
 
+def _recover_streamed_agent_tool_text(run_result: Any) -> str | None:
+    final_output = getattr(run_result, "final_output", None)
+    if isinstance(final_output, str) and final_output:
+        return final_output
+
+    new_items = getattr(run_result, "new_items", None)
+    if isinstance(new_items, list):
+        text = ItemHelpers.text_message_outputs(new_items)
+        if text:
+            return text
+
+    raw_responses = getattr(run_result, "raw_responses", None)
+    if not isinstance(raw_responses, list):
+        return None
+
+    recovered_chunks: list[str] = []
+    for raw_response in raw_responses:
+        outputs = getattr(raw_response, "output", None)
+        if not isinstance(outputs, list):
+            continue
+        for output_item in outputs:
+            text = ItemHelpers.extract_last_text(output_item)
+            if text:
+                recovered_chunks.append(text)
+
+    if not recovered_chunks:
+        return None
+    return "\n\n".join(recovered_chunks)
+
+
 class AgentToolStreamEvent(TypedDict):
     """Streaming event emitted when an agent is invoked as a tool."""
 
@@ -794,16 +825,24 @@ async def dispatch_stream_events() -> None:
                         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:
+                            recovered_text = _recover_streamed_agent_tool_text(
+                                run_result_streaming
+                            )
+                            if not recovered_text:
+                                raise
+                            run_result_streaming.final_output = recovered_text
                     finally:
                         await event_queue.put(None)
                         await event_queue.join()

tests/test_agent_as_tool.py

@@ -17,6 +17,7 @@
     FunctionTool,
     MessageOutputItem,
     ModelBehaviorError,
+    ModelResponse,
     RunConfig,
     RunContextWrapper,
     RunHooks,
@@ -27,6 +28,7 @@
     SessionSettings,
     ToolApprovalItem,
     TResponseInputItem,
+    Usage,
     tool_namespace,
 )
 from agents.agent_tool_input import StructuredToolInputBuilderOptions
@@ -39,6 +41,7 @@
 from agents.run_state import _build_agent_map
 from agents.stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent
 from agents.tool_context import ToolContext
+from tests.test_responses import get_text_message
 from tests.utils.hitl import make_function_tool_call
 
 
@@ -1671,6 +1674,85 @@ async def on_stream(payload: AgentToolStreamEvent) -> None:
     assert callbacks == stream_events
 
 
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_recovers_text_from_raw_responses_on_cancellation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    agent = Agent(name="streamer")
+    stream_event = RawResponsesStreamEvent(data=cast(Any, {"type": "response_started"}))
+
+    class DummyStreamingResult:
+        def __init__(self) -> None:
+            self.final_output = ""
+            self.current_agent = agent
+            self.new_items: list[Any] = []
+            self.raw_responses = [
+                ModelResponse(
+                    output=[get_text_message("Recovered nested summary")],
+                    usage=Usage(),
+                    response_id="resp_nested",
+                )
+            ]
+
+        async def stream_events(self):
+            yield stream_event
+            raise asyncio.CancelledError("stream-cancelled")
+
+    def fake_run_streamed(
+        cls,
+        starting_agent,
+        input,
+        *,
+        context,
+        max_turns,
+        hooks,
+        run_config,
+        previous_response_id,
+        auto_previous_response_id=False,
+        conversation_id,
+        session,
+    ):
+        return DummyStreamingResult()
+
+    async def unexpected_run(*args: Any, **kwargs: Any) -> None:
+        raise AssertionError("Runner.run should not be called when on_stream is provided.")
+
+    monkeypatch.setattr(Runner, "run_streamed", classmethod(fake_run_streamed))
+    monkeypatch.setattr(Runner, "run", classmethod(unexpected_run))
+
+    received_events: list[AgentToolStreamEvent] = []
+
+    async def on_stream(payload: AgentToolStreamEvent) -> None:
+        received_events.append(payload)
+
+    tool_call = ResponseFunctionToolCall(
+        id="call_cancelled",
+        arguments='{"input": "recover"}',
+        call_id="call-cancelled",
+        name="stream_tool",
+        type="function_call",
+    )
+
+    tool = agent.as_tool(
+        tool_name="stream_tool",
+        tool_description="Streams events",
+        on_stream=on_stream,
+    )
+
+    tool_context = ToolContext(
+        context=None,
+        tool_name="stream_tool",
+        tool_call_id=tool_call.call_id,
+        tool_arguments=tool_call.arguments,
+        tool_call=tool_call,
+    )
+    output = await tool.on_invoke_tool(tool_context, '{"input": "recover"}')
+
+    assert output == "Recovered nested summary"
+    assert len(received_events) == 1
+    assert received_events[0]["event"] == stream_event
+
+
 @pytest.mark.asyncio
 async def test_agent_as_tool_streaming_extractor_can_access_agent_tool_invocation(
     monkeypatch: pytest.MonkeyPatch,