From b22f85c6864f4792d14a69b87a0f3a6a9cbc0bfd Mon Sep 17 00:00:00 2001 From: March-77 <132431415+March-77@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:10:24 +0800 Subject: [PATCH] fix(sandbox): reconcile E2B sandboxes safely (#4443) * fix(sandbox): reconcile E2B sandboxes safely * fix(sandbox): clear failed E2B adoption intent --- README.md | 2 +- backend/AGENTS.md | 10 +- backend/docs/CONFIGURATION.md | 26 +- .../community/e2b_sandbox/__init__.py | 9 + .../e2b_sandbox/e2b_sandbox_provider.py | 649 +++++++++++++++--- .../harness/deerflow/config/sandbox_config.py | 10 +- backend/tests/test_e2b_sandbox_provider.py | 503 ++++++++++++++ 7 files changed, 1105 insertions(+), 104 deletions(-) diff --git a/README.md b/README.md index 94857b64e..82f923394 100644 --- a/README.md +++ b/README.md @@ -288,7 +288,7 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser > > With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation. > -> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. +> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. When multiple Gateway workers share the Docker/AIO or E2B sandbox backend, also configure `sandbox.ownership.type: redis`; E2B uses the leases during background startup and periodic reconciliation so duplicate/orphan cleanup cannot terminate a live peer's sandbox. See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 7960776ed..1cfcc2c37 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -547,7 +547,15 @@ copying a raw checkpoint because delta state is not self-contained in one tuple. failed kill through a new client. An unconfirmed remote ID stays tracked. `reset()` uses full shutdown semantics. It destroys tracked active and warm VMs. It wakes capacity waiters. Callers cannot reuse the old provider - instance. + instance. A background startup pass and periodic reconciliation list + provider-tagged remote sandboxes within page/item/time budgets, probe every + candidate until a healthy canonical sandbox is found, adopt only through the + shared ownership store, and reap duplicates/orphans only after their + configured grace/TTL and an atomic `del:` claim. Lease renewal is independent + of reconciliation. Failed canonical adoption clears both its capacity + reservation and acquire-intent marker even if a peer takes ownership between + the initial claim and bootstrap cleanup. Shutdown kills only IDs whose leases + are owned by this provider instance; peer-owned clients are merely closed. - **Cross-instance ownership store** (`aio_sandbox/ownership/`, #4206): gateway instances sharing a container backend coordinate container ownership through a pluggable lease store, selected by `sandbox.ownership.type` (`memory` | `redis`) and resolved like `stream_bridge` (`factory.py`, lazy per-branch import, `redis` optional extra, `DEER_FLOW_SANDBOX_OWNERSHIP_REDIS_URL` env escape hatch; a set `DEER_FLOW_STREAM_BRIDGE_REDIS_URL` implies a multi-instance deployment and infers `redis`). `memory` is single-instance only and declares `supports_cross_process = False`. - **A lease answers "who reaps this container", not "who may use it".** That splits the interface in two: `take()` transfers ownership on the **acquire** path (a container is deterministic per user/thread, so consecutive turns legitimately land on different instances — a conditional claim there would strand the thread until the previous lease expired), while `claim()` succeeds only if the container is unowned or already ours and gates every **adopt/reap** path. `release()` never clears a peer's lease. - **A lease carries a state, and that is what makes the destroy window safe.** `own:` = responsible for this container; `del:` = tearing it down (`claim(..., for_destroy=True)`). `take()` is refused against a `del:` lease, so a container cannot be re-acquired between a destroy path's claim and its container stop. Without the two states an unconditional `take()` would silently overwrite the destroyer's claim and the peer's stop would land on a container the new owner had already handed to an agent — i.e. #4206 again. That pairing is what replaced the previous same-host `flock` guard, which is gone; Redis makes the scope genuinely multi-instance instead of same-host. A destroyer that dies mid-stop leaves a `del:` marker that lapses with the TTL. On the acquire path a refused take raises `SandboxBeingDestroyedError`: the reuse/reclaim paths drop the container and cold-start, and the discover path propagates (falling through to create would collide with the not-yet-removed container name). diff --git a/backend/docs/CONFIGURATION.md b/backend/docs/CONFIGURATION.md index e90355195..f3c0c02a5 100644 --- a/backend/docs/CONFIGURATION.md +++ b/backend/docs/CONFIGURATION.md @@ -426,6 +426,15 @@ sandbox: home_dir: /home/user # /mnt/user-data is remapped under this directory idle_timeout: 600 # forwarded to e2b's server-side set_timeout() replicas: 3 # max concurrent sandboxes per gateway process + ownership: # use Redis when more than one gateway shares E2B + type: redis + redis_url: $REDIS_URL + reconciliation_interval_seconds: 60 + reconciliation_grace_seconds: 120 + reconciliation_orphan_ttl_seconds: 3600 + reconciliation_max_pages: 10 + reconciliation_max_items: 200 + reconciliation_max_seconds: 15 mounts: # one-shot upload of host files at sandbox start - host_path: /path/on/host container_path: /home/user/shared @@ -440,10 +449,19 @@ provider in `config.yaml`. Notes specific to `E2BSandboxProvider`: -- Each DeerFlow thread is bound to its e2b sandbox via metadata - (`deer_flow_user`, `deer_flow_thread`), so the same thread reuses the same - sandbox across gateway restarts and across processes — no cross-process - file lock is needed because the e2b control plane is the source of truth. +- Each DeerFlow thread is bound to its E2B sandbox via metadata + (`deer_flow_user`, `deer_flow_thread`). Startup and periodic reconciliation + probe every bounded candidate, adopt one healthy canonical sandbox, and reap + duplicates after a grace period. Provider-tagged entries without a complete + user/thread identity are reaped only after the orphan TTL. +- Ownership leases prevent one gateway from adopting or destroying a sandbox + another live gateway is responsible for. The default in-memory store is safe + only for one gateway process. Multi-worker/load-balanced deployments must use + `sandbox.ownership.type: redis`; an existing Redis stream bridge configuration + is inferred automatically. +- Reconciliation is bounded by page, item, and wall-clock limits. Its summary log + exposes discovered, adopted, duplicate, deferred, killed, dead, and budget-exhausted + counts for operational monitoring. - Idle expiry is enforced server-side by e2b's `set_timeout()`. The provider refreshes the timeout on every release so warm sandboxes stay alive long enough for the next acquire. diff --git a/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py b/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py index 1a92f8eb6..289500a51 100644 --- a/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py +++ b/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py @@ -13,6 +13,15 @@ Configuration example (``config.yaml``):: domain: e2b.dev # optional e2b domain (e.g. self-hosted) idle_timeout: 600 # forwarded to e2b ``set_timeout`` (seconds) replicas: 3 # max concurrent sandboxes (LRU eviction beyond) + ownership: # required for safe multi-worker reconciliation + type: redis + redis_url: $REDIS_URL + reconciliation_interval_seconds: 60 + reconciliation_grace_seconds: 120 + reconciliation_orphan_ttl_seconds: 3600 + reconciliation_max_pages: 10 + reconciliation_max_items: 200 + reconciliation_max_seconds: 15 mounts: # one-shot upload of host files into the sandbox - host_path: /path/on/host container_path: /path/in/sandbox 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 9197707ff..d4124b652 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 @@ -38,6 +38,7 @@ import time import uuid from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from decimal import Decimal, InvalidOperation from functools import partial from pathlib import Path @@ -51,6 +52,14 @@ from deerflow.sandbox.exceptions import SandboxCapacityExceededError from deerflow.sandbox.sandbox import Sandbox from deerflow.sandbox.sandbox_provider import SandboxProvider +from ..aio_sandbox.ownership import ( + OwnershipBackendError, + RenewOutcome, + SandboxOwnershipStore, + generate_owner_id, + make_sandbox_ownership_store, + resolve_ownership_config, +) from .e2b_sandbox import DEFAULT_E2B_HOME_DIR, E2BSandbox, _is_sandbox_gone_error logger = logging.getLogger(__name__) @@ -62,6 +71,12 @@ DEFAULT_IDLE_TIMEOUT = 1800 # 30 minutes; passed to ``Sandbox.set_timeout``. DEFAULT_REPLICAS = 3 DEFAULT_OVERFLOW_POLICY = "wait" # wait | reject | burst DEFAULT_ACQUIRE_TIMEOUT = 30 # seconds for wait policy +DEFAULT_RECONCILIATION_INTERVAL_SECONDS = 60.0 +DEFAULT_RECONCILIATION_GRACE_SECONDS = 120.0 +DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS = 3600.0 +DEFAULT_RECONCILIATION_MAX_PAGES = 10 +DEFAULT_RECONCILIATION_MAX_ITEMS = 200 +DEFAULT_RECONCILIATION_MAX_SECONDS = 15.0 # Hard upper bound for ``set_timeout`` (e2b currently caps at 24h on the # free plan; passing an excessive value is rejected by the control-plane). MAX_E2B_TIMEOUT = 24 * 60 * 60 @@ -71,10 +86,25 @@ MAX_E2B_TIMEOUT = 24 * 60 * 60 META_KEY_USER = "deer_flow_user" META_KEY_THREAD = "deer_flow_thread" META_KEY_PROVIDER = "deer_flow_provider" +META_KEY_GATEWAY = "deer_flow_gateway" +META_KEY_CREATED_AT = "deer_flow_created_at" META_VAL_PROVIDER = "e2b_sandbox_provider" E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"}) +@dataclass +class ReconciliationStats: + """Bounded reconciliation outcome, also suitable for metrics/logging.""" + + discovered: int = 0 + adopted: int = 0 + duplicates: int = 0 + deferred: int = 0 + killed: int = 0 + dead: int = 0 + budget_exhausted: bool = False + + class E2BSandboxProvider(SandboxProvider): """Sandbox provider backed by the e2b code-interpreter cloud SDK.""" @@ -116,6 +146,13 @@ class E2BSandboxProvider(SandboxProvider): self._transitioning_slots = 0 self._capacity_cond = threading.Condition(self._lock) self._shutdown_called = False + self._owned_sandbox_ids: set[str] = set() + self._acquire_inflight: set[str] = set() + self._orphan_first_seen: dict[str, float] = {} + self._maintenance_stop = threading.Event() + self._lease_thread: threading.Thread | None = None + self._reconcile_thread: threading.Thread | None = None + self._owner_id = generate_owner_id() self._config = self._load_config() acquire_workers = max(4, min(32, self._capacity_limit() + 1)) @@ -123,9 +160,20 @@ class E2BSandboxProvider(SandboxProvider): max_workers=acquire_workers, thread_name_prefix="e2b-sandbox-acquire", ) + self._ownership_config = resolve_ownership_config( + self._config.get("ownership"), + stream_bridge=self._config.get("stream_bridge"), + ) + self._ownership: SandboxOwnershipStore = make_sandbox_ownership_store( + self._ownership_config, + owner_id=self._owner_id, + ) + 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) self._register_signal_handlers() + self._start_maintenance_threads() def _load_config(self) -> dict[str, Any]: """Read e2b options off ``SandboxConfig`` (``extra="allow"``).""" @@ -181,6 +229,32 @@ class E2BSandboxProvider(SandboxProvider): "burst_limit": burst_limit, "mounts": _opt("mounts") or [], "environment": self._resolve_env_vars(_opt("environment") or {}), + "ownership": _opt("ownership"), + "stream_bridge": getattr(get_app_config(), "stream_bridge", None), + "reconciliation_interval_seconds": max( + 1.0, + float(_opt("reconciliation_interval_seconds", DEFAULT_RECONCILIATION_INTERVAL_SECONDS)), + ), + "reconciliation_grace_seconds": max( + 0.0, + float(_opt("reconciliation_grace_seconds", DEFAULT_RECONCILIATION_GRACE_SECONDS)), + ), + "reconciliation_orphan_ttl_seconds": max( + 0.0, + float(_opt("reconciliation_orphan_ttl_seconds", DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS)), + ), + "reconciliation_max_pages": max( + 1, + int(_opt("reconciliation_max_pages", DEFAULT_RECONCILIATION_MAX_PAGES)), + ), + "reconciliation_max_items": max( + 1, + int(_opt("reconciliation_max_items", DEFAULT_RECONCILIATION_MAX_ITEMS)), + ), + "reconciliation_max_seconds": max( + 0.1, + float(_opt("reconciliation_max_seconds", DEFAULT_RECONCILIATION_MAX_SECONDS)), + ), } @staticmethod @@ -322,6 +396,9 @@ 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) + 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", @@ -370,13 +447,17 @@ class E2BSandboxProvider(SandboxProvider): self._complete_transition_remote_op(target_id, remote_destroyed=True) return None + try: + self._publish_ownership(target_id) + except Exception: + self._complete_transition_remote_op(target_id, remote_destroyed=False) + self._safe_close_client(client) + raise + self._refresh_remote_timeout(client) bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id) if bootstrap_error is not None: - if remote_destroyed: - self._complete_transition_remote_op(target_id, remote_destroyed=True) - else: - self._complete_transition_remote_op(target_id, remote_destroyed=False) + self._complete_transition_remote_op(target_id, remote_destroyed=remote_destroyed) return None discard_after_shutdown = False @@ -393,10 +474,11 @@ class E2BSandboxProvider(SandboxProvider): self._end_transition_locked() if discard_after_shutdown: - self._kill_client(client) + if self._claim_ownership(target_id, for_destroy=True): + self._kill_client(client) + self._release_ownership(target_id) self._safe_close_client(client) return None - logger.info( "Reclaimed warm-pool e2b sandbox %s for user/thread %s/%s", target_id, @@ -414,81 +496,39 @@ class E2BSandboxProvider(SandboxProvider): """ sandbox_cls = self._get_sandbox_cls() seed = self._stable_seed(thread_id, user_id) - list_kwargs = self._common_kwargs() - try: - running = sandbox_cls.list( # type: ignore[attr-defined] - query={ - "metadata": { - META_KEY_PROVIDER: META_VAL_PROVIDER, - META_KEY_USER: user_id, - META_KEY_THREAD: thread_id, - } - }, - **list_kwargs, + entries, _ = self._list_remote_entries( + { + META_KEY_PROVIDER: META_VAL_PROVIDER, + META_KEY_USER: user_id, + META_KEY_THREAD: thread_id, + } + ) + candidates = sorted( + ((sandbox_id, metadata) for entry in entries if (sandbox_id := self._entry_id(entry)) and (metadata := self._entry_metadata(entry)).get(META_KEY_USER) == user_id and metadata.get(META_KEY_THREAD) == thread_id), + key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]), + ) + for target_id, _metadata in candidates: + adopted = self._adopt_remote_candidate( + sandbox_cls, + target_id, + thread_id=thread_id, + user_id=user_id, + seed=seed, ) - except TypeError: - try: - running = sandbox_cls.list( - metadata={ - META_KEY_PROVIDER: META_VAL_PROVIDER, - META_KEY_USER: user_id, - META_KEY_THREAD: thread_id, - }, - **list_kwargs, - ) - except Exception as e: - logger.debug("e2b Sandbox.list() unavailable, skipping discovery: %s", e) - return None - except Exception as e: - logger.debug( - "e2b Sandbox.list() raised while discovering thread %s: %s", - thread_id, - e, - ) - return None + if adopted is not None: + return adopted + return None - # Pick the first matching candidate; tolerate either ``SandboxInfo`` - # objects with ``sandbox_id`` or plain dicts. - # Normalise the return value of ``Sandbox.list()``: - # * Older SDKs (<= 1.x) returned a plain ``list[SandboxInfo]`` — directly iterable. - # * e2b-code-interpreter >= 2.x returns a ``SandboxPaginator`` exposing - # ``has_next: bool`` and ``next_items() -> list[SandboxInfo]`` instead - # of being iterable. Walking pages keeps discovery correct when the - # org has more sandboxes than fit in a single page. - def _iter_running(obj): - if obj is None: - return - if hasattr(obj, "next_items") and hasattr(obj, "has_next"): - for _ in range(50): - try: - page = obj.next_items() - except Exception as exc: - logger.debug("SandboxPaginator.next_items() failed: %s", exc) - return - if not page: - return - yield from page - if not getattr(obj, "has_next", False): - return - return - try: - yield from obj - except TypeError: - logger.debug("Sandbox.list() returned non-iterable %s; ignoring", type(obj).__name__) - - target_id: str | None = None - for entry in _iter_running(running): - sid = getattr(entry, "sandbox_id", None) or (entry.get("sandbox_id") if isinstance(entry, dict) else None) - metadata = getattr(entry, "metadata", None) or (entry.get("metadata") if isinstance(entry, dict) else {}) or {} - if metadata.get(META_KEY_USER) != user_id: - continue - if metadata.get(META_KEY_THREAD) != thread_id: - continue - target_id = sid - break - - if not target_id: - return None + def _adopt_remote_candidate( + self, + sandbox_cls: type[E2BClientSandbox], + target_id: str, + *, + thread_id: str, + user_id: str, + seed: str, + ) -> str | None: + """Try to adopt one discovered candidate without harming peer-owned VMs.""" try: client = self._reconnect_live_client(sandbox_cls, target_id) @@ -528,6 +568,20 @@ class E2BSandboxProvider(SandboxProvider): self._safe_close_client(client) raise + with self._lock: + shutdown_before_ownership = self._shutdown_called + if shutdown_before_ownership: + self._complete_reserved_remote_op(target_id, remote_destroyed=False) + self._safe_close_client(client) + return None + + try: + self._publish_ownership(target_id) + except Exception: + self._complete_reserved_remote_op(target_id, remote_destroyed=False) + self._safe_close_client(client) + raise + self._refresh_remote_timeout(client) bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id) if bootstrap_error is not None: @@ -543,6 +597,9 @@ class E2BSandboxProvider(SandboxProvider): self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) self._commit_capacity() if discard_after_shutdown: + if self._claim_ownership(target_id, for_destroy=True): + self._kill_client(client) + self._release_ownership(target_id) self._safe_close_client(client) return None logger.info( @@ -787,6 +844,8 @@ class E2BSandboxProvider(SandboxProvider): sandbox_cls = self._get_sandbox_cls() metadata: dict[str, str] = { META_KEY_PROVIDER: META_VAL_PROVIDER, + META_KEY_GATEWAY: self._owner_id, + META_KEY_CREATED_AT: str(time.time()), } if thread_id: metadata[META_KEY_USER] = user_id @@ -849,6 +908,17 @@ class E2BSandboxProvider(SandboxProvider): reason="shutdown", ) + try: + self._publish_ownership(sandbox_id) + except Exception: + remote_destroyed = False + if self._claim_ownership(sandbox_id, for_destroy=True): + remote_destroyed = self._kill_client(client) is None + self._release_ownership(sandbox_id) + self._safe_close_client(client) + self._complete_reserved_remote_op(sandbox_id, remote_destroyed=remote_destroyed) + raise + # Materialise DeerFlow's virtual path layout (/mnt/user-data/...) inside # the e2b VM. Without this step shell commands the agent emits — which # use the same /mnt/user-data prefix as LocalSandbox / AioSandbox — fail @@ -883,7 +953,9 @@ class E2BSandboxProvider(SandboxProvider): self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id if should_kill: - self._kill_client(client) + if self._claim_ownership(sandbox_id, for_destroy=True): + self._kill_client(client) + self._release_ownership(sandbox_id) self._safe_close_client(client) raise SandboxCapacityExceededError( f"Sandbox provider shut down during sandbox creation; killed remote sandbox {sandbox_id}", @@ -912,6 +984,339 @@ class E2BSandboxProvider(SandboxProvider): kwargs["domain"] = self._config["domain"] return kwargs + @staticmethod + def _entry_id(entry: Any) -> str | None: + value = getattr(entry, "sandbox_id", None) + if value is None and isinstance(entry, dict): + value = entry.get("sandbox_id") + return str(value) if value else None + + @staticmethod + def _entry_metadata(entry: Any) -> dict[str, Any]: + value = getattr(entry, "metadata", None) + if value is None and isinstance(entry, dict): + value = entry.get("metadata") + return value if isinstance(value, dict) else {} + + def _list_remote_entries(self, metadata: dict[str, str]) -> tuple[list[Any], bool]: + """List matching E2B entries within configured page/item/time budgets.""" + sandbox_cls = self._get_sandbox_cls() + try: + result = sandbox_cls.list(query={"metadata": metadata}, **self._common_kwargs()) # type: ignore[attr-defined] + except TypeError: + try: + result = sandbox_cls.list(metadata=metadata, **self._common_kwargs()) # type: ignore[attr-defined] + except Exception as e: + logger.warning("E2B reconciliation list failed: %s", e) + return [], False + except Exception as e: + logger.warning("E2B reconciliation list failed: %s", e) + return [], False + + max_pages = int(self._config["reconciliation_max_pages"]) + max_items = int(self._config["reconciliation_max_items"]) + deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"]) + entries: list[Any] = [] + exhausted = False + + if hasattr(result, "next_items") and hasattr(result, "has_next"): + for page_number in range(max_pages): + if time.monotonic() >= deadline or len(entries) >= max_items: + exhausted = True + break + try: + page = result.next_items() + except Exception as e: + logger.warning("E2B reconciliation paginator failed: %s", e) + break + if not page: + break + room = max_items - len(entries) + entries.extend(list(page)[:room]) + if len(page) > room: + exhausted = True + break + if not getattr(result, "has_next", False): + break + if page_number + 1 >= max_pages: + exhausted = True + else: + try: + all_entries = list(result or []) + except TypeError: + logger.warning("E2B Sandbox.list returned non-iterable %s", type(result).__name__) + return [], False + entries = all_entries[:max_items] + exhausted = len(all_entries) > max_items + + if time.monotonic() >= deadline: + exhausted = True + return entries, exhausted + + def _publish_ownership(self, sandbox_id: str) -> None: + """Publish acquire-side ownership before exposing a sandbox locally.""" + with self._lock: + self._acquire_inflight.add(sandbox_id) + try: + if not self._ownership.take(sandbox_id): + raise RuntimeError(f"E2B sandbox {sandbox_id} is being destroyed") + except Exception: + with self._lock: + self._acquire_inflight.discard(sandbox_id) + raise + with self._lock: + self._owned_sandbox_ids.add(sandbox_id) + + def _claim_ownership(self, sandbox_id: str, *, for_destroy: bool = False) -> bool: + """Exclusively claim an unowned sandbox, failing closed on store errors.""" + try: + claimed = self._ownership.claim(sandbox_id, for_destroy=for_destroy) + except OwnershipBackendError as e: + logger.warning("E2B ownership claim failed for %s: %s", sandbox_id, e) + return False + if claimed: + with self._lock: + self._owned_sandbox_ids.add(sandbox_id) + return claimed + + def _release_ownership(self, sandbox_id: str) -> None: + try: + self._ownership.release(sandbox_id) + except OwnershipBackendError as e: + logger.warning("Failed to release E2B ownership for %s: %s", sandbox_id, e) + with self._lock: + self._owned_sandbox_ids.discard(sandbox_id) + self._acquire_inflight.discard(sandbox_id) + + def _forget_local_sandbox(self, sandbox_id: str) -> None: + """Forget a lease taken by a peer without touching the remote VM.""" + with self._lock: + if sandbox_id in self._acquire_inflight: + return + sandbox = self._sandboxes.pop(sandbox_id, None) + self._warm_pool.pop(sandbox_id, None) + self._owned_sandbox_ids.discard(sandbox_id) + for key, sid in list(self._thread_sandboxes.items()): + if sid == sandbox_id: + self._thread_sandboxes.pop(key, None) + if sandbox is not None: + try: + sandbox.close() + except Exception: + pass + + def _refresh_owned_leases(self) -> None: + with self._lock: + sandbox_ids = list(self._owned_sandbox_ids) + for sandbox_id in sandbox_ids: + try: + outcome = self._ownership.renew(sandbox_id) + except OwnershipBackendError as e: + logger.warning("Could not renew E2B ownership for %s; will retry: %s", sandbox_id, e) + continue + if outcome is RenewOutcome.RENEWED: + continue + if outcome is RenewOutcome.LAPSED: + try: + if self._ownership.claim(sandbox_id): + continue + except OwnershipBackendError as e: + logger.warning("Could not re-establish E2B ownership for %s: %s", sandbox_id, e) + continue + logger.info("E2B sandbox %s ownership moved to a peer; forgetting local client", sandbox_id) + self._forget_local_sandbox(sandbox_id) + + def _start_maintenance_threads(self) -> None: + def renew() -> None: + interval = self._ownership_config.renewal_interval_seconds + while not self._maintenance_stop.wait(interval): + self._refresh_owned_leases() + + def reconcile() -> None: + interval = float(self._config["reconciliation_interval_seconds"]) + while not self._maintenance_stop.is_set(): + try: + self._reconcile_remote_sandboxes() + except Exception: + logger.exception("Periodic E2B sandbox reconciliation failed") + if self._maintenance_stop.wait(interval): + break + + self._lease_thread = threading.Thread(target=renew, name="e2b-lease-renewal", daemon=True) + self._reconcile_thread = threading.Thread(target=reconcile, name="e2b-reconciliation", daemon=True) + self._lease_thread.start() + self._reconcile_thread.start() + + def _reserve_reconciliation_capacity(self, sandbox_id: str) -> bool: + """Reserve one local slot for adoption without blocking maintenance.""" + with self._lock: + if self._shutdown_called or self._total_capacity_used_locked() >= self._capacity_limit(): + return False + self._reserved_slots += 1 + self._unowned_remote_ops_in_progress.add(sandbox_id) + self._acquire_inflight.add(sandbox_id) + return True + + def _reconcile_remote_sandboxes(self, *, now: float | None = None) -> ReconciliationStats: + """Adopt canonical E2B sandboxes and safely reap duplicates/orphans.""" + observed_at = time.monotonic() if now is None else now + deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"]) + stats = ReconciliationStats() + entries, stats.budget_exhausted = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER}) + stats.discovered = len(entries) + groups: dict[tuple[str, str], list[tuple[str, dict[str, Any]]]] = {} + orphans: list[tuple[str, dict[str, Any]]] = [] + present_ids: set[str] = set() + + for entry in entries: + sandbox_id = self._entry_id(entry) + if not sandbox_id: + continue + present_ids.add(sandbox_id) + metadata = self._entry_metadata(entry) + user_id = metadata.get(META_KEY_USER) + thread_id = metadata.get(META_KEY_THREAD) + if isinstance(user_id, str) and user_id and isinstance(thread_id, str) and thread_id: + groups.setdefault((user_id, thread_id), []).append((sandbox_id, metadata)) + else: + orphans.append((sandbox_id, metadata)) + + for (user_id, thread_id), candidates in groups.items(): + candidates.sort(key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0])) + with self._lock: + local_id = self._thread_sandboxes.get((user_id, thread_id)) + if local_id: + candidates.sort(key=lambda item: item[0] != local_id) + + live: list[tuple[str, dict[str, Any], E2BClientSandbox]] = [] + for sandbox_id, metadata in candidates: + if time.monotonic() >= deadline: + stats.budget_exhausted = True + break + try: + client = self._reconnect_live_client(self._get_sandbox_cls(), sandbox_id) + except Exception as e: + logger.debug("Could not probe E2B sandbox %s during reconciliation: %s", sandbox_id, e) + continue + if client is None: + stats.dead += 1 + continue + live.append((sandbox_id, metadata, client)) + + if not live: + continue + stats.duplicates += max(0, len(live) - 1) + canonical_id, _metadata, canonical_client = live[0] + with self._lock: + already_local = canonical_id in self._sandboxes + if already_local: + self._safe_close_client(canonical_client) + elif not self._reserve_reconciliation_capacity(canonical_id): + self._safe_close_client(canonical_client) + stats.deferred += 1 + elif not self._claim_ownership(canonical_id): + self._complete_reserved_remote_op(canonical_id, remote_destroyed=False) + with self._lock: + self._acquire_inflight.discard(canonical_id) + self._safe_close_client(canonical_client) + stats.deferred += 1 + else: + bootstrap_error, remote_destroyed = self._bootstrap_or_discard(canonical_client, canonical_id) + if bootstrap_error is not None: + self._complete_reserved_remote_op(canonical_id, remote_destroyed=remote_destroyed) + with self._lock: + self._acquire_inflight.discard(canonical_id) + stats.deferred += 1 + else: + discard_after_shutdown = False + with self._lock: + if self._shutdown_called: + discard_after_shutdown = True + else: + self._owned_sandbox_ids.add(canonical_id) + self._unowned_remote_ops_in_progress.discard(canonical_id) + self._register_connected_sandbox( + canonical_id, + canonical_client, + thread_id=thread_id, + user_id=user_id, + ) + self._commit_capacity() + if discard_after_shutdown: + if self._claim_ownership(canonical_id, for_destroy=True): + self._kill_client(canonical_client) + self._release_ownership(canonical_id) + self._safe_close_client(canonical_client) + else: + stats.adopted += 1 + + for sandbox_id, _metadata, client in live[1:]: + first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at) + if observed_at - first_seen < float(self._config["reconciliation_grace_seconds"]): + self._safe_close_client(client) + stats.deferred += 1 + continue + if not self._claim_ownership(sandbox_id, for_destroy=True): + self._safe_close_client(client) + stats.deferred += 1 + continue + if error := self._kill_client(client): + logger.warning("Failed to kill duplicate E2B sandbox %s: %s", sandbox_id, error) + self._release_ownership(sandbox_id) + stats.deferred += 1 + else: + stats.killed += 1 + self._orphan_first_seen.pop(sandbox_id, None) + self._forget_local_sandbox(sandbox_id) + self._release_ownership(sandbox_id) + self._safe_close_client(client) + + for sandbox_id, metadata in orphans: + if time.monotonic() >= deadline: + stats.budget_exhausted = True + break + first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at) + created_at = metadata.get(META_KEY_CREATED_AT) + try: + age = time.time() - float(created_at) if created_at is not None else observed_at - first_seen + except (TypeError, ValueError): + age = observed_at - first_seen + if age < float(self._config["reconciliation_orphan_ttl_seconds"]): + stats.deferred += 1 + continue + if not self._claim_ownership(sandbox_id, for_destroy=True): + stats.deferred += 1 + continue + try: + client = self._reconnect_client(self._get_sandbox_cls(), sandbox_id) + except Exception: + self._release_ownership(sandbox_id) + continue + if error := self._kill_client(client): + logger.warning("Failed to kill orphan E2B sandbox %s: %s", sandbox_id, error) + self._release_ownership(sandbox_id) + stats.deferred += 1 + else: + stats.killed += 1 + self._orphan_first_seen.pop(sandbox_id, None) + self._forget_local_sandbox(sandbox_id) + self._release_ownership(sandbox_id) + self._safe_close_client(client) + + for sandbox_id in set(self._orphan_first_seen) - present_ids: + self._orphan_first_seen.pop(sandbox_id, None) + logger.info( + "E2B reconciliation: discovered=%d adopted=%d duplicates=%d deferred=%d killed=%d dead=%d budget_exhausted=%s", + stats.discovered, + stats.adopted, + stats.duplicates, + stats.deferred, + stats.killed, + stats.dead, + stats.budget_exhausted, + ) + return stats + def _reconnect_client(self, sandbox_cls: type[E2BClientSandbox], sandbox_id: str) -> E2BClientSandbox: """Connect to an existing e2b sandbox by id, with consistent kwargs.""" return sandbox_cls.connect(sandbox_id, **self._common_kwargs()) # type: ignore[attr-defined] @@ -938,7 +1343,7 @@ class E2BSandboxProvider(SandboxProvider): sandbox_id: str, client: E2BClientSandbox, *, - thread_id: str, + thread_id: str | None, user_id: str, ) -> None: """Track a live reconnected sandbox under its thread ownership. @@ -947,7 +1352,10 @@ class E2BSandboxProvider(SandboxProvider): """ sandbox = E2BSandbox(id=sandbox_id, client=client, home_dir=self._config["home_dir"]) self._sandboxes[sandbox_id] = sandbox - self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id + self._warm_pool.pop(sandbox_id, None) + if thread_id: + self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id + self._acquire_inflight.discard(sandbox_id) def _refresh_remote_timeout(self, client: E2BClientSandbox) -> None: """Push the configured idle timeout to the e2b control plane.""" @@ -1016,11 +1424,20 @@ class E2BSandboxProvider(SandboxProvider): self._bootstrap_sandbox_paths(client) except Exception as e: logger.exception("Failed to bootstrap e2b sandbox %s. Discarding the unusable sandbox.", sandbox_id) - kill_error = self._kill_client(client) - if kill_error: - logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error) + remote_destroyed = False + if self._claim_ownership(sandbox_id, for_destroy=True): + kill_error = self._kill_client(client) + remote_destroyed = kill_error is None + if kill_error: + logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error) + self._release_ownership(sandbox_id) + else: + logger.info( + "Not killing E2B sandbox %s after bootstrap failure because a peer owns it", + sandbox_id, + ) self._safe_close_client(client) - return e, kill_error is None + return e, remote_destroyed return None, True def _bootstrap_sandbox_paths(self, client: E2BClientSandbox) -> None: @@ -1459,6 +1876,15 @@ class E2BSandboxProvider(SandboxProvider): else: return None + if not self._claim_ownership(evict_id, for_destroy=True): + logger.info("Deferring warm-pool eviction for %s because a peer owns it", evict_id) + self._forget_local_sandbox(evict_id) + with self._lock: + if evict_id in self._evictions_in_progress: + self._evictions_in_progress.discard(evict_id) + self._eviction_tombstones.discard(evict_id) + self._end_transition_locked() + return evict_id try: client = self._reconnect_live_client(self._get_sandbox_cls(), evict_id) except Exception as e: @@ -1472,6 +1898,7 @@ class E2BSandboxProvider(SandboxProvider): self._evictions_in_progress.discard(evict_id) if not self._shutdown_called: self._eviction_tombstones.add(evict_id) + self._release_ownership(evict_id) return None if client is None: @@ -1480,6 +1907,7 @@ class E2BSandboxProvider(SandboxProvider): self._evictions_in_progress.discard(evict_id) self._eviction_tombstones.discard(evict_id) self._end_transition_locked() + self._release_ownership(evict_id) logger.info("Evicted warm-pool e2b sandbox %s was already gone", evict_id) return evict_id @@ -1491,6 +1919,7 @@ class E2BSandboxProvider(SandboxProvider): self._evictions_in_progress.discard(evict_id) if not self._shutdown_called: self._eviction_tombstones.add(evict_id) + self._release_ownership(evict_id) return None self._safe_close_client(client) @@ -1499,6 +1928,7 @@ class E2BSandboxProvider(SandboxProvider): self._evictions_in_progress.discard(evict_id) self._eviction_tombstones.discard(evict_id) self._end_transition_locked() + self._release_ownership(evict_id) logger.info("Evicted warm-pool e2b sandbox %s", evict_id) return evict_id @@ -1608,8 +2038,10 @@ class E2BSandboxProvider(SandboxProvider): "Provider shut down during release of sandbox %s; killing instead of parking in warm pool", sandbox_id, ) - if error := self._kill_client(client): - logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error) + if self._claim_ownership(sandbox_id, for_destroy=True): + if error := self._kill_client(client): + logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error) + self._release_ownership(sandbox_id) self._safe_close_client(client) return @@ -1622,12 +2054,20 @@ class E2BSandboxProvider(SandboxProvider): self._free_transitioning_slot() def _kill_and_close(self, sandbox: E2BSandbox) -> None: + if not self._claim_ownership(sandbox.id, for_destroy=True): + logger.info("Not killing E2B sandbox %s because a peer owns it", sandbox.id) + try: + sandbox.close() + except Exception: + pass + return if error := self._kill_client(getattr(sandbox, "_client", None)): logger.debug( "kill() on e2b sandbox %s raised (probably already gone): %s", sandbox.id, error, ) + self._release_ownership(sandbox.id) try: sandbox.close() except Exception: @@ -1658,8 +2098,14 @@ class E2BSandboxProvider(SandboxProvider): if self._shutdown_called: return self._shutdown_called = True + 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)) + with self._lock: active = list(self._sandboxes.items()) warm_ids = list(self._warm_pool.keys() | self._eviction_tombstones | self._remote_ops_in_progress) + owned_ids = set(self._owned_sandbox_ids) self._sandboxes.clear() self._warm_pool.clear() self._eviction_tombstones.clear() @@ -1668,6 +2114,9 @@ class E2BSandboxProvider(SandboxProvider): self._unowned_remote_ops_in_progress.clear() self._thread_sandboxes.clear() self._thread_locks.clear() + self._owned_sandbox_ids.clear() + self._acquire_inflight.clear() + self._orphan_first_seen.clear() self._reserved_slots = 0 self._transitioning_slots = 0 self._capacity_cond.notify_all() @@ -1683,12 +2132,14 @@ class E2BSandboxProvider(SandboxProvider): ) for sandbox_id, sandbox in active: - if error := self._kill_client(sandbox.client): - logger.warning( - "Failed to kill active e2b sandbox %s during shutdown: %s", - sandbox_id, - error, - ) + if sandbox_id in owned_ids and self._claim_ownership(sandbox_id, for_destroy=True): + if error := self._kill_client(sandbox.client): + logger.warning( + "Failed to kill active e2b sandbox %s during shutdown: %s", + sandbox_id, + error, + ) + self._release_ownership(sandbox_id) try: sandbox.close() except Exception: @@ -1696,6 +2147,10 @@ class E2BSandboxProvider(SandboxProvider): sandbox_cls = self._get_sandbox_cls() for sandbox_id in warm_ids: + if sandbox_id not in owned_ids: + continue + if not self._claim_ownership(sandbox_id, for_destroy=True): + continue try: client = self._reconnect_client(sandbox_cls, sandbox_id) except Exception as e: @@ -1704,6 +2159,7 @@ class E2BSandboxProvider(SandboxProvider): sandbox_id, e, ) + self._release_ownership(sandbox_id) continue if error := self._kill_client(client): logger.warning( @@ -1711,9 +2167,14 @@ class E2BSandboxProvider(SandboxProvider): sandbox_id, error, ) + self._release_ownership(sandbox_id) close = getattr(client, "close", None) if callable(close): try: close() except Exception: pass + try: + self._ownership.close() + except Exception as e: + logger.warning("Failed to close E2B ownership store: %s", e) diff --git a/backend/packages/harness/deerflow/config/sandbox_config.py b/backend/packages/harness/deerflow/config/sandbox_config.py index 876d7aa8a..2c9982b96 100644 --- a/backend/packages/harness/deerflow/config/sandbox_config.py +++ b/backend/packages/harness/deerflow/config/sandbox_config.py @@ -88,8 +88,10 @@ class SandboxConfig(BaseModel): port: Base port for sandbox containers (default: 8080) container_prefix: Prefix for container names (default: deer-flow-sandbox) mounts: List of volume mounts to share directories with the container - ownership: Cross-instance container ownership store (memory | redis). Multi-instance - deployments sharing a container backend need redis; see SandboxOwnershipConfig. + + AioSandboxProvider and E2BSandboxProvider shared options: + ownership: Cross-instance sandbox ownership store (memory | redis). Multi-instance + deployments sharing a sandbox backend need redis; see SandboxOwnershipConfig. """ use: str = Field( @@ -143,8 +145,8 @@ class SandboxConfig(BaseModel): ownership: SandboxOwnershipConfig | None = Field( default=None, description=( - "AioSandboxProvider-only: where cross-instance container ownership is tracked (#4206). Omitted = memory (single-instance). " - "Multi-worker / load-balanced gateways sharing one container backend must set type: redis, or peers will adopt and idle-destroy each other's live sandboxes." + "AioSandboxProvider/E2BSandboxProvider: where cross-instance sandbox ownership is tracked (#4206, #4341). Omitted = memory (single-instance). " + "Multi-worker / load-balanced gateways sharing one sandbox backend must set type: redis, or peers can adopt and destroy each other's live sandboxes." ), ) mounts: list[VolumeMountConfig] = Field( diff --git a/backend/tests/test_e2b_sandbox_provider.py b/backend/tests/test_e2b_sandbox_provider.py index a965ce697..b654dffe4 100644 --- a/backend/tests/test_e2b_sandbox_provider.py +++ b/backend/tests/test_e2b_sandbox_provider.py @@ -175,6 +175,65 @@ class FakeSandboxClass: return self.list_return +class FakeOwnershipStore: + """Shared lease state with per-provider identities for reconciliation tests.""" + + supports_cross_process = True + + def __init__( + self, + leases: dict[str, tuple[str, str]], + *, + owner_id: str, + lock: threading.Lock | None = None, + ) -> None: + self._leases = leases + self._lock = lock or threading.Lock() + self.owner_id = owner_id + + def take(self, sandbox_id: str) -> bool: + with self._lock: + current = self._leases.get(sandbox_id) + if current is not None and current[1] == "del": + return False + self._leases[sandbox_id] = (self.owner_id, "own") + return True + + def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool: + with self._lock: + current = self._leases.get(sandbox_id) + if current is not None and current[0] != self.owner_id: + return False + if current is not None and current[1] == "del" and not for_destroy: + return False + self._leases[sandbox_id] = (self.owner_id, "del" if for_destroy else "own") + return True + + def renew(self, sandbox_id: str): + from deerflow.community.aio_sandbox.ownership import RenewOutcome + + with self._lock: + current = self._leases.get(sandbox_id) + if current is None: + return RenewOutcome.LAPSED + if current == (self.owner_id, "own"): + return RenewOutcome.RENEWED + return RenewOutcome.LOST + + def release(self, sandbox_id: str) -> None: + with self._lock: + if self._leases.get(sandbox_id, (None,))[0] == self.owner_id: + self._leases.pop(sandbox_id, None) + + def owner(self, sandbox_id: str) -> str | None: + with self._lock: + current = self._leases.get(sandbox_id) + return current[0] if current is not None else None + + def close(self) -> None: + return None + + def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_policy: str = "wait", acquire_timeout: int = 30, burst_limit: int = 0) -> Any: """Build a ``E2BSandboxProvider`` instance bypassing ``__init__``.""" mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider") @@ -192,6 +251,15 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli provider._transitioning_slots = 0 provider._capacity_cond = threading.Condition(provider._lock) provider._shutdown_called = False + provider._owner_id = "owner-a" + provider._ownership = FakeOwnershipStore({}, owner_id=provider._owner_id) + provider._ownership_config = SimpleNamespace(renewal_interval_seconds=60.0) + provider._owned_sandbox_ids = set() + provider._acquire_inflight = set() + provider._orphan_first_seen = {} + provider._maintenance_stop = threading.Event() + provider._lease_thread = None + provider._reconcile_thread = None provider._config = { "api_key": "test-key", "template": "code-interpreter-v1", @@ -204,6 +272,12 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli "burst_limit": burst_limit, "mounts": [], "environment": {}, + "reconciliation_interval_seconds": 60.0, + "reconciliation_grace_seconds": 30.0, + "reconciliation_orphan_ttl_seconds": 3600.0, + "reconciliation_max_pages": 10, + "reconciliation_max_items": 100, + "reconciliation_max_seconds": 10.0, } return provider @@ -412,6 +486,84 @@ def test_kill_and_close_swallows_kill_exceptions(): client.kill = explode p._kill_and_close(sb) + assert "fake-sb-1" not in p._owned_sandbox_ids + assert p._ownership.owner("fake-sb-1") is None + + +def test_kill_and_close_skips_peer_owned_sandbox(): + p = _make_provider() + client = FakeClient(sandbox_id="peer-owned") + sandbox = _make_sandbox(client, sandbox_id="peer-owned") + p._ownership = FakeOwnershipStore( + {"peer-owned": ("owner-peer", "own")}, + owner_id=p._owner_id, + ) + + p._kill_and_close(sandbox) + + assert client.killed is False + assert client.closed is True + assert p._ownership.owner("peer-owned") == "owner-peer" + + +def test_startup_reconciliation_runs_in_background_without_blocking_caller(monkeypatch): + p = _make_provider() + entered = threading.Event() + release = threading.Event() + + def reconcile(): + entered.set() + assert release.wait(timeout=2) + + monkeypatch.setattr(p, "_reconcile_remote_sandboxes", reconcile) + + started = time.monotonic() + p._start_maintenance_threads() + elapsed = time.monotonic() - started + + assert entered.wait(timeout=1) + assert elapsed < 0.5 + + release.set() + p._maintenance_stop.set() + for thread in (p._lease_thread, p._reconcile_thread): + assert thread is not None + thread.join(timeout=2) + + +def test_refresh_owned_leases_reclaims_lapsed_lease(): + p = _make_provider() + client = FakeClient(sandbox_id="sb-lapsed") + sandbox = _make_sandbox(client, sandbox_id="sb-lapsed") + p._sandboxes["sb-lapsed"] = sandbox + p._thread_sandboxes[("u1", "t1")] = "sb-lapsed" + p._owned_sandbox_ids.add("sb-lapsed") + + p._refresh_owned_leases() + + assert p._ownership.owner("sb-lapsed") == p._owner_id + assert p.get("sb-lapsed") is sandbox + assert client.closed is False + + +def test_refresh_owned_leases_forgets_peer_owned_sandbox(): + p = _make_provider() + client = FakeClient(sandbox_id="sb-lost") + sandbox = _make_sandbox(client, sandbox_id="sb-lost") + p._sandboxes["sb-lost"] = sandbox + p._thread_sandboxes[("u1", "t1")] = "sb-lost" + p._owned_sandbox_ids.add("sb-lost") + p._ownership = FakeOwnershipStore( + {"sb-lost": ("owner-peer", "own")}, + owner_id=p._owner_id, + ) + + p._refresh_owned_leases() + + assert p.get("sb-lost") is None + assert ("u1", "t1") not in p._thread_sandboxes + assert "sb-lost" not in p._owned_sandbox_ids + assert client.closed is True def test_reuse_in_process_sandbox_returns_cached_id_on_healthy_reuse(): @@ -614,6 +766,239 @@ def test_discover_remote_sandbox_skips_dead_candidate(monkeypatch): assert client.closed is True +def test_discover_remote_sandbox_tries_later_candidate_when_first_is_dead(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [ + _info("sb-a-dead", "u1", "t1"), + _info("sb-b-live", "u1", "t1"), + ] + dead = FakeClient( + sandbox_id="sb-a-dead", + commands=FakeCommandsAPI([FakeCommandsAPI.GONE]), + ) + live = FakeClient(sandbox_id="sb-b-live") + fake_cls.connect_factory = lambda sid, **_kw: dead if sid == "sb-a-dead" else live + + assert p._discover_remote_sandbox("t1", user_id="u1") == "sb-b-live" + assert dead.closed is True + assert p._thread_sandboxes[("u1", "t1")] == "sb-b-live" + assert "sb-b-live" in p._owned_sandbox_ids + + +def test_reconcile_defers_duplicate_with_live_peer_lease(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [ + _info("sb-canonical", "u1", "t1"), + _info("sb-duplicate", "u1", "t1"), + ] + clients = { + "sb-canonical": FakeClient(sandbox_id="sb-canonical"), + "sb-duplicate": FakeClient(sandbox_id="sb-duplicate"), + } + fake_cls.connect_factory = lambda sid, **_kw: clients[sid] + leases = {"sb-duplicate": ("owner-peer", "own")} + p._ownership = FakeOwnershipStore(leases, owner_id=p._owner_id) + p._config["reconciliation_grace_seconds"] = 0.0 + + stats = p._reconcile_remote_sandboxes(now=time.monotonic()) + + assert stats.adopted == 1 + assert stats.deferred == 1 + assert stats.killed == 0 + assert clients["sb-duplicate"].killed is False + + +def test_reconcile_kills_unowned_duplicate_after_grace(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [ + _info("sb-canonical", "u1", "t1"), + _info("sb-duplicate", "u1", "t1"), + ] + clients = { + "sb-canonical": FakeClient(sandbox_id="sb-canonical"), + "sb-duplicate": FakeClient(sandbox_id="sb-duplicate"), + } + fake_cls.connect_factory = lambda sid, **_kw: clients[sid] + p._config["reconciliation_grace_seconds"] = 5.0 + + first = p._reconcile_remote_sandboxes(now=100.0) + second = p._reconcile_remote_sandboxes(now=106.0) + + assert first.deferred == 1 + assert first.killed == 0 + assert second.killed == 1 + assert clients["sb-duplicate"].killed is True + + +def test_competing_reconcilers_only_one_kills_duplicate(monkeypatch): + shared_leases: dict[str, tuple[str, str]] = {} + shared_lock = threading.Lock() + providers = [_make_provider(), _make_provider()] + providers[0]._owner_id = "owner-a" + providers[1]._owner_id = "owner-b" + clients = { + "sb-canonical": FakeClient(sandbox_id="sb-canonical"), + "sb-duplicate": FakeClient(sandbox_id="sb-duplicate"), + } + kill_count = 0 + + def kill_duplicate() -> None: + nonlocal kill_count + kill_count += 1 + clients["sb-duplicate"].killed = True + clients["sb-duplicate"].commands._responses.append(FakeCommandsAPI.GONE) + + clients["sb-duplicate"].kill = kill_duplicate + + for provider in providers: + fake_cls = _install_fake_sdk(monkeypatch, provider) + fake_cls.list_return = [ + _info("sb-canonical", "u1", "t1"), + _info("sb-duplicate", "u1", "t1"), + ] + fake_cls.connect_factory = lambda sid, **_kw: clients[sid] + provider._ownership = FakeOwnershipStore( + shared_leases, + owner_id=provider._owner_id, + lock=shared_lock, + ) + provider._config["reconciliation_grace_seconds"] = 0.0 + + results = [provider._reconcile_remote_sandboxes(now=100.0) for provider in providers] + + assert sum(result.killed for result in results) == 1 + assert kill_count == 1 + + +def test_reconcile_honors_item_budget(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [_info(f"sb-{index}", "u1", f"t-{index}") for index in range(5)] + p._config["reconciliation_max_items"] = 2 + + stats = p._reconcile_remote_sandboxes(now=100.0) + + assert stats.discovered == 2 + assert len(fake_cls.connect_calls) == 2 + + +def test_reconcile_honors_page_budget(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + paginator = _FakePaginator( + [ + [_info("sb-first", "u1", "t1")], + [_info("sb-never-read", "u2", "t2")], + ] + ) + fake_cls.list_return = paginator + p._config["reconciliation_max_pages"] = 1 + + stats = p._reconcile_remote_sandboxes(now=100.0) + + assert stats.discovered == 1 + assert stats.budget_exhausted is True + assert paginator.calls == 1 + assert [call[0] for call in fake_cls.connect_calls] == ["sb-first"] + + +def test_reconcile_honors_wall_clock_budget(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [_info("sb-never-probed", "u1", "t1")] + p._config["reconciliation_max_seconds"] = 0.5 + provider_mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider") + ticks = iter([0.0, 0.0, 1.0, 1.0]) + monkeypatch.setattr(provider_mod.time, "monotonic", lambda: next(ticks)) + + stats = p._reconcile_remote_sandboxes(now=100.0) + + assert stats.budget_exhausted is True + assert fake_cls.connect_calls == [] + + +def test_reconcile_adopts_canonical_after_restart_loses_local_state(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [_info("sb-existing", "u1", "t1")] + + stats = p._reconcile_remote_sandboxes(now=100.0) + + assert stats.adopted == 1 + assert p._thread_sandboxes[("u1", "t1")] == "sb-existing" + assert "sb-existing" in p._owned_sandbox_ids + + +def test_reconcile_bootstrap_failure_clears_inflight_after_peer_take(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + sandbox_id = "sb-racy-bootstrap" + fake_cls.list_return = [_info(sandbox_id, "u1", "t1")] + shared_leases: dict[str, tuple[str, str]] = {} + shared_lock = threading.Lock() + p._ownership = FakeOwnershipStore( + shared_leases, + owner_id=p._owner_id, + lock=shared_lock, + ) + peer_ownership = FakeOwnershipStore( + shared_leases, + owner_id="owner-peer", + lock=shared_lock, + ) + + def peer_takes_before_bootstrap_fails(_command: str) -> SimpleNamespace: + assert peer_ownership.take(sandbox_id) is True + return SimpleNamespace(stdout="", stderr="permission denied", exit_code=1) + + client = FakeClient( + sandbox_id=sandbox_id, + commands=FakeCommandsAPI( + [ + SimpleNamespace(stdout="ok", stderr="", exit_code=0), + peer_takes_before_bootstrap_fails, + ] + ), + ) + fake_cls.connect_factory = lambda _sid, **_kw: client + + stats = p._reconcile_remote_sandboxes(now=100.0) + + assert stats.adopted == 0 + assert stats.deferred == 1 + assert sandbox_id not in p._acquire_inflight + assert sandbox_id not in p._unowned_remote_ops_in_progress + assert p._reserved_slots == 0 + assert p._ownership.owner(sandbox_id) == "owner-peer" + assert client.killed is False + assert client.closed is True + + +def test_reconcile_kills_metadata_orphan_only_after_ttl(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id="sb-orphan", + metadata={"deer_flow_provider": "e2b_sandbox_provider"}, + ) + ] + client = FakeClient(sandbox_id="sb-orphan") + fake_cls.connect_factory = lambda _sid, **_kw: client + p._config["reconciliation_orphan_ttl_seconds"] = 5.0 + + first = p._reconcile_remote_sandboxes(now=100.0) + second = p._reconcile_remote_sandboxes(now=106.0) + + assert first.deferred == 1 + assert first.killed == 0 + assert second.killed == 1 + assert client.killed is True + + def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeypatch): p = _make_provider() fake_cls = _install_fake_sdk(monkeypatch, p) @@ -635,6 +1020,65 @@ def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeyp assert ("u1", "t1") not in p._thread_sandboxes +def test_discovery_claims_ownership_before_bootstrap_cleanup(monkeypatch): + events: list[str] = [] + + class RecordingOwnershipStore(FakeOwnershipStore): + def take(self, sandbox_id: str) -> bool: + events.append("take") + return super().take(sandbox_id) + + def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool: + events.append("claim-destroy" if for_destroy else "claim") + return super().claim(sandbox_id, for_destroy=for_destroy) + + def release(self, sandbox_id: str) -> None: + events.append("release") + super().release(sandbox_id) + + p = _make_provider() + p._ownership = RecordingOwnershipStore( + {"sb-broken": ("owner-peer", "own")}, + owner_id=p._owner_id, + ) + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.list_return = [_info("sb-broken", "u1", "t1")] + client = FakeClient( + sandbox_id="sb-broken", + commands=FakeCommandsAPI( + [ + SimpleNamespace(stdout="ok", stderr="", exit_code=0), + SimpleNamespace(stdout="", stderr="permission denied", exit_code=1), + ] + ), + ) + client.kill = lambda: events.append("kill") + fake_cls.connect_factory = lambda _sid, **_kw: client + + assert p._discover_remote_sandbox("t1", user_id="u1") is None + assert events == ["take", "claim-destroy", "kill", "release"] + + +def test_bootstrap_failure_does_not_kill_without_destroy_lease(monkeypatch): + p = _make_provider() + client = FakeClient( + sandbox_id="sb-peer", + commands=FakeCommandsAPI([SimpleNamespace(stdout="", stderr="permission denied", exit_code=1)]), + ) + p._ownership = FakeOwnershipStore( + {"sb-peer": ("owner-peer", "own")}, + owner_id=p._owner_id, + ) + + error, remote_destroyed = p._bootstrap_or_discard(client, "sb-peer") + + assert error is not None + assert remote_destroyed is False + assert client.killed is False + assert client.closed is True + assert p._ownership.owner("sb-peer") == "owner-peer" + + def test_kill_client_returns_exception_without_raising(): p = _make_provider() client = FakeClient() @@ -728,6 +1172,48 @@ def test_evict_oldest_warm_keeps_slot_when_kill_lookup_raises(monkeypatch): assert client.closed is True assert p._eviction_tombstones == {"sb-warm"} assert p._transitioning_slots == 1 + assert "sb-warm" not in p._warm_pool + assert "sb-warm" not in p._owned_sandbox_ids + assert p._ownership.owner("sb-warm") is None + + +def test_evict_oldest_warm_defers_peer_owned_sandbox(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + p._warm_pool["sb-peer"] = ("seed", 12345.0) + p._ownership = FakeOwnershipStore( + {"sb-peer": ("owner-peer", "own")}, + owner_id=p._owner_id, + ) + + assert p._evict_oldest_warm() == "sb-peer" + + assert fake_cls.connect_calls == [] + assert "sb-peer" not in p._warm_pool + assert p._eviction_tombstones == set() + assert p._evictions_in_progress == set() + assert p._transitioning_slots == 0 + assert p._ownership.owner("sb-peer") == "owner-peer" + + +def test_evict_oldest_warm_releases_destroy_lease_after_kill_failure(monkeypatch): + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="sb-warm") + client.kill = MagicMock(side_effect=RuntimeError("transient kill failure")) + fake_cls.connect_factory = lambda _sid, **_kw: client + p._warm_pool["sb-warm"] = ("seed", 12345.0) + p._owned_sandbox_ids.add("sb-warm") + p._ownership.take("sb-warm") + + assert p._evict_oldest_warm() is None + + assert client.closed is True + assert "sb-warm" not in p._warm_pool + assert p._eviction_tombstones == {"sb-warm"} + assert p._transitioning_slots == 1 + assert "sb-warm" not in p._owned_sandbox_ids + assert p._ownership.owner("sb-warm") is None def test_evict_oldest_warm_uses_kill_helper_and_closes_client(monkeypatch): @@ -937,6 +1423,23 @@ def test_release_skips_warm_pool_when_sync_reveals_dead_vm(monkeypatch, tmp_path assert client.killed is True +def test_shutdown_only_kills_sandboxes_owned_by_current_instance(monkeypatch): + p = _make_provider() + owned_client = FakeClient(sandbox_id="sb-owned") + peer_client = FakeClient(sandbox_id="sb-peer") + p._sandboxes = { + "sb-owned": _make_sandbox(owned_client), + "sb-peer": _make_sandbox(peer_client), + } + p._owned_sandbox_ids = {"sb-owned"} + + p.shutdown() + + assert owned_client.killed is True + assert peer_client.killed is False + assert peer_client.closed is True + + def _setup_paths(monkeypatch, tmp_path): paths_mod = importlib.import_module("deerflow.config.paths") monkeypatch.setattr(paths_mod, "get_paths", lambda: Paths(base_dir=tmp_path), raising=False)