fix(docker): cleanup_vm() default honors persist mode (don't kill container on session close)

Commit 4 made cleanup_vm() default to force_remove=True, which was wrong:
cleanup_vm() is called from AIAgent.close() (TUI session close at
tui_gateway/server.py:2991, gateway session teardown at gateway/run.py:3569)
and from per-turn cleanup (agent/chat_completion_helpers.py:1517). All
three are session-lifecycle events that should honor persist mode, not
explicit user-initiated teardown.

Ben reported the symptom: container shared between multiple TUI sessions
(good) but killed as soon as any session closed (bad). With force_remove=True
as the default, every `session.close` JSON-RPC tore down the container.

The fix is to flip cleanup_vm()'s force_remove default back to False.
The kwarg still exists for future explicit-teardown paths (`/reset`-style
flows, "destroy my sandbox" commands) that haven't been wired up yet.

Two new unit tests pin the behavior:

* `test_cleanup_vm_default_honors_persist_mode` — asserts
  `cleanup_vm(task_id)` does neither docker stop nor docker rm on a
  persist-mode container (the regression Ben caught).
* `test_cleanup_vm_force_remove_tears_down_persist_container` —
  asserts the kwarg still flows through the runtime-signature-inspection
  plumbing to the backend's cleanup().

E2E verified against real Docker (in addition to all 17 existing checks):

  ✓ Default cleanup_vm() leaves persist-mode container running
  ✓ cleanup_vm(force_remove=True) removed the container

Refs #20561
This commit is contained in:
Ben
2026-05-29 11:49:54 +10:00
committed by Ben Barclay
parent 5c2170a7c6
commit 2f0f03c40d
2 changed files with 101 additions and 10 deletions
+16 -10
View File
@@ -1426,20 +1426,26 @@ def cleanup_all_environments():
return cleaned
def cleanup_vm(task_id: str, *, force_remove: bool = True):
def cleanup_vm(task_id: str, *, force_remove: bool = False):
"""Manually clean up a specific environment by task_id.
*force_remove* (default True) is forwarded to backends that accept it
— currently only ``DockerEnvironment``. The default of True matches the
semantics callers expect from this function: ``cleanup_vm(task_id)`` is
the explicit-teardown path (called from ``/reset``-style flows,
``AIAgent.close()``, and the idle reaper after a long inactivity window),
so the container should be removed regardless of persist mode.
*force_remove* (default False) is forwarded to backends that accept it
— currently only ``DockerEnvironment``. The default of False matches
session-lifecycle semantics: this function is called from
``AIAgent.close()`` (TUI session close, gateway session teardown) and the
per-turn cleanup branch for non-persistent envs, both of which should
honor the user's persist-mode preference. Stopping the container here
would defeat the "ONE long-lived container shared across sessions"
contract — exactly the bug Ben reported when the container was killed
on every TUI session close.
Pass ``force_remove=True`` for actual user-initiated teardown
(e.g. ``/reset``-style flows that haven't been wired yet, or future
"destroy my sandbox" commands).
The idle reaper passes the env through ``env.cleanup()`` directly (not
via this function), so persist-mode idle envs are NOT removed by idle
reaping — only by explicit ``cleanup_vm()`` or the orphan reaper at next
startup. See ``_cleanup_inactive_envs()`` for the idle path.
via this function), so persist-mode idle envs are similarly no-op'd —
only the orphan reaper at next startup reclaims them.
"""
# Remove from tracking dicts while holding the lock, but defer the
# actual (potentially slow) env.cleanup() call to outside the lock