fix(agent): reap browser/VM task resources on aborted turns

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.

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")

tools/process_registry.py

@@ -483,19 +483,49 @@ 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):
             try:
                 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 -----