Add Agent to ToolContext

Author: rm-openai

src/agents/tool_context.py

@@ -56,8 +56,8 @@ def __init__(
         tool_call_id: str | object = _MISSING,
         tool_arguments: str | object = _MISSING,
         tool_call: ResponseFunctionToolCall | None = None,
-        agent: AgentBase[Any] | None = None,
         *,
+        agent: AgentBase[Any] | None = None,
         turn_input: list[TResponseInputItem] | None = None,
         _approvals: dict[str, _ApprovalRecord] | None = None,
         tool_input: Any | None = None,

tests/test_source_compat_constructors.py

@@ -86,10 +86,10 @@ def test_tool_context_v070_positional_constructor_still_works() -> None:
     assert context.agent is None
 
 
-def test_tool_context_supports_appended_agent_positional_argument() -> None:
+def test_tool_context_supports_agent_keyword_argument() -> None:
     usage = Usage()
     agent = Agent(name="agent")
-    context = ToolContext(None, usage, "tool_name", "call_id", '{"x":1}', None, agent)
+    context = ToolContext(None, usage, "tool_name", "call_id", '{"x":1}', None, agent=agent)
 
     assert context.usage is usage
     assert context.tool_name == "tool_name"

tests/test_tool_context.py

@@ -44,3 +44,62 @@ def test_tool_context_from_agent_context_populates_fields() -> None:
     assert tool_ctx.tool_call_id == "call-123"
     assert tool_ctx.tool_arguments == '{"a": 1}'
     assert tool_ctx.agent is agent
+
+
+def test_tool_context_agent_none_by_default() -> None:
+    tool_call = ResponseFunctionToolCall(
+        type="function_call",
+        name="test_tool",
+        call_id="call-1",
+        arguments="{}",
+    )
+    ctx = make_context_wrapper()
+
+    tool_ctx = ToolContext.from_agent_context(ctx, tool_call_id="call-1", tool_call=tool_call)
+
+    assert tool_ctx.agent is None
+
+
+def test_tool_context_constructor_accepts_agent_keyword() -> None:
+    agent = Agent(name="direct-agent")
+    tool_ctx: ToolContext[dict[str, object]] = ToolContext(
+        context={},
+        tool_name="my_tool",
+        tool_call_id="call-2",
+        tool_arguments="{}",
+        agent=agent,
+    )
+
+    assert tool_ctx.agent is agent
+
+
+def test_tool_context_from_tool_context_inherits_agent() -> None:
+    original_call = ResponseFunctionToolCall(
+        type="function_call",
+        name="test_tool",
+        call_id="call-3",
+        arguments="{}",
+    )
+    derived_call = ResponseFunctionToolCall(
+        type="function_call",
+        name="test_tool",
+        call_id="call-4",
+        arguments="{}",
+    )
+    agent = Agent(name="origin-agent")
+    parent_context: ToolContext[dict[str, object]] = ToolContext(
+        context={},
+        tool_name="test_tool",
+        tool_call_id="call-3",
+        tool_arguments="{}",
+        tool_call=original_call,
+        agent=agent,
+    )
+
+    derived_context = ToolContext.from_agent_context(
+        parent_context,
+        tool_call_id="call-4",
+        tool_call=derived_call,
+    )
+
+    assert derived_context.agent is agent