From 0fcd380ca4905f24344a83231d9f77dd50655c4b Mon Sep 17 00:00:00 2001 From: tim404x Date: Fri, 5 Jun 2026 14:51:56 +0000 Subject: [PATCH] fix(agent): reap browser/VM task resources on aborted turns The gateway leaked memory slowly under real use: every conversation turn that aborted with an exception or was interrupted left its per-task resources behind. The agent-browser daemon and its Chromium/Xvfb process tree (plus any sandbox VM) would orphan to init and accumulate. Two root causes, both fixed: 1. run_conversation only cleaned up task resources on its normal return paths. On the exception / interrupt path it returned without reaping. This splits the function into a thin public run_conversation() wrapper and _run_conversation_impl(): the wrapper runs _cleanup_task_resources() in a finally that fires ONLY when the impl did not complete normally, so the success path is unchanged (no double cleanup, no timing change) while aborts are always reaped. Cleanup is idempotent and persistence-aware, so it is safe to run on abort. 2. _terminate_host_pid only sent SIGTERM to the process tree snapshot taken at call time. Processes that ignore SIGTERM (chrome_crashpad_handler, renderers under load) survived, and children reparented to init as their parents exited fell off children(recursive=True) entirely. It now snapshots the whole tree up front, SIGTERMs it (letting a live browser flush its already-persisted cookies), waits, then SIGKILLs anything still alive and reaps zombies. Tests: tests/tools/test_run_conversation_abort_cleanup.py covers the wrapper (normal return does not double-reap; exception and interrupt both reap with the correct task id; a generated id is reused for the reap; a failing reaper does not mask the original turn error). Regression-validated. The existing browser-orphan-reaper suite (18 tests) still passes. --- agent/conversation_loop.py | 51 +++++++++++ .../test_run_conversation_abort_cleanup.py | 89 +++++++++++++++++++ tools/process_registry.py | 42 +++++++-- 3 files changed, 176 insertions(+), 6 deletions(-) create mode 100644 tests/tools/test_run_conversation_abort_cleanup.py diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index d01d5d4a84485..68dd817d7d662 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -356,6 +356,57 @@ def run_conversation( task_id: str = None, stream_callback: Optional[callable] = None, persist_user_message: Optional[str] = None, +) -> Dict[str, Any]: + """Public entry point — see :func:`_run_conversation_impl` for the body. + + Thin wrapper that guarantees per-task resource cleanup (browser/VM) runs + even when the turn aborts with an exception or is interrupted + (``CancelledError``). The impl already cleans up on every *normal* + return path; this finally only covers the error/interrupt paths, where + the ``agent-browser`` daemon and its Chromium/Xvfb tree would otherwise + orphan and accumulate — the cause of the gateway's slow memory growth. + + The success path is unchanged: cleanup runs here only when the impl + raised (``completed_normally`` stays False), so there is no extra work + and no timing change on a normal turn. ``_cleanup_task_resources`` is + idempotent and persistence-aware (it skips persistent sandboxes and does + a cookie-flushing graceful close), so the reap is safe to run on abort. + """ + # Bind task_id here so the finally reaps exactly the task the impl ran + # under — the impl derives the same value from this argument. + effective_task_id = task_id or str(uuid.uuid4()) + completed_normally = False + try: + result = _run_conversation_impl( + agent, + user_message, + system_message=system_message, + conversation_history=conversation_history, + task_id=effective_task_id, + stream_callback=stream_callback, + persist_user_message=persist_user_message, + ) + completed_normally = True + return result + finally: + if not completed_normally: + try: + agent._cleanup_task_resources(effective_task_id) + except Exception: + logger.warning( + "post-abort cleanup_task_resources failed for task %s", + effective_task_id, + ) + + +def _run_conversation_impl( + agent, + user_message: str, + system_message: str = None, + conversation_history: List[Dict[str, Any]] = None, + task_id: str = None, + stream_callback: Optional[callable] = None, + persist_user_message: Optional[str] = None, ) -> Dict[str, Any]: """ Run a complete conversation with tool calling until completion. diff --git a/tests/tools/test_run_conversation_abort_cleanup.py b/tests/tools/test_run_conversation_abort_cleanup.py new file mode 100644 index 0000000000000..ce2a5f796d89b --- /dev/null +++ b/tests/tools/test_run_conversation_abort_cleanup.py @@ -0,0 +1,89 @@ +"""Regression tests for the run_conversation() abort-cleanup wrapper. + +The public ``run_conversation`` is a thin wrapper around +``_run_conversation_impl``. Its only job is to guarantee per-task resource +cleanup (the agent-browser daemon + its Chromium/Xvfb tree, sandbox VMs) runs +when the turn aborts with an exception or is interrupted — the impl already +cleans up on every normal return path, so the wrapper's ``finally`` must fire +ONLY on the error/interrupt path. Without this, an aborted browser turn orphans +its process tree and the gateway leaks memory over time. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +import agent.conversation_loop as cl + + +def _make_agent(): + agent = MagicMock() + agent._cleanup_task_resources = MagicMock() + return agent + + +def test_normal_return_does_not_trigger_wrapper_cleanup(): + """On success the impl owns cleanup; the wrapper must NOT double-reap.""" + agent = _make_agent() + sentinel = {"completed": True} + + with patch.object(cl, "_run_conversation_impl", return_value=sentinel) as impl: + result = cl.run_conversation(agent, "hi", task_id="t-1") + + assert result is sentinel + impl.assert_called_once() + agent._cleanup_task_resources.assert_not_called() + + +def test_exception_triggers_cleanup_and_repropagates(): + """An impl exception must reap the task, then re-raise unchanged.""" + agent = _make_agent() + boom = RuntimeError("kaboom") + + with patch.object(cl, "_run_conversation_impl", side_effect=boom): + with pytest.raises(RuntimeError, match="kaboom"): + cl.run_conversation(agent, "hi", task_id="t-2") + + agent._cleanup_task_resources.assert_called_once_with("t-2") + + +def test_cancelled_interrupt_triggers_cleanup(): + """Interrupt (KeyboardInterrupt / CancelledError-like) also reaps.""" + agent = _make_agent() + + with patch.object(cl, "_run_conversation_impl", side_effect=KeyboardInterrupt()): + with pytest.raises(KeyboardInterrupt): + cl.run_conversation(agent, "hi", task_id="t-3") + + agent._cleanup_task_resources.assert_called_once_with("t-3") + + +def test_cleanup_uses_generated_task_id_when_none_given(): + """When no task_id is passed, the wrapper reaps the SAME id it generated + and handed to the impl (so the reap targets exactly the task that ran).""" + agent = _make_agent() + captured = {} + + def _impl(_agent, _msg, **kwargs): + captured["task_id"] = kwargs.get("task_id") + raise RuntimeError("fail") + + with patch.object(cl, "_run_conversation_impl", side_effect=_impl): + with pytest.raises(RuntimeError): + cl.run_conversation(agent, "hi") + + assert captured["task_id"] # a uuid was generated + agent._cleanup_task_resources.assert_called_once_with(captured["task_id"]) + + +def test_cleanup_failure_is_swallowed_not_masking_original_error(): + """If the cleanup itself raises, the ORIGINAL turn error must still win + (a broken reaper can't escalate into a different exception).""" + agent = _make_agent() + agent._cleanup_task_resources.side_effect = OSError("reaper broke") + + with patch.object(cl, "_run_conversation_impl", side_effect=ValueError("original")): + with pytest.raises(ValueError, match="original"): + cl.run_conversation(agent, "hi", task_id="t-4") + + agent._cleanup_task_resources.assert_called_once_with("t-4") diff --git a/tools/process_registry.py b/tools/process_registry.py index d9eb02a4ab863..e464a31ad68ca 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -483,12 +483,6 @@ def _terminate_host_pid(pid: int) -> None: import psutil try: parent = psutil.Process(pid) - for child in parent.children(recursive=True): - try: - child.terminate() - except psutil.NoSuchProcess: - pass - parent.terminate() except psutil.NoSuchProcess: return except (OSError, PermissionError): @@ -496,6 +490,42 @@ def _terminate_host_pid(pid: int) -> None: os.kill(pid, signal.SIGTERM) except (OSError, ProcessLookupError, PermissionError): pass + return + + # Snapshot the whole tree (descendants + parent) up front, before any + # of them exit and reparent grandchildren to init — once a process is + # reparented it falls off ``children(recursive=True)`` and survives. + try: + procs = parent.children(recursive=True) + except (psutil.NoSuchProcess, OSError): + procs = [] + procs.append(parent) + + # Graceful first: SIGTERM the full tree so a live browser can flush + # (cookies are already persisted by the agent-browser ``close`` + # command before this runs on the cleanup path). + for proc in procs: + try: + proc.terminate() + except (psutil.NoSuchProcess, OSError, PermissionError): + pass + + # Then SIGKILL anything still alive. ``chrome_crashpad_handler`` (and + # renderers under load) routinely ignore SIGTERM; without escalation + # they orphan to init and accumulate — the root of the leak. + try: + _gone, alive = psutil.wait_procs(procs, timeout=5) + except (psutil.NoSuchProcess, OSError): + alive = procs + for proc in alive: + try: + proc.kill() + except (psutil.NoSuchProcess, OSError, PermissionError): + pass + try: + psutil.wait_procs(alive, timeout=3) + except (psutil.NoSuchProcess, OSError): + pass # ----- Spawn -----