feat: #864 support streaming nested tool events in Agent.as_tool

src/agents/agent.py

@@ -27,7 +27,7 @@
 if TYPE_CHECKING:
     from .lifecycle import AgentHooks
     from .mcp import MCPServer
-    from .result import RunResult
+    from .result import RunResult, RunResultStreaming
 
 
 @dataclass
@@ -199,7 +199,9 @@ def as_tool(
         self,
         tool_name: str | None,
         tool_description: str | None,
+        *,
         custom_output_extractor: Callable[[RunResult], Awaitable[str]] | None = None,
+        stream_inner_events: bool = False,
     ) -> Tool:
         """Transform this agent into a tool, callable by other agents.
 
@@ -224,17 +226,36 @@ def as_tool(
         async def run_agent(context: RunContextWrapper, input: str) -> str:
             from .run import Runner
 
-            output = await Runner.run(
-                starting_agent=self,
-                input=input,
-                context=context.context,
-            )
+            output_run: RunResult | RunResultStreaming
+            if stream_inner_events:
+                from .stream_events import RunItemStreamEvent
+
+                sub_run = Runner.run_streamed(
+                    self,
+                    input=input,
+                    context=context.context,
+                )
+                parent_queue = getattr(context, "_event_queue", None)
+                async for ev in sub_run.stream_events():
+                    if parent_queue is not None and isinstance(ev, RunItemStreamEvent):
+                        if ev.name in ("tool_called", "tool_output"):
+                            parent_queue.put_nowait(ev)
+                output_run = sub_run
+            else:
+                output_run = await Runner.run(
+                    starting_agent=self,
+                    input=input,
+                    context=context.context,
+                )
+
             if custom_output_extractor:
-                return await custom_output_extractor(output)
+                return await custom_output_extractor(cast(Any, output_run))
 
-            return ItemHelpers.text_message_outputs(output.new_items)
+            return ItemHelpers.text_message_outputs(output_run.new_items)
 
-        return run_agent
+        tool = run_agent
+        tool.stream_inner_events = stream_inner_events
+        return tool
 
     async def get_system_prompt(self, run_context: RunContextWrapper[TContext]) -> str | None:
         """Get the system prompt for the agent."""

src/agents/run.py

@@ -576,6 +576,7 @@ def run_streamed(
             trace=new_trace,
             context_wrapper=context_wrapper,
         )
+        context_wrapper._event_queue = streamed_result._event_queue
 
         # Kick off the actual agent loop in the background and return the streamed result object.
         streamed_result._run_impl_task = asyncio.create_task(

src/agents/run_context.py

@@ -1,5 +1,6 @@
+import asyncio
 from dataclasses import dataclass, field
-from typing import Any, Generic
+from typing import Any, Generic, Optional
 
 from typing_extensions import TypeVar
 
@@ -24,3 +25,5 @@ class RunContextWrapper(Generic[TContext]):
     """The usage of the agent run so far. For streamed responses, the usage will be stale until the
     last chunk of the stream is processed.
     """
+
+    _event_queue: Optional[asyncio.Queue[Any]] = field(default=None, init=False, repr=False)

src/agents/tool.py

@@ -92,6 +92,9 @@ class FunctionTool:
     and returns whether the tool is enabled. You can use this to dynamically enable/disable a tool
     based on your context/state."""
 
+    stream_inner_events: bool = False
+    """Whether to stream inner events when used as an agent tool."""
+
 
 @dataclass
 class FileSearchTool:

src/agents/tool_context.py

@@ -1,5 +1,6 @@
+import asyncio
 from dataclasses import dataclass, field, fields
-from typing import Any
+from typing import Any, Optional
 
 from .run_context import RunContextWrapper, TContext
 
@@ -15,6 +16,8 @@ class ToolContext(RunContextWrapper[TContext]):
     tool_call_id: str = field(default_factory=_assert_must_pass_tool_call_id)
     """The ID of the tool call."""
 
+    _event_queue: Optional[asyncio.Queue[Any]] = field(default=None, init=False, repr=False)
+
     @classmethod
     def from_agent_context(
         cls, context: RunContextWrapper[TContext], tool_call_id: str
@@ -26,4 +29,7 @@ def from_agent_context(
         base_values: dict[str, Any] = {
             f.name: getattr(context, f.name) for f in fields(RunContextWrapper) if f.init
         }
-        return cls(tool_call_id=tool_call_id, **base_values)
+        obj = cls(tool_call_id=tool_call_id, **base_values)
+        if hasattr(context, "_event_queue"):
+            obj._event_queue = context._event_queue
+        return obj