fix(e2b): keep shutdown retryable while maintenance workers are live (#6244)

* fix(e2b): retain teardown ownership across maintenance timeout

* fix(e2b): fence acquisition and preserve signal forwarding

* test(e2b): cover shutdown admission and signal forwarding

* docs(e2b): pin shutdown admission and signal invariants

* fix(e2b): scope deferred-cleanup signal handling

* test(e2b): exercise dedicated deferred-cleanup signal error

* style(tests): format E2B shutdown regression tests

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

backend/packages/harness/deerflow/community/e2b_sandbox/AGENTS.md

@@ -26,3 +26,13 @@ Release only needs the VM lock while leaving active state; do not hold it
 during output sync, which must not prevent ownership heartbeats.
 Track that release in `_remote_ops_in_progress` until it completes so
 reconciliation cannot probe or re-adopt a VM between active and warm states.
+
+Shutdown owns both maintenance workers as well as sandbox registries. After
+signalling the lease-renewal and reconciliation threads, a bounded join is only
+a wait budget: verify each worker actually exited before clearing registries or
+tearing down resources. If either worker is still alive, keep admission fenced
+and preserve all tracked state so a later `shutdown()` can retry cleanup.
+That fence covers cached sync/async acquisition too: recheck after serializer
+waits and immediately before exposing a reused client. Signal-triggered shutdown
+may report deferred cleanup, but it must still forward the process's original
+SIGTERM/SIGINT/SIGHUP action while retaining the pending E2B state.

backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py

@@ -119,6 +119,11 @@ _MAX_MOUNT_PASS_FILES = 2000
 # Deadline checks stop preflight work and new writes. Active SDK writes finish.
 _MOUNT_PASS_DEADLINE_SECONDS = 120
 
+
+class _E2BMaintenanceShutdownTimeout(RuntimeError):
+    """A bounded maintenance-worker join did not prove worker termination."""
+
+
 # Recursive skill projection replacement must never target an operating-system
 # tree. The configured E2B home is handled separately: an isolated descendant
 # such as /home/user/skills is supported, while the home directory itself and
@@ -329,6 +334,9 @@ class E2BSandboxProvider(SandboxProvider):
         self._transitioning_slots = 0
         self._capacity_cond = threading.Condition(self._lock)
         self._shutdown_called = False
+        # Keep admission fenced if shutdown has to retry after a bounded
+        # maintenance-thread join times out.
+        self._shutdown_cleanup_pending = False
         self._owned_sandbox_ids: set[str] = set()
         self._acquire_inflight: set[str] = set()
         self._orphan_first_seen: dict[str, float] = {}
@@ -358,7 +366,7 @@ class E2BSandboxProvider(SandboxProvider):
         if not self._ownership.supports_cross_process:
             logger.warning("E2B sandbox ownership is process-local. Multi-worker gateways must configure sandbox.ownership.type: redis for safe reconciliation.")
 
-        atexit.register(self.shutdown)
+        atexit.register(self._shutdown_at_exit)
         self._register_signal_handlers()
         self._start_maintenance_threads()
 
@@ -528,6 +536,16 @@ class E2BSandboxProvider(SandboxProvider):
 
     # ── Signal / shutdown handling ───────────────────────────────────────
 
+    def _shutdown_at_exit(self) -> None:
+        """Best-effort process-exit cleanup without turning a retryable timeout into an atexit error."""
+        try:
+            self.shutdown()
+        except _E2BMaintenanceShutdownTimeout as exc:
+            logger.warning(
+                "E2B shutdown cleanup is still pending at interpreter exit: %s",
+                exc,
+            )
+
     def _register_signal_handlers(self) -> None:
         try:
             self._original_sigterm = signal.getsignal(signal.SIGTERM)
@@ -537,7 +555,14 @@ class E2BSandboxProvider(SandboxProvider):
             return
 
         def _handler(signum, frame):
-            self.shutdown()
+            try:
+                self.shutdown()
+            except _E2BMaintenanceShutdownTimeout:
+                logger.warning(
+                    "E2B shutdown cleanup is still pending while handling signal %s; forwarding the signal action",
+                    signum,
+                )
+
             if signum == signal.SIGTERM:
                 original = self._original_sigterm
             elif hasattr(signal, "SIGHUP") and signum == signal.SIGHUP:
@@ -562,7 +587,22 @@ class E2BSandboxProvider(SandboxProvider):
                     sig_name,
                 )
 
