Python: fix(python): add release_session API to prevent BackgroundAgentsProvider memory leaks (#7450)
* fix: add release_session API to prevent BackgroundAgentsProvider memory leaks * fix: address Copilot review comments on release_session * fix(harness): make background agent session release race-safe and bounded * fix (harness): address release_session and review feedback
This commit is contained in:
pratik wayase
@@ -10,6 +10,7 @@ and retrieve results. Each background task runs in its own session concurrently.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Awaitable, MutableMapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
@@ -23,6 +24,8 @@ from .._telemetry import FeatureIndex, mark_feature_used
|
||||
from .._tools import tool
|
||||
from .._types import AgentResponse, Message
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_BACKGROUND_AGENTS_SOURCE_ID = "background_agents"
|
||||
|
||||
DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS = """\
|
||||
@@ -114,6 +117,15 @@ class _RuntimeState:
|
||||
|
||||
in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {})
|
||||
background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {})
|
||||
closed: bool = False
|
||||
|
||||
def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]]) -> None:
|
||||
"""Track a background task if this runtime is still open."""
|
||||
if self.closed:
|
||||
task.cancel()
|
||||
raise RuntimeError("Session runtime is closed; cannot start background task.")
|
||||
|
||||
self.in_flight_tasks[task_id] = task
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -126,6 +138,20 @@ async def _run_agent(awaitable: Awaitable[AgentResponse[Any]]) -> AgentResponse[
|
||||
return await awaitable
|
||||
|
||||
|
||||
def _log_abandoned_background_task(task: asyncio.Task[Any]) -> None:
|
||||
"""Retrieve exception from an abandoned task to avoid asyncio warnings."""
|
||||
if task.cancelled():
|
||||
return
|
||||
|
||||
try:
|
||||
exception = task.exception()
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
|
||||
if exception is not None:
|
||||
logger.debug("Abandoned background task raised: %s", exception)
|
||||
|
||||
|
||||
def _validate_and_build_agent_dict(agents: Sequence[SupportsAgentRun]) -> dict[str, SupportsAgentRun]:
|
||||
"""Validate agents and build a case-insensitive lookup dict.
|
||||
|
||||
@@ -308,9 +334,121 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
def _get_runtime(self, session: AgentSession) -> _RuntimeState:
|
||||
"""Get or create runtime state for a session."""
|
||||
session_id = session.session_id
|
||||
if session_id not in self._runtime:
|
||||
self._runtime[session_id] = _RuntimeState()
|
||||
return self._runtime[session_id]
|
||||
runtime = self._runtime.get(session_id)
|
||||
|
||||
if runtime is None or runtime.closed:
|
||||
runtime = _RuntimeState()
|
||||
self._runtime[session_id] = runtime
|
||||
|
||||
return runtime
|
||||
|
||||
async def release_session(
|
||||
self,
|
||||
session: AgentSession,
|
||||
*,
|
||||
cancel_running: bool = True,
|
||||
timeout: float | None = 30.0,
|
||||
) -> None:
|
||||
"""Release all runtime state for a session to prevent runtime leaks.
|
||||
|
||||
Args:
|
||||
session: The agent session whose runtime state should be released.
|
||||
cancel_running: If True, cancel pending asyncio.Tasks safely.
|
||||
timeout: Maximum seconds to wait for tasks to finish cancellation.
|
||||
If None, wait indefinitely. The default is bounded so a buggy
|
||||
task cannot wedge host eviction or shutdown.
|
||||
"""
|
||||
session_id = session.session_id
|
||||
runtime = self._runtime.get(session_id)
|
||||
|
||||
if runtime is None or runtime.closed:
|
||||
return
|
||||
|
||||
pending = [
|
||||
task
|
||||
for task in list(runtime.in_flight_tasks.values())
|
||||
if not task.done()
|
||||
]
|
||||
|
||||
if pending and not cancel_running:
|
||||
raise RuntimeError(
|
||||
f"Cannot release session {session_id}: {len(pending)} tasks still running."
|
||||
)
|
||||
|
||||
runtime.closed = True
|
||||
|
||||
try:
|
||||
if pending:
|
||||
await self._drain_runtime(
|
||||
runtime,
|
||||
cancel_running=cancel_running,
|
||||
timeout=timeout,
|
||||
)
|
||||
else:
|
||||
completed = list(runtime.in_flight_tasks.values())
|
||||
if completed:
|
||||
await asyncio.gather(*completed, return_exceptions=True)
|
||||
finally:
|
||||
runtime.in_flight_tasks.clear()
|
||||
runtime.background_sessions.clear()
|
||||
|
||||
if self._runtime.get(session_id) is runtime:
|
||||
self._runtime.pop(session_id, None)
|
||||
|
||||
async def _drain_runtime(
|
||||
self,
|
||||
runtime: _RuntimeState,
|
||||
*,
|
||||
cancel_running: bool,
|
||||
timeout: float | None,
|
||||
) -> None:
|
||||
"""Cancel and await tracked tasks, bounded by timeout."""
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = None if timeout is None else loop.time() + float(timeout)
|
||||
|
||||
while True:
|
||||
tasks = list(runtime.in_flight_tasks.values())
|
||||
pending = [task for task in tasks if not task.done()]
|
||||
|
||||
if not pending:
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
return
|
||||
|
||||
if not cancel_running:
|
||||
raise RuntimeError(f"Cannot release session: {len(pending)} tasks still running.")
|
||||
|
||||
for task in pending:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
|
||||
remaining = None
|
||||
if deadline is not None:
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
logger.warning(
|
||||
"Session release timed out before all tasks finished. Abandoning %s task(s).",
|
||||
len(pending),
|
||||
)
|
||||
for task in pending:
|
||||
if not task.done():
|
||||
task.add_done_callback(_log_abandoned_background_task)
|
||||
return
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*pending, return_exceptions=True),
|
||||
timeout=remaining,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
not_done = [task for task in pending if not task.done()]
|
||||
logger.warning(
|
||||
"Session release timed out waiting for %s task(s). They will be abandoned.",
|
||||
len(not_done),
|
||||
)
|
||||
for task in not_done:
|
||||
task.add_done_callback(_log_abandoned_background_task)
|
||||
return
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
@@ -331,6 +469,9 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
@tool(name="background_agents_start_task", approval_mode="never_require")
|
||||
def background_agents_start_task(agent_name: str, input: str, description: str) -> str:
|
||||
"""Start a background task on a named agent. Returns a confirmation with the task ID."""
|
||||
if runtime.closed:
|
||||
return "Error: Session is being released; cannot start a new background task."
|
||||
|
||||
key = agent_name.lower()
|
||||
if key not in self._agents:
|
||||
available = ", ".join(a.name or "" for a in self._agents.values())
|
||||
@@ -338,6 +479,17 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
|
||||
bg_agent = self._agents[key]
|
||||
task_id = provider_state.get("next_task_id", 1)
|
||||
|
||||
sub_session = bg_agent.create_session()
|
||||
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session)))
|
||||
try:
|
||||
runtime.track_task(task_id, async_task)
|
||||
except RuntimeError as exc:
|
||||
return f"Error: {exc}"
|
||||
|
||||
runtime.background_sessions[task_id] = sub_session
|
||||
|
||||
provider_state["next_task_id"] = task_id + 1
|
||||
|
||||
task_info = BackgroundTaskInfo(
|
||||
@@ -349,14 +501,6 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
tasks.append(task_info)
|
||||
_save_tasks(provider_state, tasks)
|
||||
|
||||
# Create a dedicated session for this background task.
|
||||
sub_session = bg_agent.create_session()
|
||||
|
||||
# Start the task concurrently.
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session)))
|
||||
runtime.in_flight_tasks[task_id] = async_task
|
||||
runtime.background_sessions[task_id] = sub_session
|
||||
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
return f"Background task {task_id} started on agent '{agent_name}'."
|
||||
|
||||
@@ -365,6 +509,9 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
@tool(name="background_agents_wait_for_first_completion", approval_mode="never_require")
|
||||
async def background_agents_wait_for_first_completion(task_ids: list[int]) -> str:
|
||||
"""Block until the first of the specified background tasks completes. Returns the completed task's ID."""
|
||||
if runtime.closed:
|
||||
return "Error: Session is being released; cannot wait for background tasks."
|
||||
|
||||
if not task_ids:
|
||||
return "Error: No task IDs provided."
|
||||
|
||||
@@ -448,6 +595,9 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
@tool(name="background_agents_continue_task", approval_mode="never_require")
|
||||
def background_agents_continue_task(task_id: int, text: str) -> str:
|
||||
"""Send follow-up input to a completed or failed task to resume its work."""
|
||||
if runtime.closed:
|
||||
return "Error: Session is being released; cannot continue a background task."
|
||||
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
task_info = next((t for t in tasks if t.id == task_id), None)
|
||||
|
||||
@@ -472,15 +622,16 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
|
||||
bg_agent = self._agents[key]
|
||||
|
||||
# Reset task state and start a new run on the existing session.
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session)))
|
||||
try:
|
||||
runtime.track_task(task_id, async_task)
|
||||
except RuntimeError as exc:
|
||||
return f"Error: {exc}"
|
||||
task_info.status = BackgroundTaskStatus.RUNNING
|
||||
task_info.result_text = None
|
||||
task_info.error_text = None
|
||||
_save_tasks(provider_state, tasks)
|
||||
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session)))
|
||||
runtime.in_flight_tasks[task_id] = async_task
|
||||
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
return f"Task {task_id} continued with new input."
|
||||
|
||||
@@ -489,6 +640,9 @@ class BackgroundAgentsProvider(ContextProvider):
|
||||
@tool(name="background_agents_clear_completed_task", approval_mode="never_require")
|
||||
def background_agents_clear_completed_task(task_id: int) -> str:
|
||||
"""Remove a completed or failed task and release its session to free memory."""
|
||||
if runtime.closed:
|
||||
return "Error: Session is being released; cannot clear tasks."
|
||||
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
task_info = next((t for t in tasks if t.id == task_id), None)
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@@ -542,3 +543,303 @@ def test_task_status_enum_values() -> None:
|
||||
assert BackgroundTaskStatus.COMPLETED == "completed"
|
||||
assert BackgroundTaskStatus.FAILED == "failed"
|
||||
assert BackgroundTaskStatus.LOST == "lost"
|
||||
|
||||
|
||||
async def test_release_session_cancels_and_clears() -> None:
|
||||
"""Should cancel pending tasks and clear runtime state."""
|
||||
provider = _make_provider(_FakeAgent("Slow", delay=10.0))
|
||||
session = _make_session()
|
||||
tools = await _get_tools(provider, session)
|
||||
|
||||
await _invoke_tool(
|
||||
tools["background_agents_start_task"],
|
||||
agent_name="Slow",
|
||||
input="task",
|
||||
description="long running",
|
||||
)
|
||||
|
||||
runtime = provider._runtime.get(session.session_id)
|
||||
assert runtime is not None
|
||||
assert len(runtime.in_flight_tasks) == 1
|
||||
|
||||
task = next(iter(runtime.in_flight_tasks.values()))
|
||||
|
||||
await provider.release_session(session, cancel_running=True)
|
||||
|
||||
assert task.done()
|
||||
assert task.cancelled()
|
||||
assert runtime.in_flight_tasks == {}
|
||||
assert runtime.background_sessions == {}
|
||||
|
||||
assert session.session_id not in provider._runtime
|
||||
|
||||
|
||||
async def test_release_session_raises_if_cancel_running_false() -> None:
|
||||
"""Should raise RuntimeError if cancel_running=False and tasks are pending."""
|
||||
provider = _make_provider(_FakeAgent("Slow", delay=10.0))
|
||||
session = _make_session()
|
||||
tools = await _get_tools(provider, session)
|
||||
|
||||
await _invoke_tool(
|
||||
tools["background_agents_start_task"],
|
||||
agent_name="Slow",
|
||||
input="task",
|
||||
description="long running",
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="tasks still running"):
|
||||
await provider.release_session(session, cancel_running=False)
|
||||
|
||||
assert session.session_id in provider._runtime
|
||||
await provider.release_session(session, cancel_running=True)
|
||||
|
||||
|
||||
async def test_release_session_idempotent() -> None:
|
||||
"""Should not raise when releasing an unknown or already released session."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
await provider.release_session(AgentSession(session_id="non_existent_session"))
|
||||
|
||||
await provider.release_session(session)
|
||||
await provider.release_session(session)
|
||||
|
||||
|
||||
async def test_release_session_isolation() -> None:
|
||||
"""Releasing one session should not affect another."""
|
||||
provider = _make_provider(_FakeAgent("Worker", delay=10.0))
|
||||
session_a = AgentSession(session_id="session_a")
|
||||
session_b = AgentSession(session_id="session_b")
|
||||
|
||||
tools_a = await _get_tools(provider, session_a)
|
||||
tools_b = await _get_tools(provider, session_b)
|
||||
|
||||
await _invoke_tool(
|
||||
tools_a["background_agents_start_task"],
|
||||
agent_name="Worker",
|
||||
input="A",
|
||||
description="A",
|
||||
)
|
||||
await _invoke_tool(
|
||||
tools_b["background_agents_start_task"],
|
||||
agent_name="Worker",
|
||||
input="B",
|
||||
description="B",
|
||||
)
|
||||
|
||||
await provider.release_session(session_a, cancel_running=True)
|
||||
|
||||
assert "session_a" not in provider._runtime
|
||||
assert "session_b" in provider._runtime
|
||||
|
||||
await provider.release_session(session_b, cancel_running=True)
|
||||
|
||||
|
||||
async def test_get_runtime_replaces_closed_runtime() -> None:
|
||||
"""A closed runtime should be replaced by a new runtime instance."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
old_runtime = provider._get_runtime(session)
|
||||
old_runtime.closed = True
|
||||
|
||||
new_runtime = provider._get_runtime(session)
|
||||
|
||||
assert new_runtime is not old_runtime
|
||||
assert provider._runtime.get(session.session_id) is new_runtime
|
||||
|
||||
|
||||
async def test_track_task_rejects_when_runtime_closed() -> None:
|
||||
"""Closed runtime should not accept new background tasks."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
runtime = provider._get_runtime(session)
|
||||
runtime.closed = True
|
||||
|
||||
async def _dummy() -> Any:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
task = asyncio.create_task(_dummy())
|
||||
|
||||
with pytest.raises(RuntimeError, match="closed"):
|
||||
runtime.track_task(1, task)
|
||||
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert task.cancelled()
|
||||
|
||||
|
||||
async def test_start_task_returns_error_when_runtime_closed() -> None:
|
||||
"""background_agents_start_task should refuse to run on a closed runtime."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
tools = await _get_tools(provider, session)
|
||||
|
||||
runtime = provider._get_runtime(session)
|
||||
runtime.closed = True
|
||||
|
||||
result = await _invoke_tool(
|
||||
tools["background_agents_start_task"],
|
||||
agent_name="Worker",
|
||||
input="task",
|
||||
description="should not start",
|
||||
)
|
||||
|
||||
assert "being released" in result
|
||||
assert runtime.in_flight_tasks == {}
|
||||
|
||||
|
||||
async def test_tools_return_error_when_runtime_closed() -> None:
|
||||
"""Mutating/background tools should refuse to run on a closed runtime."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
tools = await _get_tools(provider, session)
|
||||
|
||||
runtime = provider._get_runtime(session)
|
||||
runtime.closed = True
|
||||
|
||||
wait_result = await _invoke_tool(
|
||||
tools["background_agents_wait_for_first_completion"],
|
||||
task_ids=[1],
|
||||
)
|
||||
assert "being released" in wait_result
|
||||
|
||||
continue_result = await _invoke_tool(
|
||||
tools["background_agents_continue_task"],
|
||||
task_id=1,
|
||||
text="continue",
|
||||
)
|
||||
assert "being released" in continue_result
|
||||
|
||||
clear_result = await _invoke_tool(
|
||||
tools["background_agents_clear_completed_task"],
|
||||
task_id=1,
|
||||
)
|
||||
assert "being released" in clear_result
|
||||
|
||||
|
||||
async def test_release_session_times_out_if_task_ignores_cancellation() -> None:
|
||||
"""release_session should return within bounded time even if task ignores cancel."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
runtime = provider._get_runtime(session)
|
||||
unblock = asyncio.Event()
|
||||
|
||||
async def _ignore_cancel() -> Any:
|
||||
try:
|
||||
await unblock.wait()
|
||||
except asyncio.CancelledError:
|
||||
await unblock.wait()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(_ignore_cancel())
|
||||
runtime.in_flight_tasks[1] = task
|
||||
|
||||
start = asyncio.get_running_loop().time()
|
||||
|
||||
await asyncio.wait_for(
|
||||
provider.release_session(
|
||||
session,
|
||||
cancel_running=True,
|
||||
timeout=0.05,
|
||||
),
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
elapsed = asyncio.get_running_loop().time() - start
|
||||
|
||||
assert elapsed < 1.0
|
||||
assert session.session_id not in provider._runtime
|
||||
|
||||
unblock.set()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1.0)
|
||||
|
||||
|
||||
async def test_release_session_does_not_pop_replacement_runtime() -> None:
|
||||
"""A release of an old runtime should not remove a replacement runtime."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
|
||||
old_runtime = provider._get_runtime(session)
|
||||
unblock = asyncio.Event()
|
||||
|
||||
async def _blocked_task() -> Any:
|
||||
try:
|
||||
await unblock.wait()
|
||||
except asyncio.CancelledError:
|
||||
await unblock.wait()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(_blocked_task())
|
||||
old_runtime.in_flight_tasks[1] = task
|
||||
|
||||
release_task = asyncio.create_task(
|
||||
provider.release_session(
|
||||
session,
|
||||
cancel_running=True,
|
||||
timeout=5.0,
|
||||
)
|
||||
)
|
||||
|
||||
for _ in range(100):
|
||||
if old_runtime.closed:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert old_runtime.closed
|
||||
|
||||
new_runtime = provider._get_runtime(session)
|
||||
assert new_runtime is not old_runtime
|
||||
|
||||
unblock.set()
|
||||
|
||||
await asyncio.wait_for(release_task, timeout=1.0)
|
||||
|
||||
assert provider._runtime.get(session.session_id) is new_runtime
|
||||
|
||||
with suppress(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1.0)
|
||||
|
||||
await provider.release_session(
|
||||
session,
|
||||
cancel_running=True,
|
||||
timeout=1.0,
|
||||
)
|
||||
|
||||
|
||||
async def test_release_session_skips_drain_when_no_pending_tasks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Should not invoke the drain path when there are no pending tasks."""
|
||||
provider = _make_provider(_FakeAgent("Worker"))
|
||||
session = _make_session()
|
||||
runtime = provider._get_runtime(session)
|
||||
|
||||
async def completed_task() -> Any:
|
||||
return None
|
||||
|
||||
task = asyncio.create_task(completed_task())
|
||||
await task
|
||||
|
||||
runtime.in_flight_tasks[1] = task
|
||||
|
||||
drain_called = False
|
||||
original_drain = provider._drain_runtime
|
||||
|
||||
async def fake_drain(*args: Any, **kwargs: Any) -> None:
|
||||
nonlocal drain_called
|
||||
drain_called = True
|
||||
await original_drain(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(provider, "_drain_runtime", fake_drain)
|
||||
|
||||
await provider.release_session(session)
|
||||
|
||||
assert drain_called is False
|
||||
assert session.session_id not in provider._runtime
|
||||
|
||||
Reference in New Issue
Block a user