feat: Add on_stream to agents as tools (#2169)

Author: seratch

examples/agent_patterns/README.md

@@ -28,6 +28,7 @@ The mental model for handoffs is that the new agent "takes over". It sees the pr
 For example, you could model the translation task above as tool calls instead: rather than handing over to the language-specific agent, you could call the agent as a tool, and then use the result in the next step. This enables things like translating multiple languages at once.
 
 See the [`agents_as_tools.py`](./agents_as_tools.py) file for an example of this.
+See the [`agents_as_tools_streaming.py`](./agents_as_tools_streaming.py) file for a streaming variant that taps into nested agent events via `on_stream`.
 
 ## LLM-as-a-judge
 

examples/agent_patterns/agents_as_tools_streaming.py

@@ -0,0 +1,59 @@
+import asyncio
+
+from agents import Agent, AgentToolStreamEvent, ModelSettings, Runner, function_tool, trace
+
+
+@function_tool(
+    name_override="billing_status_checker",
+    description_override="Answer questions about customer billing status.",
+)
+def billing_status_checker(customer_id: str | None = None, question: str = "") -> str:
+    """Return a canned billing answer or a fallback when the question is unrelated."""
+    normalized = question.lower()
+    if "bill" in normalized or "billing" in normalized:
+        return f"This customer (ID: {customer_id})'s bill is $100"
+    return "I can only answer questions about billing."
+
+
+def handle_stream(event: AgentToolStreamEvent) -> None:
+    """Print streaming events emitted by the nested billing agent."""
+    stream = event["event"]
+    tool_call = event.get("tool_call")
+    tool_call_info = tool_call.call_id if tool_call is not None else "unknown"
+    print(f"[stream] agent={event['agent'].name} call={tool_call_info} type={stream.type} {stream}")
+
+
+async def main() -> None:
+    with trace("Agents as tools streaming example"):
+        billing_agent = Agent(
+            name="Billing Agent",
+            instructions="You are a billing agent that answers billing questions.",
+            model_settings=ModelSettings(tool_choice="required"),
+            tools=[billing_status_checker],
+        )
+
+        billing_agent_tool = billing_agent.as_tool(
+            tool_name="billing_agent",
+            tool_description="You are a billing agent that answers billing questions.",
+            on_stream=handle_stream,
+        )
+
+        main_agent = Agent(
+            name="Customer Support Agent",
+            instructions=(
+                "You are a customer support agent. Always call the billing agent to answer billing "
+                "questions and return the billing agent response to the user."
+            ),
+            tools=[billing_agent_tool],
+        )
+
+        result = await Runner.run(
+            main_agent,
+            "Hello, my customer ID is ABC123. How much is my bill for this month?",
+        )
+
+    print(f"\nFinal response:\n{result.final_output}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/financial_research_agent/manager.py

@@ -6,7 +6,7 @@
 
 from rich.console import Console
 
-from agents import Runner, RunResult, custom_span, gen_trace_id, trace
+from agents import Runner, RunResult, RunResultStreaming, custom_span, gen_trace_id, trace
 
 from .agents.financials_agent import financials_agent
 from .agents.planner_agent import FinancialSearchItem, FinancialSearchPlan, planner_agent
@@ -17,7 +17,7 @@
 from .printer import Printer
 
 
-async def _summary_extractor(run_result: RunResult) -> str:
+async def _summary_extractor(run_result: RunResult | RunResultStreaming) -> str:
     """Custom output extractor for sub‑agents that return an AnalysisSummary."""
     # The financial/risk analyst agents emit an AnalysisSummary with a `summary` field.
     # We want the tool call to return just that summary text so the writer can drop it inline.

src/agents/__init__.py

@@ -8,6 +8,7 @@
 from .agent import (
     Agent,
     AgentBase,
+    AgentToolStreamEvent,
     StopAtTools,
     ToolsToFinalOutputFunction,
     ToolsToFinalOutputResult,
@@ -214,6 +215,7 @@ def enable_verbose_stdout_logging():
 __all__ = [
     "Agent",
     "AgentBase",
+    "AgentToolStreamEvent",
     "StopAtTools",
     "ToolsToFinalOutputFunction",
     "ToolsToFinalOutputResult",

src/agents/agent.py

@@ -25,15 +25,19 @@
 from .prompts import DynamicPromptFunction, Prompt, PromptUtil
 from .run_context import RunContextWrapper, TContext
 from .tool import FunctionTool, FunctionToolResult, Tool, function_tool
+from .tool_context import ToolContext
 from .util import _transforms
 from .util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
+    from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
+
     from .lifecycle import AgentHooks, RunHooks
     from .mcp import MCPServer
     from .memory.session import Session
-    from .result import RunResult
+    from .result import RunResult, RunResultStreaming
     from .run import RunConfig
+    from .stream_events import StreamEvent
 
 
 @dataclass
@@ -58,6 +62,19 @@ class ToolsToFinalOutputResult:
 """
 
 
+class AgentToolStreamEvent(TypedDict):
+    """Streaming event emitted when an agent is invoked as a tool."""
+
+    event: StreamEvent
+    """The streaming event from the nested agent run."""
+
+    agent: Agent[Any]
+    """The nested agent emitting the event."""
+
+    tool_call: ResponseFunctionToolCall | None
+    """The originating tool call, if available."""
+
+
 class StopAtTools(TypedDict):
     stop_at_tool_names: list[str]
     """A list of tool names, any of which will stop the agent from running further."""
@@ -382,9 +399,12 @@ def as_tool(
         self,
         tool_name: str | None,
         tool_description: str | None,
-        custom_output_extractor: Callable[[RunResult], Awaitable[str]] | None = None,
+        custom_output_extractor: (
+            Callable[[RunResult | RunResultStreaming], Awaitable[str]] | None
+        ) = None,
         is_enabled: bool
         | Callable[[RunContextWrapper[Any], AgentBase[Any]], MaybeAwaitable[bool]] = True,
+        on_stream: Callable[[AgentToolStreamEvent], MaybeAwaitable[None]] | None = None,
         run_config: RunConfig | None = None,
         max_turns: int | None = None,
         hooks: RunHooks[TContext] | None = None,
@@ -409,33 +429,100 @@ def as_tool(
             is_enabled: Whether the tool is enabled. Can be a bool or a callable that takes the run
                 context and agent and returns whether the tool is enabled. Disabled tools are hidden
                 from the LLM at runtime.
+            on_stream: Optional callback (sync or async) to receive streaming events from the nested
+                agent run. The callback receives an `AgentToolStreamEvent` containing the nested
+                agent, the originating tool call (when available), and each stream event. When
+                provided, the nested agent is executed in streaming mode.
         """
 
         @function_tool(
             name_override=tool_name or _transforms.transform_string_function_style(self.name),
             description_override=tool_description or "",
             is_enabled=is_enabled,
         )
-        async def run_agent(context: RunContextWrapper, input: str) -> Any:
+        async def run_agent(context: ToolContext, input: str) -> Any:
             from .run import DEFAULT_MAX_TURNS, Runner
 
             resolved_max_turns = max_turns if max_turns is not None else DEFAULT_MAX_TURNS
-
-            output = await Runner.run(
-                starting_agent=self,
-                input=input,
-                context=context.context,
-                run_config=run_config,
-                max_turns=resolved_max_turns,
-                hooks=hooks,
-                previous_response_id=previous_response_id,
-                conversation_id=conversation_id,
-                session=session,
-            )
+            run_result: RunResult | RunResultStreaming
+
+            if on_stream is not None:
+                run_result = Runner.run_streamed(
+                    starting_agent=self,
+                    input=input,
+                    context=context.context,
+                    run_config=run_config,
+                    max_turns=resolved_max_turns,
+                    hooks=hooks,
+                    previous_response_id=previous_response_id,
+                    conversation_id=conversation_id,
+                    session=session,
+                )
+                # Dispatch callbacks in the background so slow handlers do not block
+                # event consumption.
+                event_queue: asyncio.Queue[AgentToolStreamEvent | None] = asyncio.Queue()
+
+                async def _run_handler(payload: AgentToolStreamEvent) -> None:
+                    """Execute the user callback while capturing exceptions."""
+                    try:
+                        maybe_result = on_stream(payload)
+                        if inspect.isawaitable(maybe_result):
+                            await maybe_result
+                    except Exception:
+                        logger.exception(
+                            "Error while handling on_stream event for agent tool %s.",
+                            self.name,
+                        )
+
+                async def dispatch_stream_events() -> None:
+                    while True:
+                        payload = await event_queue.get()
+                        is_sentinel = payload is None  # None marks the end of the stream.
+                        try:
+                            if payload is not None:
+                                await _run_handler(payload)
+                        finally:
+                            event_queue.task_done()
+
+                        if is_sentinel:
+                            break
+
+                dispatch_task = asyncio.create_task(dispatch_stream_events())
+
+                try:
+                    from .stream_events import AgentUpdatedStreamEvent
+
+                    current_agent = run_result.current_agent
+                    async for event in run_result.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)
+                finally:
+                    await event_queue.put(None)
+                    await event_queue.join()
+                    await dispatch_task
+            else:
+                run_result = await Runner.run(
+                    starting_agent=self,
+                    input=input,
+                    context=context.context,
+                    run_config=run_config,
+                    max_turns=resolved_max_turns,
+                    hooks=hooks,
+                    previous_response_id=previous_response_id,
+                    conversation_id=conversation_id,
+                    session=session,
+                )
             if custom_output_extractor:
-                return await custom_output_extractor(output)
+                return await custom_output_extractor(run_result)
 
-            return output.final_output
+            return run_result.final_output
 
         return run_agent
 

src/agents/tool_context.py

@@ -31,6 +31,9 @@ class ToolContext(RunContextWrapper[TContext]):
     tool_arguments: str = field(default_factory=_assert_must_pass_tool_arguments)
     """The raw arguments string of the tool call."""
 
+    tool_call: Optional[ResponseFunctionToolCall] = None
+    """The tool call object associated with this invocation."""
+
     @classmethod
     def from_agent_context(
         cls,
@@ -50,6 +53,11 @@ def from_agent_context(
             tool_call.arguments if tool_call is not None else _assert_must_pass_tool_arguments()
         )
 
-        return cls(
-            tool_name=tool_name, tool_call_id=tool_call_id, tool_arguments=tool_args, **base_values
+        tool_context = cls(
+            tool_name=tool_name,
+            tool_call_id=tool_call_id,
+            tool_arguments=tool_args,
+            tool_call=tool_call,
+            **base_values,
         )
+        return tool_context