address code review comments

Author: habema

src/agents/agent.py

@@ -537,10 +537,10 @@ async def dispatch_stream_events() -> None:
 
             return run_result.final_output
 
-        # Set origin tracking on the FunctionTool created by @function_tool
+        # Set origin tracking on run_agent (the FunctionTool returned by @function_tool)
         run_agent._tool_origin = ToolOrigin(
             type=ToolOriginType.AGENT_AS_TOOL,
-            agent_as_tool_name=self.name,
+            agent_as_tool=self,
         )
         return run_agent
 

src/agents/mcp/util.py

@@ -179,7 +179,7 @@ def to_function_tool(
         )
         function_tool._tool_origin = ToolOrigin(
             type=ToolOriginType.MCP,
-            mcp_server_name=server.name,
+            mcp_server=server,
         )
         return function_tool
 

src/agents/tool.py

@@ -48,6 +48,7 @@
 if TYPE_CHECKING:
     from .agent import Agent, AgentBase
     from .items import RunItem
+    from .mcp.server import MCPServer
 
 
 ToolParams = ParamSpec("ToolParams")
@@ -200,19 +201,19 @@ class ToolOrigin:
     type: ToolOriginType
     """The type of tool origin."""
 
-    mcp_server_name: str | None = None
-    """The name of the MCP server. Only set when type is MCP."""
+    mcp_server: MCPServer | None = None
+    """The MCP server object. Only set when type is MCP."""
 
-    agent_as_tool_name: str | None = None
-    """The name of the agent. Only set when type is AGENT_AS_TOOL."""
+    agent_as_tool: Agent[Any] | None = None
+    """The agent object. Only set when type is AGENT_AS_TOOL."""
 
     def __repr__(self) -> str:
         """Custom repr that only includes relevant fields."""
         parts = [f"type={self.type.value!r}"]
-        if self.mcp_server_name is not None:
-            parts.append(f"mcp_server_name={self.mcp_server_name!r}")
-        if self.agent_as_tool_name is not None:
-            parts.append(f"agent_as_tool_name={self.agent_as_tool_name!r}")
+        if self.mcp_server is not None:
+            parts.append(f"mcp_server_name={self.mcp_server.name!r}")
+        if self.agent_as_tool is not None:
+            parts.append(f"agent_as_tool_name={self.agent_as_tool.name!r}")
         return f"ToolOrigin({', '.join(parts)})"
 
 

tests/test_tool_origin.py

@@ -44,14 +44,14 @@ def test_tool(x: int) -> str:
     assert len(tool_call_items) == 1
     assert tool_call_items[0].tool_origin is not None
     assert tool_call_items[0].tool_origin.type == ToolOriginType.FUNCTION
-    assert tool_call_items[0].tool_origin.mcp_server_name is None
-    assert tool_call_items[0].tool_origin.agent_as_tool_name is None
+    assert tool_call_items[0].tool_origin.mcp_server is None
+    assert tool_call_items[0].tool_origin.agent_as_tool is None
 
     assert len(tool_output_items) == 1
     assert tool_output_items[0].tool_origin is not None
     assert tool_output_items[0].tool_origin.type == ToolOriginType.FUNCTION
-    assert tool_output_items[0].tool_origin.mcp_server_name is None
-    assert tool_output_items[0].tool_origin.agent_as_tool_name is None
+    assert tool_output_items[0].tool_origin.mcp_server is None
+    assert tool_output_items[0].tool_origin.agent_as_tool is None
 
 
 @pytest.mark.asyncio
@@ -78,14 +78,16 @@ async def test_mcp_tool_origin():
     assert len(tool_call_items) == 1
     assert tool_call_items[0].tool_origin is not None
     assert tool_call_items[0].tool_origin.type == ToolOriginType.MCP
-    assert tool_call_items[0].tool_origin.mcp_server_name == "test_mcp_server"
-    assert tool_call_items[0].tool_origin.agent_as_tool_name is None
+    assert tool_call_items[0].tool_origin.mcp_server is not None
+    assert tool_call_items[0].tool_origin.mcp_server.name == "test_mcp_server"
+    assert tool_call_items[0].tool_origin.agent_as_tool is None
 
     assert len(tool_output_items) == 1
     assert tool_output_items[0].tool_origin is not None
     assert tool_output_items[0].tool_origin.type == ToolOriginType.MCP
-    assert tool_output_items[0].tool_origin.mcp_server_name == "test_mcp_server"
-    assert tool_output_items[0].tool_origin.agent_as_tool_name is None
+    assert tool_output_items[0].tool_origin.mcp_server is not None
+    assert tool_output_items[0].tool_origin.mcp_server.name == "test_mcp_server"
+    assert tool_output_items[0].tool_origin.agent_as_tool is None
 
 
 @pytest.mark.asyncio
