diff --git a/.github/workflows/backend-unit-tests.yml b/.github/workflows/backend-unit-tests.yml index f6b48effd..29779d00a 100644 --- a/.github/workflows/backend-unit-tests.yml +++ b/.github/workflows/backend-unit-tests.yml @@ -61,6 +61,15 @@ jobs: --health-interval 10s --health-timeout 5s --health-retries 5 + redis: + image: redis:7-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 10s + --health-timeout 5s + --health-retries 5 env: TEST_POSTGRES_URI: postgresql://deerflow:deerflow@localhost:5432/deerflow_test?sslmode=disable diff --git a/README.md b/README.md index 642eecea7..62b56e362 100644 --- a/README.md +++ b/README.md @@ -907,7 +907,17 @@ 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. +With in-memory ownership, `replicas` limits one Gateway process. With Redis +ownership, E2B shares one capacity Hash between workers using the same +`sandbox.ownership.key_prefix`; `replicas` (plus a configured burst) is then a +deployment-wide hard limit. Use one unique prefix and the same effective limit +per deployment. To change the limit, stop its Gateways, delete the capacity +Hash, and restart; mismatched workers fail closed. + +The Hash counts remote VMs and in-flight creates, repairs interrupted creates +from E2B metadata, grace-protects stale inventory omissions, and blocks new +creates while Redis or initial inventory is unavailable. Run Redis with persistence, non-evicting memory, and HA. + E2B acquisition uses a bounded executor. Waiting acquisitions do not use the default asyncio executor. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 7b947ed96..2ce7a1e60 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -616,8 +616,13 @@ that cannot tell sibling branches apart. 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. + policy can evict one warm VM before it returns an error. With memory + ownership, `replicas` limits one Gateway process. Redis ownership shares one + `:e2b-capacity` Hash, making the limit (plus a bounded + burst) deployment-wide. Lua atomically manages VM and in-flight-create + entries; missing or unavailable state fails closed. E2B reservation metadata + repairs interrupted creates. Inventory replacement is revision-CAS guarded, + and incomplete inventories never remove entries; complete omissions get a grace period. 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 @@ -656,7 +661,7 @@ that cannot tell sibling branches apart. - **Teardown join budget covers refresh plus release.** Redis bounds each ownership operation at five seconds, and context exit can catch the heartbeat in one final refresh before its `finally` performs the final release. `_TEARDOWN_JOIN_TIMEOUT_SECONDS` is therefore 12 seconds — greater than both sequential operation bounds — so a normal pair of socket timeouts does not emit the deferred-release warning; a still-running heartbeat continues to own the release safely. - **An absent lease means the same thing on both paths, and reconciliation must say so too.** The `LAPSED` rule above only covers an owner renewing its *own* lease; on its own it does not make state loss safe, because reconciliation reads the same absent key as "orphan, adopt". After a Redis flush (restart without persistence, or eviction under `maxmemory`) every owner is alive and merely pre-renewal-tick, so whichever instance reconciles first would adopt every live container, each real owner's next renewal would report `LOST`, and it would drop a sandbox mid-turn for the adopter to idle-destroy — #4206 through the back door. `_adoptable_after_grace` closes it: an untracked container must be seen unowned (`owner()`, a read-only peek — the atomic `claim()` is still what actually gates adoption) across a full lease TTL before it can be adopted, tracked per container in `_unowned_since`. That rebuilds the delay the flush erased — a live owner republishes within one renewal interval, shorter than the TTL by construction (`ttl_multiplier >= 2`) — while a genuinely crashed owner never republishes, so its containers are still adopted one grace later rather than leaking. A republished lease **resets** the grace; a pausing-only timer would still expire over a live owner's lease. The grace is skipped when `supports_cross_process` is `False`: no peer can hold a lease such a store would show us, so single-instance deployments keep instant orphan cleanup, and a grace could not help a multi-worker gateway on `memory` anyway (peers are invisible to each other's leases with or without it). - **The `memory` store is single-instance only** and says so via `supports_cross_process = False`; the provider logs a warning at startup when the configured store cannot see peers. A multi-worker gateway on `memory` has no cross-process coordination at all — same contract as `stream_bridge`'s memory backend. This is why the redis inference matters: it reads `app_config.stream_bridge` **and** the env var, in the same order the bridge's own resolver does, so any deployment already pointing the bridge at Redis (i.e. every multi-instance one) gets a redis ownership store without extra config. - - `get()` stays a pure in-memory lookup and must never call the store (that is blocking filesystem/network IO on the event loop); anchored by `tests/blocking_io/test_aio_sandbox_get.py`, which injects a deliberately-blocking probe store so the anchor keeps its teeth regardless of the configured backend. Tests: `tests/test_sandbox_ownership_store.py` (store contract, defined once for **both** backends — but the redis tier is `@pytest.mark.integration` + opt-in via `DEER_FLOW_TEST_REDIS_URL` and self-skips, and **CI provisions no redis**, so the merge gate runs the memory tier only and the Lua scripts never execute there; drift between the backends is caught only when the suite runs against a live redis. There is no fake-redis tier because the redis exclusion lives in Lua a fake would not execute) and `tests/test_sandbox_orphan_reconciliation.py` (provider behaviour, two providers sharing one store). + - `get()` stays a pure in-memory lookup and must never call the store (that is blocking filesystem/network IO on the event loop); anchored by `tests/blocking_io/test_aio_sandbox_get.py`, which injects a deliberately-blocking probe store so the anchor keeps its teeth regardless of the configured backend. Tests: `tests/test_sandbox_ownership_store.py` (store contract, defined once for **both** backends — the redis tier is `@pytest.mark.integration`, uses `DEER_FLOW_TEST_REDIS_URL` when set, and otherwise self-skips without a reachable Redis. Backend CI provisions Redis, so the merge gate executes the real Lua tier; there is no fake-redis tier because a fake would not execute the Lua exclusions) and `tests/test_sandbox_orphan_reconciliation.py` (provider behaviour, two providers sharing one store). - `BoxliteProvider` (`packages/harness/deerflow/community/boxlite/`) - BoxLite micro-VM isolation. The `boxlite` runtime is optional (`deerflow-harness[boxlite]`) and lazy-imported only when this provider is selected. The provider owns one private asyncio event loop on a daemon thread because BoxLite handles are loop-affine; sync `Sandbox` calls marshal onto that loop with `run_coroutine_threadsafe`. Boxes are named deterministically from `user_id:thread_id`, released into an in-process warm pool after each agent turn, and reclaimed only by the same user/thread. Warm-pool health checks use a short explicit timeout and forward that timeout through both BoxLite `exec(timeout=...)` and the private-loop `.result(timeout)` bridge so a hung VM cannot pin the per-thread acquire lock indefinitely. `sandbox.replicas` caps active + warm VMs per gateway process; if capacity is exhausted, only warm-pool VMs are evicted. `sandbox.idle_timeout` stops idle warm VMs after the configured seconds. `reset()` is intentionally a lightweight registry clear for `reset_sandbox_provider()` and does not close boxes, stop the idle reaper, or close the private loop; full teardown remains `shutdown()`. diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/ownership/factory.py b/backend/packages/harness/deerflow/community/aio_sandbox/ownership/factory.py index 3070f8c21..0be236c88 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/ownership/factory.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/ownership/factory.py @@ -61,7 +61,10 @@ def resolve_ownership_config(config: SandboxOwnershipConfig | None, *, stream_br return SandboxOwnershipConfig() -def _resolve_redis_url(config: SandboxOwnershipConfig) -> str: +def resolve_ownership_redis_url( + config: SandboxOwnershipConfig, +) -> str: + """Resolve the Redis endpoint shared by ownership-adjacent stores.""" return config.redis_url or os.getenv(_ENV_OWNERSHIP_REDIS_URL) or os.getenv(_ENV_STREAM_BRIDGE_REDIS_URL) or os.getenv("REDIS_URL") or "redis://localhost:6379/0" @@ -96,7 +99,7 @@ def make_sandbox_ownership_store(config: SandboxOwnershipConfig | None, *, owner if resolved.type == "redis": from .redis import RedisOwnershipStore - redis_url = _resolve_redis_url(resolved) + redis_url = resolve_ownership_redis_url(resolved) logger.info("Sandbox ownership store: redis (ttl=%.1fs, renewal=%.1fs)", ttl, resolved.renewal_interval_seconds) return RedisOwnershipStore( owner_id=effective_owner_id, diff --git a/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py b/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py index 289500a51..dd9e36ef7 100644 --- a/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py +++ b/backend/packages/harness/deerflow/community/e2b_sandbox/__init__.py @@ -12,8 +12,8 @@ Configuration example (``config.yaml``):: template: code-interpreter-v1 # e2b template id; defaults to e2b code-interpreter domain: e2b.dev # optional e2b domain (e.g. self-hosted) idle_timeout: 600 # forwarded to e2b ``set_timeout`` (seconds) - replicas: 3 # max concurrent sandboxes (LRU eviction beyond) - ownership: # required for safe multi-worker reconciliation + replicas: 3 # hard capacity shared when ownership is Redis + ownership: # multi-worker ownership + capacity coordination type: redis redis_url: $REDIS_URL reconciliation_interval_seconds: 60 diff --git a/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/__init__.py b/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/__init__.py new file mode 100644 index 000000000..e69750783 --- /dev/null +++ b/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/__init__.py @@ -0,0 +1,6 @@ +"""Redis-backed deployment-wide E2B capacity.""" + +from .redis import CapacityBackendError as CapacityBackendError +from .redis import RedisE2BCapacityStore as RedisE2BCapacityStore +from .redis import ReserveStatus as ReserveStatus +from .redis import make_e2b_capacity_store as make_e2b_capacity_store diff --git a/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/redis.py b/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/redis.py new file mode 100644 index 000000000..22795639b --- /dev/null +++ b/backend/packages/harness/deerflow/community/e2b_sandbox/capacity/redis.py @@ -0,0 +1,279 @@ +"""Atomic Redis Hash ledger for deployment-wide E2B capacity.""" + +from __future__ import annotations + +import enum +import logging + +from deerflow.community.aio_sandbox.ownership.factory import resolve_ownership_redis_url +from deerflow.config.sandbox_config import SandboxOwnershipConfig + +logger = logging.getLogger(__name__) + +_SOCKET_TIMEOUT_SECONDS = 5.0 + +_LEDGER_SCRIPT = """ +local function now_ms() + local current = redis.call('TIME') + return tonumber(current[1]) * 1000 + math.floor(tonumber(current[2]) / 1000) +end + +local function initialize(hard_limit) + redis.call('HSET', KEYS[1], + 'meta:state', 'initializing', + 'meta:hard_limit', hard_limit, + 'meta:revision', '0') +end + +local function mark_present(sandbox_id) + local field = 's:' .. sandbox_id + if redis.call('HGET', KEYS[1], field) == '1' then + return false + end + redis.call('HSET', KEYS[1], field, '1') + return true +end + +local operation = ARGV[1] +local hard_limit = ARGV[2] +local state = redis.call('HGET', KEYS[1], 'meta:state') +local stored_limit = redis.call('HGET', KEYS[1], 'meta:hard_limit') +if state ~= false and stored_limit ~= hard_limit then + return redis.error_reply( + 'E2B capacity ledger configuration mismatch: configured hard_limit=' + .. hard_limit .. ', ledger hard_limit=' .. (stored_limit or '') + ) +end + +if operation == 'revision' then + return tonumber(redis.call('HGET', KEYS[1], 'meta:revision') or '0') +end + +if operation == 'reserve' then + if state ~= 'ready' then + return 'NOT_READY' + end + local field = 'r:' .. ARGV[3] + if redis.call('HEXISTS', KEYS[1], field) == 1 then + return 'GRANTED' + end + if redis.call('HLEN', KEYS[1]) - 3 >= tonumber(hard_limit) then + return 'FULL' + end + redis.call('HSET', KEYS[1], field, tostring(now_ms())) + redis.call('HINCRBY', KEYS[1], 'meta:revision', 1) + return 'GRANTED' +end + +if operation == 'release' then + if state ~= false and redis.call('HDEL', KEYS[1], 's:' .. ARGV[3]) == 1 then + redis.call('HINCRBY', KEYS[1], 'meta:revision', 1) + end + return 'OK' +end + +if state == false then + if operation == 'reconcile' and tonumber(ARGV[3]) ~= 0 then + return 'STALE' + end + initialize(hard_limit) + state = 'initializing' +end + +if operation == 'track' then + local changed = false + if ARGV[3] ~= '' and redis.call('HDEL', KEYS[1], 'r:' .. ARGV[3]) == 1 then + changed = true + end + if mark_present(ARGV[4]) then + changed = true + end + if changed then + redis.call('HINCRBY', KEYS[1], 'meta:revision', 1) + end + return 'OK' +end + +if operation ~= 'reconcile' then + return redis.error_reply('unknown E2B capacity operation: ' .. operation) +end + +local expected_revision = tonumber(ARGV[3]) +local revision = tonumber(redis.call('HGET', KEYS[1], 'meta:revision') or '0') +if revision ~= expected_revision then + return 'STALE' +end + +local complete = ARGV[4] == '1' +local remote_ids = {} +local changed = false +for index = 6, #ARGV, 2 do + local sandbox_id = ARGV[index] + local token = ARGV[index + 1] + remote_ids[sandbox_id] = true + if token ~= '' and redis.call('HDEL', KEYS[1], 'r:' .. token) == 1 then + changed = true + end + if mark_present(sandbox_id) then + changed = true + end +end + +if complete then + local stale_before_ms = now_ms() - tonumber(ARGV[5]) + for _, field in ipairs(redis.call('HKEYS', KEYS[1])) do + local prefix = string.sub(field, 1, 2) + local identifier = string.sub(field, 3) + if prefix == 's:' and remote_ids[identifier] ~= true then + local missing_since_ms = tonumber(string.match(redis.call('HGET', KEYS[1], field) or '', '^m:(%d+)$')) + if missing_since_ms == nil then + redis.call('HSET', KEYS[1], field, 'm:' .. now_ms()) + changed = true + elseif missing_since_ms <= stale_before_ms then + redis.call('HDEL', KEYS[1], field) + changed = true + end + elseif prefix == 'r:' then + local created_ms = tonumber(redis.call('HGET', KEYS[1], field)) + if created_ms ~= nil and created_ms <= stale_before_ms then + redis.call('HDEL', KEYS[1], field) + changed = true + end + end + end + if state ~= 'ready' then + redis.call('HSET', KEYS[1], 'meta:state', 'ready') + changed = true + end +end + +if changed then + redis.call('HINCRBY', KEYS[1], 'meta:revision', 1) +end +return 'APPLIED' +""" + + +class CapacityBackendError(RuntimeError): + """Redis could not return a definitive capacity decision.""" + + +class ReserveStatus(enum.StrEnum): + GRANTED = "GRANTED" + FULL = "FULL" + NOT_READY = "NOT_READY" + + +def _text(value: object) -> str: + return value.decode() if isinstance(value, bytes) else str(value) + + +class RedisE2BCapacityStore: + """One capacity scope stored in one Redis Hash.""" + + def __init__( + self, + *, + redis_url: str, + hard_limit: int, + key_prefix: str = "deerflow:sandbox:owner", + ) -> None: + if hard_limit < 1: + raise ValueError("hard_limit must be at least 1") + try: + from redis import Redis + from redis.exceptions import RedisError + except ImportError: # pragma: no cover - optional extra + raise ImportError("Redis E2B capacity requires: cd backend && uv sync --extra redis") from None + + self._hard_limit = hard_limit + self._key = f"{key_prefix.rstrip(':')}:e2b-capacity" + self._redis_error = RedisError + self._redis = Redis.from_url( + redis_url, + decode_responses=True, + socket_timeout=_SOCKET_TIMEOUT_SECONDS, + socket_connect_timeout=_SOCKET_TIMEOUT_SECONDS, + ) + self._script = self._redis.register_script(_LEDGER_SCRIPT) + + @property + def key(self) -> str: + return self._key + + def _run(self, operation: str, *args: object) -> object: + try: + return self._script(keys=[self._key], args=[operation, self._hard_limit, *args]) + except self._redis_error as error: + raise CapacityBackendError(f"failed to {operation} E2B capacity in Redis: {error}") from error + + def revision(self) -> int: + return int(self._run("revision")) + + def reserve(self, token: str) -> ReserveStatus: + if not token: + raise ValueError("token must not be empty") + try: + return ReserveStatus(_text(self._run("reserve", token))) + except ValueError as error: + raise CapacityBackendError("unexpected E2B capacity reserve result") from error + + def track( + self, + sandbox_id: str, + *, + reservation_token: str | None = None, + ) -> None: + if not sandbox_id: + raise ValueError("sandbox_id must not be empty") + self._run("track", reservation_token or "", sandbox_id) + + def release(self, sandbox_id: str) -> None: + if sandbox_id: + self._run("release", sandbox_id) + + def reconcile( + self, + *, + expected_revision: int, + remote_sandboxes: dict[str, str | None], + complete: bool, + reservation_max_age_ms: int, + ) -> bool: + remote_args = [item for sandbox_id, token in remote_sandboxes.items() for item in (sandbox_id, token or "")] + status = _text( + self._run( + "reconcile", + expected_revision, + "1" if complete else "0", + reservation_max_age_ms, + *remote_args, + ) + ) + if status not in {"APPLIED", "STALE"}: + raise CapacityBackendError(f"unexpected E2B capacity reconciliation result: {status}") + return status == "APPLIED" + + def close(self) -> None: + try: + self._redis.close() + except Exception as error: # pragma: no cover - teardown best effort + logger.warning("Error closing E2B capacity Redis client: %s", error) + + +def make_e2b_capacity_store( + ownership: SandboxOwnershipConfig, + *, + hard_limit: int, +) -> RedisE2BCapacityStore | None: + """Enable the shared ledger only with Redis ownership.""" + if ownership.type == "memory": + return None + if ownership.type != "redis": + raise ValueError(f"Unknown sandbox ownership type: {ownership.type!r}") + logger.info("E2B deployment capacity: redis (key_prefix=%s, hard_limit=%d)", ownership.key_prefix, hard_limit) + return RedisE2BCapacityStore( + redis_url=resolve_ownership_redis_url(ownership), + hard_limit=hard_limit, + key_prefix=ownership.key_prefix, + ) 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 d4124b652..e0cfd4d91 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 @@ -15,6 +15,9 @@ provider fields during startup. 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) + ownership: + type: redis # shares ownership and capacity across Gateways + redis_url: redis://redis:6379/0 mounts: # one-shot uploads on sandbox start - host_path: /data/skills container_path: /home/user/skills @@ -44,6 +47,7 @@ from functools import partial from pathlib import Path from typing import Any +from e2b import SandboxQuery from e2b_code_interpreter import Sandbox as E2BClientSandbox from deerflow.config import get_app_config @@ -56,10 +60,16 @@ from ..aio_sandbox.ownership import ( OwnershipBackendError, RenewOutcome, SandboxOwnershipStore, + compute_lease_ttl, generate_owner_id, make_sandbox_ownership_store, resolve_ownership_config, ) +from .capacity import ( + CapacityBackendError, + ReserveStatus, + make_e2b_capacity_store, +) from .e2b_sandbox import DEFAULT_E2B_HOME_DIR, E2BSandbox, _is_sandbox_gone_error logger = logging.getLogger(__name__) @@ -77,6 +87,9 @@ DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS = 3600.0 DEFAULT_RECONCILIATION_MAX_PAGES = 10 DEFAULT_RECONCILIATION_MAX_ITEMS = 200 DEFAULT_RECONCILIATION_MAX_SECONDS = 15.0 +# Twice the E2B SDK's 60-second create request timeout. Short ownership lease +# settings must not make an in-flight create look abandoned. +MIN_CAPACITY_RESERVATION_SECONDS = 120.0 # Hard upper bound for ``set_timeout`` (e2b currently caps at 24h on the # free plan; passing an excessive value is rejected by the control-plane). MAX_E2B_TIMEOUT = 24 * 60 * 60 @@ -88,6 +101,8 @@ META_KEY_THREAD = "deer_flow_thread" META_KEY_PROVIDER = "deer_flow_provider" META_KEY_GATEWAY = "deer_flow_gateway" META_KEY_CREATED_AT = "deer_flow_created_at" +META_KEY_CAPACITY_LEDGER = "deer_flow_capacity_ledger" +META_KEY_CAPACITY_RESERVATION = "deer_flow_capacity_reservation" META_VAL_PROVIDER = "e2b_sandbox_provider" E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"}) @@ -168,6 +183,10 @@ class E2BSandboxProvider(SandboxProvider): self._ownership_config, owner_id=self._owner_id, ) + self._deployment_capacity = make_e2b_capacity_store( + self._ownership_config, + hard_limit=self._capacity_limit(), + ) if not self._ownership.supports_cross_process: logger.warning("E2B sandbox ownership is process-local. Multi-worker gateways must configure sandbox.ownership.type: redis for safe reconciliation.") @@ -285,6 +304,24 @@ class E2BSandboxProvider(SandboxProvider): def _stable_seed(thread_id: str, user_id: str) -> str: return hashlib.sha256(f"{user_id}:{thread_id}".encode()).hexdigest()[:16] + def _metadata_matches_capacity_ledger( + self, + metadata: dict[str, Any], + ) -> bool: + """Include this ledger and legacy sandboxes that predate the tag.""" + store = self._deployment_capacity + if store is None: + return True + remote_ledger = metadata.get(META_KEY_CAPACITY_LEDGER) + return remote_ledger in (None, "", store.key) + + @staticmethod + def _capacity_reservation_from_metadata( + metadata: dict[str, Any], + ) -> str | None: + token = metadata.get(META_KEY_CAPACITY_RESERVATION) + return token if isinstance(token, str) and token else None + # ── Signal / shutdown handling ─────────────────────────────────────── def _register_signal_handlers(self) -> None: @@ -390,6 +427,7 @@ class E2BSandboxProvider(SandboxProvider): sandbox.close() except Exception: pass + self._release_deployment_sandbox(sid) return None try: @@ -496,7 +534,7 @@ class E2BSandboxProvider(SandboxProvider): """ sandbox_cls = self._get_sandbox_cls() seed = self._stable_seed(thread_id, user_id) - entries, _ = self._list_remote_entries( + entries, _, _ = self._list_remote_entries( { META_KEY_PROVIDER: META_VAL_PROVIDER, META_KEY_USER: user_id, @@ -504,16 +542,21 @@ class E2BSandboxProvider(SandboxProvider): } ) candidates = sorted( - ((sandbox_id, metadata) for entry in entries if (sandbox_id := self._entry_id(entry)) and (metadata := self._entry_metadata(entry)).get(META_KEY_USER) == user_id and metadata.get(META_KEY_THREAD) == thread_id), + ( + (sandbox_id, metadata) + for entry in entries + if (sandbox_id := self._entry_id(entry)) and (metadata := self._entry_metadata(entry)).get(META_KEY_USER) == user_id and metadata.get(META_KEY_THREAD) == thread_id and self._metadata_matches_capacity_ledger(metadata) + ), key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]), ) - for target_id, _metadata in candidates: + for target_id, metadata in candidates: adopted = self._adopt_remote_candidate( sandbox_cls, target_id, thread_id=thread_id, user_id=user_id, seed=seed, + capacity_reservation=(self._capacity_reservation_from_metadata(metadata)), ) if adopted is not None: return adopted @@ -527,6 +570,7 @@ class E2BSandboxProvider(SandboxProvider): thread_id: str, user_id: str, seed: str, + capacity_reservation: str | None = None, ) -> str | None: """Try to adopt one discovered candidate without harming peer-owned VMs.""" @@ -548,6 +592,10 @@ class E2BSandboxProvider(SandboxProvider): return None try: + self._track_deployment_sandbox( + target_id, + reservation_token=capacity_reservation, + ) self._reserve_capacity( thread_id, user_id, @@ -664,6 +712,66 @@ class E2BSandboxProvider(SandboxProvider): with self._lock: self._end_transition_locked() + def _capacity_reservation_max_age_ms(self) -> int: + configured = compute_lease_ttl(self._ownership_config) + float(self._config["reconciliation_grace_seconds"]) + return int(max(configured, MIN_CAPACITY_RESERVATION_SECONDS) * 1_000) + + def _capacity_error( + self, + message: str, + *, + reason: str = "capacity", + ) -> SandboxCapacityExceededError: + with self._lock: + return SandboxCapacityExceededError( + message, + active=len(self._sandboxes), + warm=len(self._warm_pool), + reserved=self._reserved_slots, + replicas=int(self._config["replicas"]), + reason=reason, + ) + + def _track_deployment_sandbox( + self, + sandbox_id: str, + *, + reservation_token: str | None = None, + required: bool = True, + ) -> None: + store = self._deployment_capacity + if store is None: + return + try: + store.track( + sandbox_id, + reservation_token=reservation_token, + ) + except CapacityBackendError as error: + if required: + raise self._capacity_error( + f"Deployment-wide E2B capacity is unavailable; cannot safely track sandbox {sandbox_id}", + reason="capacity_backend", + ) from error + logger.warning( + "Could not track E2B sandbox %s in deployment capacity; reconciliation will retry: %s", + sandbox_id, + error, + ) + + def _release_deployment_sandbox(self, sandbox_id: str) -> None: + store = self._deployment_capacity + if store is None: + return + try: + store.release(sandbox_id) + except CapacityBackendError as error: + logger.warning( + "Could not release deployment capacity for destroyed E2B sandbox %s; reconciliation will retry: %s", + sandbox_id, + error, + ) + def _reserve_capacity( self, thread_id: str | None, @@ -671,23 +779,71 @@ class E2BSandboxProvider(SandboxProvider): *, 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. - """ + ) -> str | None: + """Reserve local capacity, then deployment capacity for a new VM.""" + store = self._deployment_capacity policy = self._config["overflow_policy"] timeout = float(self._config["acquire_timeout"]) deadline = time.monotonic() + timeout + token = uuid.uuid4().hex if store is not None and remote_id is None else None + + while True: + self._reserve_local_capacity( + thread_id, + user_id, + remote_id=remote_id, + remote_owned=remote_owned, + deadline=deadline, + ) + if store is None or remote_id is not None: + return token + + assert token is not None + backend_error = None + try: + status = store.reserve(token) + except CapacityBackendError as error: + backend_error = error + status = None + if status is ReserveStatus.GRANTED: + return token + self._release_capacity() + if status is ReserveStatus.FULL and self._evict_oldest_warm() is not None: + continue + + if policy != "wait": + if backend_error is not None: + raise self._capacity_error( + "Deployment-wide E2B capacity is unavailable", + reason="capacity_backend", + ) from backend_error + if status is ReserveStatus.NOT_READY: + raise self._capacity_error( + "Deployment-wide E2B capacity is initializing", + reason="capacity_initializing", + ) + raise self._capacity_error("Deployment-wide E2B capacity is full") + remaining = deadline - time.monotonic() + if remaining <= 0: + raise self._capacity_error(f"Timed out after {timeout}s waiting for deployment-wide E2B capacity") + with self._capacity_cond: + self._capacity_cond.wait(timeout=min(remaining, 1.0)) + + def _reserve_local_capacity( + self, + thread_id: str | None, + user_id: str, + *, + remote_id: str | None = None, + remote_owned: bool = True, + deadline: float | None = None, + ) -> None: + """Acquire the existing process-local lifecycle slot.""" + policy = self._config["overflow_policy"] + timeout = float(self._config["acquire_timeout"]) + deadline = deadline or time.monotonic() + timeout while True: - # Reject immediately if the provider is shutting down. with self._lock: if self._shutdown_called: raise SandboxCapacityExceededError( @@ -697,7 +853,6 @@ class E2BSandboxProvider(SandboxProvider): reason="shutdown", ) - # 1. Try immediate atomic reservation. with self._lock: if self._shutdown_called: raise SandboxCapacityExceededError( @@ -714,7 +869,6 @@ class E2BSandboxProvider(SandboxProvider): remote_ops.add(remote_id) return - # 2. Try evicting a warm entry to free a slot. evicted = self._evict_oldest_warm() if evicted is not None: with self._lock: @@ -731,9 +885,7 @@ class E2BSandboxProvider(SandboxProvider): 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( @@ -764,7 +916,6 @@ class E2BSandboxProvider(SandboxProvider): replicas=int(self._config["replicas"]), ) - # policy == "wait": block until a slot frees or timeout. remaining = deadline - time.monotonic() if remaining <= 0: raise SandboxCapacityExceededError( @@ -839,7 +990,7 @@ class E2BSandboxProvider(SandboxProvider): Capacity is enforced atomically via :meth:`_reserve_capacity`. """ - self._reserve_capacity(thread_id, user_id) + reservation_token = self._reserve_capacity(thread_id, user_id) sandbox_cls = self._get_sandbox_cls() metadata: dict[str, str] = { @@ -847,6 +998,10 @@ class E2BSandboxProvider(SandboxProvider): META_KEY_GATEWAY: self._owner_id, META_KEY_CREATED_AT: str(time.time()), } + if self._deployment_capacity is not None: + metadata[META_KEY_CAPACITY_LEDGER] = self._deployment_capacity.key + if reservation_token is not None: + metadata[META_KEY_CAPACITY_RESERVATION] = reservation_token if thread_id: metadata[META_KEY_USER] = user_id metadata[META_KEY_THREAD] = thread_id @@ -869,6 +1024,11 @@ class E2BSandboxProvider(SandboxProvider): raise sandbox_id: str = getattr(client, "sandbox_id", None) or str(uuid.uuid4())[:8] + self._track_deployment_sandbox( + sandbox_id, + reservation_token=reservation_token, + required=False, + ) if not self._track_reserved_remote_op(sandbox_id): kill_error = self._kill_client(client) cleanup_confirmed = kill_error is None @@ -998,26 +1158,25 @@ class E2BSandboxProvider(SandboxProvider): value = entry.get("metadata") return value if isinstance(value, dict) else {} - def _list_remote_entries(self, metadata: dict[str, str]) -> tuple[list[Any], bool]: - """List matching E2B entries within configured page/item/time budgets.""" + def _list_remote_entries( + self, + metadata: dict[str, str], + ) -> tuple[list[Any], bool, bool]: + """List E2B entries and report budget exhaustion and completeness.""" sandbox_cls = self._get_sandbox_cls() try: - result = sandbox_cls.list(query={"metadata": metadata}, **self._common_kwargs()) # type: ignore[attr-defined] - except TypeError: - try: - result = sandbox_cls.list(metadata=metadata, **self._common_kwargs()) # type: ignore[attr-defined] - except Exception as e: - logger.warning("E2B reconciliation list failed: %s", e) - return [], False + query = SandboxQuery(metadata=metadata) + result = sandbox_cls.list(query=query, **self._common_kwargs()) # type: ignore[attr-defined] except Exception as e: logger.warning("E2B reconciliation list failed: %s", e) - return [], False + return [], False, False max_pages = int(self._config["reconciliation_max_pages"]) max_items = int(self._config["reconciliation_max_items"]) deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"]) entries: list[Any] = [] exhausted = False + complete = True if hasattr(result, "next_items") and hasattr(result, "has_next"): for page_number in range(max_pages): @@ -1028,6 +1187,7 @@ class E2BSandboxProvider(SandboxProvider): page = result.next_items() except Exception as e: logger.warning("E2B reconciliation paginator failed: %s", e) + complete = False break if not page: break @@ -1045,13 +1205,15 @@ class E2BSandboxProvider(SandboxProvider): all_entries = list(result or []) except TypeError: logger.warning("E2B Sandbox.list returned non-iterable %s", type(result).__name__) - return [], False + return [], False, False entries = all_entries[:max_items] exhausted = len(all_entries) > max_items if time.monotonic() >= deadline: exhausted = True - return entries, exhausted + if exhausted: + complete = False + return entries, exhausted, complete def _publish_ownership(self, sandbox_id: str) -> None: """Publish acquire-side ownership before exposing a sandbox locally.""" @@ -1147,8 +1309,21 @@ class E2BSandboxProvider(SandboxProvider): self._lease_thread.start() self._reconcile_thread.start() - def _reserve_reconciliation_capacity(self, sandbox_id: str) -> bool: + def _reserve_reconciliation_capacity( + self, + sandbox_id: str, + *, + reservation_token: str | None = None, + ) -> bool: """Reserve one local slot for adoption without blocking maintenance.""" + try: + self._track_deployment_sandbox( + sandbox_id, + reservation_token=reservation_token, + ) + except SandboxCapacityExceededError as error: + logger.warning("Could not track discovered E2B sandbox %s: %s", sandbox_id, error) + return False with self._lock: if self._shutdown_called or self._total_capacity_used_locked() >= self._capacity_limit(): return False @@ -1162,8 +1337,38 @@ class E2BSandboxProvider(SandboxProvider): observed_at = time.monotonic() if now is None else now deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"]) stats = ReconciliationStats() - entries, stats.budget_exhausted = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER}) + capacity_revision = None + capacity_store = self._deployment_capacity + if capacity_store is not None: + try: + capacity_revision = capacity_store.revision() + except CapacityBackendError as error: + logger.warning( + "Could not read E2B capacity before reconciliation: %s", + error, + ) + + entries, stats.budget_exhausted, inventory_complete = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER}) + entries = [entry for entry in entries if self._metadata_matches_capacity_ledger(self._entry_metadata(entry))] stats.discovered = len(entries) + + if capacity_store is not None and capacity_revision is not None: + records = {sandbox_id: self._capacity_reservation_from_metadata(self._entry_metadata(entry)) for entry in entries if (sandbox_id := self._entry_id(entry))} + try: + applied = capacity_store.reconcile( + expected_revision=capacity_revision, + remote_sandboxes=records, + complete=inventory_complete, + reservation_max_age_ms=self._capacity_reservation_max_age_ms(), + ) + if not applied: + logger.debug("E2B capacity inventory became stale during reconciliation; retrying on the next pass") + except CapacityBackendError as error: + logger.warning( + "Could not apply E2B capacity inventory: %s", + error, + ) + groups: dict[tuple[str, str], list[tuple[str, dict[str, Any]]]] = {} orphans: list[tuple[str, dict[str, Any]]] = [] present_ids: set[str] = set() @@ -1206,12 +1411,15 @@ class E2BSandboxProvider(SandboxProvider): if not live: continue stats.duplicates += max(0, len(live) - 1) - canonical_id, _metadata, canonical_client = live[0] + canonical_id, canonical_metadata, canonical_client = live[0] with self._lock: already_local = canonical_id in self._sandboxes if already_local: self._safe_close_client(canonical_client) - elif not self._reserve_reconciliation_capacity(canonical_id): + elif not self._reserve_reconciliation_capacity( + canonical_id, + reservation_token=self._capacity_reservation_from_metadata(canonical_metadata), + ): self._safe_close_client(canonical_client) stats.deferred += 1 elif not self._claim_ownership(canonical_id): @@ -1332,10 +1540,16 @@ class E2BSandboxProvider(SandboxProvider): reaped the VM. Closing that host-side client before returning ``None`` keeps both acquire paths from leaking a connection. """ - client = self._reconnect_client(sandbox_cls, sandbox_id) + try: + client = self._reconnect_client(sandbox_cls, sandbox_id) + except Exception as error: + if _is_sandbox_gone_error(error): + self._release_deployment_sandbox(sandbox_id) + raise if self._client_alive(client): return client self._safe_close_client(client) + self._release_deployment_sandbox(sandbox_id) return None def _register_connected_sandbox( @@ -2073,13 +2287,14 @@ class E2BSandboxProvider(SandboxProvider): except Exception: pass - @staticmethod def _kill_client( + self, client: E2BClientSandbox | None, ) -> Exception | None: """Kill a remote VM and return an exception for the caller to log.""" if client is None: return RuntimeError("Cannot confirm remote VM destruction without a client") + sandbox_id = getattr(client, "sandbox_id", None) try: kill = getattr(client, "kill", None) if not callable(kill): @@ -2087,6 +2302,8 @@ class E2BSandboxProvider(SandboxProvider): kill() except Exception as e: return e + if sandbox_id is not None: + self._release_deployment_sandbox(sandbox_id) return None def reset(self) -> None: @@ -2178,3 +2395,11 @@ class E2BSandboxProvider(SandboxProvider): self._ownership.close() except Exception as e: logger.warning("Failed to close E2B ownership store: %s", e) + if self._deployment_capacity is not None: + try: + self._deployment_capacity.close() + except Exception as e: + logger.warning( + "Failed to close E2B deployment capacity store: %s", + e, + ) diff --git a/backend/packages/harness/deerflow/config/sandbox_config.py b/backend/packages/harness/deerflow/config/sandbox_config.py index 732c47621..f6c8f10f8 100644 --- a/backend/packages/harness/deerflow/config/sandbox_config.py +++ b/backend/packages/harness/deerflow/config/sandbox_config.py @@ -76,8 +76,9 @@ class SandboxConfig(BaseModel): AioSandboxProvider, BoxliteProvider, and E2BSandboxProvider shared options: image: Sandbox image to use (Docker/AIO image or BoxLite OCI image) - replicas: Positive provider capacity per gateway process. Each provider - defines which lifecycle states count toward this limit. + replicas: Positive provider capacity. E2B shares it across Gateway + workers when ownership uses Redis; other modes/providers keep + process-local accounting. 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) @@ -115,7 +116,7 @@ class SandboxConfig(BaseModel): replicas: int | None = Field( default=None, gt=0, - description="Positive provider capacity per gateway process. Each provider defines which lifecycle states count toward this limit.", + description=("Positive provider capacity. E2B enforces it deployment-wide when sandbox ownership uses Redis; otherwise accounting is per Gateway process. Each provider defines which lifecycle states count."), ) overflow_policy: SandboxOverflowPolicy = Field( default="wait", diff --git a/backend/tests/test_e2b_capacity_store_redis.py b/backend/tests/test_e2b_capacity_store_redis.py new file mode 100644 index 000000000..65e59ae06 --- /dev/null +++ b/backend/tests/test_e2b_capacity_store_redis.py @@ -0,0 +1,178 @@ +"""Integration tests for deployment-wide E2B admission.""" + +from __future__ import annotations + +import os +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from deerflow.community.e2b_sandbox.capacity import ( + CapacityBackendError, + RedisE2BCapacityStore, + ReserveStatus, + make_e2b_capacity_store, +) +from deerflow.config.sandbox_config import SandboxOwnershipConfig + +REDIS_URL = os.environ.get("DEER_FLOW_TEST_REDIS_URL", "redis://localhost:6379/15") +pytestmark = pytest.mark.integration + + +@pytest.fixture +def make_store(): + redis = pytest.importorskip("redis") + probe = redis.Redis.from_url(REDIS_URL, decode_responses=True, socket_connect_timeout=0.5) + try: + probe.ping() + except Exception: + probe.close() + pytest.skip(f"Redis not reachable at {REDIS_URL}") + + prefix = f"deerflow:test:{uuid.uuid4().hex}" + stores = [] + + def make(hard_limit: int = 1): + store = RedisE2BCapacityStore( + redis_url=REDIS_URL, + hard_limit=hard_limit, + key_prefix=prefix, + ) + stores.append(store) + return store + + try: + yield make + finally: + probe.delete(f"{prefix}:e2b-capacity") + probe.close() + for store in stores: + store.close() + + +def _initialize(store) -> None: + assert store.reconcile( + expected_revision=store.revision(), + remote_sandboxes={}, + complete=True, + reservation_max_age_ms=0, + ) + + +def _counts(store) -> tuple[int, int]: + fields = store._redis.hkeys(store.key) + return ( + sum(field.startswith("s:") for field in fields), + sum(field.startswith("r:") for field in fields), + ) + + +def test_factory_is_lazy_and_backend_errors_fail_closed() -> None: + assert make_e2b_capacity_store(SandboxOwnershipConfig(type="memory"), hard_limit=3) is None + store = make_e2b_capacity_store( + SandboxOwnershipConfig(type="redis", redis_url="redis://127.0.0.1:1/0", key_prefix="test"), + hard_limit=3, + ) + assert store is not None and store.key == "test:e2b-capacity" + try: + with pytest.raises(CapacityBackendError): + store.reserve("reservation") + finally: + store.close() + + +def test_two_gateways_atomically_share_one_hash(make_store) -> None: + gateway_a, gateway_b = make_store(), make_store() + assert gateway_a.reserve("not-ready") is ReserveStatus.NOT_READY + _initialize(gateway_a) + barrier = threading.Barrier(2) + + def reserve(args): + store, token = args + barrier.wait() + return store.reserve(token) + + tokens = ["reservation-a", "reservation-b"] + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(reserve, zip((gateway_a, gateway_b), tokens))) + + assert results.count(ReserveStatus.GRANTED) == 1 + assert results.count(ReserveStatus.FULL) == 1 + assert list(gateway_a._redis.scan_iter(f"{gateway_a.key}")) == [gateway_a.key] + winner = tokens[results.index(ReserveStatus.GRANTED)] + gateway_a.track("sandbox-a", reservation_token=winner) + gateway_a.track("sandbox-a", reservation_token=winner) + assert _counts(gateway_a) == (1, 0) + # A successful but stale list must not release a just-tracked slot. + assert gateway_b.reconcile( + expected_revision=gateway_b.revision(), + remote_sandboxes={}, + complete=True, + reservation_max_age_ms=120_000, + ) + assert gateway_b.reserve("stale-inventory") is ReserveStatus.FULL + gateway_b.release("sandbox-a") + gateway_b.release("sandbox-a") + assert _counts(gateway_a) == (0, 0) + + +def test_reconcile_repairs_crashes_without_erasing_concurrent_changes(make_store) -> None: + gateway_a, gateway_b = make_store(2), make_store(2) + _initialize(gateway_a) + assert gateway_a.reserve("crashed-create") is ReserveStatus.GRANTED + assert gateway_b.reconcile( + expected_revision=gateway_b.revision(), + remote_sandboxes={"sandbox-a": "crashed-create"}, + complete=True, + reservation_max_age_ms=0, + ) + stale_revision = gateway_a.revision() + assert gateway_b.reserve("concurrent") is ReserveStatus.GRANTED + assert not gateway_a.reconcile( + expected_revision=stale_revision, + remote_sandboxes={}, + complete=True, + reservation_max_age_ms=0, + ) + assert _counts(gateway_a) == (1, 1) + + +def test_reconcile_keeps_incomplete_inventory_and_fresh_reservations(make_store) -> None: + store = make_store(2) + _initialize(store) + store.track("sandbox-a") + assert store.reserve("creating") is ReserveStatus.GRANTED + + assert store.reconcile( + expected_revision=store.revision(), + remote_sandboxes={}, + complete=False, + reservation_max_age_ms=0, + ) + assert _counts(store) == (1, 1) + assert store.reconcile( + expected_revision=store.revision(), + remote_sandboxes={"sandbox-a": None}, + complete=True, + reservation_max_age_ms=120_000, + ) + assert _counts(store) == (1, 1) + store._redis.hset(store.key, "s:sandbox-a", "m:0") + assert store.reconcile( + expected_revision=store.revision(), + remote_sandboxes={}, + complete=True, + reservation_max_age_ms=0, + ) + assert _counts(store) == (0, 0) + + +def test_mismatched_hard_limits_fail_closed(make_store) -> None: + gateway_a, gateway_b = make_store(), make_store(2) + _initialize(gateway_a) + with pytest.raises(CapacityBackendError, match="configuration mismatch"): + gateway_b.revision() + with pytest.raises(CapacityBackendError, match="configuration mismatch"): + gateway_b.reserve("reservation") diff --git a/backend/tests/test_e2b_sandbox_provider.py b/backend/tests/test_e2b_sandbox_provider.py index cd3aa865f..e49a5f435 100644 --- a/backend/tests/test_e2b_sandbox_provider.py +++ b/backend/tests/test_e2b_sandbox_provider.py @@ -17,6 +17,10 @@ from unittest.mock import MagicMock import pytest from pydantic import ValidationError +from deerflow.community.e2b_sandbox.capacity import ( + CapacityBackendError, + ReserveStatus, +) from deerflow.config.paths import Paths from deerflow.config.sandbox_config import SandboxConfig from deerflow.sandbox.exceptions import SandboxCapacityExceededError @@ -253,7 +257,12 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli provider._shutdown_called = False provider._owner_id = "owner-a" provider._ownership = FakeOwnershipStore({}, owner_id=provider._owner_id) - provider._ownership_config = SimpleNamespace(renewal_interval_seconds=60.0) + provider._ownership_config = SimpleNamespace( + renewal_interval_seconds=60.0, + ttl_multiplier=4.0, + key_prefix="deerflow:test", + ) + provider._deployment_capacity = None provider._owned_sandbox_ids = set() provider._acquire_inflight = set() provider._orphan_first_seen = {} @@ -282,6 +291,22 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli return provider +def _install_shared_deployment_capacity( + *providers, + reserve_results: list[ReserveStatus] | None = None, +) -> MagicMock: + store = MagicMock() + store.key = "deerflow:test:e2b-capacity" + store.revision.return_value = 0 + store.reserve.return_value = ReserveStatus.GRANTED + store.reconcile.return_value = True + if reserve_results is not None: + store.reserve.side_effect = reserve_results + for provider in providers: + provider._deployment_capacity = store + return store + + def _install_fake_sdk(monkeypatch, provider) -> FakeSandboxClass: fake_cls = FakeSandboxClass() monkeypatch.setattr(provider, "_get_sandbox_cls", lambda: fake_cls) @@ -1081,11 +1106,17 @@ def test_bootstrap_failure_does_not_kill_without_destroy_lease(monkeypatch): def test_kill_client_returns_exception_without_raising(): p = _make_provider() - client = FakeClient() + store = _install_shared_deployment_capacity(p) + failed_client = FakeClient() error = RuntimeError("already gone") - client.kill = MagicMock(side_effect=error) + failed_client.kill = MagicMock(side_effect=error) - assert p._kill_client(client) is error + assert p._kill_client(failed_client) is error + store.release.assert_not_called() + + client = FakeClient() + assert p._kill_client(client) is None + store.release.assert_called_once_with(client.sandbox_id) def test_kill_client_reports_uncertain_cleanup_without_callable_kill(): @@ -2107,9 +2138,132 @@ def test_grep_single_file_path_with_matching_glob(): assert truncated is False -# ────────────────────────────────────────────────────────────────────────────── # Capacity enforcement tests (#4339) -# ────────────────────────────────────────────────────────────────────────────── + + +def test_deployment_capacity_reserves_commits_and_rejects_globally(monkeypatch) -> None: + gateway_a = _make_provider(replicas=1, overflow_policy="reject") + gateway_b = _make_provider(replicas=1, overflow_policy="reject") + store = _install_shared_deployment_capacity( + gateway_a, + gateway_b, + reserve_results=[ReserveStatus.GRANTED, ReserveStatus.FULL], + ) + sdk_a = _install_fake_sdk(monkeypatch, gateway_a) + sdk_b = FakeSandboxClass() + monkeypatch.setattr(gateway_b, "_get_sandbox_cls", lambda: sdk_b) + + sandbox_id = gateway_a.acquire("thread-a", user_id="user-a") + with pytest.raises(SandboxCapacityExceededError): + gateway_b.acquire("thread-b", user_id="user-b") + + metadata = sdk_a.create_calls[0]["metadata"] + assert metadata["deer_flow_capacity_ledger"] == store.key + assert metadata["deer_flow_capacity_reservation"] + store.track.assert_called_once_with( + sandbox_id, + reservation_token=metadata["deer_flow_capacity_reservation"], + ) + assert len(sdk_a.create_calls) == 1 + assert sdk_b.create_calls == [] + assert store.reserve.call_count == 2 + + +def test_ambiguous_create_failure_retains_deployment_reservation(monkeypatch) -> None: + provider = _make_provider(replicas=1, overflow_policy="reject") + store = _install_shared_deployment_capacity(provider) + sdk = _install_fake_sdk(monkeypatch, provider) + sdk.create_factory = lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("control-plane timeout")) + + with pytest.raises(RuntimeError, match="control-plane timeout"): + provider.acquire("thread-a", user_id="user-a") + + assert provider._reserved_slots == 0 + store.reserve.assert_called_once() + store.track.assert_not_called() + store.release.assert_not_called() + + +def test_discovery_uses_sdk_query_and_tracks_without_reserving(monkeypatch) -> None: + provider = _make_provider(replicas=1, overflow_policy="reject") + store = _install_shared_deployment_capacity(provider) + sdk = _install_fake_sdk(monkeypatch, provider) + entry = SimpleNamespace( + sandbox_id="sandbox-existing", + metadata={ + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_user": "user-a", + "deer_flow_thread": "thread-a", + "deer_flow_capacity_ledger": store.key, + }, + ) + expected_query = {key: entry.metadata[key] for key in ("deer_flow_provider", "deer_flow_user", "deer_flow_thread")} + sdk.list_return = SimpleNamespace( + has_next=False, + next_items=lambda: [entry] if sdk.list_calls[-1]["query"].metadata == expected_query else [], + ) + assert provider.acquire("thread-a", user_id="user-a") == entry.sandbox_id + assert sdk.create_calls == [] + store.reserve.assert_not_called() + store.track.assert_called_once_with(entry.sandbox_id, reservation_token=None) + + +def test_reconciliation_repairs_crash_and_uses_safe_reservation_age(monkeypatch) -> None: + provider = _make_provider(replicas=1, overflow_policy="reject") + provider._ownership_config.renewal_interval_seconds = 1.0 + provider._ownership_config.ttl_multiplier = 2.0 + provider._config["reconciliation_grace_seconds"] = 0.0 + store = _install_shared_deployment_capacity(provider) + sdk = _install_fake_sdk(monkeypatch, provider) + sdk.list_return = [ + { + "sandbox_id": "sandbox-existing", + "metadata": { + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_capacity_ledger": store.key, + "deer_flow_capacity_reservation": "reservation-crashed", + }, + }, + { + "sandbox_id": "sandbox-other-deployment", + "metadata": { + "deer_flow_provider": "e2b_sandbox_provider", + "deer_flow_capacity_ledger": "deerflow:other:e2b-capacity", + }, + }, + ] + + stats = provider._reconcile_remote_sandboxes(now=100.0) + args = store.reconcile.call_args.kwargs + assert stats.discovered == 1 + assert args["remote_sandboxes"] == {"sandbox-existing": "reservation-crashed"} + assert args["complete"] is True + assert args["reservation_max_age_ms"] == 120_000 + + +def test_failed_inventory_and_redis_error_both_prevent_create(monkeypatch) -> None: + provider = _make_provider(replicas=1, overflow_policy="reject") + store = _install_shared_deployment_capacity( + provider, + reserve_results=[ReserveStatus.NOT_READY], + ) + sdk = _install_fake_sdk(monkeypatch, provider) + sdk.list = MagicMock(side_effect=RuntimeError("E2B unavailable")) + + provider._reconcile_remote_sandboxes(now=100.0) + + reconcile_args = store.reconcile.call_args.kwargs + assert reconcile_args["complete"] is False + assert reconcile_args["remote_sandboxes"] == {} + with pytest.raises(SandboxCapacityExceededError): + provider._create_sandbox("thread-a", user_id="user-a") + store.reserve.side_effect = CapacityBackendError("Redis unavailable") + with pytest.raises(SandboxCapacityExceededError) as error: + provider._create_sandbox("thread-a", user_id="user-a") + + assert error.value.reason == "capacity_backend" + assert sdk.create_calls == [] + assert provider._reserved_slots == 0 def test_capacity_reject_policy_raises_when_full(monkeypatch): @@ -2146,7 +2300,11 @@ def test_capacity_reject_frees_slot_on_release(monkeypatch): 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") + p = _make_provider(replicas=3, overflow_policy="reject") + store = _install_shared_deployment_capacity( + p, + reserve_results=[ReserveStatus.GRANTED, ReserveStatus.FULL, ReserveStatus.GRANTED], + ) fake_cls = _install_fake_sdk(monkeypatch, p) sid1 = p.acquire("t1", user_id="u1") @@ -2158,6 +2316,7 @@ def test_capacity_reject_evicts_other_thread_warm_entry_before_create(monkeypatc assert sid2 != sid1 assert len(p._warm_pool) == 0 assert len(fake_cls.create_calls) == 2 + store.release.assert_called_once_with(sid1) def test_capacity_wait_policy_times_out(monkeypatch): @@ -3330,8 +3489,14 @@ def test_shutdown_during_discovery_does_not_kill_unowned_vm(monkeypatch): 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( + def pause_after_reserve( + thread_id, + user_id, + *, + remote_id=None, + remote_owned=True, + ): + reservation = reserve_capacity( thread_id, user_id, remote_id=remote_id, @@ -3339,6 +3504,7 @@ def test_shutdown_during_discovery_does_not_kill_unowned_vm(monkeypatch): ) reserved.set() assert allow_commit.wait(timeout=2) + return reservation monkeypatch.setattr(p, "_reserve_capacity", pause_after_reserve) result: list[str | None] = []