Add agent to ToolContext calls

Author: rm-openai

src/agents/agent.py

@@ -599,6 +599,7 @@ async def _run_agent_impl(context: ToolContext, input_json: str) -> Any:
                     tool_call_id=context.tool_call_id,
                     tool_arguments=context.tool_arguments,
                     tool_call=context.tool_call,
+                    agent=context.agent,
                 )
                 if should_capture_tool_input:
                     nested_context.tool_input = params_data

src/agents/realtime/session.py

@@ -600,6 +600,7 @@ async def _handle_tool_call(
                 tool_name=event.name,
                 tool_call_id=event.call_id,
                 tool_arguments=event.arguments,
+                agent=agent,
             )
             result = await func_tool.on_invoke_tool(tool_context, event.arguments)
 
@@ -626,6 +627,7 @@ async def _handle_tool_call(
                 tool_name=event.name,
                 tool_call_id=event.call_id,
                 tool_arguments=event.arguments,
+                agent=agent,
             )
 
             # Execute the handoff to get the new agent

src/agents/run_internal/tool_execution.py

@@ -820,6 +820,7 @@ async def run_single_tool(func_tool: FunctionTool, tool_call: ResponseFunctionTo
                 context_wrapper,
                 tool_call.call_id,
                 tool_call=tool_call,
+                agent=agent,
             )
             agent_hooks = agent.hooks
             if config.trace_include_sensitive_data:

src/agents/tool_context.py

@@ -9,6 +9,7 @@
 from .usage import Usage
 
 if TYPE_CHECKING:
+    from .agent import AgentBase
     from .items import TResponseInputItem
     from .run_context import _ApprovalRecord
 
@@ -44,6 +45,9 @@ class ToolContext(RunContextWrapper[TContext]):
     tool_call: ResponseFunctionToolCall | None = None
     """The tool call object associated with this invocation."""
 
+    agent: AgentBase[Any] | None = None
+    """The active agent for this tool call, when available."""
+
     def __init__(
         self,
         context: TContext,
@@ -52,6 +56,7 @@ def __init__(
         tool_call_id: str | object = _MISSING,
         tool_arguments: str | object = _MISSING,
         tool_call: ResponseFunctionToolCall | None = None,
+        agent: AgentBase[Any] | None = None,
         *,
         turn_input: list[TResponseInputItem] | None = None,
         _approvals: dict[str, _ApprovalRecord] | None = None,
@@ -80,13 +85,15 @@ def __init__(
             else cast(str, tool_call_id)
         )
         self.tool_call = tool_call
+        self.agent = agent
 
     @classmethod
     def from_agent_context(
         cls,
         context: RunContextWrapper[TContext],
         tool_call_id: str,
         tool_call: ResponseFunctionToolCall | None = None,
+        agent: AgentBase[Any] | None = None,
     ) -> ToolContext:
         """
         Create a ToolContext from a RunContextWrapper.
@@ -99,12 +106,16 @@ def from_agent_context(
         tool_args = (
             tool_call.arguments if tool_call is not None else _assert_must_pass_tool_arguments()
         )
+        tool_agent = agent
+        if tool_agent is None and isinstance(context, ToolContext):
+            tool_agent = context.agent
 
         tool_context = cls(
             tool_name=tool_name,
             tool_call_id=tool_call_id,
             tool_arguments=tool_args,
             tool_call=tool_call,
+            agent=tool_agent,
             **base_values,
         )
         return tool_context

tests/realtime/test_session.py

@@ -989,6 +989,7 @@ async def test_function_tool_execution_success(
         call_args = mock_function_tool.on_invoke_tool.call_args
         tool_context = call_args[0][0]
         assert isinstance(tool_context, ToolContext)
+        assert tool_context.agent == mock_agent
         assert call_args[0][1] == '{"param": "value"}'
 
         # Verify tool output was sent to model

tests/test_agent_runner.py

@@ -63,6 +63,7 @@
 from agents.run_internal.tool_use_tracker import AgentToolUseTracker
 from agents.run_state import RunState
 from agents.tool import ComputerTool, FunctionToolResult, function_tool
+from agents.tool_context import ToolContext
 from agents.usage import Usage
 
 from .fake_model import FakeModel
@@ -437,6 +438,36 @@ async def test_tool_call_runs():
     )
 
 
+@pytest.mark.asyncio
+async def test_tool_call_context_includes_current_agent() -> None:
+    model = FakeModel()
+    captured_contexts: list[ToolContext[Any]] = []
+
+    @function_tool(name_override="foo")
+    def foo(context: ToolContext[Any]) -> str:
+        captured_contexts.append(context)
+        return "tool_result"
+
+    agent = Agent(
+        name="test",
+        model=model,
+        tools=[foo],
+    )
+
+    model.add_multiple_turn_outputs(
+        [
+            [get_function_tool_call("foo", "{}")],
+            [get_text_message("done")],
+        ]
+    )
+
+    result = await Runner.run(agent, input="user_message")
+
+    assert result.final_output == "done"
+    assert len(captured_contexts) == 1
+    assert captured_contexts[0].agent is agent
+
+
 @pytest.mark.asyncio
 async def test_handoffs():
     model = FakeModel()

tests/test_source_compat_constructors.py

@@ -83,6 +83,19 @@ def test_tool_context_v070_positional_constructor_still_works() -> None:
     assert context.tool_name == "tool_name"
     assert context.tool_call_id == "call_id"
     assert context.tool_arguments == '{"x":1}'
+    assert context.agent is None
+
+
+def test_tool_context_supports_appended_agent_positional_argument() -> None:
+    usage = Usage()
+    agent = Agent(name="agent")
+    context = ToolContext(None, usage, "tool_name", "call_id", '{"x":1}', None, agent)
+
+    assert context.usage is usage
+    assert context.tool_name == "tool_name"
+    assert context.tool_call_id == "call_id"
+    assert context.tool_arguments == '{"x":1}'
+    assert context.agent is agent
 
 
 def test_run_result_v070_positional_constructor_still_works() -> None:

tests/test_tool_context.py

@@ -1,6 +1,7 @@
 import pytest
 from openai.types.responses import ResponseFunctionToolCall
 
+from agents import Agent
 from agents.run_context import RunContextWrapper
 from agents.tool_context import ToolContext
 from tests.utils.hitl import make_context_wrapper
@@ -30,9 +31,16 @@ def test_tool_context_from_agent_context_populates_fields() -> None:
         arguments='{"a": 1}',
     )
     ctx = make_context_wrapper()
+    agent = Agent(name="agent")
 
-    tool_ctx = ToolContext.from_agent_context(ctx, tool_call_id="call-123", tool_call=tool_call)
+    tool_ctx = ToolContext.from_agent_context(
+        ctx,
+        tool_call_id="call-123",
+        tool_call=tool_call,
+        agent=agent,
+    )
 
     assert tool_ctx.tool_name == "test_tool"
     assert tool_ctx.tool_call_id == "call-123"
     assert tool_ctx.tool_arguments == '{"a": 1}'
+    assert tool_ctx.agent is agent