+    def _raise_if_shutting_down(self) -> None:
+        with self._lock:
+            if not self._shutdown_called:
+                return
+            raise SandboxCapacityExceededError(
+                "Sandbox provider is shutting down; cannot acquire sandbox",
+                active=len(self._sandboxes),
+                warm=len(self._warm_pool),
+                reserved=self._reserved_slots,
+                replicas=int(self._config["replicas"]),
+                retry_after_seconds=30.0,
+                reason="shutdown",
+            )
+
     def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
+        self._raise_if_shutting_down()
         effective_user_id = self._effective_acquire_user_id(user_id)
         if thread_id:
             with self._acquire_serializer.hold(self._thread_key(thread_id, effective_user_id)):
@@ -576,6 +616,9 @@ class E2BSandboxProvider(SandboxProvider):
         return await loop.run_in_executor(self._acquire_executor, acquire)
 
     def _acquire_internal(self, thread_id: str | None, *, user_id: str) -> str:
+        # The serializer may have queued this caller before shutdown started.
+        # Recheck admission after the wait and before any cache/reclaim path.
+        self._raise_if_shutting_down()
         if thread_id:
             cached = self._reuse_in_process_sandbox(thread_id, user_id=user_id)
             if cached is not None:
@@ -630,17 +673,30 @@ class E2BSandboxProvider(SandboxProvider):
             self._refresh_remote_timeout(sandbox.client)
         except Exception as e:  # pragma: no cover - defensive
             logger.debug("Failed to refresh timeout on reuse: %s", e)
+        # Shutdown may have started while the remote ping/timeout refresh was
+        # in flight. Do not publish or expose this cached client after admission
+        # closes.
+        self._raise_if_shutting_down()
         self._publish_ownership(sid)
         with self._lock:
             self._acquire_inflight.discard(sid)
