From bc352ccef085cef39280ce0b19ffbf7678199348 Mon Sep 17 00:00:00 2001 From: Daniele Morotti <58258368+DanieleMorotti@users.noreply.github.com> Date: Thu, 30 Oct 2025 12:16:09 +0100 Subject: [PATCH] Fix agent memory leak with weakref --- src/agents/items.py | 45 ++++++++++++++++++++++++++++++++- tests/test_agent_memory_leak.py | 34 +++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) create mode 100644 tests/test_agent_memory_leak.py diff --git a/src/agents/items.py b/src/agents/items.py index 8e7d1cfc3..96f9577a9 100644 --- a/src/agents/items.py +++ b/src/agents/items.py @@ -1,7 +1,8 @@ from __future__ import annotations import abc -from dataclasses import dataclass +import weakref +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, Union import pydantic @@ -84,6 +85,22 @@ class RunItemBase(Generic[T], abc.ABC): (i.e. `openai.types.responses.ResponseInputItemParam`). """ + _agent_ref: weakref.ReferenceType[Agent[Any]] | None = field( + init=False, + repr=False, + default=None, + ) + + def __post_init__(self) -> None: + # Store the producing agent weakly to avoid keeping it alive after the run. + self._agent_ref = weakref.ref(self.agent) + object.__delattr__(self, "agent") + + def __getattr__(self, name: str) -> Any: + if name == "agent": + return self._agent_ref() if self._agent_ref else None + raise AttributeError(name) + def to_input_item(self) -> TResponseInputItem: """Converts this item into an input item suitable for passing to the model.""" if isinstance(self.raw_item, dict): @@ -131,6 +148,32 @@ class HandoffOutputItem(RunItemBase[TResponseInputItem]): type: Literal["handoff_output_item"] = "handoff_output_item" + _source_agent_ref: weakref.ReferenceType[Agent[Any]] | None = field( + init=False, + repr=False, + default=None, + ) + _target_agent_ref: weakref.ReferenceType[Agent[Any]] | None = field( + init=False, + repr=False, + default=None, + ) + + def __post_init__(self) -> None: + super().__post_init__() + # Handoff metadata should not hold strong references to the agents either. + self._source_agent_ref = weakref.ref(self.source_agent) + self._target_agent_ref = weakref.ref(self.target_agent) + object.__delattr__(self, "source_agent") + object.__delattr__(self, "target_agent") + + def __getattr__(self, name: str) -> Any: + if name == "source_agent": + return self._source_agent_ref() if self._source_agent_ref else None + if name == "target_agent": + return self._target_agent_ref() if self._target_agent_ref else None + return super().__getattr__(name) + ToolCallItemTypes: TypeAlias = Union[ ResponseFunctionToolCall, diff --git a/tests/test_agent_memory_leak.py b/tests/test_agent_memory_leak.py new file mode 100644 index 000000000..6690d2189 --- /dev/null +++ b/tests/test_agent_memory_leak.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import gc +import weakref + +import pytest +from openai.types.responses import ResponseOutputMessage, ResponseOutputText + +from agents import Agent, Runner +from tests.fake_model import FakeModel + + +def _make_message(text: str) -> ResponseOutputMessage: + return ResponseOutputMessage( + id="msg-1", + content=[ResponseOutputText(annotations=[], text=text, type="output_text")], + role="assistant", + status="completed", + type="message", + ) + + +@pytest.mark.asyncio +async def test_agent_is_released_after_run() -> None: + fake_model = FakeModel(initial_output=[_make_message("Paris")]) + agent = Agent(name="leaker", instructions="Answer questions.", model=fake_model) + agent_ref = weakref.ref(agent) + + await Runner.run(agent, "What is the capital of France?") + + del agent + gc.collect() + + assert agent_ref() is None