fix memory leak in code review and add test

Author: habema

src/agents/items.py

@@ -252,6 +252,12 @@ class ToolCallItem(RunItemBase[Any]):
     tool_origin: ToolOrigin | None = field(default=None, repr=False)
     """Information about the origin/source of the tool call. Only set for FunctionTool calls."""
 
+    def release_agent(self) -> None:
+        """Release agent references including tool_origin.agent_as_tool."""
+        super().release_agent()
+        if self.tool_origin is not None:
+            self.tool_origin.release_agent()
+
 
 ToolCallOutputTypes: TypeAlias = Union[
     FunctionCallOutput,
@@ -278,6 +284,12 @@ class ToolCallOutputItem(RunItemBase[Any]):
     tool_origin: ToolOrigin | None = field(default=None, repr=False)
     """Information about the origin/source of the tool call. Only set for FunctionTool calls."""
 
+    def release_agent(self) -> None:
+        """Release agent references including tool_origin.agent_as_tool."""
+        super().release_agent()
+        if self.tool_origin is not None:
+            self.tool_origin.release_agent()
+
     def to_input_item(self) -> TResponseInputItem:
         """Converts the tool output into an input item for the next model turn.
 

src/agents/tool.py

@@ -210,13 +210,50 @@ class ToolOrigin:
     agent_as_tool: Agent[Any] | None = None
     """The agent object. Only set when type is AGENT_AS_TOOL."""
 
+    _agent_as_tool_ref: weakref.ReferenceType[Agent[Any]] | None = field(
+        default=None, init=False, repr=False
+    )
+    """Weak reference to agent_as_tool for memory management."""
+
+    def __post_init__(self) -> None:
+        """Initialize weak reference for agent_as_tool."""
+        if self.agent_as_tool is not None:
+            self._agent_as_tool_ref = weakref.ref(self.agent_as_tool)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Lazily resolve agent_as_tool via weakref when strong ref is cleared."""
+        if name == "agent_as_tool":
+            # Check if strong reference still exists
+            value = object.__getattribute__(self, "__dict__").get("agent_as_tool")
+            if value is not None:
+                return value
+            # Try to resolve via weakref
+            ref = object.__getattribute__(self, "_agent_as_tool_ref")
+            if ref is not None:
+                agent = ref()
+                if agent is not None:
+                    return agent
+            return None
+        return super().__getattribute__(name)
+
+    def release_agent(self) -> None:
+        """Release the strong reference to agent_as_tool while keeping a weak reference."""
+        if "agent_as_tool" not in self.__dict__:
+            return
+        agent = self.__dict__.get("agent_as_tool")
+        if agent is not None:
+            self._agent_as_tool_ref = weakref.ref(agent)
+        # Set to None instead of deleting so dataclass repr/asdict keep working.
+        self.__dict__["agent_as_tool"] = None
+
     def __repr__(self) -> str:
         """Custom repr that only includes relevant fields."""
         parts = [f"type={self.type.value!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}")
+        agent = self.agent_as_tool
+        if agent is not None:
+            parts.append(f"agent_as_tool_name={agent.name!r}")
         return f"ToolOrigin({', '.join(parts)})"
 
 

tests/test_tool_origin.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import gc
 import sys
+import weakref
 from typing import cast
 
 import pytest
@@ -331,3 +333,54 @@ def func_tool() -> str:
     )
 
     assert item.tool_origin is None
+
+
+def test_tool_origin_release_agent_clears_strong_reference():
+    """Test that release_agent() clears strong reference to agent_as_tool."""
+    # Create a ToolOrigin with an agent_as_tool
+    nested_agent = Agent(
+        name="nested_agent",
+        model=FakeModel(),
+        instructions="You are a nested agent.",
+    )
+
+    tool_origin = ToolOrigin(
+        type=ToolOriginType.AGENT_AS_TOOL,
+        agent_as_tool=nested_agent,
+    )
+
+    # Create a ToolCallItem with this tool_origin
+    tool_call_item = ToolCallItem(
+        raw_item=cast(
+            ToolCallItemTypes,
+            {
+                "type": "function_call",
+                "name": "test_tool",
+                "call_id": "call_123",
+                "arguments": "{}",
+            },
+        ),
+        agent=nested_agent,
+        tool_origin=tool_origin,
+    )
+
+    # Verify agent_as_tool is set
+    assert tool_call_item.tool_origin is not None
+    assert tool_call_item.tool_origin.agent_as_tool is nested_agent
+
+    # Create weak reference to verify GC behavior
+    nested_agent_ref = weakref.ref(nested_agent)
+
+    # Release agent - this should clear strong reference in tool_origin
+    tool_call_item.release_agent()
+
+    # After release, agent_as_tool should still be accessible via weakref
+    assert tool_call_item.tool_origin.agent_as_tool is nested_agent
+
+    # Delete the agent and force GC
+    del nested_agent
+    gc.collect()
+
+    # After GC, agent_as_tool should be None since strong refs were cleared
+    assert nested_agent_ref() is None
+    assert tool_call_item.tool_origin.agent_as_tool is None