DRU-380 - Add sandbox hold

backend/druks/workflows.py

@@ -2,7 +2,7 @@
 from collections.abc import Awaitable, Callable
 from contextlib import nullcontext, suppress
 from contextvars import ContextVar
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 from functools import partial
 from typing import (
     TYPE_CHECKING,
@@ -290,13 +290,20 @@ async def answer(cls, subject: Subject | StoredSubject, **reply: Any) -> None:
 
     @classmethod
     async def wait(
-        cls, *, input_request: dict[str, Any] | None = None, ttl_seconds: float = GATE_TTL_SECONDS
+        cls,
+        *,
+        input_request: dict[str, Any] | None = None,
+        ttl_seconds: float = GATE_TTL_SECONDS,
+        hold_sandbox: bool | timedelta | None = False,
     ) -> Self:
         # Suspend the running workflow until its gate is answered. A gate is a
         # run-level state — the read surfaces "needs you" straight off the parked run.
         # ``input_request`` is the plain-dict ask (at least a ``label`` and
         # ``presentation``), stored on the run beside ``input_gate`` and cleared on
         # resume — so an app declares the ask here, beside on_wait, not at read time.
+        # ``hold_sandbox`` keeps the warm VM across the park (see ``_hold_host``):
+        # ``True`` holds it for as long as its lease could still cover one more
+        # worst-case agent call, a timedelta for at most that long.
         workflow = current_workflow.get()
         if not workflow._subject and cls.on_wait.__func__ is Gate.on_wait.__func__:
             # No subject means no feed surface; if on_wait wasn't overridden
@@ -311,7 +318,9 @@ async def _on_wait() -> None:
                 await cls.on_wait(workflow)
 
         await DBOS.run_step_async(StepOptions(name=f"{cls.name}._on_wait"), _on_wait)
-        payload = await _park(workflow, cls.name, input_request, ttl_seconds)
+        payload = await _park(
+            workflow, cls.name, input_request, ttl_seconds, hold_sandbox=hold_sandbox
+        )
         reply = cls.model_validate(payload)
         workflow.journal.add(reply)
         return reply
@@ -333,10 +342,15 @@ async def _park(
     gate: str,
     input_request: dict[str, Any] | None,
     ttl_seconds: float,
+    hold_sandbox: bool | timedelta | None = False,
 ) -> dict[str, Any]:
     # Shared park core: a park lasts days, so reap the warm VM, then suspend on the
-    # gate's channel until Run.resume answers it.
-    await workflow._reap_run()
+    # gate's channel until Run.resume answers it. A caller that expects a quick
+    # answer can hold the VM instead — the clipped lease is what ends the hold.
+    if hold_sandbox:
+        await workflow._hold_host(hold_sandbox)
+    else:
+        await workflow._reap_run()
     await _emit_run_event(
         workflow.workflow_id,
         RunState.PARKED,
@@ -899,6 +913,25 @@ async def _reap_run(self) -> None:
         host, self._host = self._host, None
         await sandbox_client.release(host_id=host.id)
 
+    async def _hold_host(self, hold: bool | timedelta) -> None:
+        # Keep the warm VM across a park instead of reaping it, by clipping its lease
+        # down: drukbox reaps at the new expiry, so a hold nobody ever answers still
+        # frees the VM with no druks-side sweep. ``True`` holds it for as long as the
+        # lease could still cover one more worst-case call — past that the next call
+        # would rotate anyway. Never extends: the lease drukbox already granted is the
+        # ceiling. A run with no warm host has nothing to hold.
+        if not self._host:
+            return
+        span = (
+            hold
+            if isinstance(hold, timedelta)
+            else timedelta(seconds=SANDBOX_HOST_ROTATE_BEFORE_SECONDS)
+        )
+        expires_at = datetime.now(UTC) + span
+        if self._host.expires_at:
+            expires_at = min(self._host.expires_at, expires_at)
+        await sandbox_client.set_expiry(host_id=self._host.id, expires_at=expires_at)
+
     @property
     def workflow_id(self) -> str:
         return self._workflow_id

backend/tests/test_warm_host_rotation.py

@@ -37,6 +37,7 @@ def __init__(self, *, lease: timedelta) -> None:
         self.secrets: list[dict[str, Secret]] = []
         self.released: list[str] = []
         self.reattached: list[str] = []
+        self.expiry_sets: list[tuple[str, datetime]] = []
 
     async def provision(
         self,
@@ -60,6 +61,9 @@ async def reattach(self, *, host_id: str) -> _FakeSandbox:
         self.reattached.append(host_id)
         return _FakeSandbox(id=host_id, expires_at=datetime.now(UTC) + self.lease)
 
+    async def set_expiry(self, *, host_id: str, expires_at: datetime) -> None:
+        self.expiry_sets.append((host_id, expires_at))
+
 
 def _warm_workflow(*, reuse: bool = True) -> Workflow:
     # __new__ skips __init__/__init_subclass__ so the host logic can be exercised
@@ -73,6 +77,19 @@ def _warm_workflow(*, reuse: bool = True) -> Workflow:
     return flow
 
 
+def _park_without_dbos(monkeypatch) -> None:
+    # The park's durable surroundings — the run event it emits and the channel it
+    # suspends on — say nothing about the hold, so the gate answers immediately.
+    async def _emit(*args, **kwargs) -> None:
+        return
+
+    async def _answer(gate, timeout_seconds=None) -> dict[str, str]:
+        return {"action": "approve"}
+
+    monkeypatch.setattr(sdk, "_emit_run_event", _emit)
+    monkeypatch.setattr(sdk.DBOS, "recv_async", _answer)
+
+
 @pytest.mark.asyncio
 async def test_warm_host_reused_while_lease_covers_another_call(monkeypatch):
     """A warm host with lease to spare is reused across calls, never re-provisioned."""
@@ -198,3 +215,134 @@ async def test_a_replay_finds_the_warm_box_through_its_identity(
 
     assert client.reattached == ["host-crashed"]
     assert client.provisions == []
+
+
+@pytest.mark.asyncio
+async def test_park_without_hold_releases_the_warm_host(monkeypatch):
+    """A park with no hold is today's park: the VM goes, nothing is clipped."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    _park_without_dbos(monkeypatch)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+
+    await sdk._park(flow, "review", None, 60.0)
+
+    assert fake.released == ["host-1"]
+    assert fake.expiry_sets == []
+    assert flow._host is None
+
+
+@pytest.mark.asyncio
+async def test_park_with_hold_clips_the_lease_and_keeps_the_host(monkeypatch):
+    """A held park clips the lease instead of deleting the VM, and the run keeps
+    the handle so a same-worker resume reattaches warm."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    _park_without_dbos(monkeypatch)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+
+    await sdk._park(flow, "review", None, 60.0, hold_sandbox=True)
+
+    assert [host_id for host_id, _ in fake.expiry_sets] == ["host-1"]
+    assert fake.released == []
+    assert flow._host is not None
+    assert flow._host.id == "host-1"
+
+
+@pytest.mark.asyncio
+async def test_hold_true_clips_to_one_more_worst_case_call(monkeypatch):
+    """``True`` holds the VM for as long as its lease could still cover one more
+    worst-case call — past that the next call would rotate anyway."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+
+    await flow._hold_host(True)
+
+    ((_, expires_at),) = fake.expiry_sets
+    clip = datetime.now(UTC) + timedelta(seconds=SANDBOX_HOST_ROTATE_BEFORE_SECONDS)
+    assert clip - timedelta(seconds=5) <= expires_at <= clip
+    assert expires_at <= flow._host.expires_at
+
+
+@pytest.mark.asyncio
+async def test_hold_never_outlasts_the_lease_drukbox_granted(monkeypatch):
+    """The clip is a floor, never an extension: a lease shorter than the hold
+    stands as it is."""
+    fake = _FakeSandboxClient(lease=timedelta(minutes=20))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+
+    await flow._hold_host(timedelta(hours=1))
+
+    assert fake.expiry_sets == [("host-1", flow._host.expires_at)]
+
+
+@pytest.mark.asyncio
+async def test_hold_timedelta_clips_to_the_requested_span(monkeypatch):
+    """A timedelta hold ends at ``now + hold`` when the lease outlasts it."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+
+    await flow._hold_host(timedelta(minutes=30))
+
+    ((_, expires_at),) = fake.expiry_sets
+    clip = datetime.now(UTC) + timedelta(minutes=30)
+    assert clip - timedelta(seconds=5) <= expires_at <= clip
+
+
+@pytest.mark.asyncio
+async def test_hold_without_a_warm_host_touches_nothing(monkeypatch):
+    """Without steps_reuse_sandbox there is no warm host to hold, so a held park
+    neither clips nor deletes."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    _park_without_dbos(monkeypatch)
+    flow = _warm_workflow(reuse=False)
+    await flow._lease_host(_NONE)
+
+    await sdk._park(flow, "review", None, 60.0, hold_sandbox=True)
+
+    assert fake.expiry_sets == []
+    assert fake.released == []
+    assert fake.provisions == []
+
+
+@pytest.mark.asyncio
+async def test_resume_after_a_hold_reuses_the_held_host(monkeypatch):
+    """The worker that survived the recv still holds the handle, so the first
+    call after the resume lands on the same VM."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    _park_without_dbos(monkeypatch)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+    await sdk._park(flow, "review", None, 60.0, hold_sandbox=True)
+
+    assert await flow._lease_host(_NONE) == "host-1"
+    assert fake.provisions == ["wf-1:workflow"]
+
+
+@pytest.mark.asyncio
+async def test_resume_on_a_restarted_worker_re_leases_under_the_run_key(monkeypatch):
+    """A worker that died over the park has no handle: the resume goes back
+    through the run's idempotency key exactly once — warm if the clipped lease
+    still stands, cold if drukbox already reaped it."""
+    fake = _FakeSandboxClient(lease=timedelta(hours=2))
+    monkeypatch.setattr(sdk, "sandbox_client", fake)
+    _park_without_dbos(monkeypatch)
+    flow = _warm_workflow()
+    await flow._lease_host(_NONE)
+    await sdk._park(flow, "review", None, 60.0, hold_sandbox=timedelta(minutes=30))
+    flow._host = None
+
+    await flow._lease_host(_NONE)
+
+    assert fake.provisions == ["wf-1:workflow", "wf-1:workflow"]
+    assert fake.released == []