diff --git a/apps/api/alerting.py b/apps/api/alerting.py index 48ae207f3..a9e60dd0b 100644 --- a/apps/api/alerting.py +++ b/apps/api/alerting.py @@ -102,6 +102,7 @@ "run_claimed_without_dispatch", "run_execution_exception", "sandbox_crash", + "sandbox_liveness_unconfirmed", "sandbox_transport_retry_exhausted", "schedule_missed", "scheduler_missed", diff --git a/apps/api/run_service.py b/apps/api/run_service.py index d9a45bde1..f50987dd4 100644 --- a/apps/api/run_service.py +++ b/apps/api/run_service.py @@ -601,6 +601,7 @@ def _schedule_retry( "output_token_limit", "output_too_large", "quality_gate_failed", + "sandbox_liveness_unconfirmed", "schema_violation", "spend_cap_exceeded", "token_cap_exceeded", @@ -634,6 +635,15 @@ def _schedule_retry( "transient_network_retry_exhausted", } +# Safety terminal codes that must NEVER be auto-retried, not even when a worker +# manifest names them in retry.on. sandbox_liveness_unconfirmed means we could +# not confirm the original sandbox stopped after a transport drop; re-running +# would risk duplicating real side effects (emails, CRM writes, sends), so a +# manifest opt-in must not be able to force it. +_NEVER_RETRY_ERROR_CODES = { + "sandbox_liveness_unconfirmed", +} + _PERMANENT_RETRY_CATEGORIES = { "auth", "cancelled", @@ -695,6 +705,10 @@ def _classify_retry_failure( from services import run_metrics category = run_metrics.classify_failure(error_code=code, error=error) + if code in _NEVER_RETRY_ERROR_CODES: + # Hard safety veto: overrides manifest retry.on and retryable=True so a + # liveness-unconfirmed run can never be re-dispatched. + return _RetryDecision(False, True, category, "never_retry_safety") if code in _RETRY_EXHAUSTED_ERROR_CODES: return _RetryDecision(False, True, category, "retry_exhausted") if _retry_config_exactly_allows_error(retry_cfg, code): @@ -2444,6 +2458,12 @@ def update_run_status( _RUN_REAPER_DEFAULT_GRACE_SECONDS = 60 _RUN_REAPER_DEFAULT_INTERVAL_SECONDS = 180 _RESTART_RETRY_BACKOFF_SECONDS = 60 +# Per-sandbox control-plane kill timeout used inside the cancel loop. The loop +# runs serially over every active run, so the driver default (60s) could make a +# single hung kill blow the whole shutdown/cancel budget. Bound each kill and +# stop issuing kills once the overall cancel budget is spent (leftover sandboxes +# are handled by thread-join + startup recovery). +_CANCEL_LOOP_SANDBOX_KILL_REQUEST_TIMEOUT_SECONDS = 10.0 @dataclass @@ -2454,6 +2474,19 @@ class _ActiveRun: thread: threading.Thread started_monotonic: float = field(default_factory=time.monotonic) stage: str = "claimed" + # Set once the sandbox worker command (run.py) has started for this run. Once + # True, the run must NOT be requeued/re-dispatched on shutdown: re-running a + # started worker could duplicate side effects (emails, CRM writes, sends). + worker_command_started: bool = False + + +def mark_run_worker_command_started(run_id: str) -> None: + """Called by the sandbox driver when the worker command starts, so the + graceful-shutdown requeue can skip runs whose worker already began.""" + with _active_runs_lock: + active = _active_runs.get(run_id) + if active is not None: + active.worker_command_started = True _active_runs: dict[str, _ActiveRun] = {} @@ -3305,6 +3338,12 @@ def _cancel_active_runs( cancel_sandbox = None cancelled_at = _now_iso() + # Single shared deadline for BOTH the kill pass and the join pass so the + # whole cancel operation stays within timeout_seconds (previously the kill + # pass and join pass each got a full budget, so total could reach ~2x, and a + # hung control-plane kill could blow it). Stragglers are handled by startup + # recovery. + overall_deadline = time.monotonic() + max(0.0, timeout_seconds) for run in active: if run.user_id: try: @@ -3324,14 +3363,28 @@ def _cancel_active_runs( except Exception as exc: logger.warning("Failed to mark run %s cancelled: %s", run.run_id, exc) if cancel_sandbox is not None: - try: - cancel_sandbox(run.run_id, reason=reason) - except Exception: - logger.debug("E2B cancel failed for run %s", run.run_id, exc_info=True) + remaining = overall_deadline - time.monotonic() + if remaining <= 0: + logger.warning( + "Cancel loop kill budget exhausted; leaving sandbox for run %s to " + "thread-join and startup recovery", + run.run_id, + ) + else: + try: + cancel_sandbox( + run.run_id, + reason=reason, + request_timeout=min( + _CANCEL_LOOP_SANDBOX_KILL_REQUEST_TIMEOUT_SECONDS, + remaining, + ), + ) + except Exception: + logger.debug("E2B cancel failed for run %s", run.run_id, exc_info=True) - deadline = time.monotonic() + max(0.0, timeout_seconds) for run in active: - remaining = deadline - time.monotonic() + remaining = overall_deadline - time.monotonic() if remaining <= 0: break run.thread.join(timeout=remaining) @@ -3438,6 +3491,18 @@ def request_active_run_shutdown( for run in active: if run.run_id in remaining_ids: continue + # Double-execute guard: never requeue a run whose sandbox worker command + # had already started. Re-running it would re-execute committed side + # effects (emails, CRM writes, sends). Leave it cancelled (from the + # cancel pass) for operator action instead. Pre-command runs (setup / + # queued-but-not-yet-executing) are safe to requeue. + if run.worker_command_started: + logger.warning( + "Not requeuing run %s on shutdown: its worker command had started; " + "re-running could duplicate side effects", + run.run_id, + ) + continue if _requeue_interrupted_run_in_place(repos_obj, run.run_id, run.user_id): requeued += 1 if requeued: diff --git a/apps/api/runner_sandbox/e2b_driver.py b/apps/api/runner_sandbox/e2b_driver.py index 34955fb4a..889e6e8da 100644 --- a/apps/api/runner_sandbox/e2b_driver.py +++ b/apps/api/runner_sandbox/e2b_driver.py @@ -12,6 +12,7 @@ import re import shlex import shutil +import sys import tarfile import threading import time @@ -65,6 +66,19 @@ E2B_INSTALL_COMMAND_MIN_REQUEST_TIMEOUT_SECONDS = 300 E2B_CREATE_REQUEST_TIMEOUT_SECONDS = 120 E2B_CONTROL_REQUEST_TIMEOUT_SECONDS = 60 +# Bounded kill-and-verify budget used before a transport-drop retry. A dropped +# worker-command transport must NOT be retried until the ORIGINAL sandbox is +# CONFIRMED terminated, otherwise the original keeps running (real side effects: +# emails, CRM writes, sends) while the retry runs a second time -> duplicate. +# These bound the confirmation so we never block on the unbounded 60s control +# timeout over an already-dead transport. +E2B_KILL_VERIFY_REQUEST_TIMEOUT_SECONDS = 15 +E2B_KILL_VERIFY_BUDGET_SECONDS = 30.0 +E2B_KILL_VERIFY_POLL_INTERVAL_SECONDS = 1.0 +# Terminal, operator-action error code: a transport dropped AFTER the worker +# command started and the original sandbox could not be confirmed dead. We +# refuse to re-run rather than risk double-executing side effects. +SANDBOX_LIVENESS_UNCONFIRMED_ERROR_CODE = "sandbox_liveness_unconfirmed" MAX_WORKER_ERROR_OUTPUT_CHARS = 1000 _OOM_EXIT_CODES = {137, -9} _OOM_MARKERS = ( @@ -133,6 +147,81 @@ def _kill_sandbox_quietly(sandbox: Any) -> None: logger.debug("E2B sandbox kill suppressed", exc_info=True) +def _e2b_kill_verify_budget_seconds() -> float: + raw = os.environ.get("WORKEROS_E2B_KILL_VERIFY_BUDGET_SECONDS", "") + if not raw: + return E2B_KILL_VERIFY_BUDGET_SECONDS + try: + return max(0.0, float(raw)) + except ValueError: + return E2B_KILL_VERIFY_BUDGET_SECONDS + + +def _e2b_kill_verify_request_timeout_seconds() -> float: + raw = os.environ.get("WORKEROS_E2B_KILL_VERIFY_REQUEST_TIMEOUT_SECONDS", "") + if not raw: + return float(E2B_KILL_VERIFY_REQUEST_TIMEOUT_SECONDS) + try: + return max(0.1, float(raw)) + except ValueError: + return float(E2B_KILL_VERIFY_REQUEST_TIMEOUT_SECONDS) + + +def _sandbox_confirmed_not_running(sandbox: Any, *, request_timeout: float) -> bool: + """Positive liveness probe: return True ONLY when the sandbox is provably + not running. e2b ``is_running()`` hits the sandbox's own envd health route + and returns False when it is gone (502) or raises on any ambiguity + (timeout / connection error). Any exception -> False, so an unreachable + sandbox is never mistaken for a dead one.""" + try: + return sandbox.is_running(request_timeout=request_timeout) is False + except Exception: + logger.debug("E2B is_running probe inconclusive", exc_info=True) + return False + + +def _terminate_and_confirm_dead( + sandbox: Any, *, log_fn: Callable[[str, str], None] +) -> bool: + """Kill the sandbox and return True ONLY when termination is CONFIRMED. + + e2b ``kill()`` is a DELETE against the control plane (api.e2b.dev), which + travels a DIFFERENT transport than the worker-command envd stream, so it + usually succeeds even when the command stream dropped. A ``kill()`` that + returns (killed, or 404 already-gone) WITHOUT raising is authoritative proof + the sandbox is terminated. If ``kill()`` keeps raising (control plane also + unreachable) we fall back to polling ``is_running()`` for a bounded budget. + + Returns False when termination cannot be POSITIVELY confirmed within the + budget. The caller MUST NOT re-dispatch on a False result: a still-running + original would double-execute the worker's side effects. The budget is + bounded (never the unbounded 60s control timeout on a dead transport).""" + budget = _e2b_kill_verify_budget_seconds() + request_timeout = _e2b_kill_verify_request_timeout_seconds() + deadline = time.monotonic() + budget + attempted = False + while True: + if attempted and time.monotonic() >= deadline: + return False + attempted = True + try: + # Control-plane terminate. Returns True (killed) or False (404, + # already gone); either is confirmed-dead. Only a raise is ambiguous. + sandbox.kill(request_timeout=request_timeout) + return True + except Exception: + logger.debug( + "E2B bounded kill did not confirm termination; probing liveness", + exc_info=True, + ) + if _sandbox_confirmed_not_running(sandbox, request_timeout=request_timeout): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(E2B_KILL_VERIFY_POLL_INTERVAL_SECONDS, remaining)) + + def _warm_pool_lease(key: str, *, log_fn: Callable[[str, str], None]) -> _WarmSandboxEntry | None: if not _warm_pool_enabled() or not key: return None @@ -1532,13 +1621,18 @@ def active_sandbox_count() -> int: return len(_active_sandboxes) -def cancel_sandbox(run_id: str, *, reason: str | None = None) -> bool: +def cancel_sandbox( + run_id: str, + *, + reason: str | None = None, + request_timeout: float = E2B_CONTROL_REQUEST_TIMEOUT_SECONDS, +) -> bool: with _active_sandboxes_lock: sandbox = _active_sandboxes.get(run_id) if sandbox is None: return False try: - sandbox.kill(request_timeout=E2B_CONTROL_REQUEST_TIMEOUT_SECONDS) + sandbox.kill(request_timeout=request_timeout) logger.warning("Killed active E2B sandbox for run %s: %s", run_id, reason or "cancel requested") return True except Exception as exc: @@ -2181,6 +2275,56 @@ def run( error="Cancelled by user", error_code="user_cancel", ) + # Double-execute guard: only retry when it is SAFE. A retry + # re-runs the worker (run.py) from scratch in a fresh sandbox. + # Once the worker command has STARTED, its side effects + # (emails, CRM writes, sends) may already be committed: + # partially if it dropped mid-run, fully if it completed. The + # bounded kill-and-verify below stops the ORIGINAL sandbox, but + # killing cannot UNDO an action that already committed, and the + # platform does not yet enforce action-level idempotency at the + # connector boundary. So a post-command-start retry could + # DUPLICATE a committed action even when the original is + # confirmed dead. We therefore refuse to auto-retry any run + # whose command started, and surface an operator-action + # terminal instead. Retry stays enabled ONLY for create/upload + # phase drops (command never started -> no worker code ran -> + # no side effects possible), which covers the common + # sandbox-create / HPACK / header-block transport drops. + # (Re-enabling safe post-start retry requires durable + # run_id+operation idempotency keys threaded into connector + # calls; tracked as the documented residual.) + command_started = bool(getattr(exc, "_e2b_worker_command_started", False)) + command_completed = bool(getattr(exc, "_e2b_worker_command_completed", False)) + sandbox_confirmed_dead = bool(getattr(exc, "_e2b_sandbox_confirmed_dead", False)) + if command_started: + logger.error( + "E2B transport dropped after the worker command %s for worker %s run %s; " + "refusing to re-run to avoid duplicate side effects " + "(command_completed=%s, sandbox_confirmed_dead=%s)", + "completed" if command_completed else "started", + worker_id, + run_id, + command_completed, + sandbox_confirmed_dead, + ) + log_fn( + "[e2b] Sandbox transport dropped after the worker started; NOT retrying " + "to avoid duplicating actions (emails, CRM writes, sends). Operator " + "action required.", + "error", + ) + return WorkerResult( + status="error", + error=( + "The sandbox connection dropped after the worker started. Floom did " + "not re-run it to avoid duplicating side effects (emails, CRM writes, " + "sends). Check whether the run's actions completed, then re-run " + "manually if needed." + ), + error_code=SANDBOX_LIVENESS_UNCONFIRMED_ERROR_CODE, + retryable=False, + ) detail = ( _sanitize_sandbox_exception_detail(str(transport_exc)) or transport_exc.__class__.__name__ @@ -2257,6 +2401,38 @@ def run( error_code="sandbox_oom", retryable=False, ) + # Double-execute guard (outer handler): an exception that is NOT a + # recognized transient transport drop still reaches here after the + # worker command started (e.g. an opaque error during result + # processing). _sandbox_exception_result would classify it as a + # RETRYABLE code (e2b_sandbox_error / upstream_http_5xx), and the + # run_service retry scheduler would then re-dispatch the whole run in + # a fresh sandbox, re-running the worker -> duplicate side effects. If + # the command started, force a non-retryable safety terminal so no + # layer re-runs it. + if getattr(exc, "_e2b_worker_command_started", False): + logger.error( + "E2B sandbox failed for worker %s run %s AFTER the worker command started; " + "surfacing a non-retryable terminal to avoid duplicate side effects: %s", + worker_id, + run_id, + exc, + ) + log_fn( + "[e2b] Sandbox failed after the worker started; NOT retrying to avoid " + "duplicating actions. Operator action required.", + "error", + ) + return WorkerResult( + status="error", + error=( + "The sandbox failed after the worker started. Floom did not re-run it to " + "avoid duplicating side effects (emails, CRM writes, sends). Check whether " + "the run's actions completed, then re-run manually if needed." + ), + error_code=SANDBOX_LIVENESS_UNCONFIRMED_ERROR_CODE, + retryable=False, + ) logger.exception( "E2B sandbox failed for worker %s run %s: %s", worker_id, run_id, exc ) @@ -2427,6 +2603,20 @@ def _run_in_sandbox( error_code="context_mount_failed", retryable=True, ) + # Post-cancel create race: a shutdown or user-cancel pass sweeps the + # registered sandboxes and sets cancel_requested, but a thread that is + # still in pre-sandbox setup here has no registered sandbox yet, so the + # sweep misses it and it would spawn an ORPHAN the pass believes it + # terminated. Re-check cancel as close to the spawn as possible and + # refuse to create/lease a sandbox for an already-cancelled run. + if run_cancel_requested(run_id): + logger.info("E2B run %s cancel-requested before sandbox spawn; not creating a sandbox", run_id) + log_fn("[e2b] Run cancelled before sandbox spawn; not creating a sandbox", "info") + return WorkerResult( + status="cancelled", + error="Cancelled by user", + error_code="user_cancel", + ) warm_entry = _warm_pool_lease(warm_key or "", log_fn=log_fn) if warm_key else None perf.mark("warm_lease") sandbox_prepared = warm_entry is not None @@ -2448,6 +2638,14 @@ def _run_in_sandbox( _register_sandbox(run_id, sandbox) perf.mark("register_sandbox") keep_warm = False + # Tracks whether the worker command (run.py) has started. Once True, a + # transport drop carries side-effect risk, so the finally block below + # must CONFIRM this sandbox dead before the outer run() loop may retry. + worker_command_started = False + # Tracks whether the worker command RETURNED (ran to completion). Once + # True, side effects are done and a retry would duplicate them, so run() + # never retries a completed run regardless of sandbox liveness. + worker_command_completed = False try: workdir = "/home/user/worker" @@ -2823,6 +3021,34 @@ def _finish_from_result_data(result_data: Dict[str, Any]) -> WorkerResult: ) perf.log(log_fn, "e2b.before_worker_command") + # Second cancel gate (closes the create-race window): the sandbox is + # now registered, but a cancel may have landed AFTER the pre-spawn + # check and DURING Sandbox.create()/setup, when the cancel sweep saw + # no registered sandbox to kill. Re-check here, immediately before the + # worker runs. Together with registration this closes the window: once + # registered, the cancel sweep can find + kill the sandbox; a cancel + # set before registration is caught here. Return through the finally, + # which kills the (still unused) sandbox. + if run_cancel_requested(run_id): + logger.info("E2B run %s cancel-requested before worker command; aborting without execution", run_id) + log_fn("[e2b] Run cancelled before worker command; not executing", "info") + return WorkerResult( + status="cancelled", + error="Cancelled by user", + error_code="user_cancel", + ) + # From here on the worker's own code may run and cause side effects, + # so a transport drop must not be blindly retried (see run() guard). + worker_command_started = True + # Signal run_service so the graceful-shutdown requeue never re-runs a + # started worker. Lazy + guarded: the OSS/single-tenant path and unit + # tests may run the driver without the run_service executor loaded. + try: + from run_service import mark_run_worker_command_started # noqa: PLC0415 + + mark_run_worker_command_started(run_id) + except Exception: + logger.debug("Could not mark worker_command_started for run %s", run_id, exc_info=True) try: proc = sandbox.commands.run( command, @@ -2833,6 +3059,11 @@ def _finish_from_result_data(result_data: Dict[str, Any]) -> WorkerResult: timeout=float(effective_timeout_seconds), request_timeout=_e2b_command_request_timeout(effective_timeout_seconds), ) + # The command RETURNED (exit code known): the worker ran to + # completion, so its side effects are already done. A later + # transient failure (e.g. result.json read) must NOT trigger a + # retry, which would re-run those side effects. See run() guard. + worker_command_completed = True except Exception as exc: exc_stdout = _coerce_output_text(getattr(exc, "stdout", None)) exc_stderr = _coerce_output_text(getattr(exc, "stderr", None)) @@ -2998,8 +3229,23 @@ def _finish_from_result_data(result_data: Dict[str, Any]) -> WorkerResult: finally: _unregister_sandbox(run_id, sandbox) + pending_exc = sys.exc_info()[1] + unwinding_after_command_start = pending_exc is not None and worker_command_started + # Stamp liveness intent on EVERY post-command-start unwind, BEFORE the + # warm-pool branch. A warm-pool return (pooled=True) must never be able + # to bypass this: if the stamp were skipped, run()'s outer handler + # would fall back to a RETRYABLE code and the run_service scheduler + # would re-dispatch the worker -> duplicate side effects. + if unwinding_after_command_start: + try: + pending_exc._e2b_worker_command_started = True + pending_exc._e2b_worker_command_completed = worker_command_completed + except Exception: + logger.debug("Could not stamp command-start state on exception", exc_info=True) pooled = False - if keep_warm and warm_key: + # Never return a sandbox to the warm pool while unwinding a failure; + # only pool a cleanly-completed run's sandbox. + if keep_warm and warm_key and not unwinding_after_command_start: entry = warm_entry or _WarmSandboxEntry( key=warm_key, sandbox=sandbox, @@ -3009,14 +3255,33 @@ def _finish_from_result_data(result_data: Dict[str, Any]) -> WorkerResult: ) pooled = _warm_pool_return(entry, log_fn=log_fn) if not pooled: - try: - # e2b 2.x: kill() may raise if the sandbox already exited. - # We attempt gracefully; any exception is a warning, not a failure. - sandbox.kill(request_timeout=E2B_CONTROL_REQUEST_TIMEOUT_SECONDS) - log_fn("[e2b] Sandbox killed", "debug") - except Exception as close_exc: - # Sandbox may have self-terminated (timeout, OOM) — not an error. - logger.debug("E2B sandbox already gone (kill suppressed): %s", close_exc) + if unwinding_after_command_start: + # Bounded kill-and-verify to STOP the original promptly (it may + # still be mid-run). The confirmed-dead result is informational + # now that run() refuses any post-command-start retry. + confirmed_dead = _terminate_and_confirm_dead(sandbox, log_fn=log_fn) + try: + pending_exc._e2b_sandbox_confirmed_dead = confirmed_dead + except Exception: + logger.debug("Could not stamp liveness state on exception", exc_info=True) + if confirmed_dead: + log_fn("[e2b] Sandbox confirmed terminated after transport drop", "debug") + else: + log_fn( + "[e2b] Sandbox termination could NOT be confirmed after transport drop", + "warning", + ) + else: + # Normal completion, or a pre-command failure (no worker side + # effects yet): best-effort quiet kill, as before. + try: + # e2b 2.x: kill() may raise if the sandbox already exited. + # We attempt gracefully; any exception is a warning, not a failure. + sandbox.kill(request_timeout=E2B_CONTROL_REQUEST_TIMEOUT_SECONDS) + log_fn("[e2b] Sandbox killed", "debug") + except Exception as close_exc: + # Sandbox may have self-terminated (timeout, OOM); not an error. + logger.debug("E2B sandbox already gone (kill suppressed): %s", close_exc) def _upload_contexts_to_sandbox( self, diff --git a/apps/api/services/public_view.py b/apps/api/services/public_view.py index 855aa8abf..47f0edede 100644 --- a/apps/api/services/public_view.py +++ b/apps/api/services/public_view.py @@ -109,6 +109,11 @@ def _json_noindex(payload: Dict[str, Any], *, status_code: int = 200) -> JSONRes # Sandbox / timeout / resource. "e2b_sandbox_error": _SANDBOX_HEADLINE, "sandbox_transport_retry_exhausted": _SANDBOX_HEADLINE, + "sandbox_liveness_unconfirmed": ( + "The sandbox connection dropped after this worker started, and Floom could not confirm " + "it stopped. To avoid duplicate actions, the run was not retried automatically. Check " + "whether its actions completed, then re-run it manually if needed." + ), "transient_network_error": "A temporary network connection dropped. Floom is retrying this run automatically.", "transient_network_retry_exhausted": ( "The network connection dropped repeatedly and Floom could not complete the run after automatic retries. " diff --git a/apps/api/services/run_metrics.py b/apps/api/services/run_metrics.py index d4ef1b8b7..bc7768fbb 100644 --- a/apps/api/services/run_metrics.py +++ b/apps/api/services/run_metrics.py @@ -76,6 +76,7 @@ # a reworded message. "e2b_sandbox_error": "crash", "sandbox_transport_retry_exhausted": "crash", + "sandbox_liveness_unconfirmed": "crash", "sandbox_crash": "crash", "invalid_result_json": "crash", "missing_result": "crash", diff --git a/apps/api/tests/test_cancel_loop_kill_bounded.py b/apps/api/tests/test_cancel_loop_kill_bounded.py new file mode 100644 index 000000000..e6445787b --- /dev/null +++ b/apps/api/tests/test_cancel_loop_kill_bounded.py @@ -0,0 +1,183 @@ +"""The shutdown/worker-deletion cancel loop must stay within budget. + +``_cancel_active_runs`` kills every active run's sandbox serially. The driver's +default control-plane kill timeout is 60s, so a single hung kill could blow the +whole shutdown budget (N runs x 60s). The loop now bounds each kill and stops +issuing kills once the overall budget is spent. +""" +from __future__ import annotations + +import sys +import threading +import time +from pathlib import Path +from unittest import mock + +API_DIR = Path(__file__).resolve().parents[1] +if str(API_DIR) not in sys.path: + sys.path.insert(0, str(API_DIR)) + +import run_service +from runner_sandbox import e2b_driver + + +class _FakeRuns: + def cancel(self, **_kwargs): + return None + + def add_log(self, **_kwargs): + return None + + +class _FakeRepos: + def __init__(self): + self.runs = _FakeRuns() + + +def _finished_thread() -> threading.Thread: + t = threading.Thread(target=lambda: None) + t.start() + t.join() + return t + + +def test_liveness_unconfirmed_is_never_retried_even_with_manifest_optin(): + """sandbox_liveness_unconfirmed is a safety terminal: a worker manifest that + names it in retry.on (or retryable=True) must NOT force a re-dispatch, which + would risk duplicating side effects.""" + + class _RetryCfg: + on = ["sandbox_liveness_unconfirmed"] + max_attempts = 5 + + decision = run_service._classify_retry_failure( + error_code="sandbox_liveness_unconfirmed", + error="sandbox could not be confirmed stopped", + result_retryable=True, + retry_cfg=_RetryCfg(), + ) + assert decision.retryable is False + assert decision.permanent is True + assert decision.reason == "never_retry_safety" + + +def test_cancel_sandbox_forwards_bounded_request_timeout(): + kills: list[float] = [] + + class _Sandbox: + def kill(self, request_timeout=None, **_kwargs): + kills.append(request_timeout) + + e2b_driver._register_sandbox("run-forward", _Sandbox()) + try: + assert e2b_driver.cancel_sandbox("run-forward", request_timeout=7.5) is True + finally: + e2b_driver._active_sandboxes.pop("run-forward", None) + assert kills == [7.5] + + +def test_shutdown_does_not_requeue_started_worker_runs(monkeypatch): + """Graceful shutdown must not requeue a run whose worker command already + started (re-running it would duplicate side effects). A pre-command run is + still requeued.""" + # Each run thread waits until its sandbox is cancelled, then (like the real + # executor) unregisters itself so the shutdown join sees it stop. + stop_events = {"run-started": threading.Event(), "run-pre": threading.Event()} + + def fake_cancel(run_id, *, reason=None, request_timeout=None): + stop_events[run_id].set() + return True + + monkeypatch.setattr(e2b_driver, "cancel_sandbox", fake_cancel) + + requeued: list[str] = [] + monkeypatch.setattr( + run_service, + "_requeue_interrupted_run_in_place", + lambda _repos, run_id, _user_id: (requeued.append(run_id) or True), + ) + + def make_thread(run_id: str) -> threading.Thread: + def target(): + stop_events[run_id].wait(timeout=3) + run_service._unregister_active_run(run_id) + + return threading.Thread(target=target) + + started = run_service._ActiveRun( + run_id="run-started", worker_id="w", user_id="u", thread=make_thread("run-started") + ) + started.worker_command_started = True + pre_command = run_service._ActiveRun( + run_id="run-pre", worker_id="w", user_id="u", thread=make_thread("run-pre") + ) + + run_service._register_active_run(started) + run_service._register_active_run(pre_command) + started.thread.start() + pre_command.thread.start() + try: + run_service.request_active_run_shutdown(repos=_FakeRepos(), timeout_seconds=5.0) + finally: + run_service._unregister_active_run("run-started") + run_service._unregister_active_run("run-pre") + + assert "run-started" not in requeued # started worker never re-run + assert "run-pre" in requeued # pre-command run safely requeued + + +def test_mark_run_worker_command_started_sets_flag(): + run = run_service._ActiveRun( + run_id="run-mark", worker_id="w", user_id="u", thread=_finished_thread() + ) + run_service._register_active_run(run) + try: + assert run.worker_command_started is False + run_service.mark_run_worker_command_started("run-mark") + assert run.worker_command_started is True + finally: + run_service._unregister_active_run("run-mark") + + +def test_cancel_loop_is_time_bounded_on_hung_kills(): + # Use REAL wall-clock (no global time.monotonic patch, which is pollution + # prone). The first kill "hangs" long enough to exceed the small budget, so + # the loop must stop issuing kills instead of running one-per-active-run. + timeout_seconds = 0.2 + hang_seconds = 0.5 + + calls: list[tuple[str, float]] = [] + + def fake_cancel(run_id, *, reason=None, request_timeout=None): + calls.append((run_id, request_timeout)) + time.sleep(hang_seconds) # a hung/slow kill consumes real wall time + return True + + with mock.patch.object(e2b_driver, "cancel_sandbox", fake_cancel): + active = [ + run_service._ActiveRun( + run_id=f"run-{i}", + worker_id="w", + user_id="u", + thread=_finished_thread(), + ) + for i in range(4) + ] + run_service._cancel_active_runs( + active, + repos=_FakeRepos(), + timeout_seconds=timeout_seconds, + reason="shutdown", + mark_shutdown_cancelled=False, + ) + + # Only the first run is killed before the budget is exhausted; the rest are + # skipped (handled by thread-join + startup recovery), so the loop does NOT + # run for 4 x 60s. + assert len(calls) == 1 + run_id, request_timeout = calls[0] + assert run_id == "run-0" + # Each kill timeout is bounded by both the per-call cap and remaining budget. + assert request_timeout is not None + assert request_timeout <= run_service._CANCEL_LOOP_SANDBOX_KILL_REQUEST_TIMEOUT_SECONDS + assert request_timeout <= timeout_seconds diff --git a/apps/api/tests/test_e2b_double_execute_prevention.py b/apps/api/tests/test_e2b_double_execute_prevention.py new file mode 100644 index 000000000..25343d7f5 --- /dev/null +++ b/apps/api/tests/test_e2b_double_execute_prevention.py @@ -0,0 +1,506 @@ +"""Double-execute prevention on the E2B transport-retry path. + +The failure this guards against: a transport drop (Server disconnected / +ConnectionTerminated / HPACK / reset / 503) after the worker command has +STARTED. The retry loop would re-run run.py from scratch in a fresh sandbox. +If the original sandbox is still running, its side effects (emails, CRM writes, +sends) run once in the original AND once in the retry -> a duplicate. + +The fix: before ANY re-dispatch after a post-command-start drop, require PROOF +the original sandbox is terminated (bounded kill-and-verify). If it cannot be +confirmed dead, the run ends with the terminal ``sandbox_liveness_unconfirmed`` +code instead of blind-retrying. Plus: refuse to spawn a sandbox for a run that +was cancel-requested while it sat in pre-sandbox setup (post-cancel create +race). +""" +from __future__ import annotations + +import shutil +import sys +import types +from pathlib import Path + +API_DIR = Path(__file__).resolve().parents[1] +if str(API_DIR) not in sys.path: + sys.path.insert(0, str(API_DIR)) + +from models import WorkerConfig, WorkerRuntime, WorkerTrigger +from runner_sandbox import e2b_driver +from runner_sandbox.e2b_driver import E2BSandboxDriver + + +# --- fake e2b sandbox ------------------------------------------------------- + + +class _Files: + def __init__(self, host_root: Path): + self.host_root = host_root + self._files: dict[str, bytes] = {} + + def _host_path(self, sandbox_path: str) -> Path: + return self.host_root / sandbox_path.removeprefix("/") + + def make_dir(self, sandbox_path: str, **_kwargs): + self._host_path(sandbox_path).mkdir(parents=True, exist_ok=True) + + def write(self, sandbox_path: str, content, **_kwargs): + if isinstance(content, str): + content = content.encode("utf-8") + data = bytes(content) + self._files[sandbox_path] = data + host_path = self._host_path(sandbox_path) + host_path.parent.mkdir(parents=True, exist_ok=True) + host_path.write_bytes(data) + + def exists(self, sandbox_path: str, **_kwargs): + return sandbox_path in self._files or self._host_path(sandbox_path).exists() + + def read(self, sandbox_path: str, format="text", **_kwargs): + if _ConfigurableSandbox.state.get("result_read_raises") and sandbox_path.endswith( + "result.json" + ): + raise _TransientDrop("Request timed out") + data = self._files.get(sandbox_path) + if data is None: + data = self._host_path(sandbox_path).read_bytes() # raises if missing + if format == "bytes": + return bytearray(data) + return data.decode("utf-8") + + +class _TransientDrop(RuntimeError): + """A transport error the driver classifies as a transient drop.""" + + +class _Commands: + def __init__(self, sandbox: "_ConfigurableSandbox"): + self.sandbox = sandbox + self.run_calls: list[str] = [] + + def run(self, command: str, **kwargs): + self.run_calls.append(command) + if "run.py" in command: + state = self.sandbox.__class__.state + state["worker_attempts"] += 1 + behaviour = state["command_behaviour"](state["worker_attempts"]) + if behaviour == "drop": + raise _TransientDrop("Server disconnected") + if behaviour == "boom": + # Opaque, NON-transient failure (not a recognized transport drop + # and not OOM) raised after the command started. + raise RuntimeError("totally unexpected worker sandbox failure") + # success: write result.json exactly as a real worker would + self.sandbox.files.write( + "/home/user/worker/result.json", + '{"status": "success", "outputs": {"attempt": ' + + str(state["worker_attempts"]) + + '}, "artifacts": []}', + ) + return types.SimpleNamespace(exit_code=0, stdout="", stderr="") + # any non-worker command (installs, diagnostics) succeeds trivially + return types.SimpleNamespace(exit_code=0, stdout="", stderr="") + + +class _ConfigurableSandbox: + # class-level shared state so behaviour spans the per-attempt fresh sandboxes + state: dict = {} + + def __init__(self): + assert self.state.get("host_root") is not None + self.files = _Files(self.state["host_root"]) + self.commands = _Commands(self) + self.kill_calls: list[float | None] = [] + self.is_running_calls = 0 + self.state["instances"].append(self) + + @classmethod + def create(cls, **_kwargs): + cls.state["creates"] += 1 + behaviour = cls.state.get("create_behaviour") + if behaviour is not None and behaviour(cls.state["creates"]) == "drop": + # Transport drop DURING Sandbox.create (before any worker code runs). + raise _TransientDrop("Server disconnected") + return cls() + + def set_timeout(self, *_args, **_kwargs): + return None + + def kill(self, request_timeout=None, **_kwargs): + self.kill_calls.append(request_timeout) + if self.state["kill_raises"]: + raise _TransientDrop("Server disconnected") + return None + + def is_running(self, request_timeout=None, **_kwargs): + self.is_running_calls += 1 + return self.state["is_running"] + + +def _install( + monkeypatch, + tmp_path, + *, + command_behaviour, + kill_raises, + is_running=True, + result_read_raises=False, + create_behaviour=None, +): + monkeypatch.setenv("E2B_API_KEY", "e2b-test") + monkeypatch.setenv("WORKEROS_E2B_TRANSPORT_MAX_ATTEMPTS", "3") + monkeypatch.setenv("WORKEROS_E2B_TRANSPORT_RETRY_BASE_SECONDS", "0") + monkeypatch.setenv("WORKEROS_E2B_KILL_VERIFY_BUDGET_SECONDS", "0") + monkeypatch.setenv("WORKEROS_E2B_CREATE_MIN_INTERVAL_SECONDS", "0") + monkeypatch.setitem(sys.modules, "e2b", types.SimpleNamespace(Sandbox=_ConfigurableSandbox)) + monkeypatch.setattr(e2b_driver, "WORKERS_DIR", tmp_path / "workers") + monkeypatch.setattr(e2b_driver, "run_cancel_requested", lambda _run_id: False) + _ConfigurableSandbox.state = { + "host_root": tmp_path / "sandbox", + "instances": [], + "creates": 0, + "worker_attempts": 0, + "command_behaviour": command_behaviour, + "kill_raises": kill_raises, + "is_running": is_running, + "result_read_raises": result_read_raises, + "create_behaviour": create_behaviour, + } + + +def _worker_config(tmp_path) -> WorkerConfig: + worker_dir = tmp_path / "worker" + worker_dir.mkdir(exist_ok=True) + (worker_dir / "requirements.txt").write_text("", encoding="utf-8") + (worker_dir / "run.py").write_text("print('placeholder')\n", encoding="utf-8") + return WorkerConfig( + id="side-effect-worker", + name="Side Effect Worker", + trigger=WorkerTrigger(type="manual"), + runtime=WorkerRuntime( + type="python311", + command="python3 run.py", + mode="pure-script", + bundle_path=str(worker_dir), + ), + secrets=[], + memory=False, + outputs=[], + ) + + +def _cleanup(): + host_root = _ConfigurableSandbox.state.get("host_root") + if host_root: + shutil.rmtree(host_root, ignore_errors=True) + + +# --- tests ------------------------------------------------------------------ + + +def test_drop_after_command_start_with_unconfirmed_kill_does_not_double_execute(tmp_path, monkeypatch): + """Transport drops mid-command and the original sandbox CANNOT be confirmed + dead -> NO second sandbox is spawned; the run ends terminally with + ``sandbox_liveness_unconfirmed`` rather than blind-retrying.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "drop", # every attempt drops + kill_raises=True, # control plane also unreachable -> unconfirmed + is_running=True, # sandbox still reports running + ) + config = _worker_config(tmp_path) + logs: list[tuple[str, str]] = [] + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-unconfirmed", + inputs={}, + secrets={}, + log_fn=lambda msg, level="info": logs.append((msg, level)), + trace_id="trace-unconfirmed", + timeout_seconds=30, + config=config, + ) + + # The core assertion: the worker command ran exactly ONCE across the whole + # lifecycle, and only ONE sandbox was ever created. No duplicate execution. + assert _ConfigurableSandbox.state["creates"] == 1 + assert _ConfigurableSandbox.state["worker_attempts"] == 1 + assert result.status == "error" + assert result.error_code == "sandbox_liveness_unconfirmed" + assert result.retryable is False + assert any( + "not retrying" in msg.lower() and level == "error" + for msg, level in logs + ) + _cleanup() + + +def test_drop_after_command_start_even_confirmed_dead_does_not_retry(tmp_path, monkeypatch): + """Transport drops mid-command and the original sandbox IS confirmed dead + (kill succeeds). Confirmed-dead stops concurrent execution, but it cannot + undo a side effect the worker may have already committed before the drop, + and the platform does not enforce action-level idempotency. So we still + refuse to auto-retry a run whose command started: NO second sandbox, and the + run ends terminally with sandbox_liveness_unconfirmed.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "drop", + kill_raises=False, # control-plane kill confirms termination... + ) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-confirmed-dead", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-confirmed-dead", + timeout_seconds=30, + config=config, + ) + + # ...yet a completed/in-flight worker is never re-run. + assert _ConfigurableSandbox.state["creates"] == 1 + assert _ConfigurableSandbox.state["worker_attempts"] == 1 + assert result.status == "error" + assert result.error_code == "sandbox_liveness_unconfirmed" + assert result.retryable is False + _cleanup() + + +def test_opaque_failure_after_command_start_is_non_retryable(tmp_path, monkeypatch): + """An opaque, non-transient failure raised AFTER the worker command started + must NOT come back as a retryable code (which run_service's retry scheduler + would re-dispatch, re-running the worker). The driver's outer handler forces + the non-retryable sandbox_liveness_unconfirmed terminal.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "boom", # opaque non-transient failure + kill_raises=False, + ) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-opaque-boom", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-opaque-boom", + timeout_seconds=30, + config=config, + ) + + assert _ConfigurableSandbox.state["creates"] == 1 + assert _ConfigurableSandbox.state["worker_attempts"] == 1 + assert result.status == "error" + assert result.error_code == "sandbox_liveness_unconfirmed" + assert result.retryable is False + _cleanup() + + +def test_drop_before_command_start_retries(tmp_path, monkeypatch): + """A transport drop DURING Sandbox.create (before any worker code runs) has + no side-effect risk, so retry is safe and preserved: a fresh sandbox is + created and the worker runs to success.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "success", + kill_raises=False, + create_behaviour=lambda create_count: "drop" if create_count == 1 else "ok", + ) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-create-drop", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-create-drop", + timeout_seconds=30, + config=config, + ) + + # First create dropped (no worker ran); retry created a second sandbox and + # the worker ran exactly once. + assert _ConfigurableSandbox.state["creates"] == 2 + assert _ConfigurableSandbox.state["worker_attempts"] == 1 + assert result.status == "success" + assert result.outputs == {"attempt": 1} + _cleanup() + + +def test_completed_worker_then_result_read_drop_does_not_retry(tmp_path, monkeypatch): + """The worker command RETURNS successfully (exit 0, side effects done), then + the result.json read drops transiently. Retrying would re-run the completed + worker -> duplicate. So NO second sandbox is spawned and the run ends + terminally rather than re-executing.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "success", # command returns exit 0 + kill_raises=False, # kill confirms dead - but that must NOT greenlight retry + result_read_raises=True, # the post-completion result.json read drops + ) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-completed-read-drop", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-completed-read-drop", + timeout_seconds=30, + config=config, + ) + + # The worker ran exactly ONCE and no retry sandbox was created, even though + # the sandbox was confirmed dead. A completed worker is never re-run. + assert _ConfigurableSandbox.state["creates"] == 1 + assert _ConfigurableSandbox.state["worker_attempts"] == 1 + assert result.status == "error" + assert result.error_code == "sandbox_liveness_unconfirmed" + assert result.retryable is False + _cleanup() + + +def test_cancel_during_setup_after_registration_aborts_before_command(tmp_path, monkeypatch): + """Create-race second gate: cancel lands AFTER the pre-spawn check (e.g. + during Sandbox.create/setup). The sandbox is created and registered but the + worker command must NOT run.""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "success", + kill_raises=False, + ) + # False on the pre-spawn check, True on the pre-command re-check. + calls = {"n": 0} + + def _cancel(_run_id): + calls["n"] += 1 + return calls["n"] >= 2 + + monkeypatch.setattr(e2b_driver, "run_cancel_requested", _cancel) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-cancel-during-setup", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-cancel-during-setup", + timeout_seconds=30, + config=config, + ) + + # Sandbox was created (cancel landed after the pre-spawn check) but the + # worker command never ran, and the created sandbox was killed via finally. + assert _ConfigurableSandbox.state["creates"] == 1 + assert _ConfigurableSandbox.state["worker_attempts"] == 0 + assert result.status == "cancelled" + assert result.error_code == "user_cancel" + assert _ConfigurableSandbox.state["instances"][0].kill_calls # killed on the way out + _cleanup() + + +def test_warm_pool_does_not_bypass_post_command_start_guard(tmp_path, monkeypatch): + """If a failure unwinds AFTER the worker command started, the warm-pool + return must NOT bypass the double-execute guard: the exception is still + stamped (so run() returns a non-retryable terminal, not a retryable code the + scheduler would re-dispatch) and the sandbox is NOT returned to the pool.""" + from models import WorkerConfig, WorkerRuntime, WorkerTrigger + + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "success", # command succeeds, keep_warm set + kill_raises=False, + ) + monkeypatch.setenv("WORKEROS_E2B_WARM_POOL_ENABLED", "1") + e2b_driver.clear_warm_pool() + + worker_dir = tmp_path / "worker" + worker_dir.mkdir(exist_ok=True) + (worker_dir / "requirements.txt").write_text("", encoding="utf-8") + (worker_dir / "run.py").write_text("print('x')\n", encoding="utf-8") + config = WorkerConfig( + id="side-effect-worker", + name="Side Effect Worker", + trigger=WorkerTrigger(type="manual"), + runtime=WorkerRuntime( + type="python311", command="python3 run.py", mode="pure-script", + bundle_path=str(worker_dir), + ), + secrets=["OPENAI_API_KEY"], # qualifies for a warm-pool key + memory=False, + outputs=[], + ) + + # Raise AFTER the worker completed and keep_warm was set (the exact window + # where warm-pool return could skip stamping). + raised = {"done": False} + + def log_fn(msg, level="info"): + if not raised["done"] and "Run completed successfully" in msg: + raised["done"] = True + raise RuntimeError("post-completion logging failure") + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-warm-bypass", + inputs={}, + secrets={"OPENAI_API_KEY": "sk-test"}, + log_fn=log_fn, + trace_id="trace-warm-bypass", + timeout_seconds=30, + config=config, + ) + + assert raised["done"] is True # the injected failure fired + assert result.status == "error" + assert result.error_code == "sandbox_liveness_unconfirmed" + assert result.retryable is False + assert e2b_driver.warm_pool_size() == 0 # sandbox NOT pooled while unwinding + _cleanup() + + +def test_cancel_requested_before_spawn_does_not_create_sandbox(tmp_path, monkeypatch): + """Post-cancel create race: a run cancel-requested while it sat in + pre-sandbox setup must NOT spawn a sandbox (which the cancel sweep already + believes it terminated).""" + _install( + monkeypatch, + tmp_path, + command_behaviour=lambda _attempt: "success", + kill_raises=False, + ) + # Simulate the cancel sweep having set cancel_requested before this thread + # reached the spawn point. + monkeypatch.setattr(e2b_driver, "run_cancel_requested", lambda _run_id: True) + config = _worker_config(tmp_path) + + result = E2BSandboxDriver().run( + worker_id="side-effect-worker", + run_id="run-cancel-race", + inputs={}, + secrets={}, + log_fn=lambda *_a, **_k: None, + trace_id="trace-cancel-race", + timeout_seconds=30, + config=config, + ) + + assert _ConfigurableSandbox.state["creates"] == 0 # never spawned + assert _ConfigurableSandbox.state["worker_attempts"] == 0 + assert result.status == "cancelled" + assert result.error_code == "user_cancel" + _cleanup() diff --git a/tests/test_e2b_artifact_collection.py b/tests/test_e2b_artifact_collection.py index e2b7c824b..db64f00f5 100644 --- a/tests/test_e2b_artifact_collection.py +++ b/tests/test_e2b_artifact_collection.py @@ -242,16 +242,28 @@ def kill(self, **_kwargs): class FakeCancelledCommandRunner: + def __init__(self, sandbox): + self._sandbox = sandbox + def run(self, _command, **_kwargs): + # Simulate the sandbox being killed by a concurrent cancel WHILE the + # worker command is running (a post-spawn cancellation). The optional + # class hook lets a test flip the run's cancel flag at this exact moment. + hook = getattr(self._sandbox.__class__, "on_command_run", None) + if hook is not None: + hook() raise RuntimeError("sandbox was killed") class FakeCancelledSandbox: instances = [] + # Optional callable invoked when the worker command runs (post-spawn). Set by + # tests to move a cancellation to after the sandbox has spawned. + on_command_run = None def __init__(self): self.files = FakeWritableFiles({}) - self.commands = FakeCancelledCommandRunner() + self.commands = FakeCancelledCommandRunner(self) self.killed = False FakeCancelledSandbox.instances.append(self) @@ -1586,15 +1598,27 @@ def test_e2b_cancelled_sandbox_exception_reads_active_repository(tmp_path, monke with e2b_driver._active_sandboxes_lock: e2b_driver._active_sandboxes.clear() + # The run is cancelled AFTER the sandbox spawns (the sandbox is killed while + # the worker command runs), so the driver must terminate and read the active + # repository to detect the cancel. Not cancelled at spawn time, so the + # pre-spawn create-race guard does not fire. + cancel_state = {"requested": False} + class Runs: def get_any(self, *, run_id: str): assert run_id == "run_cancelled" - return {"run_id": run_id, "cancel_requested": True} + return {"run_id": run_id, "cancel_requested": cancel_state["requested"]} class Repos: runs = Runs() monkeypatch.setattr(db_module, "get_repositories", lambda: Repos()) + # Flip the cancel flag the moment the worker command runs (post-spawn). + monkeypatch.setattr( + FakeCancelledSandbox, + "on_command_run", + lambda: cancel_state.__setitem__("requested", True), + ) worker_dir = tmp_path / "workers" / "cancelled-worker" worker_dir.mkdir(parents=True) (worker_dir / "run.py").write_text("print('unused')\n") @@ -1610,6 +1634,9 @@ class Repos: timeout_seconds=30, ) + # The sandbox DID spawn and the worker command ran (post-spawn path), then the + # cancel was detected via the active repository. + assert len(FakeCancelledSandbox.instances) == 1 assert result.status == "cancelled" assert result.error_code == "user_cancel" assert ("info", "[e2b] Sandbox terminated - run cancelled by user") in logs diff --git a/tests/test_s35_shutdown.py b/tests/test_s35_shutdown.py index 396364a2a..0d179e379 100644 --- a/tests/test_s35_shutdown.py +++ b/tests/test_s35_shutdown.py @@ -51,7 +51,9 @@ def __init__(self): monkeypatch.setattr( e2b_driver, "cancel_sandbox", - lambda run_id, reason=None: killed.append((run_id, reason)) or True, + # cancel_sandbox now takes a bounded request_timeout (the cancel loop + # bounds each kill so a hung kill cannot blow the shutdown budget). + lambda run_id, reason=None, request_timeout=None: killed.append((run_id, reason)) or True, ) try: