furyhawk/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-07-21 02:05:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(sandbox): consolidate E2B client lifecycle helpers (#4262)

Browse Source 
* refactor(sandbox): consolidate E2B client lifecycle

* test(sandbox): cover E2B lifecycle cleanup
This commit is contained in:
luo jiyin
2026-07-19 16:20:57 +08:00
committed by GitHub
parent 0f088033fe
commit 90d511f3d2
2 changed files with 84 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py
+63 -46
View File
@@ -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):
backend/tests/test_e2b_sandbox_provider.py
+21 -2
View File
@@ -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):