workflow: recover stranded local runs on restart

src/vercel/_internal/workflow/worlds/local.py

@@ -1,8 +1,11 @@
+import asyncio
 import hashlib
 import json
+import logging
 import math
 import os
 import pathlib
+import threading
 import traceback
 from datetime import datetime
 from typing import Any, TypeVar
@@ -11,11 +14,18 @@
 import pydantic
 
 from vercel._internal.polyfills import UTC
+from vercel._internal.retryloop import RetryLoop, exp_backoff
 from vercel.workers import client as vqs_client
 
 from .. import world as w
 from ..ulid import monotonic_factory
 
+logger = logging.getLogger("vercel.workflow")
+
+# Local dev only: how long to keep retrying recovery enqueues while the dev
+# server's queue finishes coming up before giving up.
+RECOVERY_TIMEOUT_SECONDS = float(os.getenv("WORKFLOW_LOCAL_RECOVERY_TIMEOUT", "30"))
+
 MAX_DELAY_SECONDS = float(
     os.getenv("VERCEL_QUEUE_MAX_DELAY_SECONDS", "82800")
 )  # 23 hours - leave 1h buffer before 24h retention limit
@@ -71,6 +81,7 @@ class LocalWorld(w.World):
     def __init__(self) -> None:
         self.monotonic_ulid = monotonic_factory()
         self.data_dir = pathlib.Path(os.getenv("WORKFLOW_LOCAL_DATA_DIR", ".workflow-data"))
+        self._recovery_started = False
 
     def delete_all_hooks_for_run(self, run_id: str) -> None:
         hooks_dir = self.data_dir / "hooks"
@@ -207,8 +218,77 @@ async def http_handler(request: w.HTTPRequest) -> w.HTTPResponse:
             except Exception as error:
                 return w.HTTPResponse.json({"error": str(error)}, status=500)
 
+        # When the workflow worker comes up, recover any run stranded by a dev
+        # server restart (see _resume_pending_runs). Triggered off the workflow
+        # topic only, once per process.
+        if queue_name_prefix == "__wkf_workflow_":
+            self._start_recovery_once()
+
         return http_handler
 
+    def _start_recovery_once(self) -> None:
+        if self._recovery_started or os.getenv("WORKFLOW_LOCAL_DISABLE_RECOVERY"):
+            return
+        self._recovery_started = True
+
+        # No event loop runs at worker import time, so drive the async sweep on
+        # a throwaway loop in a daemon thread; it never blocks worker startup.
+        def _run() -> None:
+            try:
+                asyncio.run(self._resume_pending_runs())
+            except Exception as e:  # never break worker startup
+                logger.warning("Local run recovery failed: %r", e)
+
+        threading.Thread(target=_run, name="wkf-local-recovery", daemon=True).start()
+
+    async def _resume_pending_runs(self) -> None:
+        """Re-enqueue runs left mid-flight by a dev server restart (local dev only).
+
+        ``vercel dev``'s queue is in-memory, so a ``sleep``'s delayed wake-up
+        message is lost if the server restarts mid-run, stranding the run in
+        ``running`` with no one to resume it. Re-invoking the run lets
+        ``workflow_handler`` turn any elapsed wait into a ``wait_completed`` and
+        continue; replay is idempotent, so re-invoking a healthy run is harmless.
+        Production uses ``VercelWorld`` (durable queue), so this never runs there.
+        """
+        runs_dir = self.data_dir / "runs"
+        if not runs_dir.exists():
+            return
+
+        pending: list[tuple[str, str]] = []
+        for run_file in runs_dir.glob("*.json"):
+            try:
+                run = await self.runs_get(run_file.stem)
+            except Exception:
+                continue
+            if run.status in ("pending", "running"):
+                pending.append((run.workflow_name, run.run_id))
+        if not pending:
+            return
+
+        # The dev server's queue may not be reachable the instant the worker
+        # imports. Retry until it accepts the messages (popping each run once
+        # enqueued, so none is sent twice) or the timeout elapses.
+        try:
+            async for iteration in RetryLoop(
+                backoff=exp_backoff(), timeout=RECOVERY_TIMEOUT_SECONDS, ignore=Exception
+            ):
+                async with iteration:
+                    while pending:
+                        workflow_name, run_id = pending[-1]
+                        await self.queue(
+                            f"__wkf_workflow_{workflow_name}",
+                            w.WorkflowInvokePayload(runId=run_id),
+                        )
+                        pending.pop()
+        except Exception as e:
+            logger.warning(
+                "Local run recovery timed out after %ss; %d run(s) still stranded: %r",
+                RECOVERY_TIMEOUT_SECONDS,
+                len(pending),
+                e,
+            )
+
     async def runs_get(self, run_id: str) -> w.WorkflowRun:
         run_path = self.data_dir / "runs" / f"{run_id}.json"
         run = read_json(run_path, w.WorkflowRunAdaptor)

tests/unit/test_workflow_local_world.py

@@ -0,0 +1,51 @@
+"""Tests for LocalWorld dev-restart recovery."""
+
+import pytest
+
+from vercel._internal.workflow import world as w
+
+# LocalWorld imports vercel.workers, which is only installed on Python >= 3.12.
+pytest.importorskip("vercel.workers")
+
+from vercel._internal.workflow.worlds import local as local_mod  # noqa: E402
+
+
+@pytest.fixture
+def world(tmp_path, monkeypatch) -> local_mod.LocalWorld:
+    monkeypatch.setenv("WORKFLOW_LOCAL_DATA_DIR", str(tmp_path))
+    return local_mod.LocalWorld()
+
+
+async def _create_run(world: local_mod.LocalWorld) -> str:
+    """Create a run and move it to 'running'; returns the run id."""
+    created = await world.events_create(
+        None,
+        w.RunCreatedEventData(
+            deploymentId="", workflowName="workflow//tests.wf", input=[b"json[[], {}]"]
+        ).into_event(),
+    )
+    assert created.run is not None
+    run_id = created.run.run_id
+    await world.events_create(run_id, w.RunStartedEvent())
+    return run_id
+
+
+async def test_resume_pending_runs_reenqueues_only_nonterminal(world) -> None:
+    """Recovery sweep re-enqueues runs left 'running' (e.g. a dev restart lost a
+    sleep's wake-up), but leaves terminal runs alone."""
+    captured: list[tuple[str, str]] = []
+
+    async def fake_queue(queue_name, message, **kwargs):  # type: ignore[no-untyped-def]
+        captured.append((queue_name, message.run_id))
+        return "msg_x"
+
+    world.queue = fake_queue  # type: ignore[method-assign]
+
+    running_id = await _create_run(world)
+
+    completed_id = await _create_run(world)
+    await world.events_create(completed_id, w.RunCompletedEventData(output=[b"json1"]).into_event())
+
+    await world._resume_pending_runs()
+
+    assert captured == [("__wkf_workflow_workflow//tests.wf", running_id)]