fix review comment

Author: seratch

src/agents/agent.py

@@ -120,10 +120,6 @@ 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)
@@ -140,7 +136,7 @@ def _recover_streamed_agent_tool_text(run_result: Any) -> str | None:
         if not isinstance(outputs, list):
             continue
         for output_item in outputs:
-            chunk = ItemHelpers.extract_last_text(output_item)
+            chunk = ItemHelpers.extract_text(output_item)
             if chunk:
                 recovered_chunks.append(chunk)
 
@@ -149,6 +145,20 @@ def _recover_streamed_agent_tool_text(run_result: Any) -> str | None:
     return "\n\n".join(recovered_chunks)
 
 
+def _finalize_cancelled_streamed_agent_tool_recovery(run_result: Any) -> bool:
+    """Preserve completed outputs and only backfill plain-text results when missing."""
+    final_output = getattr(run_result, "final_output", None)
+    if final_output is not None and not (isinstance(final_output, str) and not final_output):
+        return True
+
+    recovered_text = _recover_streamed_agent_tool_text(run_result)
+    if recovered_text is None:
+        return final_output is not None
+
+    run_result.final_output = recovered_text
+    return True
+
+
 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."""
     run_loop_task = getattr(run_result, "run_loop_task", None)
@@ -852,10 +862,10 @@ async def dispatch_stream_events() -> None:
                         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)
-                            if not recovered_text:
+                            if not _finalize_cancelled_streamed_agent_tool_recovery(
+                                run_result_streaming
+                            ):
                                 raise
-                            run_result_streaming.final_output = recovered_text
                     finally:
                         await event_queue.put(None)
                         await event_queue.join()

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]

tests/test_agent_as_tool.py

@@ -1758,6 +1758,189 @@ async def on_stream(payload: AgentToolStreamEvent) -> None:
     assert received_events[0]["event"] == stream_event
 
 
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_recovers_all_text_chunks_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=[
+                        ResponseOutputMessage(
+                            id="msg_recover",
+                            role="assistant",
+                            status="completed",
+                            type="message",
+                            content=[
+                                ResponseOutputText(
+                                    annotations=[],
+                                    text="first ",
+                                    type="output_text",
+                                    logprobs=[],
+                                ),
+                                ResponseOutputText(
+                                    annotations=[],
+                                    text="second",
+                                    type="output_text",
+                                    logprobs=[],
+                                ),
+                            ],
+                        )
+                    ],
+                    usage=Usage(),
+                    response_id="resp_nested",
+                )
+            ]
+            self.run_loop_task = asyncio.create_task(asyncio.sleep(0))
+
+        async def stream_events(self):
+            yield stream_event
+            raise asyncio.CancelledError("stream-cancelled")
+
+    streaming_result = DummyStreamingResult()
+    await streaming_result.run_loop_task
+
+    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_multi_chunk",
+        arguments='{"input": "recover"}',
+        call_id="call-cancelled-multi-chunk",
+        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 == "first second"
+
+
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_preserves_structured_output_on_cancellation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class StructuredOutput(BaseModel):
+        answer: str
+
+    agent = Agent(name="streamer", output_type=StructuredOutput)
+    stream_event = RawResponsesStreamEvent(data=cast(Any, {"type": "response_started"}))
+
+    class DummyStreamingResult:
+        def __init__(self) -> None:
+            self.final_output = StructuredOutput(answer="structured")
+            self.current_agent = agent
+            self.new_items: list[Any] = []
+            self.raw_responses = [
+                ModelResponse(
+                    output=[get_text_message('{"answer":"text-json"}')],
+                    usage=Usage(),
+                    response_id="resp_nested",
+                )
+            ]
+            self.run_loop_task = asyncio.create_task(asyncio.sleep(0))
+
+        async def stream_events(self):
+            yield stream_event
+            raise asyncio.CancelledError("stream-cancelled")
+
+    streaming_result = DummyStreamingResult()
+    await streaming_result.run_loop_task
+
+    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_structured",
+        arguments='{"input": "recover"}',
+        call_id="call-cancelled-structured",
+        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 == StructuredOutput(answer="structured")
+    assert not isinstance(output, str)
+
+
 @pytest.mark.asyncio
 async def test_agent_as_tool_streaming_reraises_caller_cancellation_with_live_run(
     monkeypatch: pytest.MonkeyPatch,

tests/test_items_helpers.py

@@ -107,6 +107,27 @@ def test_extract_last_text_returns_text_only() -> None:
     assert ItemHelpers.extract_last_text(message2) is None
 
 
+def test_extract_text_concatenates_all_text_segments() -> None:
+    first_text = ResponseOutputText(annotations=[], text="part1", type="output_text", logprobs=[])
+    second_text = ResponseOutputText(annotations=[], text="part2", type="output_text", logprobs=[])
+    refusal = ResponseOutputRefusal(refusal="no", type="refusal")
+    message = make_message([first_text, refusal, second_text])
+
+    assert ItemHelpers.extract_text(message) == "part1part2"
+    assert (
+        ItemHelpers.extract_text(
+            ResponseFunctionToolCall(
+                id="tool123",
+                arguments="{}",
+                call_id="call123",
+                name="func",
+                type="function_call",
+            )
+        )
+        is None
+    )
+
+
 def test_input_to_new_input_list_from_string() -> None:
     result = ItemHelpers.input_to_new_input_list("hi")
     # Should wrap the string into a list with a single dict containing content and user role.