From 3b77a7401b549fa6da4c8e1f8c2c0081d56e3d7a Mon Sep 17 00:00:00 2001 From: luo jiyin Date: Sat, 25 Jul 2026 10:54:14 +0800 Subject: [PATCH] fix(sandbox): enforce E2B replica capacity limits (#4391) * fix(sandbox): enforce E2B replica capacity limits (in-process) Add SandboxCapacityExceededError with diagnostic fields. Add overflow_policy (wait/reject/burst), acquire_timeout, and burst_limit config options. Implement atomic capacity reservation with a four-slot model: reserved / active / warm / transitioning. Transitioning slots close the window where active-to-warm or warm-to-active transitions appear to have zero occupied slots, which would let concurrent acquires exceed the configured replica ceiling. Re-route release, reclaim, and evict through transitioning counters. Add shutdown guard: reject waiters, kill VMs created during shutdown. Add 14 tests: policy enforcement, release+acquire race, warm-reclaim race, shutdown-waiter interaction, shutdown-during-create, and concurrent different-thread capacity assertion. Related: #4339 * fix: harden e2b sandbox capacity lifecycle * fix: retain e2b capacity during uncertain eviction * fix: serialize e2b tombstone eviction * fix: retain capacity after uncertain e2b cleanup * fix: track e2b remote operations during shutdown * fix(sandbox): validate E2B capacity config * fix(sandbox): classify capacity errors * fix(sandbox): harden E2B capacity lifecycle * test(sandbox): cover E2B review findings * docs(changelog): note E2B capacity behavior * docs(readme): explain E2B overflow handling * docs(backend): record E2B lifecycle rules * docs(sandbox): clarify destructive E2B reset * fix(sandbox): close E2B capacity race gaps --------- Co-authored-by: Willem Jiang --- CHANGELOG.md | 5 + README.md | 20 + backend/AGENTS.md | 18 +- .../e2b_sandbox/e2b_sandbox_provider.py | 683 +++++++-- .../harness/deerflow/config/sandbox_config.py | 23 +- .../harness/deerflow/sandbox/exceptions.py | 42 + .../deerflow/sandbox/sandbox_provider.py | 11 +- backend/tests/test_e2b_sandbox_provider.py | 1337 ++++++++++++++++- 8 files changed, 2024 insertions(+), 115 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9239dfaea..382bb17d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,11 @@ This section accumulates work toward the **2.1.0** milestone ### ⚠ Breaking changes +- **sandbox:** E2B now enforces `sandbox.replicas` as a process-local capacity + limit. The default `wait` policy waits for `acquire_timeout`, then fails the + agent turn. DeerFlow does not retry the turn automatically. Use `burst` with + `burst_limit` to permit bounded extra VMs. The `reject` policy can remove one + warm VM before it returns a capacity error. ([#4391]) - **skills:** A directory containing `SKILL.md` is now a runtime package boundary. Nested `SKILL.md` files inside that package are supporting data and are no longer registered as independent skills; unusual custom layouts must diff --git a/README.md b/README.md index 5ad7acd11..27e3363c3 100644 --- a/README.md +++ b/README.md @@ -771,6 +771,26 @@ This is how DeerFlow handles tasks that take minutes to hours: a research task m ### Sandbox & File System +`E2BSandboxProvider` uses `wait` as its default overflow policy. It waits for +`acquire_timeout`, then fails the agent turn. DeerFlow does not retry the turn +automatically. Clients can use the structured error to schedule a retry. + +Use `burst` with `burst_limit` to permit bounded extra VMs. The `wait` and +`reject` policies use only `replicas`. The `reject` policy can remove one warm +VM before it returns an error. + +`replicas` limits one Gateway process. It does not limit all Gateway processes. +E2B acquisition uses a bounded executor. Waiting acquisitions do not use the +default asyncio executor. + +An E2B VM keeps its slot until E2B confirms destruction. This rule covers +create and reclaim operations. Discovery can find a VM from another Gateway. +Shutdown closes an unowned discovery client without destroying its VM. +Release stops counting a transition when the VM enters the warm pool. +Shutdown races retry remote cleanup after a transient kill failure. +Reset destroys tracked active and warm E2B VMs. The old provider instance +cannot accept new acquisitions. + DeerFlow doesn't just *talk* about doing things. It has its own computer. Each task gets its own execution environment with a full filesystem view — skills, workspace, uploads, outputs. The agent reads, writes, and edits files. It can view images and, when configured safely, execute shell commands. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 923651d61..417ebc455 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -470,7 +470,23 @@ copying a raw checkpoint because delta state is not self-contained in one tuple. **Implementations**: - `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-thread `LocalSandbox` (id `local:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Legacy global-custom mounts are gated by the same user-scoped skill discovery rule used for prompt/list visibility; providers must not infer visibility from raw directory presence alone. - `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker-based isolation. Active-cache and warm-pool entries are checked with the backend during acquire/reuse; definitively dead containers are dropped from all in-process maps so the thread can discover or create a fresh sandbox instead of reusing a stale client. Backend health-check failures are treated as unknown, not dead; local discovery likewise treats an unverifiable container as not adoptable and falls through to create rather than failing acquire. `get()` remains an in-memory lookup for event-loop-safe tool paths — it never touches the ownership store (that would be blocking IO on the event loop); ownership is published on acquire/reclaim and refreshed off the event loop by the dedicated renewal thread (`_renew_owned_leases`). Legacy global-custom mounts follow the same shared visibility helper as local and remote providers. -- `E2BSandboxProvider` (`packages/harness/deerflow/community/e2b_sandbox/`) - E2B-backed remote isolation. Acquire and the complete release transition (output sync, timeout refresh, client close, and warm-pool publication) share a per-user/thread lock, so a same-process acquire cannot discover or create a replacement while release is between active and warm states. The provider-wide registry lock is not held across remote IO. +- `E2BSandboxProvider` (`packages/harness/deerflow/community/e2b_sandbox/`) provides E2B remote isolation. + Acquire and release share a per-user and thread lock. The provider lock does + not cover remote IO. `burst_limit` adds capacity only for the `burst` policy. + The `wait` policy fails the turn after `acquire_timeout`. The runtime does not + retry the turn automatically. E2B acquisition uses a bounded executor. + Waiting calls do not consume the default asyncio executor. The `reject` + policy can evict one warm VM before it returns an error. `replicas` limits + one Gateway process. It does not provide multi-process capacity control. + Uncertain cleanup keeps a tombstone slot. Shutdown tracks owned remote + operation IDs. Discovery can find a VM from another Gateway. Shutdown closes + an unowned discovery client without destroying its VM. Release ends its + transition count when the VM enters the warm pool. Local client cleanup does + not consume a second slot. A create that returns after shutdown retries one + 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. - **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/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py b/backend/packages/harness/deerflow/community/e2b_sandbox/e2b_sandbox_provider.py index 92dce7f23..9197707ff 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 @@ -1,8 +1,7 @@ """``E2BSandboxProvider`` — DeerFlow :class:`SandboxProvider` for e2b cloud. -Configuration is read from :class:`SandboxConfig` (which has -``extra="allow"``), so any keys below can appear under ``sandbox:`` in -``config.yaml`` even though they are not declared on the model: +Configuration is read from :class:`SandboxConfig`. E2B reports unknown +provider fields during startup. .. code-block:: yaml @@ -11,8 +10,11 @@ Configuration is read from :class:`SandboxConfig` (which has api_key: $E2B_API_KEY # required (or via E2B_API_KEY env var) template: code-interpreter-v1 # default: e2b code-interpreter template domain: e2b.dev # optional; for self-hosted e2b - idle_timeout: 600 # forwarded to ``set_timeout`` - replicas: 3 # max concurrent sandboxes + idle_timeout: 1800 # forwarded to ``set_timeout`` + replicas: 3 # max capacity (active + warm) + overflow_policy: wait # wait | reject | burst (default: wait) + acquire_timeout: 30 # seconds for ``wait`` policy (default: 30) + burst_limit: 2 # extra slots for ``burst`` policy (default: 0) mounts: # one-shot uploads on sandbox start - host_path: /data/skills container_path: /home/user/skills @@ -35,7 +37,9 @@ import threading import time import uuid from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from decimal import Decimal, InvalidOperation +from functools import partial from pathlib import Path from typing import Any @@ -43,6 +47,7 @@ from e2b_code_interpreter import Sandbox as E2BClientSandbox from deerflow.config import get_app_config from deerflow.runtime.user_context import get_effective_user_id +from deerflow.sandbox.exceptions import SandboxCapacityExceededError from deerflow.sandbox.sandbox import Sandbox from deerflow.sandbox.sandbox_provider import SandboxProvider @@ -55,6 +60,8 @@ logger = logging.getLogger(__name__) DEFAULT_TEMPLATE = "code-interpreter-v1" # the public e2b code-interpreter template 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 # 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 @@ -65,6 +72,7 @@ META_KEY_USER = "deer_flow_user" META_KEY_THREAD = "deer_flow_thread" META_KEY_PROVIDER = "deer_flow_provider" META_VAL_PROVIDER = "e2b_sandbox_provider" +E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"}) class E2BSandboxProvider(SandboxProvider): @@ -90,9 +98,31 @@ class E2BSandboxProvider(SandboxProvider): # Warm pool: released sandboxes whose remote micro-VM is still alive. # ``OrderedDict`` maintains insertion / move_to_end order for LRU. self._warm_pool: OrderedDict[str, tuple[str, float]] = OrderedDict() + # Evictions with unknown remote state. Each id keeps its transition + # slot until a later eviction attempt confirms destruction. + self._eviction_tombstones: set[str] = set() + # IDs currently reconnecting or stopping. This grants one retry owner + # to each tombstone and prevents a second call from freeing its slot. + self._evictions_in_progress: set[str] = set() + # Remote IDs that are between tracked lifecycle states. Shutdown uses + # this set to retain cleanup visibility during remote calls. + self._remote_ops_in_progress: set[str] = set() + # Discovery can find a VM that another Gateway still uses. + # Shutdown closes these clients but does not destroy their VMs. + self._unowned_remote_ops_in_progress: set[str] = set() + # In-flight creates that have reserved capacity but not yet committed + # to ``_sandboxes``. Guarded by ``_lock``. + self._reserved_slots = 0 + self._transitioning_slots = 0 + self._capacity_cond = threading.Condition(self._lock) self._shutdown_called = False self._config = self._load_config() + acquire_workers = max(4, min(32, self._capacity_limit() + 1)) + self._acquire_executor = ThreadPoolExecutor( + max_workers=acquire_workers, + thread_name_prefix="e2b-sandbox-acquire", + ) atexit.register(self.shutdown) self._register_signal_handlers() @@ -100,6 +130,12 @@ class E2BSandboxProvider(SandboxProvider): def _load_config(self) -> dict[str, Any]: """Read e2b options off ``SandboxConfig`` (``extra="allow"``).""" sandbox_config = get_app_config().sandbox + unknown_keys = sorted(set(getattr(sandbox_config, "model_extra", None) or {}) - E2B_EXTRA_CONFIG_KEYS) + if unknown_keys: + logger.warning( + "E2BSandboxProvider: unknown sandbox config fields: %s", + ", ".join(unknown_keys), + ) def _opt(name: str, default: Any = None) -> Any: return getattr(sandbox_config, name, default) @@ -114,7 +150,24 @@ class E2BSandboxProvider(SandboxProvider): idle_timeout = max(0, min(int(idle_timeout), MAX_E2B_TIMEOUT)) replicas = _opt("replicas") - replicas = DEFAULT_REPLICAS if replicas is None else max(1, int(replicas)) + replicas = DEFAULT_REPLICAS if replicas is None else int(replicas) + + overflow_policy = _opt("overflow_policy") or DEFAULT_OVERFLOW_POLICY + if overflow_policy not in ("wait", "reject", "burst"): + logger.warning("E2BSandboxProvider: invalid overflow_policy %r; falling back to %r", overflow_policy, DEFAULT_OVERFLOW_POLICY) + overflow_policy = DEFAULT_OVERFLOW_POLICY + + acquire_timeout = _opt("acquire_timeout") + if acquire_timeout is None: + acquire_timeout = DEFAULT_ACQUIRE_TIMEOUT + else: + acquire_timeout = max(1, int(acquire_timeout)) + + burst_limit_raw = _opt("burst_limit") + burst_limit = max(0, int(burst_limit_raw)) if burst_limit_raw is not None else 0 + if overflow_policy == "burst" and burst_limit == 0: + logger.warning("E2BSandboxProvider: overflow_policy is 'burst' but burst_limit is 0; falling back to 'reject'") + overflow_policy = "reject" return { "api_key": api_key, @@ -123,6 +176,9 @@ class E2BSandboxProvider(SandboxProvider): "home_dir": _opt("home_dir") or DEFAULT_E2B_HOME_DIR, "idle_timeout": idle_timeout, "replicas": replicas, + "overflow_policy": overflow_policy, + "acquire_timeout": acquire_timeout, + "burst_limit": burst_limit, "mounts": _opt("mounts") or [], "environment": self._resolve_env_vars(_opt("environment") or {}), } @@ -209,7 +265,9 @@ class E2BSandboxProvider(SandboxProvider): async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: effective_user_id = self._effective_acquire_user_id(user_id) - return await asyncio.to_thread(self.acquire, thread_id, user_id=effective_user_id) + loop = asyncio.get_running_loop() + acquire = partial(self.acquire, thread_id, user_id=effective_user_id) + return await loop.run_in_executor(self._acquire_executor, acquire) def _acquire_internal(self, thread_id: str | None, *, user_id: str) -> str: if thread_id: @@ -274,6 +332,13 @@ class E2BSandboxProvider(SandboxProvider): 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. + + The warm-pool entry is popped and a transitioning slot is taken + immediately. The slot is committed to active when the sandbox is + registered in ``_sandboxes``, or freed if the reclaim fails. + If the provider shut down during the transition, the VM is killed. + """ seed = self._stable_seed(thread_id, user_id) with self._lock: target_id = next( @@ -283,6 +348,8 @@ class E2BSandboxProvider(SandboxProvider): if target_id is None: return None self._warm_pool.pop(target_id) + self._begin_transition_locked() + self._remote_ops_in_progress.add(target_id) try: client = self._reconnect_live_client(self._get_sandbox_cls(), target_id) @@ -292,6 +359,7 @@ class E2BSandboxProvider(SandboxProvider): target_id, e, ) + self._complete_transition_remote_op(target_id, remote_destroyed=False) return None if client is None: @@ -299,12 +367,36 @@ class E2BSandboxProvider(SandboxProvider): "Warm-pool e2b sandbox %s is no longer alive (reaped by control plane); dropping and falling back to create", target_id, ) + self._complete_transition_remote_op(target_id, remote_destroyed=True) return None self._refresh_remote_timeout(client) - if self._bootstrap_or_discard(client, target_id) is not None: + 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) return None - self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) + + discard_after_shutdown = False + with self._lock: + if self._shutdown_called: + logger.info( + "Provider shut down during reclaim of sandbox %s; killing VM", + target_id, + ) + discard_after_shutdown = True + else: + self._remote_ops_in_progress.discard(target_id) + self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) + self._end_transition_locked() + + if discard_after_shutdown: + self._kill_client(client) + self._safe_close_client(client) + return None + logger.info( "Reclaimed warm-pool e2b sandbox %s for user/thread %s/%s", target_id, @@ -415,10 +507,44 @@ class E2BSandboxProvider(SandboxProvider): ) return None + try: + self._reserve_capacity( + thread_id, + user_id, + remote_id=target_id, + remote_owned=False, + ) + except SandboxCapacityExceededError as error: + if error.reason == "shutdown": + logger.info( + "Discovered e2b sandbox %s while the provider is shutting down; not adopting it", + target_id, + ) + else: + logger.info( + "Discovered e2b sandbox %s, but capacity is full; not adopting it", + target_id, + ) + self._safe_close_client(client) + raise + self._refresh_remote_timeout(client) - if self._bootstrap_or_discard(client, target_id) is not None: + bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id) + if bootstrap_error is not None: + self._complete_reserved_remote_op(target_id, remote_destroyed=remote_destroyed) + return None + + discard_after_shutdown = False + with self._lock: + if self._shutdown_called: + discard_after_shutdown = True + else: + self._unowned_remote_ops_in_progress.discard(target_id) + self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id) + self._commit_capacity() + if discard_after_shutdown: + self._safe_close_client(client) return None - 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, @@ -428,20 +554,235 @@ class E2BSandboxProvider(SandboxProvider): ) return target_id - def _create_sandbox(self, thread_id: str | None, *, user_id: str) -> str: - """Allocate a fresh e2b sandbox and hydrate it with configured mounts.""" + # ── Capacity reservation ────────────────────────────────────────────── + # + # Every sandbox holds one *slot*. A slot is in exactly one of four states: + # + # reserved _reserved_slots create in flight, not yet in _sandboxes + # active _sandboxes serving a thread + # warm _warm_pool released, parked for reuse + # transitioning _transitioning_slots being moved between states + # + # Total = _reserved_slots + len(_sandboxes) + len(_warm_pool) + _transitioning_slots + # + # The ``transitioning`` bucket closes the window where a sandbox has been + # removed from ``_sandboxes`` (or ``_warm_pool``) but not yet parked in its + # destination. Without it, ``_release_internal`` (active → warm), + # ``_reclaim_warm_pool_sandbox`` (warm → active), and ``_evict_oldest_warm`` + # (warm → destroyed) all temporarily appear to have *zero* slots occupied, + # letting a concurrent acquire reserve a new slot and exceed the configured + # ``replicas``. + + def _total_capacity_used_locked(self) -> int: + """Return reserved + active + warm + transitioning (``_lock`` must be held).""" + return self._reserved_slots + len(self._sandboxes) + len(self._warm_pool) + self._transitioning_slots + + def _capacity_limit(self) -> int: + """Hard ceiling for reserved + active + warm + transitioning slots.""" replicas = int(self._config["replicas"]) + if self._config["overflow_policy"] == "burst": + return replicas + int(self._config["burst_limit"]) + return replicas + + def _begin_transition_locked(self) -> None: + """Increment the transitioning counter (``_lock`` must be held). + + Call before removing a sandbox from ``_sandboxes`` or ``_warm_pool`` + so the slot is still counted while the transition is in flight. + """ + self._transitioning_slots += 1 + + def _end_transition_locked(self) -> None: + """Decrement the transitioning counter (``_lock`` must be held). + + Call after the transition completes — either the slot has been parked + in its destination dict or the sandbox has been destroyed. + """ + if self._transitioning_slots > 0: + self._transitioning_slots -= 1 + self._capacity_cond.notify_all() + + def _free_transitioning_slot(self) -> None: + """Release a transitioning slot after the sandbox was destroyed.""" with self._lock: - in_use = len(self._sandboxes) + len(self._warm_pool) - if in_use >= replicas: + self._end_transition_locked() + + def _reserve_capacity( + self, + thread_id: str | None, + user_id: str, + *, + remote_id: str | None = None, + remote_owned: bool = True, + ) -> None: + """Acquire a capacity slot, blocking or raising as configured. + + Must be called before ``Sandbox.create()``. The caller MUST call + ``_commit_capacity()`` on success or ``_release_capacity()`` on + failure — otherwise the reserved slot is leaked until shutdown. + + Raises: + SandboxCapacityExceededError: when the overflow policy is + ``reject`` or the wait timeout expires. + """ + policy = self._config["overflow_policy"] + timeout = float(self._config["acquire_timeout"]) + deadline = time.monotonic() + timeout + + while True: + # Reject immediately if the provider is shutting down. + with self._lock: + if self._shutdown_called: + raise SandboxCapacityExceededError( + "Sandbox provider is shutting down; cannot acquire capacity", + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) + + # 1. Try immediate atomic reservation. + with self._lock: + if self._shutdown_called: + raise SandboxCapacityExceededError( + "Sandbox provider is shutting down; cannot acquire capacity", + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) + cap = self._capacity_limit() + if self._total_capacity_used_locked() < cap: + self._reserved_slots += 1 + if remote_id is not None: + remote_ops = self._remote_ops_in_progress if remote_owned else self._unowned_remote_ops_in_progress + remote_ops.add(remote_id) + return + + # 2. Try evicting a warm entry to free a slot. evicted = self._evict_oldest_warm() - if evicted is None: - logger.warning( - "All %d e2b replica slots are in active use; creating a new sandbox beyond the soft limit (active=%d, warm=%d)", - replicas, - len(self._sandboxes), - len(self._warm_pool), - ) + if evicted is not None: + with self._lock: + if self._shutdown_called: + raise SandboxCapacityExceededError( + "Sandbox provider shut down while acquiring capacity", + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) + if self._total_capacity_used_locked() < cap: + self._reserved_slots += 1 + if remote_id is not None: + remote_ops = self._remote_ops_in_progress if remote_owned else self._unowned_remote_ops_in_progress + remote_ops.add(remote_id) + return + # Slot was stolen; fall through to policy / wait. + + # 3. Apply overflow policy. + with self._lock: + if self._shutdown_called: + raise SandboxCapacityExceededError( + "Sandbox provider is shutting down; cannot acquire capacity", + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) + used = self._total_capacity_used_locked() + cap = self._capacity_limit() + + if used >= cap: + if policy == "reject": + raise SandboxCapacityExceededError( + f"All {cap} sandbox capacity slots are in use and overflow_policy is 'reject'", + active=len(self._sandboxes), + warm=len(self._warm_pool), + reserved=self._reserved_slots, + replicas=int(self._config["replicas"]), + ) + + if policy == "burst": + raise SandboxCapacityExceededError( + f"All {cap} sandbox capacity slots are in use (replicas={self._config['replicas']}, burst={self._config['burst_limit']})", + active=len(self._sandboxes), + warm=len(self._warm_pool), + reserved=self._reserved_slots, + replicas=int(self._config["replicas"]), + ) + + # policy == "wait": block until a slot frees or timeout. + remaining = deadline - time.monotonic() + if remaining <= 0: + raise SandboxCapacityExceededError( + f"Timed out after {timeout}s waiting for a sandbox capacity slot (replicas={self._config['replicas']}, active={len(self._sandboxes)}, warm={len(self._warm_pool)}, reserved={self._reserved_slots})", + active=len(self._sandboxes), + warm=len(self._warm_pool), + reserved=self._reserved_slots, + replicas=int(self._config["replicas"]), + ) + self._capacity_cond.wait(timeout=min(remaining, 1.0)) + + def _release_capacity(self) -> None: + """Release a reserved slot (call on create failure or destroy).""" + with self._lock: + if self._reserved_slots > 0: + self._reserved_slots -= 1 + self._capacity_cond.notify_all() + + def _complete_transition_remote_op(self, sandbox_id: str, *, remote_destroyed: bool) -> None: + """Finish a remote operation that already owns a transition slot.""" + with self._lock: + if sandbox_id not in self._remote_ops_in_progress: + return + self._remote_ops_in_progress.discard(sandbox_id) + if self._shutdown_called: + return + if remote_destroyed: + self._end_transition_locked() + else: + self._eviction_tombstones.add(sandbox_id) + + def _track_reserved_remote_op(self, sandbox_id: str) -> bool: + """Make a reserved remote ID visible to shutdown.""" + with self._lock: + if self._shutdown_called: + return False + self._remote_ops_in_progress.add(sandbox_id) + return True + + def _complete_reserved_remote_op(self, sandbox_id: str, *, remote_destroyed: bool) -> None: + """Finish a remote operation that owns a reserved slot.""" + with self._lock: + tracked = sandbox_id in self._remote_ops_in_progress + tracked_unowned = sandbox_id in self._unowned_remote_ops_in_progress + if not tracked and not tracked_unowned: + return + self._remote_ops_in_progress.discard(sandbox_id) + self._unowned_remote_ops_in_progress.discard(sandbox_id) + if self._shutdown_called: + return + if self._reserved_slots > 0: + self._reserved_slots -= 1 + if remote_destroyed: + self._capacity_cond.notify_all() + else: + self._begin_transition_locked() + self._eviction_tombstones.add(sandbox_id) + self._capacity_cond.notify_all() + + def _commit_capacity(self) -> None: + """Convert a reserved slot to a committed active slot. + + The reservation is dropped and the newly-created sandbox fills the + slot. Must be called inside the same critical section that inserts + into ``_sandboxes``. + """ + if self._reserved_slots > 0: + self._reserved_slots -= 1 + + def _create_sandbox(self, thread_id: str | None, *, user_id: str) -> str: + """Allocate a fresh e2b sandbox and hydrate it with configured mounts. + + Capacity is enforced atomically via :meth:`_reserve_capacity`. + """ + self._reserve_capacity(thread_id, user_id) sandbox_cls = self._get_sandbox_cls() metadata: dict[str, str] = { @@ -465,18 +806,58 @@ class E2BSandboxProvider(SandboxProvider): client = sandbox_cls.create(**create_kwargs) # type: ignore[attr-defined] except Exception as e: logger.error("Failed to create e2b sandbox: %s", e) + self._release_capacity() raise sandbox_id: str = getattr(client, "sandbox_id", None) or str(uuid.uuid4())[:8] + if not self._track_reserved_remote_op(sandbox_id): + kill_error = self._kill_client(client) + cleanup_confirmed = kill_error is None + self._safe_close_client(client) + if kill_error is not None: + with self._lock: + self._remote_ops_in_progress.add(sandbox_id) + try: + retry_client = self._reconnect_client(sandbox_cls, sandbox_id) + except Exception as reconnect_error: + logger.warning( + "Failed to reconnect e2b sandbox %s after shutdown cleanup failed: %s", + sandbox_id, + reconnect_error, + ) + else: + retry_error = self._kill_client(retry_client) + self._safe_close_client(retry_client) + if retry_error is None: + cleanup_confirmed = True + with self._lock: + self._remote_ops_in_progress.discard(sandbox_id) + else: + logger.warning( + "Failed to kill e2b sandbox %s after reconnecting during shutdown: %s", + sandbox_id, + retry_error, + ) + if cleanup_confirmed: + message = f"Sandbox provider shut down during sandbox creation; cleaned up remote sandbox {sandbox_id}" + else: + message = f"Sandbox provider shut down during sandbox creation; could not confirm cleanup for remote sandbox {sandbox_id}" + raise SandboxCapacityExceededError( + message, + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) # 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 # with PermissionError because /mnt is owned by root in the e2b # template. See the path-mapping note in :class:`E2BSandbox`. - error = self._bootstrap_or_discard(client, sandbox_id) - if error is not None: - raise RuntimeError(f"Failed to bootstrap e2b sandbox {sandbox_id}") from error + bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, sandbox_id) + if bootstrap_error is not None: + self._complete_reserved_remote_op(sandbox_id, remote_destroyed=remote_destroyed) + raise RuntimeError(f"Failed to bootstrap e2b sandbox {sandbox_id}") from bootstrap_error # One-shot mount uploads. e2b has no host bind-mount, so we copy # files from ``host_path`` into ``container_path`` at sandbox start. @@ -486,13 +867,34 @@ class E2BSandboxProvider(SandboxProvider): logger.warning("Failed to apply some mounts to e2b sandbox %s: %s", sandbox_id, e) sandbox = E2BSandbox(id=sandbox_id, client=client, home_dir=self._config["home_dir"]) - with self._lock: - self._sandboxes[sandbox_id] = sandbox - if thread_id: - self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id + # Commit atomically. If the provider shut down during bootstrap or + # mounts, kill the VM rather than parking it under ``_sandboxes`` + # where the next shutdown won't see it. + should_kill = False + with self._lock: + if self._shutdown_called: + should_kill = True + else: + self._remote_ops_in_progress.discard(sandbox_id) + self._commit_capacity() + self._sandboxes[sandbox_id] = sandbox + if thread_id: + self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id + + if should_kill: + self._kill_client(client) + self._safe_close_client(client) + raise SandboxCapacityExceededError( + f"Sandbox provider shut down during sandbox creation; killed remote sandbox {sandbox_id}", + replicas=int(self._config["replicas"]), + retry_after_seconds=30.0, + reason="shutdown", + ) + + replicas = self._config["replicas"] logger.info( - "Created e2b sandbox %s for user/thread %s/%s (template=%s, replicas=%d)", + "Created e2b sandbox %s for user/thread %s/%s (template=%s, replicas=%s)", sandbox_id, user_id, thread_id, @@ -539,11 +941,13 @@ class E2BSandboxProvider(SandboxProvider): thread_id: str, user_id: str, ) -> None: - """Track a live reconnected sandbox under its thread ownership.""" + """Track a live reconnected sandbox under its thread ownership. + + The caller must hold ``self._lock``. + """ 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 + 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.""" @@ -606,17 +1010,18 @@ class E2BSandboxProvider(SandboxProvider): logger.debug("e2b client close raised: %s", e) return - def _bootstrap_or_discard(self, client: E2BClientSandbox, sandbox_id: str) -> Exception | None: - """Bootstrap a sandbox or return its error after cleanup.""" + def _bootstrap_or_discard(self, client: E2BClientSandbox, sandbox_id: str) -> tuple[Exception | None, bool]: + """Bootstrap a sandbox and report whether cleanup destroyed the VM.""" try: self._bootstrap_sandbox_paths(client) except Exception as e: logger.exception("Failed to bootstrap e2b sandbox %s. Discarding the unusable sandbox.", sandbox_id) - if error := self._kill_client(client): - logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, error) + 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) self._safe_close_client(client) - return e - return None + return e, kill_error is None + return None, True def _bootstrap_sandbox_paths(self, client: E2BClientSandbox) -> None: """Materialise DeerFlow's virtual path layout inside the e2b VM. @@ -1036,29 +1441,64 @@ class E2BSandboxProvider(SandboxProvider): pass def _evict_oldest_warm(self) -> str | None: + """Evict the oldest warm entry, holding a transitioning slot. + + The warm entry is popped and a transitioning slot is taken. The slot + stays occupied until the control plane confirms the VM is gone. + """ with self._lock: - if not self._warm_pool: + retryable = self._eviction_tombstones - self._evictions_in_progress + if retryable: + evict_id = next(iter(retryable)) + self._evictions_in_progress.add(evict_id) + elif self._warm_pool: + evict_id, (_, _) = self._warm_pool.popitem(last=False) + self._eviction_tombstones.add(evict_id) + self._evictions_in_progress.add(evict_id) + self._begin_transition_locked() + else: return None - evict_id, (_, _) = self._warm_pool.popitem(last=False) try: - client = self._reconnect_client(self._get_sandbox_cls(), evict_id) + client = self._reconnect_live_client(self._get_sandbox_cls(), evict_id) except Exception as e: logger.warning( "Evicted warm-pool e2b sandbox %s could not be reconnected for kill: %s", evict_id, e, ) + with self._lock: + if evict_id in self._evictions_in_progress: + self._evictions_in_progress.discard(evict_id) + if not self._shutdown_called: + self._eviction_tombstones.add(evict_id) + return None + + if client is None: + 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() + logger.info("Evicted warm-pool e2b sandbox %s was already gone", evict_id) return evict_id if error := self._kill_client(client): logger.warning("Failed to kill evicted e2b sandbox %s: %s", evict_id, error) - close = getattr(client, "close", None) - if callable(close): - try: - close() - except Exception: - pass + self._safe_close_client(client) + with self._lock: + if evict_id in self._evictions_in_progress: + self._evictions_in_progress.discard(evict_id) + if not self._shutdown_called: + self._eviction_tombstones.add(evict_id) + return None + + self._safe_close_client(client) + 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() logger.info("Evicted warm-pool e2b sandbox %s", evict_id) return evict_id @@ -1088,13 +1528,25 @@ class E2BSandboxProvider(SandboxProvider): self._release_internal(sandbox_id) def _release_internal(self, sandbox_id: str) -> None: - """Complete one release while the thread transition lock is held.""" + """Complete one release while the thread transition lock is held. + + The active slot becomes a *transitioning* slot the moment the sandbox + is removed from ``_sandboxes``. It stays counted through output sync + and timeout refresh. The transition ends when the VM enters its + destination or destruction completes. Shutdown kills the VM instead + of parking it in ``_warm_pool``. + """ sandbox: E2BSandbox | None = None seed: str | None = None + removed_keys: list[tuple[str, str]] = [] + transition_slot_held = False with self._lock: sandbox = self._sandboxes.pop(sandbox_id, None) - # Find the (user, thread) the sandbox was bound to. + if sandbox is None: + return + self._begin_transition_locked() + transition_slot_held = True removed_keys = [key for key, sid in self._thread_sandboxes.items() if sid == sandbox_id] for key in removed_keys: self._thread_sandboxes.pop(key, None) @@ -1102,53 +1554,72 @@ class E2BSandboxProvider(SandboxProvider): user_id, thread_id = removed_keys[0] seed = self._stable_seed(thread_id, user_id) - if sandbox is None: - return + # E2BSandbox.close() clears its client reference. Keep this reference + # so a shutdown that races release can still kill the remote VM. + client = sandbox.client - if sandbox.is_dead: - logger.info( - "Releasing dead e2b sandbox %s; skipping output sync and warm pool, killing remote VM", - sandbox_id, - ) - self._kill_and_close(sandbox) - return - - sync_failed_due_to_dead_vm = False - if seed is not None and removed_keys: - user_id_sync, thread_id_sync = removed_keys[0] - try: - self._sync_outputs_to_host(sandbox, thread_id=thread_id_sync, user_id=user_id_sync) - except Exception as e: # pragma: no cover - defensive - logger.warning( - "Failed to mirror e2b sandbox %s outputs to host: %s", - sandbox_id, - e, - ) + try: if sandbox.is_dead: - sync_failed_due_to_dead_vm = True + logger.info( + "Releasing dead e2b sandbox %s; skipping output sync and warm pool, killing remote VM", + sandbox_id, + ) + self._kill_and_close(sandbox) + return - if sync_failed_due_to_dead_vm: - logger.info( - "Sandbox %s was reaped during release; not parking in warm pool", - sandbox_id, - ) - self._kill_and_close(sandbox) - return + sync_failed_due_to_dead_vm = False + if seed is not None and removed_keys: + user_id_sync, thread_id_sync = removed_keys[0] + try: + self._sync_outputs_to_host(sandbox, thread_id=thread_id_sync, user_id=user_id_sync) + except Exception as e: # pragma: no cover - defensive + logger.warning( + "Failed to mirror e2b sandbox %s outputs to host: %s", + sandbox_id, + e, + ) + if sandbox.is_dead: + sync_failed_due_to_dead_vm = True - try: - self._refresh_remote_timeout(sandbox.client) - except Exception as e: - logger.debug("Failed to refresh timeout during release: %s", e) + if sync_failed_due_to_dead_vm: + logger.info( + "Sandbox %s was reaped during release; not parking in warm pool", + sandbox_id, + ) + self._kill_and_close(sandbox) + return - try: - sandbox.close() - except Exception as e: - logger.warning("Error closing e2b sandbox %s during release: %s", sandbox_id, e) + try: + self._refresh_remote_timeout(client) + except Exception as e: + logger.debug("Failed to refresh timeout during release: %s", e) - with self._lock: - self._warm_pool[sandbox_id] = (seed or "", time.time()) - self._warm_pool.move_to_end(sandbox_id) - logger.info("Released e2b sandbox %s to warm pool", sandbox_id) + with self._lock: + should_kill = self._shutdown_called + if not should_kill: + self._warm_pool[sandbox_id] = (seed or "", time.time()) + self._warm_pool.move_to_end(sandbox_id) + self._end_transition_locked() + transition_slot_held = False + logger.info("Released e2b sandbox %s to warm pool", sandbox_id) + + if should_kill: + logger.info( + "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) + self._safe_close_client(client) + return + + try: + sandbox.close() + except Exception as e: + logger.warning("Error closing e2b sandbox %s during release: %s", sandbox_id, e) + finally: + if transition_slot_held: + self._free_transitioning_slot() def _kill_and_close(self, sandbox: E2BSandbox) -> None: if error := self._kill_client(getattr(sandbox, "_client", None)): @@ -1168,21 +1639,19 @@ class E2BSandboxProvider(SandboxProvider): ) -> Exception | None: """Kill a remote VM and return an exception for the caller to log.""" if client is None: - return None + return RuntimeError("Cannot confirm remote VM destruction without a client") try: kill = getattr(client, "kill", None) - if callable(kill): - kill() + if not callable(kill): + return RuntimeError("Cannot confirm remote VM destruction without a callable kill method") + kill() except Exception as e: return e return None def reset(self) -> None: - with self._lock: - self._sandboxes.clear() - self._thread_sandboxes.clear() - self._thread_locks.clear() - self._warm_pool.clear() + """Destroy tracked E2B VMs and make this detached provider unusable.""" + self.shutdown() def shutdown(self) -> None: with self._lock: @@ -1190,10 +1659,22 @@ class E2BSandboxProvider(SandboxProvider): return self._shutdown_called = True active = list(self._sandboxes.items()) - warm_ids = list(self._warm_pool.keys()) + warm_ids = list(self._warm_pool.keys() | self._eviction_tombstones | self._remote_ops_in_progress) self._sandboxes.clear() self._warm_pool.clear() + self._eviction_tombstones.clear() + self._evictions_in_progress.clear() + self._remote_ops_in_progress.clear() + self._unowned_remote_ops_in_progress.clear() self._thread_sandboxes.clear() + self._thread_locks.clear() + self._reserved_slots = 0 + self._transitioning_slots = 0 + self._capacity_cond.notify_all() + + executor = getattr(self, "_acquire_executor", None) + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) logger.info( "Shutting down E2BSandboxProvider: %d active + %d warm sandboxes", diff --git a/backend/packages/harness/deerflow/config/sandbox_config.py b/backend/packages/harness/deerflow/config/sandbox_config.py index 218f00b73..876d7aa8a 100644 --- a/backend/packages/harness/deerflow/config/sandbox_config.py +++ b/backend/packages/harness/deerflow/config/sandbox_config.py @@ -3,6 +3,7 @@ from typing import Literal from pydantic import BaseModel, ConfigDict, Field SandboxOwnershipType = Literal["memory", "redis"] +SandboxOverflowPolicy = Literal["wait", "reject", "burst"] class SandboxOwnershipConfig(BaseModel): @@ -73,9 +74,10 @@ class SandboxConfig(BaseModel): allow_host_bash: Enable host-side bash execution for LocalSandboxProvider. Dangerous and intended only for fully trusted local workflows. - AioSandboxProvider and BoxliteProvider shared options: + AioSandboxProvider, BoxliteProvider, and E2BSandboxProvider shared options: image: Sandbox image to use (Docker/AIO image or BoxLite OCI image) - replicas: Maximum active + warm sandboxes/VMs per gateway process (default: 3). When the limit is reached, warm/least-recently-used sandboxes are evicted to make room; active sandboxes are not forcibly stopped. + replicas: Positive provider capacity per gateway process. Each provider + defines which lifecycle states count toward this limit. idle_timeout: Idle timeout in seconds before released warm sandboxes/VMs are stopped (default: 600 = 10 minutes). Set to 0 to disable. environment: Environment variables to inject into the sandbox (values starting with $ are resolved from host env) @@ -108,7 +110,22 @@ class SandboxConfig(BaseModel): ) replicas: int | None = Field( default=None, - description="Maximum active + warm sandboxes/VMs per gateway process (default: 3). Warm/least-recently-used entries are evicted to make room; active sandboxes are not forcibly stopped.", + gt=0, + description="Positive provider capacity per gateway process. Each provider defines which lifecycle states count toward this limit.", + ) + overflow_policy: SandboxOverflowPolicy = Field( + default="wait", + description="E2B capacity policy. Use wait, reject, or burst.", + ) + acquire_timeout: int = Field( + default=30, + gt=0, + description="Seconds that E2B wait policy waits for capacity.", + ) + burst_limit: int = Field( + default=0, + ge=0, + description="Extra E2B capacity slots when overflow_policy is burst.", ) container_prefix: str | None = Field( default=None, diff --git a/backend/packages/harness/deerflow/sandbox/exceptions.py b/backend/packages/harness/deerflow/sandbox/exceptions.py index bf55f7314..91dd0fb67 100644 --- a/backend/packages/harness/deerflow/sandbox/exceptions.py +++ b/backend/packages/harness/deerflow/sandbox/exceptions.py @@ -69,3 +69,45 @@ class SandboxFileNotFoundError(SandboxFileError): """Raised when a file or directory is not found.""" pass + + +class SandboxCapacityExceededError(SandboxError): + """Raised when the sandbox provider has no available capacity. + + The reason distinguishes occupied capacity from provider shutdown. + The caller controls retry scheduling. DeerFlow does not retry automatically. + """ + + CODE = "SANDBOX_CAPACITY_EXCEEDED" + + def __init__( + self, + message: str = "All sandbox replica slots are in use", + *, + active: int = 0, + warm: int = 0, + reserved: int = 0, + replicas: int = 0, + retry_after_seconds: float = 5.0, + reason: str = "capacity", + ) -> None: + details: dict[str, object] = { + "code": self.CODE, + "reason": reason, + "replicas": replicas, + "retryable": True, + "retry_after_seconds": retry_after_seconds, + } + if active: + details["active"] = active + if warm: + details["warm"] = warm + if reserved: + details["reserved"] = reserved + super().__init__(message, details) + self.active = active + self.warm = warm + self.reserved = reserved + self.replicas = replicas + self.retry_after_seconds = retry_after_seconds + self.reason = reason diff --git a/backend/packages/harness/deerflow/sandbox/sandbox_provider.py b/backend/packages/harness/deerflow/sandbox/sandbox_provider.py index e8c7b163c..58ee72604 100644 --- a/backend/packages/harness/deerflow/sandbox/sandbox_provider.py +++ b/backend/packages/harness/deerflow/sandbox/sandbox_provider.py @@ -51,7 +51,10 @@ class SandboxProvider(ABC): pass def reset(self) -> None: - """Clear cached state that survives provider instance replacement.""" + """Clear cached state that survives provider instance replacement. + + Provider overrides can release resources and make the instance unusable. + """ pass @@ -115,7 +118,7 @@ def get_sandbox_provider(**kwargs) -> SandboxProvider: def reset_sandbox_provider() -> None: """Reset the sandbox provider singleton. - This clears the cached instance without calling shutdown. + This clears the cached instance without calling shutdown directly. The next call to `get_sandbox_provider()` will create a new instance. Useful for testing or when switching configurations. @@ -124,7 +127,9 @@ def reset_sandbox_provider() -> None: `LocalSandbox` singleton). Without it, config/mount changes would not take effect on the next acquire(). - Note: If the provider has active sandboxes, they will be orphaned. + A provider override can release active sandboxes during reset. + Otherwise, active sandboxes become orphaned. + Do not reuse the detached provider after reset. Use `shutdown_sandbox_provider()` for proper cleanup. """ global _default_sandbox_provider diff --git a/backend/tests/test_e2b_sandbox_provider.py b/backend/tests/test_e2b_sandbox_provider.py index 8b17172b3..a965ce697 100644 --- a/backend/tests/test_e2b_sandbox_provider.py +++ b/backend/tests/test_e2b_sandbox_provider.py @@ -2,18 +2,24 @@ from __future__ import annotations +import asyncio import importlib import json import os import threading +import time from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock import pytest +from pydantic import ValidationError from deerflow.config.paths import Paths +from deerflow.config.sandbox_config import SandboxConfig +from deerflow.sandbox.exceptions import SandboxCapacityExceededError # ────────────────────────────────────────────────────────────────────────────── # Fakes for the e2b SDK @@ -169,7 +175,7 @@ class FakeSandboxClass: return self.list_return -def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800) -> Any: +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") provider = mod.E2BSandboxProvider.__new__(mod.E2BSandboxProvider) @@ -178,6 +184,13 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800) -> Any: provider._thread_sandboxes = {} provider._thread_locks = {} provider._warm_pool = OrderedDict() + provider._eviction_tombstones = set() + provider._evictions_in_progress = set() + provider._remote_ops_in_progress = set() + provider._unowned_remote_ops_in_progress = set() + provider._reserved_slots = 0 + provider._transitioning_slots = 0 + provider._capacity_cond = threading.Condition(provider._lock) provider._shutdown_called = False provider._config = { "api_key": "test-key", @@ -186,6 +199,9 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800) -> Any: "home_dir": "/home/user", "idle_timeout": idle_timeout, "replicas": replicas, + "overflow_policy": overflow_policy, + "acquire_timeout": acquire_timeout, + "burst_limit": burst_limit, "mounts": [], "environment": {}, } @@ -628,14 +644,67 @@ def test_kill_client_returns_exception_without_raising(): assert p._kill_client(client) is error -def test_kill_client_ignores_missing_or_uncallable_clients(): +def test_kill_client_reports_uncertain_cleanup_without_callable_kill(): p = _make_provider() - assert p._kill_client(None) is None - assert p._kill_client(SimpleNamespace()) is None + assert isinstance(p._kill_client(None), RuntimeError) + assert isinstance(p._kill_client(SimpleNamespace()), RuntimeError) -def test_evict_oldest_warm_closes_client_when_kill_lookup_raises(monkeypatch): +def test_sandbox_config_validates_e2b_capacity_fields(): + config = SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + overflow_policy="burst", + acquire_timeout=12, + burst_limit=2, + ) + + assert config.overflow_policy == "burst" + assert config.acquire_timeout == 12 + assert config.burst_limit == 2 + + with pytest.raises(ValidationError): + SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + overflow_policy="invalid", + ) + + with pytest.raises(ValidationError): + SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + acquire_timeout=0, + ) + + with pytest.raises(ValidationError): + SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + burst_limit=-1, + ) + + with pytest.raises(ValidationError): + SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + replicas=0, + ) + + +def test_e2b_config_warns_about_unknown_fields(monkeypatch, caplog): + mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider") + config = SandboxConfig( + use="deerflow.community.e2b_sandbox:E2BSandboxProvider", + api_key="test-key", + overflo_policy="reject", + ) + provider = mod.E2BSandboxProvider.__new__(mod.E2BSandboxProvider) + monkeypatch.setattr(mod, "get_app_config", lambda: SimpleNamespace(sandbox=config)) + + with caplog.at_level("WARNING"): + provider._load_config() + + assert "overflo_policy" in caplog.text + + +def test_evict_oldest_warm_keeps_slot_when_kill_lookup_raises(monkeypatch): p = _make_provider() fake_cls = _install_fake_sdk(monkeypatch, p) error = RuntimeError("kill unavailable") @@ -655,8 +724,10 @@ def test_evict_oldest_warm_closes_client_when_kill_lookup_raises(monkeypatch): fake_cls.connect_factory = lambda _sid, **_kw: client p._warm_pool["sb-warm"] = ("seed", 12345.0) - assert p._evict_oldest_warm() == "sb-warm" + assert p._evict_oldest_warm() is None assert client.closed is True + assert p._eviction_tombstones == {"sb-warm"} + assert p._transitioning_slots == 1 def test_evict_oldest_warm_uses_kill_helper_and_closes_client(monkeypatch): @@ -738,7 +809,7 @@ def test_acquire_rejects_falsey_bootstrap_error(monkeypatch): client = FakeClient(sandbox_id="bootstrap-failure") error = FalseyError("bootstrap failed") fake_cls.create_factory = lambda **kwargs: client - provider._bootstrap_or_discard = MagicMock(return_value=error) + provider._bootstrap_or_discard = MagicMock(return_value=(error, True)) with pytest.raises(RuntimeError, match="bootstrap") as caught: provider.acquire("thread-1", user_id="user-1") @@ -1496,3 +1567,1255 @@ def test_grep_without_glob_is_unaffected(): assert [m.path for m in matches] == ["/home/user/workspace/anywhere/file.txt"] assert truncated is False + + +# ────────────────────────────────────────────────────────────────────────────── +# Capacity enforcement tests (#4339) +# ────────────────────────────────────────────────────────────────────────────── + + +def test_capacity_reject_policy_raises_when_full(monkeypatch): + """With overflow_policy='reject' and replicas=1, a second acquire raises + SandboxCapacityExceededError instead of creating an unbounded sandbox.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + assert sid1 is not None + assert len(p._sandboxes) == 1 + + with pytest.raises(SandboxCapacityExceededError) as exc_info: + p.acquire("t2", user_id="u2") + assert exc_info.value.replicas == 1 + assert exc_info.value.retry_after_seconds > 0 + assert exc_info.value.details["code"] == "SANDBOX_CAPACITY_EXCEEDED" + assert exc_info.value.details["retryable"] is True + assert len(fake_cls.create_calls) == 1 # no second create + + +def test_capacity_reject_frees_slot_on_release(monkeypatch): + """Releasing a sandbox frees a capacity slot so the next acquire succeeds.""" + p = _make_provider(replicas=1, overflow_policy="reject") + _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + + sid2 = p.acquire("t2", user_id="u2") + assert sid2 is not None + assert sid2 != sid1 + + +def test_capacity_reject_evicts_other_thread_warm_entry_before_create(monkeypatch): + """Reject policy can evict one warm VM before it rejects new capacity.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + assert len(p._warm_pool) == 1 + + sid2 = p.acquire("t2", user_id="u2") + + assert sid2 != sid1 + assert len(p._warm_pool) == 0 + assert len(fake_cls.create_calls) == 2 + + +def test_capacity_wait_policy_times_out(monkeypatch): + """With overflow_policy='wait' and a short timeout, the provider raises + SandboxCapacityExceededError when no slot frees up.""" + p = _make_provider(replicas=1, overflow_policy="wait", acquire_timeout=1) + _install_fake_sdk(monkeypatch, p) + + p.acquire("t1", user_id="u1") + + with pytest.raises(SandboxCapacityExceededError) as exc_info: + p.acquire("t2", user_id="u2") + assert "Timed out" in str(exc_info.value) + + +def test_capacity_wait_policy_succeeds_when_slot_freed(monkeypatch): + """A blocked waiter proceeds once a slot is freed by another thread.""" + p = _make_provider(replicas=1, overflow_policy="wait", acquire_timeout=10) + _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + + results: list[str | Exception] = [] + barrier = threading.Barrier(2, timeout=5) + acquired = threading.Event() + + def waiter() -> None: + barrier.wait() + try: + results.append(p.acquire("t2", user_id="u2")) + except Exception as e: + results.append(e) + acquired.set() + + t = threading.Thread(target=waiter) + t.start() + barrier.wait() + + # Give the waiter a moment to block on the condition. + time.sleep(0.2) + + p.release(sid1) + t.join(timeout=5) + + assert not t.is_alive(), "waiter thread must complete" + assert len(results) == 1 + assert isinstance(results[0], str) + + +def test_capacity_burst_policy_allows_limited_overflow(monkeypatch): + """With overflow_policy='burst' and burst_limit=2, the provider allows + up to replicas + burst_limit sandboxes.""" + p = _make_provider(replicas=1, overflow_policy="burst", burst_limit=2) + _install_fake_sdk(monkeypatch, p) + + sids = [] + for i in range(3): # replicas(1) + burst(2) = 3 + sids.append(p.acquire(f"t{i}", user_id=f"u{i}")) + assert len(sids) == 3 + assert len(p._sandboxes) == 3 + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t-extra", user_id="u-extra") + + +@pytest.mark.parametrize("overflow_policy", ["reject", "wait"]) +def test_non_burst_policy_ignores_burst_limit(monkeypatch, overflow_policy): + """Only the burst policy can use slots above the replica limit.""" + p = _make_provider( + replicas=1, + overflow_policy=overflow_policy, + acquire_timeout=1, + burst_limit=2, + ) + fake_cls = _install_fake_sdk(monkeypatch, p) + + p.acquire("t1", user_id="u1") + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + assert len(fake_cls.create_calls) == 1 + + +def test_capacity_burst_with_zero_limit_falls_back_to_reject(monkeypatch): + """overflow_policy='burst' with burst_limit=0 is treated as 'reject'.""" + p = _make_provider(replicas=1, overflow_policy="burst", burst_limit=0) + _install_fake_sdk(monkeypatch, p) + + p.acquire("t1", user_id="u1") + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + +def test_capacity_release_on_create_failure(monkeypatch): + """A failed create releases the reserved slot so capacity is not leaked.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + call_count = 0 + + def flaky_create(**kw): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("API down") + return FakeClient(sandbox_id=f"sb-{call_count}") + + fake_cls.create_factory = flaky_create + + with pytest.raises(RuntimeError, match="API down"): + p.acquire("t1", user_id="u1") + + assert p._reserved_slots == 0 + sid = p.acquire("t1", user_id="u1") + assert sid is not None + + +def test_capacity_release_on_bootstrap_failure(monkeypatch): + """A failed bootstrap releases the reserved slot.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + call_count = 0 + + def flaky_create(**kw): + nonlocal call_count + call_count += 1 + if call_count == 1: + return FakeClient( + sandbox_id="sb-broken", + commands=FakeCommandsAPI([SimpleNamespace(stdout="", stderr="fail", exit_code=1)]), + ) + return FakeClient(sandbox_id=f"sb-ok-{call_count}") + + fake_cls.create_factory = flaky_create + + with pytest.raises(RuntimeError, match="bootstrap"): + p.acquire("t1", user_id="u1") + + assert p._reserved_slots == 0 + sid = p.acquire("t1", user_id="u1") + assert sid is not None + + +def test_capacity_reject_policy_does_not_leak_reserved_slots(monkeypatch): + """Repeated reject failures must not accumulate reserved slots.""" + p = _make_provider(replicas=1, overflow_policy="reject") + _install_fake_sdk(monkeypatch, p) + + p.acquire("t1", user_id="u1") + + for _ in range(5): + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t-extra", user_id="u-extra") + assert p._reserved_slots == 0 + + +def test_capacity_keeps_slot_when_warm_eviction_reconnect_fails(monkeypatch): + """An uncertain warm eviction must not make room for a new VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + fake_cls.connect_factory = lambda _sid, **_kw: (_ for _ in ()).throw(RuntimeError("gone")) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + assert len(p._warm_pool) == 1 + assert len(p._sandboxes) == 0 + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + assert len(fake_cls.create_calls) == 1 + assert "sb-1" not in p._warm_pool + assert p._transitioning_slots == 1 + + +def test_capacity_keeps_slot_when_warm_reclaim_reconnect_fails(monkeypatch): + """An uncertain warm reclaim must not make room for a new VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + sid = p.acquire("t1", user_id="u1") + p.release(sid) + fake_cls.connect_factory = lambda _sid, **_kw: (_ for _ in ()).throw(RuntimeError("network down")) + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t1", user_id="u1") + + assert len(fake_cls.create_calls) == 1 + assert p._eviction_tombstones == {sid} + assert p._reserved_slots == 0 + assert p._transitioning_slots == 1 + + destroy_client = FakeClient(sandbox_id=sid) + fake_cls.connect_factory = lambda _sid, **_kw: destroy_client + assert p.acquire("t2", user_id="u2") != sid + assert destroy_client.killed + assert p._transitioning_slots == 0 + + +def test_shutdown_during_reclaim_reconnect_failure_tracks_vm(monkeypatch): + """Shutdown must see a warm VM while reclaim reconnects.""" + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + sid = p.acquire("t1", user_id="u1") + p.release(sid) + + reconnect_started = threading.Event() + allow_failure = threading.Event() + shutdown_client = FakeClient(sandbox_id=sid) + calls = 0 + + def reconnect(_sid, **_kw): + nonlocal calls + calls += 1 + if calls == 1: + reconnect_started.set() + assert allow_failure.wait(timeout=2) + raise RuntimeError("network down") + return shutdown_client + + fake_cls.connect_factory = reconnect + result: list[str | None] = [] + reclaim = threading.Thread( + target=lambda: result.append(p._reclaim_warm_pool_sandbox("t1", user_id="u1")), + ) + reclaim.start() + assert reconnect_started.wait(timeout=1) + + p.shutdown() + allow_failure.set() + reclaim.join(timeout=5) + + assert result == [None] + assert shutdown_client.killed + assert p._warm_pool == {} + assert p._eviction_tombstones == set() + assert p._remote_ops_in_progress == set() + + +def test_capacity_keeps_slot_when_warm_reclaim_bootstrap_kill_fails(monkeypatch): + """An uncertain reclaim cleanup must not make room for a new VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + sid = p.acquire("t1", user_id="u1") + p.release(sid) + reconnect_client = FakeClient( + sandbox_id=sid, + commands=FakeCommandsAPI( + [ + SimpleNamespace(stdout="ok", stderr="", exit_code=0), + SimpleNamespace(stdout="", stderr="bootstrap failed", exit_code=1), + ] + ), + ) + reconnect_client.kill = MagicMock(side_effect=RuntimeError("kill failed")) + fake_cls.connect_factory = lambda _sid, **_kw: reconnect_client + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t1", user_id="u1") + + assert len(fake_cls.create_calls) == 1 + assert p._eviction_tombstones == {sid} + assert p._reserved_slots == 0 + assert p._transitioning_slots == 1 + + +def test_capacity_keeps_slot_when_create_bootstrap_kill_fails(monkeypatch): + """An uncertain create cleanup must not make room for a new VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient( + sandbox_id="bootstrap-uncertain", + commands=FakeCommandsAPI([SimpleNamespace(stdout="", stderr="bootstrap failed", exit_code=1)]), + ) + client.kill = MagicMock(side_effect=RuntimeError("kill failed")) + fake_cls.create_factory = lambda **_kw: client + + with pytest.raises(RuntimeError, match="bootstrap"): + p.acquire("t1", user_id="u1") + + fake_cls.connect_factory = lambda _sid, **_kw: (_ for _ in ()).throw(RuntimeError("network down")) + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + assert len(fake_cls.create_calls) == 1 + assert p._eviction_tombstones == {client.sandbox_id} + assert p._reserved_slots == 0 + assert p._transitioning_slots == 1 + + destroy_client = FakeClient(sandbox_id=client.sandbox_id) + fake_cls.connect_factory = lambda _sid, **_kw: destroy_client + fake_cls.create_factory = lambda **_kw: FakeClient(sandbox_id="bootstrap-replacement") + assert p.acquire("t2", user_id="u2") != client.sandbox_id + assert destroy_client.killed + assert p._transitioning_slots == 0 + + +def test_shutdown_during_create_bootstrap_kill_failure_tracks_vm(monkeypatch): + """Shutdown must see a created VM while bootstrap is in progress.""" + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="create-bootstrap-race") + client.kill = MagicMock(side_effect=RuntimeError("kill failed")) + fake_cls.create_factory = lambda **_kw: client + shutdown_client = FakeClient(sandbox_id=client.sandbox_id) + fake_cls.connect_factory = lambda _sid, **_kw: shutdown_client + bootstrap_started = threading.Event() + allow_failure = threading.Event() + + def slow_bootstrap(_client): + bootstrap_started.set() + assert allow_failure.wait(timeout=2) + raise RuntimeError("bootstrap failed") + + monkeypatch.setattr(p, "_bootstrap_sandbox_paths", slow_bootstrap) + result: list[Exception] = [] + + def create(): + try: + p.acquire("t1", user_id="u1") + except Exception as error: + result.append(error) + + create_thread = threading.Thread(target=create) + create_thread.start() + assert bootstrap_started.wait(timeout=1) + + p.shutdown() + allow_failure.set() + create_thread.join(timeout=5) + + assert len(result) == 1 + assert shutdown_client.killed + assert p._eviction_tombstones == set() + assert p._remote_ops_in_progress == set() + + +def test_capacity_keeps_slot_when_warm_eviction_kill_fails(monkeypatch): + """An uncertain kill must not make room for a new VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + reconnect_client = FakeClient(sandbox_id="created-1") + + def fail_kill(): + raise RuntimeError("control plane unavailable") + + reconnect_client.kill = fail_kill + fake_cls.connect_factory = lambda _sid, **_kw: reconnect_client + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + assert len(fake_cls.create_calls) == 1 + assert reconnect_client.closed + + +def test_capacity_retries_tombstone_until_warm_vm_is_destroyed(monkeypatch): + """A later confirmed eviction can release the retained capacity slot.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + fake_cls.connect_factory = lambda _sid, **_kw: (_ for _ in ()).throw(RuntimeError("network down")) + + with pytest.raises(SandboxCapacityExceededError): + p.acquire("t2", user_id="u2") + + reconnect_client = FakeClient(sandbox_id=sid1) + fake_cls.connect_factory = lambda _sid, **_kw: reconnect_client + sid2 = p.acquire("t2", user_id="u2") + + assert sid2 != sid1 + assert reconnect_client.killed + assert p._eviction_tombstones == set() + assert p._transitioning_slots == 0 + + +def test_tombstone_eviction_has_one_retry_owner(monkeypatch): + """Only one thread can retry a tombstone at a time.""" + p = _make_provider() + _install_fake_sdk(monkeypatch, p) + p._eviction_tombstones = {"a", "b"} + p._transitioning_slots = 2 + p._evictions_in_progress = {"b"} + + reconnect_started = threading.Event() + allow_reconnect = threading.Event() + + def slow_reconnect(_cls, sandbox_id): + assert sandbox_id == "a" + reconnect_started.set() + assert allow_reconnect.wait(timeout=2) + return None + + monkeypatch.setattr(p, "_reconnect_live_client", slow_reconnect) + first = threading.Thread(target=p._evict_oldest_warm) + first.start() + assert reconnect_started.wait(timeout=1) + + assert p._evict_oldest_warm() is None + allow_reconnect.set() + first.join(timeout=5) + + assert p._eviction_tombstones == {"b"} + assert p._transitioning_slots == 1 + + +def test_shutdown_during_initial_eviction_reconnect_failure_tracks_vm(monkeypatch): + """Shutdown must destroy a VM while its first eviction reconnects.""" + p = _make_provider() + fake_cls = _install_fake_sdk(monkeypatch, p) + p._warm_pool["warm-1"] = ("seed", time.time()) + + reconnect_started = threading.Event() + allow_failure = threading.Event() + shutdown_client = FakeClient(sandbox_id="warm-1") + calls = 0 + + def reconnect(_sid, **_kw): + nonlocal calls + calls += 1 + if calls == 1: + reconnect_started.set() + assert allow_failure.wait(timeout=2) + raise RuntimeError("network down") + return shutdown_client + + fake_cls.connect_factory = reconnect + eviction = threading.Thread(target=p._evict_oldest_warm) + eviction.start() + assert reconnect_started.wait(timeout=1) + + p.shutdown() + allow_failure.set() + eviction.join(timeout=5) + + assert shutdown_client.killed + assert p._warm_pool == {} + assert p._eviction_tombstones == set() + assert p._evictions_in_progress == set() + + +def test_capacity_reset_uses_destructive_shutdown_semantics(monkeypatch): + """reset() destroys tracked E2B resources and ends the provider.""" + p = _make_provider(replicas=1, overflow_policy="wait", acquire_timeout=30) + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="active") + fake_cls.create_factory = lambda **_kwargs: client + p.acquire("t1", user_id="u1") + + p._reserved_slots = 3 + p.reset() + assert p._reserved_slots == 0 + assert p._shutdown_called + assert client.killed + assert client.closed + + +def test_capacity_reset_wakes_waiter_with_shutdown_error(monkeypatch): + p = _make_provider(replicas=1, overflow_policy="wait", acquire_timeout=30) + _install_fake_sdk(monkeypatch, p) + p.acquire("t1", user_id="u1") + result: list[Exception] = [] + started = threading.Event() + + def wait_for_capacity(): + started.set() + try: + p.acquire("t2", user_id="u2") + except Exception as error: + result.append(error) + + waiter = threading.Thread(target=wait_for_capacity) + waiter.start() + assert started.wait(timeout=1) + time.sleep(0.1) + + p.reset() + waiter.join(timeout=0.5) + + assert not waiter.is_alive() + assert len(result) == 1 + assert isinstance(result[0], SandboxCapacityExceededError) + assert result[0].reason == "shutdown" + + +def test_capacity_default_config_values(): + """Default config values are backward-compatible.""" + p = _make_provider() + assert p._config["overflow_policy"] == "wait" + assert p._config["acquire_timeout"] == 30 + assert p._config["burst_limit"] == 0 + + +@pytest.mark.anyio +async def test_e2b_acquire_async_uses_dedicated_executor(monkeypatch): + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + p._acquire_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="e2b-sandbox-acquire", + ) + thread_names: list[str] = [] + + def create_client(**_kwargs): + thread_names.append(threading.current_thread().name) + return FakeClient(sandbox_id="async-e2b") + + async def fail_to_thread(*_args, **_kwargs): + raise AssertionError("E2B acquire must not use the default asyncio executor") + + fake_cls.create_factory = create_client + monkeypatch.setattr(asyncio, "to_thread", fail_to_thread) + + try: + sandbox_id = await p.acquire_async("t1", user_id="u1") + finally: + p.shutdown() + p._acquire_executor.shutdown(wait=True, cancel_futures=True) + + assert sandbox_id == "async-e2b" + assert thread_names == ["e2b-sandbox-acquire_0"] + + +# ── Race-condition regression tests ────────────────────────────────────── + + +def test_capacity_release_holds_slot_during_transition(monkeypatch): + """replicas=1. T1 releases sb1 (slow sync). T2 acquires → reject.""" + p = _make_provider(replicas=1, overflow_policy="reject") + _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + + sync_started = threading.Event() + allow_sync = threading.Event() + + def slow_sync(*_args, **_kwargs): + sync_started.set() + assert allow_sync.wait(timeout=2) + + monkeypatch.setattr(p, "_sync_outputs_to_host", slow_sync) + + result: list[str | Exception] = [] + + def do_acquire(): + try: + result.append(p.acquire("t2", user_id="u2")) + except Exception as e: + result.append(e) + + release_thread = threading.Thread(target=p.release, args=(sid1,)) + release_thread.start() + assert sync_started.wait(timeout=1), "release must enter sync" + + t = threading.Thread(target=do_acquire) + t.start() + t.join(timeout=2) + + allow_sync.set() + release_thread.join(timeout=2) + + assert isinstance(result[0], SandboxCapacityExceededError), f"expected reject, got {result[0]!r}" + assert len(p._sandboxes) + len(p._warm_pool) <= p._capacity_limit() + + +def test_capacity_reclaim_holds_slot_during_transition(monkeypatch): + """replicas=1. T1 reclaims warm sb (slow reconnect). T2 acquires → reject.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + + reconnect_started = threading.Event() + allow_reconnect = threading.Event() + original_connect = fake_cls.connect_factory + + def slow_connect(sid, **kw): + reconnect_started.set() + assert allow_reconnect.wait(timeout=2) + return original_connect(sid, **kw) + + fake_cls.connect_factory = slow_connect + + result: list[str | Exception] = [] + + def do_acquire(): + try: + result.append(p.acquire("t2", user_id="u2")) + except Exception as e: + result.append(e) + + reclaim_thread = threading.Thread(target=p._reclaim_warm_pool_sandbox, args=("t1",), kwargs={"user_id": "u1"}) + reclaim_thread.start() + assert reconnect_started.wait(timeout=1), "reclaim must enter reconnect" + + t = threading.Thread(target=do_acquire) + t.start() + t.join(timeout=2) + + allow_reconnect.set() + reclaim_thread.join(timeout=2) + + assert isinstance(result[0], SandboxCapacityExceededError), f"expected reject during reclaim, got {result[0]!r}" + + +def test_capacity_shutdown_wakes_waiter_with_error(monkeypatch): + """Waiter blocked on capacity must raise on shutdown, not create VM.""" + p = _make_provider(replicas=1, overflow_policy="wait", acquire_timeout=30) + fake_cls = _install_fake_sdk(monkeypatch, p) + + p.acquire("t1", user_id="u1") + + result: list[Exception | None] = [] + started = threading.Event() + + def waiter(): + started.set() + try: + p.acquire("t2", user_id="u2") + except Exception as e: + result.append(e) + + t = threading.Thread(target=waiter) + t.start() + assert started.wait(timeout=1) + + time.sleep(0.2) + p.shutdown() + t.join(timeout=5) + + assert len(result) == 1 + assert isinstance(result[0], SandboxCapacityExceededError) + assert "shutting down" in str(result[0]).lower() + assert len(fake_cls.create_calls) == 1 + + +def test_capacity_create_aborted_by_shutdown_kills_vm(monkeypatch): + """shutdown after create() but before commit: kill VM, raise error.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + created_client: FakeClient | None = None + create_returned = threading.Event() + allow_commit = threading.Event() + + def intercept_create(**kw): + nonlocal created_client + c = FakeClient(sandbox_id="sb-shutdown-race") + created_client = c + create_returned.set() + assert allow_commit.wait(timeout=2) + return c + + fake_cls.create_factory = intercept_create + + result: list[str | Exception] = [] + + def do_acquire(): + try: + result.append(p.acquire("t1", user_id="u1")) + except Exception as e: + result.append(e) + + t = threading.Thread(target=do_acquire) + t.start() + assert create_returned.wait(timeout=1) + + p.shutdown() + allow_commit.set() + t.join(timeout=5) + + assert len(result) == 1 + assert isinstance(result[0], SandboxCapacityExceededError) + assert created_client is not None + assert created_client.killed, "VM must be killed when shutdown aborts create" + + +def test_capacity_create_aborted_by_shutdown_retries_failed_kill(monkeypatch): + """A transient kill failure must not orphan a VM created during shutdown.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + created_client = FakeClient(sandbox_id="sb-shutdown-retry") + created_client.kill = MagicMock(side_effect=RuntimeError("control plane unavailable")) + retry_client = FakeClient(sandbox_id=created_client.sandbox_id) + fake_cls.connect_factory = lambda _sid, **_kwargs: retry_client + create_returned = threading.Event() + allow_create_return = threading.Event() + + def intercept_create(**_kwargs): + create_returned.set() + assert allow_create_return.wait(timeout=2) + return created_client + + fake_cls.create_factory = intercept_create + result: list[Exception] = [] + + def do_acquire(): + try: + p.acquire("t1", user_id="u1") + except Exception as error: + result.append(error) + + acquire_thread = threading.Thread(target=do_acquire) + acquire_thread.start() + assert create_returned.wait(timeout=1) + + p.shutdown() + allow_create_return.set() + acquire_thread.join(timeout=5) + + assert len(result) == 1 + assert isinstance(result[0], SandboxCapacityExceededError) + assert retry_client.killed + assert retry_client.closed + + +def test_capacity_create_aborted_by_shutdown_tracks_uncertain_cleanup(monkeypatch): + """A persistent cleanup failure must keep the remote VM ID visible.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + created_client = FakeClient(sandbox_id="sb-shutdown-uncertain") + created_client.kill = MagicMock(side_effect=RuntimeError("control plane unavailable")) + fake_cls.connect_factory = lambda _sid, **_kwargs: (_ for _ in ()).throw(RuntimeError("network unavailable")) + create_returned = threading.Event() + allow_create_return = threading.Event() + + def intercept_create(**_kwargs): + create_returned.set() + assert allow_create_return.wait(timeout=2) + return created_client + + fake_cls.create_factory = intercept_create + result: list[Exception] = [] + + def do_acquire(): + try: + p.acquire("t1", user_id="u1") + except Exception as error: + result.append(error) + + acquire_thread = threading.Thread(target=do_acquire) + acquire_thread.start() + assert create_returned.wait(timeout=1) + + p.shutdown() + allow_create_return.set() + acquire_thread.join(timeout=5) + + assert len(result) == 1 + assert "could not confirm cleanup" in str(result[0]) + assert p._remote_ops_in_progress == {created_client.sandbox_id} + + +def test_capacity_concurrent_different_threads_only_one_creates(monkeypatch): + """replicas=1. Two threads for different users/threads → exactly 1 create.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + create_count = 0 + create_lock = threading.Lock() + create_started = threading.Event() + allow_first_create = threading.Event() + + def counted_create(**kw): + nonlocal create_count + with create_lock: + create_count += 1 + if create_count == 1: + create_started.set() + assert allow_first_create.wait(timeout=2) + return FakeClient(sandbox_id=f"sb-{create_count}") + + fake_cls.create_factory = counted_create + + results: list[str | Exception] = [] + barrier = threading.Barrier(2, timeout=5) + + def worker_a(): + barrier.wait() + try: + results.append(p.acquire("t-a", user_id="u-a")) + except Exception as e: + results.append(e) + + def worker_b(): + barrier.wait() + time.sleep(0.1) + try: + results.append(p.acquire("t-b", user_id="u-b")) + except Exception as e: + results.append(e) + + ta = threading.Thread(target=worker_a) + tb = threading.Thread(target=worker_b) + ta.start() + tb.start() + assert create_started.wait(timeout=2) + + tb.join(timeout=5) + allow_first_create.set() + ta.join(timeout=5) + + assert create_count == 1, f"expected 1 create, got {create_count}" + assert len(results) == 2 + sids = [r for r in results if isinstance(r, str)] + errs = [r for r in results if isinstance(r, Exception)] + assert len(sids) == 1, f"expected 1 successful acquire, got {len(sids)}" + assert len(errs) == 1 + assert isinstance(errs[0], SandboxCapacityExceededError) + + +def test_shutdown_during_release_does_not_repopulate_warm_pool(monkeypatch): + """release in flight when shutdown fires must kill the VM, not park it.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="sb-release-shutdown") + fake_cls.create_factory = lambda **_kw: client + + sid1 = p.acquire("t1", user_id="u1") + + sync_entered = threading.Event() + allow_sync = threading.Event() + + def slow_sync(*_args, **_kwargs): + sync_entered.set() + assert allow_sync.wait(timeout=2) + + monkeypatch.setattr(p, "_sync_outputs_to_host", slow_sync) + + release_done = threading.Event() + + def do_release(): + p.release(sid1) + release_done.set() + + t = threading.Thread(target=do_release) + t.start() + assert sync_entered.wait(timeout=1) + + p.shutdown() + allow_sync.set() + t.join(timeout=5) + + assert release_done.is_set() + assert sid1 not in p._warm_pool, "must not park in warm pool after shutdown" + assert client.killed, "release must kill its saved client after shutdown" + + +def test_shutdown_during_release_close_kills_published_warm_vm(monkeypatch): + """Shutdown must find a released VM before its client transport closes.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="sb-release-close-race") + fake_cls.create_factory = lambda **_kw: client + sid = p.acquire("t1", user_id="u1") + + close_started = threading.Event() + allow_close = threading.Event() + + def slow_close(): + client.closed = True + close_started.set() + assert allow_close.wait(timeout=2) + + def fail_kill_after_close(): + if client.closed: + raise RuntimeError("transport closed") + client.killed = True + + client.close = slow_close + client.kill = fail_kill_after_close + shutdown_client = FakeClient(sandbox_id=sid) + fake_cls.connect_factory = lambda _sid, **_kw: shutdown_client + + release_thread = threading.Thread(target=p.release, args=(sid,)) + release_thread.start() + assert close_started.wait(timeout=1) + + shutdown_thread = threading.Thread(target=p.shutdown) + shutdown_thread.start() + shutdown_thread.join(timeout=2) + allow_close.set() + release_thread.join(timeout=5) + + assert shutdown_client.killed + + +def test_release_keeps_parked_vm_reclaimable_during_client_close(monkeypatch): + """A concurrent acquire must not evict a VM because release counts it twice.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + first_client = FakeClient(sandbox_id="sb-first") + second_client = FakeClient(sandbox_id="sb-second") + created_clients = iter((first_client, second_client)) + fake_cls.create_factory = lambda **_kwargs: next(created_clients) + + first_sid = p.acquire("t1", user_id="u1") + close_started = threading.Event() + allow_close = threading.Event() + + def slow_close(): + close_started.set() + assert allow_close.wait(timeout=2) + first_client.closed = True + + first_client.close = slow_close + release_thread = threading.Thread(target=p.release, args=(first_sid,)) + release_thread.start() + assert close_started.wait(timeout=1) + + assert p.acquire("t2", user_id="u2") == "sb-second" + allow_close.set() + release_thread.join(timeout=5) + + assert p.acquire("t1", user_id="u1") == first_sid + + +def test_shutdown_during_reclaim_does_not_register_active(monkeypatch): + """reclaim in flight when shutdown fires must kill the VM, not register.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + sid1 = p.acquire("t1", user_id="u1") + p.release(sid1) + + reconnect_entered = threading.Event() + allow_reconnect = threading.Event() + original_connect = fake_cls.connect_factory + + def slow_connect(sid, **kw): + reconnect_entered.set() + assert allow_reconnect.wait(timeout=2) + return original_connect(sid, **kw) + + fake_cls.connect_factory = slow_connect + + result: list[str | None] = [] + + def do_reclaim(): + result.append(p._reclaim_warm_pool_sandbox("t1", user_id="u1")) + + t = threading.Thread(target=do_reclaim) + t.start() + assert reconnect_entered.wait(timeout=1) + + p.shutdown() + allow_reconnect.set() + t.join(timeout=5) + + assert result == [None], "reclaim must return None after shutdown" + assert sid1 not in p._sandboxes, "must not register active after shutdown" + + +def test_shutdown_during_bootstrap_does_not_commit(monkeypatch): + """create in bootstrap when shutdown fires must kill, not commit.""" + p = _make_provider(replicas=2, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + + client = FakeClient(sandbox_id="sb-bootstrap-shutdown") + fake_cls.create_factory = lambda **kw: client + + bootstrap_entered = threading.Event() + allow_bootstrap = threading.Event() + + original_bootstrap = p._bootstrap_sandbox_paths + + def slow_bootstrap(c): + bootstrap_entered.set() + assert allow_bootstrap.wait(timeout=2) + return original_bootstrap(c) + + monkeypatch.setattr(p, "_bootstrap_sandbox_paths", slow_bootstrap) + + result: list[str | Exception] = [] + + def do_acquire(): + try: + result.append(p.acquire("t1", user_id="u1")) + except Exception as e: + result.append(e) + + t = threading.Thread(target=do_acquire) + t.start() + assert bootstrap_entered.wait(timeout=1) + + p.shutdown() + allow_bootstrap.set() + t.join(timeout=5) + + assert len(result) == 1 + assert isinstance(result[0], SandboxCapacityExceededError) + assert "sb-bootstrap-shutdown" not in p._sandboxes + assert client.killed, "VM must be killed" + + +def test_discovery_reports_busy_capacity_without_killing_remote_vm(monkeypatch): + """Discovery must report a full provider without destroying the remote VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + discovered_client = FakeClient(sandbox_id="sb-remote") + fake_cls.connect_factory = lambda _sid, **_kw: discovered_client + bootstrap = MagicMock() + monkeypatch.setattr(p, "_bootstrap_sandbox_paths", bootstrap) + + # Fill the single slot. + p.acquire("t1", user_id="u1") + bootstrap.reset_mock() + + # Discovery finds a matching sandbox. + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id="sb-remote", + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "u2", + "deer_flow_thread": "t2", + }, + ) + ] + + with pytest.raises(SandboxCapacityExceededError): + p._discover_remote_sandbox("t2", user_id="u2") + + assert "sb-remote" not in p._sandboxes + bootstrap.assert_not_called() + assert not discovered_client.killed + assert discovered_client.closed + + +def test_discovery_reports_shutdown_without_killing_remote_vm(monkeypatch, caplog): + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="sb-remote") + fake_cls.connect_factory = lambda _sid, **_kw: client + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id=client.sandbox_id, + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "u1", + "deer_flow_thread": "t1", + }, + ) + ] + p._shutdown_called = True + + with caplog.at_level("INFO"), pytest.raises(SandboxCapacityExceededError) as error: + p._discover_remote_sandbox("t1", user_id="u1") + + assert error.value.reason == "shutdown" + assert "shutting down" in caplog.text + assert "capacity is full" not in caplog.text + assert not client.killed + assert client.closed + + +def test_discovery_bootstrap_kill_failure_retains_reserved_slot(monkeypatch): + """Discovery keeps capacity when bootstrap cleanup cannot destroy the VM.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient( + sandbox_id="discovery-bootstrap-failure", + commands=FakeCommandsAPI( + [ + SimpleNamespace(stdout="ok", stderr="", exit_code=0), + SimpleNamespace(stdout="", stderr="bootstrap failed", exit_code=1), + ] + ), + ) + client.kill = MagicMock(side_effect=RuntimeError("kill failed")) + fake_cls.connect_factory = lambda _sid, **_kw: client + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id=client.sandbox_id, + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "u1", + "deer_flow_thread": "t1", + }, + ) + ] + + assert p._discover_remote_sandbox("t1", user_id="u1") is None + + assert p._eviction_tombstones == {client.sandbox_id} + assert p._reserved_slots == 0 + assert p._transitioning_slots == 1 + assert p._remote_ops_in_progress == set() + + +def test_shutdown_does_not_retry_kill_for_unowned_discovery_vm(monkeypatch): + """Shutdown does not claim an unowned discovery VM after cleanup fails.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="discovery-bootstrap-race") + client.kill = MagicMock(side_effect=RuntimeError("kill failed")) + shutdown_client = FakeClient(sandbox_id=client.sandbox_id) + connect_calls = 0 + + def reconnect(_sid, **_kw): + nonlocal connect_calls + connect_calls += 1 + return client if connect_calls == 1 else shutdown_client + + fake_cls.connect_factory = reconnect + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id=client.sandbox_id, + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "u1", + "deer_flow_thread": "t1", + }, + ) + ] + bootstrap_started = threading.Event() + allow_failure = threading.Event() + + def slow_bootstrap(_client): + bootstrap_started.set() + assert allow_failure.wait(timeout=2) + raise RuntimeError("bootstrap failed") + + monkeypatch.setattr(p, "_bootstrap_sandbox_paths", slow_bootstrap) + result: list[str | None] = [] + discovery = threading.Thread( + target=lambda: result.append(p._discover_remote_sandbox("t1", user_id="u1")), + ) + discovery.start() + assert bootstrap_started.wait(timeout=1) + + p.shutdown() + allow_failure.set() + discovery.join(timeout=5) + + assert result == [None] + assert not shutdown_client.killed + assert connect_calls == 1 + assert p._eviction_tombstones == set() + assert p._remote_ops_in_progress == set() + assert p._unowned_remote_ops_in_progress == set() + + +def test_shutdown_during_discovery_does_not_kill_unowned_vm(monkeypatch): + """Shutdown closes discovery clients without killing unowned remote VMs.""" + p = _make_provider(replicas=1, overflow_policy="reject") + fake_cls = _install_fake_sdk(monkeypatch, p) + client = FakeClient(sandbox_id="sb-discovery-shutdown") + fake_cls.connect_factory = lambda _sid, **_kw: client + fake_cls.list_return = [ + SimpleNamespace( + sandbox_id=client.sandbox_id, + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "u1", + "deer_flow_thread": "t1", + }, + ) + ] + + reserved = threading.Event() + allow_commit = threading.Event() + reserve_capacity = p._reserve_capacity + + def pause_after_reserve(thread_id, user_id, *, remote_id=None, remote_owned=True): + reserve_capacity( + thread_id, + user_id, + remote_id=remote_id, + remote_owned=remote_owned, + ) + reserved.set() + assert allow_commit.wait(timeout=2) + + monkeypatch.setattr(p, "_reserve_capacity", pause_after_reserve) + result: list[str | None] = [] + thread = threading.Thread( + target=lambda: result.append(p._discover_remote_sandbox("t1", user_id="u1")), + ) + thread.start() + assert reserved.wait(timeout=1) + + p.shutdown() + allow_commit.set() + thread.join(timeout=5) + + assert result == [None] + assert not client.killed + assert client.closed + assert p._sandboxes == {} + assert p._reserved_slots == 0