ENG-822 - Sandbox provisioning failures are transient but unclassified — the in-run retry never covers them

backend/druks/harnesses/exceptions.py

@@ -62,6 +62,15 @@ class HarnessSandboxError(HarnessError):
     retry_delays = (60, 300)
 
 
+class HarnessSandboxProvisioningError(HarnessSandboxError):
+    """Provisioning or transport failure during host creation, attach, or SSH setup.
+
+    Inherits the transient retry schedule from HarnessSandboxError; uses
+    ``sandbox_provisioning`` so exhausted retries record a named failure code."""
+
+    code = "sandbox_provisioning"
+
+
 class OAuthTokenError(Exception):
     """No usable subscription credential is available.
 

backend/druks/sandbox/client.py

@@ -5,25 +5,37 @@
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 
+import asyncssh
 from drukbox_sdk import SandboxAPI, SandboxHost
 from drukbox_sdk.exceptions import (
     SandboxAPIError,
     SandboxNotFoundError,
+    SandboxProvisioningError,
     SandboxUnavailableError,
 )
 from uuid_utils import uuid7
 
+from druks.harnesses.exceptions import HarnessSandboxProvisioningError
 from druks.settings import load_settings
 
 from .constants import SANDBOX_HOST_LEASE_SECONDS
-from .exceptions import HostGone, SandboxUnreachable
+from .exceptions import HostGone, SandboxError, SandboxUnreachable
 from .host import Sandbox
 from .layout import get_helper_script_path, get_remote_home
 
 logger = logging.getLogger(__name__)
 
 _DRUKS_SANDBOX_LOCAL_SCRIPT = Path(__file__).parent / "druks-sandbox.sh"
 
+# SSH/socket/sandbox errors that mean a fresh VM never became usable.
+# CancelledError and programming errors are excluded so they propagate unchanged.
+_ACQUIRE_SETUP_REACHABILITY_ERRORS = (
+    SandboxError,
+    asyncssh.Error,
+    OSError,
+    TimeoutError,
+)
+
 
 class Client:
     """Ambient client for the drukbox control plane.
