diff --git a/backend/druks/workflows.py b/backend/druks/workflows.py index 35ad61a8..9725aab6 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, @@ -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 diff --git a/backend/tests/test_warm_host_rotation.py b/backend/tests/test_warm_host_rotation.py index fa333ca7..038ee417 100644 --- a/backend/tests/test_warm_host_rotation.py +++ b/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 == []