-
-        logger.info(
-            "Reusing in-process e2b sandbox %s for user/thread %s/%s",
-            sid,
-            user_id,
-            thread_id,
-        )
-        return sid
+            if self._shutdown_called:
+                raise SandboxCapacityExceededError(
+                    "Sandbox provider shut down while reusing a cached sandbox",
+                    active=len(self._sandboxes),
+                    warm=len(self._warm_pool),
+                    reserved=self._reserved_slots,
+                    replicas=int(self._config["replicas"]),
+                    retry_after_seconds=30.0,
+                    reason="shutdown",
+                )
+            logger.info(
+                "Reusing in-process e2b sandbox %s for user/thread %s/%s",
+                sid,
+                user_id,
+                thread_id,
+            )
+            return sid
 
     def _reclaim_warm_pool_sandbox(self, thread_id: str, *, user_id: str) -> str | None:
         """Reclaim a warm-pool sandbox, holding a transitioning slot throughout.
@@ -2904,13 +2960,29 @@ class E2BSandboxProvider(SandboxProvider):
 
     def shutdown(self) -> None:
         with self._lock:
-            if self._shutdown_called:
+            if self._shutdown_called and not self._shutdown_cleanup_pending:
                 return
             self._shutdown_called = True
+            self._shutdown_cleanup_pending = False
+
         self._maintenance_stop.set()
-        for thread in (self._lease_thread, self._reconcile_thread):
-            if thread is not None and thread is not threading.current_thread():
-                thread.join(timeout=max(5.0, float(self._config["reconciliation_max_seconds"]) + 1.0))
+        join_timeout = max(5.0, float(self._config["reconciliation_max_seconds"]) + 1.0)
+        live_threads: list[str] = []
+        for name, thread in (
+            ("lease renewal", self._lease_thread),
+            ("reconciliation", self._reconcile_thread),
+        ):
+            if thread is None or thread is threading.current_thread():
+                continue
+            thread.join(timeout=join_timeout)
+            if thread.is_alive():
+                live_threads.append(name)
+
+        if live_threads:
+            with self._lock:
+                self._shutdown_cleanup_pending = True
+            raise _E2BMaintenanceShutdownTimeout("E2B maintenance thread shutdown timed out: " + ", ".join(live_threads))
+
         with self._lock:
             active = list(self._sandboxes.items())
             warm_ids = list(self._warm_pool.keys() | self._eviction_tombstones | self._remote_ops_in_progress)

backend/tests/test_e2b_sandbox_provider.py

@@ -8,6 +8,7 @@ import importlib
 import json
 import os
 import shutil
+import signal
 import subprocess
 import threading
 import time
@@ -278,6 +279,7 @@ def _make_provider(
     provider._transitioning_slots = 0
     provider._capacity_cond = threading.Condition(provider._lock)
     provider._shutdown_called = False
+    provider._shutdown_cleanup_pending = False
     provider._owner_id = "owner-a"
     provider._ownership = FakeOwnershipStore({}, owner_id=provider._owner_id)
     provider._ownership_config = SimpleNamespace(
@@ -3175,6 +3177,107 @@ def test_release_skips_warm_pool_when_sync_reveals_dead_vm(monkeypatch, tmp_path
     assert client.killed is True
 
 
+@pytest.mark.anyio
+async def test_shutdown_defers_teardown_and_fences_cached_acquire_while_maintenance_thread_is_alive():
+    p = _make_provider()
+    p._acquire_executor = ThreadPoolExecutor(
+        max_workers=1,
+        thread_name_prefix="e2b-shutdown-acquire-test",
+    )
+    client = FakeClient(sandbox_id="sb-owned")
+    sandbox = _make_sandbox(client, sandbox_id="sb-owned")
+    p._sandboxes = {"sb-owned": sandbox}
+    p._owned_sandbox_ids = {"sb-owned"}
+    p._thread_sandboxes[p._thread_key("thread-owned", "user-owned")] = "sb-owned"
+
+    class JoinControlledThread:
+        def __init__(self) -> None:
+            self.alive = True
+            self.join_timeout: float | None = None
+
+        def join(self, timeout: float | None = None) -> None:
+            self.join_timeout = timeout
+
+        def is_alive(self) -> bool:
+            return self.alive
+
+    lease_thread = JoinControlledThread()
+    p._lease_thread = lease_thread
+
+    with pytest.raises(RuntimeError, match="lease renewal"):
+        p.shutdown()
+
+    assert p._shutdown_called is True
+    assert p._shutdown_cleanup_pending is True
+    assert p._maintenance_stop.is_set()
+    assert p._sandboxes == {"sb-owned": sandbox}
+    assert p._thread_sandboxes[p._thread_key("thread-owned", "user-owned")] == "sb-owned"
+    assert client.killed is False
+    assert lease_thread.join_timeout == 11.0
+
+    with pytest.raises(SandboxCapacityExceededError) as sync_exc:
+        p.acquire("thread-owned", user_id="user-owned")
+    assert sync_exc.value.reason == "shutdown"
+
+    with pytest.raises(SandboxCapacityExceededError) as async_exc:
+        await p.acquire_async("thread-owned", user_id="user-owned")
+    assert async_exc.value.reason == "shutdown"
+
+    lease_thread.alive = False
+    p.shutdown()
+
+    assert p._shutdown_called is True
+    assert p._shutdown_cleanup_pending is False
+    assert p._sandboxes == {}
+    assert client.killed is True
+
+
+def test_atexit_shutdown_logs_pending_cleanup_without_raising(monkeypatch, caplog):
+    mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
+    p = _make_provider()
+
+    def blocked_shutdown() -> None:
+        with p._lock:
+            p._shutdown_called = True
+            p._shutdown_cleanup_pending = True
+        raise mod._E2BMaintenanceShutdownTimeout("E2B maintenance thread shutdown timed out: lease renewal")
+
+    monkeypatch.setattr(p, "shutdown", blocked_shutdown)
+
+    p._shutdown_at_exit()
+
+    assert p._shutdown_cleanup_pending is True
+    assert "still pending at interpreter exit" in caplog.text
+
+
+def test_signal_handler_forwards_original_action_when_shutdown_cleanup_is_pending(monkeypatch):
+    mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
+    p = _make_provider()
+    registered: dict[int, Any] = {}
+    forwarded: list[tuple[int, Any]] = []
+
+    def original_handler(signum, frame):
+        forwarded.append((signum, frame))
+
+    monkeypatch.setattr(signal, "getsignal", lambda _signum: original_handler)
+    monkeypatch.setattr(signal, "signal", lambda signum, handler: registered.__setitem__(signum, handler))
+
+    def blocked_shutdown() -> None:
+        with p._lock:
+            p._shutdown_called = True
+            p._shutdown_cleanup_pending = True
+        raise mod._E2BMaintenanceShutdownTimeout("E2B maintenance thread shutdown timed out: lease renewal")
+
+    monkeypatch.setattr(p, "shutdown", blocked_shutdown)
+    p._register_signal_handlers()
+
+    frame = object()
+    registered[signal.SIGTERM](signal.SIGTERM, frame)
+
+    assert forwarded == [(signal.SIGTERM, frame)]
+    assert p._shutdown_cleanup_pending is True
+
+
 def test_shutdown_only_kills_sandboxes_owned_by_current_instance(monkeypatch):
     p = _make_provider()
     owned_client = FakeClient(sandbox_id="sb-owned")