diff --git a/backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py b/backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py index 55b7f7873..ce0052ae4 100644 --- a/backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py +++ b/backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py @@ -271,7 +271,6 @@ class E2BSandboxProvider(SandboxProvider): return sid def _reclaim_warm_pool_sandbox(self, thread_id: str, *, user_id: str) -> str | None: - key = self._thread_key(thread_id, user_id) seed = self._stable_seed(thread_id, user_id) with self._lock: target_id = next( @@ -282,9 +281,8 @@ class E2BSandboxProvider(SandboxProvider): return None self._warm_pool.pop(target_id) - sandbox_cls = self._get_sandbox_cls() try: - client = self._reconnect_client(sandbox_cls, target_id) + client = self._reconnect_live_client(self._get_sandbox_cls(), target_id) except Exception as e: logger.warning( "Warm-pool e2b sandbox %s failed to reconnect, dropping: %s", @@ -293,18 +291,11 @@ class E2BSandboxProvider(SandboxProvider): ) return None - # Verify the reconnected client actually corresponds to a live VM. - # ``Sandbox.connect`` succeeds for paused/expired sandboxes too on - # some SDK versions, but the very next command then fails with - # "sandbox not found" mid-tool-call. Pinging here moves that failure - # into the acquire path, where we cleanly fall back to creating a - # fresh sandbox. - if not self._client_alive(client): + if client is None: logger.warning( "Warm-pool e2b sandbox %s is no longer alive (reaped by control plane); dropping and falling back to create", target_id, ) - self._safe_close_client(client) return None self._refresh_remote_timeout(client) @@ -312,10 +303,7 @@ class E2BSandboxProvider(SandboxProvider): self._bootstrap_sandbox_paths(client) except Exception as e: logger.debug("bootstrap on warm-pool reclaim failed: %s", e) - sandbox = E2BSandbox(id=target_id, client=client, home_dir=self._config["home_dir"]) - with self._lock: - self._sandboxes[target_id] = sandbox - self._thread_sandboxes[key] = target_id + self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) logger.info( "Reclaimed warm-pool e2b sandbox %s for user/thread %s/%s", target_id, @@ -410,7 +398,7 @@ class E2BSandboxProvider(SandboxProvider): return None try: - client = self._reconnect_client(sandbox_cls, target_id) + client = self._reconnect_live_client(sandbox_cls, target_id) except Exception as e: logger.warning( "Discovered e2b sandbox %s could not be reconnected: %s", @@ -419,12 +407,11 @@ class E2BSandboxProvider(SandboxProvider): ) return None - if not self._client_alive(client): + if client is None: logger.warning( "Discovered e2b sandbox %s is no longer alive; falling back to create", target_id, ) - self._safe_close_client(client) return None self._refresh_remote_timeout(client) @@ -432,10 +419,7 @@ class E2BSandboxProvider(SandboxProvider): self._bootstrap_sandbox_paths(client) except Exception as e: logger.debug("bootstrap on remote discovery failed: %s", e) - sandbox = E2BSandbox(id=target_id, client=client, home_dir=self._config["home_dir"]) - with self._lock: - self._sandboxes[target_id] = sandbox - self._thread_sandboxes[self._thread_key(thread_id, user_id)] = target_id + self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) logger.info( "Discovered remote e2b sandbox %s for user/thread %s/%s (seed=%s)", target_id, @@ -536,6 +520,37 @@ class E2BSandboxProvider(SandboxProvider): """Connect to an existing e2b sandbox by id, with consistent kwargs.""" return sandbox_cls.connect(sandbox_id, **self._common_kwargs()) # type: ignore[attr-defined] + def _reconnect_live_client( + self, + sandbox_cls: type[E2BClientSandbox], + sandbox_id: str, + ) -> E2BClientSandbox | None: + """Reconnect to *sandbox_id* and reject clients for reaped VMs. + + ``Sandbox.connect`` may succeed even after the E2B control plane has + reaped the VM. Closing that host-side client before returning ``None`` + keeps both acquire paths from leaking a connection. + """ + client = self._reconnect_client(sandbox_cls, sandbox_id) + if self._client_alive(client): + return client + self._safe_close_client(client) + return None + + def _register_connected_sandbox( + self, + sandbox_id: str, + client: E2BClientSandbox, + *, + thread_id: str, + user_id: str, + ) -> None: + """Track a live reconnected sandbox under its thread ownership.""" + sandbox = E2BSandbox(id=sandbox_id, client=client, home_dir=self._config["home_dir"]) + with self._lock: + self._sandboxes[sandbox_id] = sandbox + self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id + def _refresh_remote_timeout(self, client: E2BClientSandbox) -> None: """Push the configured idle timeout to the e2b control plane.""" idle_timeout = int(self._config["idle_timeout"]) @@ -998,23 +1013,33 @@ class E2BSandboxProvider(SandboxProvider): logger.info("Released e2b sandbox %s to warm pool", sandbox_id) def _kill_and_close(self, sandbox: E2BSandbox) -> None: - client = getattr(sandbox, "_client", None) - if client is not None: - kill = getattr(client, "kill", None) - if callable(kill): - try: - kill() - except Exception as e: - logger.debug( - "kill() on e2b sandbox %s raised (probably already gone): %s", - sandbox.id, - e, - ) + if error := self._kill_client(getattr(sandbox, "_client", None)): + logger.debug( + "kill() on e2b sandbox %s raised (probably already gone): %s", + sandbox.id, + error, + ) try: sandbox.close() except Exception: pass + @staticmethod + def _kill_client( + client: E2BClientSandbox | None, + ) -> Exception | None: + """Kill a remote VM and return an exception for the caller to log.""" + if client is None: + return + kill = getattr(client, "kill", None) + if not callable(kill): + return + try: + kill() + except Exception as e: + return e + return None + def reset(self) -> None: with self._lock: self._sandboxes.clear() @@ -1040,15 +1065,11 @@ class E2BSandboxProvider(SandboxProvider): ) for sandbox_id, sandbox in active: - try: - kill = getattr(sandbox.client, "kill", None) - if callable(kill): - kill() - except Exception as e: + if error := self._kill_client(sandbox.client): logger.warning( "Failed to kill active e2b sandbox %s during shutdown: %s", sandbox_id, - e, + error, ) try: sandbox.close() @@ -1066,15 +1087,11 @@ class E2BSandboxProvider(SandboxProvider): e, ) continue - try: - kill = getattr(client, "kill", None) - if callable(kill): - kill() - except Exception as e: + if error := self._kill_client(client): logger.warning( "Failed to kill warm-pool e2b sandbox %s during shutdown: %s", sandbox_id, - e, + error, ) close = getattr(client, "close", None) if callable(close): diff --git a/backend/tests/test_e2b_sandbox_provider.py b/backend/tests/test_e2b_sandbox_provider.py index 07b335e4c..740806a2d 100644 --- a/backend/tests/test_e2b_sandbox_provider.py +++ b/backend/tests/test_e2b_sandbox_provider.py @@ -458,7 +458,11 @@ def test_reclaim_warm_pool_sandbox_happy_path(monkeypatch): def test_reclaim_warm_pool_sandbox_drops_dead_entry(monkeypatch): p = _make_provider() fake_cls = _install_fake_sdk(monkeypatch, p) - fake_cls.connect_factory = lambda sid, **kw: FakeClient(sandbox_id=sid, commands=FakeCommandsAPI([FakeCommandsAPI.GONE])) + client = FakeClient( + sandbox_id="sb-zombie", + commands=FakeCommandsAPI([FakeCommandsAPI.GONE]), + ) + fake_cls.connect_factory = lambda _sid, **_kw: client seed = p._stable_seed("t1", "u1") p._warm_pool["sb-zombie"] = (seed, 12345.0) @@ -466,6 +470,7 @@ def test_reclaim_warm_pool_sandbox_drops_dead_entry(monkeypatch): assert sid is None assert "sb-zombie" not in p._sandboxes assert "sb-zombie" not in p._warm_pool + assert client.closed is True def test_reclaim_warm_pool_sandbox_handles_reconnect_exception(monkeypatch): @@ -550,10 +555,24 @@ def test_discover_remote_sandbox_skips_dead_candidate(monkeypatch): p = _make_provider() fake_cls = _install_fake_sdk(monkeypatch, p) fake_cls.list_return = [_info("sb-dead", "u1", "t1")] - fake_cls.connect_factory = lambda sid, **kw: FakeClient(sandbox_id=sid, commands=FakeCommandsAPI([FakeCommandsAPI.GONE])) + client = FakeClient( + sandbox_id="sb-dead", + commands=FakeCommandsAPI([FakeCommandsAPI.GONE]), + ) + fake_cls.connect_factory = lambda _sid, **_kw: client assert p._discover_remote_sandbox("t1", user_id="u1") is None assert ("u1", "t1") not in p._thread_sandboxes + assert client.closed is True + + +def test_kill_client_returns_exception_without_raising(): + p = _make_provider() + client = FakeClient() + error = RuntimeError("already gone") + client.kill = MagicMock(side_effect=error) + + assert p._kill_client(client) is error def test_discover_remote_sandbox_returns_none_when_list_raises(monkeypatch):