@@ -123,14 +125,16 @@ async def test_agent_as_tool_origin():
     assert len(tool_call_items) == 1
     assert tool_call_items[0].tool_origin is not None
     assert tool_call_items[0].tool_origin.type == ToolOriginType.AGENT_AS_TOOL
-    assert tool_call_items[0].tool_origin.mcp_server_name is None
-    assert tool_call_items[0].tool_origin.agent_as_tool_name == "nested_agent"
+    assert tool_call_items[0].tool_origin.mcp_server is None
+    assert tool_call_items[0].tool_origin.agent_as_tool is not None
+    assert tool_call_items[0].tool_origin.agent_as_tool.name == "nested_agent"
 
     assert len(tool_output_items) == 1
     assert tool_output_items[0].tool_origin is not None
     assert tool_output_items[0].tool_origin.type == ToolOriginType.AGENT_AS_TOOL
-    assert tool_output_items[0].tool_origin.mcp_server_name is None
-    assert tool_output_items[0].tool_origin.agent_as_tool_name == "nested_agent"
+    assert tool_output_items[0].tool_origin.mcp_server is None
+    assert tool_output_items[0].tool_origin.agent_as_tool is not None
+    assert tool_output_items[0].tool_origin.agent_as_tool.name == "nested_agent"
 
 
 @pytest.mark.asyncio
@@ -192,10 +196,12 @@ def func_tool(x: int) -> str:
     assert function_item.tool_origin.type == ToolOriginType.FUNCTION
     assert mcp_item.tool_origin is not None
     assert mcp_item.tool_origin.type == ToolOriginType.MCP
-    assert mcp_item.tool_origin.mcp_server_name == "mcp_server"
+    assert mcp_item.tool_origin.mcp_server is not None
+    assert mcp_item.tool_origin.mcp_server.name == "mcp_server"
     assert agent_item.tool_origin is not None
     assert agent_item.tool_origin.type == ToolOriginType.AGENT_AS_TOOL
-    assert agent_item.tool_origin.agent_as_tool_name == "nested"
+    assert agent_item.tool_origin.agent_as_tool is not None
+    assert agent_item.tool_origin.agent_as_tool.name == "nested"
 
 
 @pytest.mark.asyncio
@@ -229,12 +235,14 @@ async def test_tool_origin_streaming():
     assert len(tool_call_items) == 1
     assert tool_call_items[0].tool_origin is not None
     assert tool_call_items[0].tool_origin.type == ToolOriginType.MCP
-    assert tool_call_items[0].tool_origin.mcp_server_name == "streaming_server"
+    assert tool_call_items[0].tool_origin.mcp_server is not None
+    assert tool_call_items[0].tool_origin.mcp_server.name == "streaming_server"
 
     assert len(tool_output_items) == 1
     assert tool_output_items[0].tool_origin is not None
     assert tool_output_items[0].tool_origin.type == ToolOriginType.MCP
-    assert tool_output_items[0].tool_origin.mcp_server_name == "streaming_server"
+    assert tool_output_items[0].tool_origin.mcp_server is not None
+    assert tool_output_items[0].tool_origin.mcp_server.name == "streaming_server"
 
 
 @pytest.mark.asyncio
@@ -246,12 +254,18 @@ async def test_tool_origin_repr():
     assert "agent_as_tool_name" not in repr(function_origin)
 
     # MCP origin
-    mcp_origin = ToolOrigin(type=ToolOriginType.MCP, mcp_server_name="test_server")
-    assert "mcp_server_name='test_server'" in repr(mcp_origin)
-    assert "agent_as_tool_name" not in repr(mcp_origin)
+    if sys.version_info >= (3, 10):
+        from .mcp.helpers import FakeMCPServer
+
+        test_server = FakeMCPServer(server_name="test_server")
+        mcp_origin = ToolOrigin(type=ToolOriginType.MCP, mcp_server=test_server)
+        assert "mcp_server_name='test_server'" in repr(mcp_origin)
+        assert "agent_as_tool_name" not in repr(mcp_origin)
 
     # AGENT_AS_TOOL origin
-    agent_origin = ToolOrigin(type=ToolOriginType.AGENT_AS_TOOL, agent_as_tool_name="test_agent")
+    model = FakeModel()
+    test_agent = Agent(name="test_agent", model=model, instructions="Test agent")
+    agent_origin = ToolOrigin(type=ToolOriginType.AGENT_AS_TOOL, agent_as_tool=test_agent)
     assert "agent_as_tool_name='test_agent'" in repr(agent_origin)
     assert "mcp_server_name" not in repr(agent_origin)