From 6ddc4254e85d28498c1278f3f04c756ae99e8b87 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Fri, 31 Jul 2026 15:34:31 +0530 Subject: [PATCH 1/3] fix: add release_session API to prevent BackgroundAgentsProvider memory leaks --- .../_harness/_background_agents.py | 30 +++++++ .../core/test_harness_background_agents.py | 78 +++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index c1c70b64a34..56cc4f0108b 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -312,6 +312,36 @@ def _get_runtime(self, session: AgentSession) -> _RuntimeState: self._runtime[session_id] = _RuntimeState() return self._runtime[session_id] + async def release_session(self, session_id: str, *, cancel_running: bool = True) -> None: + """Release all runtime state for a session to prevent runtime leaks. + + Args: + session_id: The session ID to release. + cancel_running: If True, cancel pending asyncio.Tasks safely. + """ + + runtime = self._runtime.get(session_id) + if runtime is None: + return + + pending = [t for t in runtime.in_flight_tasks.values() if not t.done()] + + if pending and not cancel_running: + raise RuntimeError( + f"Cannot release session {session_id}: {len(pending)} tasks still running." + ) + + if pending: + for task in pending: + task.cancel() + + await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED) + + self._runtime.pop(session_id, None) + + runtime.in_flight_tasks.clear() + runtime.background_sessions.clear() + async def before_run( self, *, diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 98e6fa2a676..7fb9e9540dc 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -542,3 +542,81 @@ 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 + + await provider.release_session(session.session_id, cancel_running=True) + 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.session_id, cancel_running=False) + + assert session.session_id in provider._runtime + await provider.release_session(session.session_id, 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("non_existent_session") + + await provider.release_session(session.session_id) + await provider.release_session(session.session_id) + + +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) From 777b695b6c1cc449b03ed6bd0ac948990f441fbc Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Fri, 31 Jul 2026 15:55:30 +0530 Subject: [PATCH 2/3] fix: address Copilot review comments on release_session --- .../agent_framework/_harness/_background_agents.py | 11 ++++++++--- .../core/tests/core/test_harness_background_agents.py | 8 ++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 56cc4f0108b..208f4dc84be 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -319,24 +319,29 @@ async def release_session(self, session_id: str, *, cancel_running: bool = True) session_id: The session ID to release. cancel_running: If True, cancel pending asyncio.Tasks safely. """ - + + runtime = self._runtime.get(session_id) if runtime is None: return - pending = [t for t in runtime.in_flight_tasks.values() if not t.done()] + tasks = list(runtime.in_flight_tasks.values()) + pending = [t for t in tasks if not t.done()] if pending and not cancel_running: raise RuntimeError( f"Cannot release session {session_id}: {len(pending)} tasks still running." ) - if pending: + if cancel_running and pending: for task in pending: task.cancel() await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._runtime.pop(session_id, None) runtime.in_flight_tasks.clear() diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 7fb9e9540dc..a5d4bda3c93 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -560,8 +560,16 @@ async def test_release_session_cancels_and_clears() -> None: 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.session_id, 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 From 1c7e9f897b44db93a1f97e28ac6d2d9207fe0dd6 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Tue, 4 Aug 2026 00:31:39 +0530 Subject: [PATCH 3/3] fix(harness): make background agent session release race-safe and bounded --- .../_harness/_background_agents.py | 171 +++++++++++++--- .../core/test_harness_background_agents.py | 186 +++++++++++++++++- 2 files changed, 323 insertions(+), 34 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 208f4dc84be..27689d645e8 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -10,6 +10,7 @@ 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 .._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,44 +334,111 @@ def __init__( 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 - async def release_session(self, session_id: str, *, cancel_running: bool = True) -> None: + return runtime + + async def release_session( + self, + session_id: str, + *, + cancel_running: bool = True, + timeout: float | None = 30.0, + ) -> None: """Release all runtime state for a session to prevent runtime leaks. Args: session_id: The session ID to release. 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. """ - - runtime = self._runtime.get(session_id) if runtime is None: return - tasks = list(runtime.in_flight_tasks.values()) - pending = [t for t in tasks if not t.done()] + if 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." + raise RuntimeError(f"Cannot release session {session_id}: {len(pending)} tasks still running.") + + runtime.closed = True + + try: + await self._drain_runtime( + runtime, + cancel_running=cancel_running, + timeout=timeout, ) + finally: + runtime.in_flight_tasks.clear() + runtime.background_sessions.clear() - if cancel_running and pending: - for task in pending: - task.cancel() + if self._runtime.get(session_id) is runtime: + self._runtime.pop(session_id, None) - await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED) + 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) - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + while True: + tasks = list(runtime.in_flight_tasks.values()) + pending = [task for task in tasks if not task.done()] - self._runtime.pop(session_id, None) + if not pending: + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + return - runtime.in_flight_tasks.clear() - runtime.background_sessions.clear() + 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, @@ -366,6 +459,9 @@ async def before_run( @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()) @@ -373,6 +469,17 @@ def background_agents_start_task(agent_name: str, input: str, description: str) 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( @@ -384,14 +491,6 @@ def background_agents_start_task(agent_name: str, input: str, description: str) 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}'." @@ -400,6 +499,9 @@ def background_agents_start_task(agent_name: str, input: str, description: str) @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." @@ -483,6 +585,9 @@ def background_agents_get_all_tasks() -> str: @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) @@ -507,15 +612,16 @@ def background_agents_continue_task(task_id: int, text: str) -> str: 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." @@ -524,6 +630,9 @@ def background_agents_continue_task(task_id: int, text: str) -> str: @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) diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index a5d4bda3c93..5c0eb8e2b4f 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from contextlib import suppress from typing import Any import pytest @@ -560,7 +561,7 @@ async def test_release_session_cancels_and_clears() -> None: 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.session_id, cancel_running=True) @@ -615,11 +616,15 @@ async def test_release_session_isolation() -> None: await _invoke_tool( tools_a["background_agents_start_task"], - agent_name="Worker", input="A", description="A", + agent_name="Worker", + input="A", + description="A", ) await _invoke_tool( tools_b["background_agents_start_task"], - agent_name="Worker", input="B", description="B", + agent_name="Worker", + input="B", + description="B", ) await provider.release_session("session_a", cancel_running=True) @@ -628,3 +633,178 @@ async def test_release_session_isolation() -> None: 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 + + task = asyncio.create_task(asyncio.sleep(0)) + + 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() -> None: + 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.session_id, + 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() -> None: + 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.session_id, + 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.session_id, + cancel_running=True, + timeout=1.0, + )