From b591f4bb0129670174a978ed473e5d2d54ec23b3 Mon Sep 17 00:00:00 2001 From: "druks-operator-treadstone[bot]" <322217521+druks-operator-treadstone[bot]@users.noreply.github.com> Date: Fri, 4 Sep 2026 06:35:06 +0000 Subject: [PATCH] A park can hold its warm VM by clipping the lease instead of reaping it Co-Authored-By: Claude Opus 5 --- backend/druks/workflows.py | 43 ++++++- backend/tests/test_warm_host_rotation.py | 149 +++++++++++++++++++++++ 2 files changed, 187 insertions(+), 5 deletions(-) diff --git a/backend/druks/workflows.py b/backend/druks/workflows.py index 1c4ee075..4f8f0a93 100644 --- a/backend/druks/workflows.py +++ b/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, @@ -293,13 +293,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 @@ -314,7 +321,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 @@ -336,10 +345,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, @@ -867,6 +881,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 diff --git a/backend/tests/test_warm_host_rotation.py b/backend/tests/test_warm_host_rotation.py index 15812487..69e2ca83 100644 --- a/backend/tests/test_warm_host_rotation.py +++ b/backend/tests/test_warm_host_rotation.py @@ -36,6 +36,7 @@ def __init__(self, *, lease: timedelta) -> None: self.provisions: list[str] = [] self.secrets: list[dict[str, Secret]] = [] self.released: list[str] = [] + self.expiry_sets: list[tuple[str, datetime]] = [] async def provision( self, *, idempotency_key: str, secrets: dict[str, Secret], template: str | None @@ -49,6 +50,9 @@ async def provision( async def release(self, *, host_id: str) -> None: self.released.append(host_id) + 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 @@ -56,10 +60,24 @@ def _warm_workflow(*, reuse: bool = True) -> Workflow: flow = Workflow.__new__(Workflow) flow.steps_reuse_sandbox = reuse flow._host = None + flow._subject = None flow._workflow_id = "wf-1" 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.""" @@ -155,3 +173,134 @@ async def test_no_warm_host_when_reuse_disabled(monkeypatch): assert await flow._lease_host(_NONE) is None assert fake.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() + + 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() + + 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() + + 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() + + 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() + + 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() + + 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() + await sdk._park(flow, "review", None, 60.0, hold_sandbox=True) + + assert await flow._lease_host() == "host-1" + assert fake.provisions == ["wf-1:sandbox"] + + +@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() + await sdk._park(flow, "review", None, 60.0, hold_sandbox=timedelta(minutes=30)) + flow._host = None + + await flow._lease_host() + + assert fake.provisions == ["wf-1:sandbox", "wf-1:sandbox"] + assert fake.released == []