@@ -80,13 +92,25 @@ async def acquire(
             # Fixed lease: drukbox reaps the host when this lapses, so a run whose
             # worker dies frees its VM without a druks-side reconciler.
             expires_at = datetime.now(UTC) + timedelta(seconds=SANDBOX_HOST_LEASE_SECONDS)
-            record = await api.create_host(
-                expires_at=expires_at,
-                env=sandbox_env,
-                idempotency_key=key,
-                image=image or None,
-                provider=provider,
-            )
+            try:
+                record = await api.create_host(
+                    expires_at=expires_at,
+                    env=sandbox_env,
+                    idempotency_key=key,
+                    image=image or None,
+                    provider=provider,
+                )
+            except (SandboxProvisioningError, SandboxUnavailableError) as exc:
+                # Transient control-plane failures — a 502 the service raises
+                # when the provider/Tailscale/keyscan step fails, or a
+                # transport/503 SandboxUnavailableError. Classify them into the
+                # in-run retry path so a slow provider window recovers instead
+                # of dead-ending the run. Fatal SDK errors (auth, validation,
+                # conflict, not-found, generic response) are subclasses of the
+                # untouched SandboxAPIError base and fall through unretried.
+                raise HarnessSandboxProvisioningError(
+                    f"sandbox host provisioning failed: {exc}"
+                ) from exc
             logger.info("sandbox host created id=%s", record.id)
             key_path = settings.sandbox_keys_dir / record.id
             if record.private_key:
@@ -96,8 +120,15 @@ async def acquire(
             sandbox = Sandbox(record=record)
             try:
                 await _upload_helper_script(sandbox)
+            except _ACQUIRE_SETUP_REACHABILITY_ERRORS as error:
+                # Fresh VM never became usable; roll back and classify as provisioning.
+                await sandbox.aclose()
+                await self._best_effort_delete(api, record.id)
+                key_path.unlink(missing_ok=True)
+                raise HarnessSandboxProvisioningError(
+                    f"sandbox host {record.id} unreachable during setup: {error}"
+                ) from error
             except BaseException:
-                # Caller never sees this host; release rather than orphan.
                 await sandbox.aclose()
                 await self._best_effort_delete(api, record.id)
                 key_path.unlink(missing_ok=True)
@@ -139,6 +170,14 @@ async def attach(self, *, host_id: str) -> AsyncIterator[Sandbox]:
                 raise HostGone(
                     f"sandbox host {host_id} no longer exists",
                 ) from exc
+            except SandboxUnavailableError as exc:
+                # Transport/503 while looking up an existing host — transient,
+                # so classify it the same as a create-time control-plane
+                # failure and let the in-run retry re-attach once the service
+                # recovers.
+                raise HarnessSandboxProvisioningError(
+                    f"sandbox host {host_id} lookup failed: {exc}"
+                ) from exc
             sandbox = Sandbox(record=record)
             try:
                 yield sandbox

backend/tests/test_agents.py

@@ -580,6 +580,186 @@ async def run_agent(_self, **_context):
     durable_sleep.assert_not_awaited()
 
 
+def _flaky_ephemeral(sandbox, keys, *, failures):
+    """An ephemeral() that raises each error in ``failures`` in turn (recording
+    the idempotency key each attempt presents) then yields ``sandbox``. Used to
+    simulate a transient provisioning failure at acquire time that recovers."""
+    attempts = 0
+
+    @asynccontextmanager
+    async def _fake(self, *, idempotency_key=None, **_kwargs):
+        nonlocal attempts
+        keys.append(idempotency_key)
+        index = attempts
+        attempts += 1
+        if index < len(failures):
+            raise failures[index]
+        yield sandbox
+
+    return _fake
+
+
+async def test_provisioning_failure_recovers_through_durable_retry(
+    druks_db, tmp_path, monkeypatch, current_run, _inline_agent_steps
+):
+    """A transient provisioning failure at acquire time is retried by the
+    body-level durable path with backoff; a later attempt succeeds, and every
+    attempt for the one logical acquire presents the same ephemeral key."""
+    from druks.harnesses.exceptions import HarnessSandboxProvisioningError
+
+    sandbox = _patch_runtime(monkeypatch, tmp_path, {"ok": True})
+    keys: list[str | None] = []
+    monkeypatch.setattr(
+        "druks.sandbox.client.Client.ephemeral",
+        _flaky_ephemeral(
+            sandbox, keys, failures=[HarnessSandboxProvisioningError("exe.dev create timed out")]
+        ),
+    )
+    monkeypatch.setattr(agents.random, "uniform", lambda _low, _high: 1.0)
+    current_run._reap_run = AsyncMock()
+    checkpoints, sleep = _inline_agent_steps
+
+    result = await DUMMY_AGENT()
+
+    assert result == DummyOutput(ok=True)
+    # First delay off the schedule HarnessSandboxProvisioningError inherits (60, 300).
+    assert [awaited.args[0] for awaited in sleep.await_args_list] == [60.0]
+    # Each attempt re-enters _run with the same workflow/agent identity → same key.
+    assert keys == ["wf-9:dummy", "wf-9:dummy"]
+    # A 60s wait is under the reap-before threshold, so the (never-acquired) VM
+    # isn't reaped between attempts.
+    current_run._reap_run.assert_not_awaited()
+    # Both attempts are separate durable steps under the agent's step name.
+    assert [options["name"] for options in checkpoints] == ["test.agent.dummy"] * 2
+
+
+async def test_provisioning_failure_recovers_in_step_retry(
+    druks_db, tmp_path, monkeypatch, current_run
+):
+    """An agent invoked inside an enclosing @step retries the same transient
+    provisioning failure in memory (plain asyncio.sleep, no durable sleep) and
+    recovers, holding the ephemeral key stable across attempts."""
+    from druks.harnesses.exceptions import HarnessSandboxProvisioningError
+    from druks.workflows import _in_step
+
+    sandbox = _patch_runtime(monkeypatch, tmp_path, {"ok": True})
+    keys: list[str | None] = []
+    monkeypatch.setattr(
+        "druks.sandbox.client.Client.ephemeral",
+        _flaky_ephemeral(
+            sandbox, keys, failures=[HarnessSandboxProvisioningError("exe.dev create timed out")]
+        ),
+    )
+    monkeypatch.setattr(agents.random, "uniform", lambda _low, _high: 1.0)
+    memory_sleep = AsyncMock()
+    durable_sleep = AsyncMock()
+    monkeypatch.setattr(agents.asyncio, "sleep", memory_sleep)
+    monkeypatch.setattr(agents.DBOS, "sleep_async", durable_sleep)
+
+    token = _in_step.set(True)
+    try:
+        result = await DUMMY_AGENT()
+    finally:
+        _in_step.reset(token)
+
+    assert result == DummyOutput(ok=True)
+    memory_sleep.assert_awaited_once_with(60.0)
+    durable_sleep.assert_not_awaited()
+    assert keys == ["wf-9:dummy", "wf-9:dummy"]
+
+
+async def test_reused_host_retry_presents_a_stable_idempotency_key(monkeypatch, current_run):
+    """A warm workflow that reuses its sandbox re-provisions the host on the
+    same logical acquire with an unchanged idempotency key, so an ambiguous
+    create timeout resolves to the existing host instead of leaking a duplicate."""
+    from druks.harnesses.exceptions import HarnessSandboxProvisioningError
+
+    current_run.steps_reuse_sandbox = True
+    keys: list[str | None] = []
+    host = MagicMock()
+    host.id = "warm-host"
+    host.expires_at = None
+    attempts = 0
+
+    async def fake_provision(self, *, idempotency_key=None, **_kwargs):
+        nonlocal attempts
+        keys.append(idempotency_key)
+        index = attempts
+        attempts += 1
+        if index == 0:
+            raise HarnessSandboxProvisioningError("provider slow")
+        return host
+
+    monkeypatch.setattr("druks.sandbox.client.Client.provision", fake_provision)
+
+    with pytest.raises(HarnessSandboxProvisioningError):
+        await current_run._ensure_host()
+    host_id = await current_run._ensure_host()
+
+    assert host_id == "warm-host"
+    assert keys == ["wf-9:sandbox", "wf-9:sandbox"]
+
+
+async def test_provisioning_failure_exhausts_retries_with_classified_code(
+    druks_db, tmp_path, monkeypatch, current_run, _inline_agent_steps
+):
+    """When transient provisioning retries exhaust, the propagated error still
+    carries ``code=sandbox_provisioning`` — the classification survives to the
+    run's failure record instead of falling back to an empty code."""
+    from druks.harnesses.exceptions import HarnessSandboxProvisioningError
+
+    _patch_runtime(monkeypatch, tmp_path, {"ok": True})
+    failures = [HarnessSandboxProvisioningError(f"provider down {i}") for i in range(3)]
+    monkeypatch.setattr(
+        "druks.sandbox.client.Client.ephemeral", _flaky_ephemeral(None, [], failures=failures)
+    )
+    monkeypatch.setattr(agents.random, "uniform", lambda _low, _high: 1.0)
+    current_run._reap_run = AsyncMock()
+    _, sleep = _inline_agent_steps
+
+    with pytest.raises(HarnessSandboxProvisioningError) as excinfo:
+        await DUMMY_AGENT()
+
+    assert excinfo.value is failures[-1]
+    assert excinfo.value.code == "sandbox_provisioning"
+    # One sleep per inherited delay; past the schedule the failure stands.
+    assert [awaited.args[0] for awaited in sleep.await_args_list] == [60.0, 300.0]
+    # The 300s wait is over the reap-before threshold, so the run is reaped once.
+    assert current_run._reap_run.await_count == 1
+
+
+async def test_fatal_sdk_error_does_not_trigger_agent_retry(
+    druks_db, tmp_path, monkeypatch, current_run, _inline_agent_steps
+):
+    """A non-transient SDK failure surfacing from acquire is not a HarnessError,
+    so it propagates on the first attempt with no retry sleep — auth/validation
+    failures can't recover without changed inputs."""
+    from drukbox_sdk.exceptions import SandboxValidationError
+
+    _patch_runtime(monkeypatch, tmp_path, {"ok": True})
+    fatal = SandboxValidationError("bad image")
+    attempts = 0
+
+    @asynccontextmanager
+    async def fatal_ephemeral(self, *, idempotency_key=None, **_kwargs):
+        nonlocal attempts
+        attempts += 1
+        raise fatal
+        yield  # pragma: no cover
+
+    monkeypatch.setattr("druks.sandbox.client.Client.ephemeral", fatal_ephemeral)
+    current_run._reap_run = AsyncMock()
+    _, sleep = _inline_agent_steps
+
+    with pytest.raises(SandboxValidationError) as excinfo:
+        await DUMMY_AGENT()
+
+    assert excinfo.value is fatal
+    assert attempts == 1
+    sleep.assert_not_awaited()
+    current_run._reap_run.assert_not_awaited()
+
+
 async def test_recovery_supersedes_the_orphaned_running_call(druks_db):
     """A worker crash leaves a RUNNING row; the recovered step re-runs with a
     fresh id and abandons the orphan, so the timeline shows one live step."""

backend/tests/test_run_state.py

@@ -289,6 +289,25 @@ async def body() -> None:
     assert Run.get(run.id).failure_code == "overloaded"
 
 
+@pytest.mark.asyncio
+async def test_an_exhausted_provisioning_failure_stamps_its_code(druks_db, _inline_steps):
+    # An exhausted transient provisioning failure records the classified
+    # ``sandbox_provisioning`` code rather than the empty string a raw drukbox
+    # SDK exception used to leave behind — so the dashboard/taxonomy can name it.
+    from druks.harnesses.exceptions import HarnessSandboxProvisioningError
+
+    item, run = _item_and_run(druks_db, "running")
+
+    async def body() -> None:
+        raise HarnessSandboxProvisioningError("exe.dev VM creation timed out")
+
+    with pytest.raises(HarnessSandboxProvisioningError):
+        await _execute_run(run.id, run.kind, {"type": "work_item", "id": item.id}, None, body)
+
+    ambient_session().expire_all()
+    assert Run.get(run.id).failure_code == "sandbox_provisioning"
+
+
 @pytest.mark.asyncio
 async def test_a_foreign_code_never_becomes_the_failure_code(druks_db, _inline_steps):
     """``code`` is a common attribute name — asyncssh's is an int — so only