Preserve caller cancellation in agent tool streaming

Only recover streamed nested-agent text once the nested run is already locally terminal, and add a regression that caller cancellation still re-raises even if raw responses already contain text.

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

Author: 2 people

src/agents/agent.py

@@ -149,6 +149,18 @@ def _recover_streamed_agent_tool_text(run_result: Any) -> str | None:
     return "\n\n".join(recovered_chunks)
 
 
+def _can_recover_cancelled_streamed_agent_tool(run_result: Any) -> bool:
+    """Recover text only after the nested streamed run has already reached local terminal state."""
+    if getattr(run_result, "is_complete", False):
+        return True
+
+    run_loop_task = getattr(run_result, "run_loop_task", None)
+    if isinstance(run_loop_task, asyncio.Task):
+        return run_loop_task.done()
+
+    return False
+
+
 class AgentToolStreamEvent(TypedDict):
     """Streaming event emitted when an agent is invoked as a tool."""
 
@@ -837,6 +849,8 @@ async def dispatch_stream_events() -> None:
                                 }
                                 await event_queue.put(payload)
                         except asyncio.CancelledError:
+                            if not _can_recover_cancelled_streamed_agent_tool(run_result_streaming):
+                                raise
                             recovered_text = _recover_streamed_agent_tool_text(
                                 run_result_streaming
                             )

tests/test_agent_as_tool.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import dataclasses
 import json
 from typing import Any, cast
@@ -1693,6 +1694,7 @@ def __init__(self) -> None:
                     response_id="resp_nested",
                 )
             ]
+            self.is_complete = True
 
         async def stream_events(self):
             yield stream_event
@@ -1753,6 +1755,89 @@ async def on_stream(payload: AgentToolStreamEvent) -> None:
     assert received_events[0]["event"] == stream_event
 
 
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_reraises_caller_cancellation_with_live_run(
+    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",
+                )
+            ]
+            self.is_complete = False
+            self.run_loop_task = asyncio.create_task(asyncio.sleep(60))
+
+        async def stream_events(self):
+            yield stream_event
+            raise asyncio.CancelledError("stream-cancelled")
+
+    streaming_result = DummyStreamingResult()
+
+    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 streaming_result
+
+    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))
+
+    async def on_stream(payload: AgentToolStreamEvent) -> None:
+        del payload
+
+    tool_call = ResponseFunctionToolCall(
+        id="call_cancelled_live",
+        arguments='{"input": "recover"}',
+        call_id="call-cancelled-live",
+        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,
+    )
+    try:
+        with pytest.raises(asyncio.CancelledError, match="stream-cancelled"):
+            await tool.on_invoke_tool(tool_context, '{"input": "recover"}')
+    finally:
+        streaming_result.run_loop_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await streaming_result.run_loop_task
+
+
 @pytest.mark.asyncio
 async def test_agent_as_tool_streaming_extractor_can_access_agent_tool_invocation(
     monkeypatch: pytest.MonkeyPatch,