diff --git a/src/vercel/_internal/workflow/worlds/local.py b/src/vercel/_internal/workflow/worlds/local.py index 0fe6a731..bd63574a 100644 --- a/src/vercel/_internal/workflow/worlds/local.py +++ b/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) diff --git a/tests/unit/test_workflow_local_world.py b/tests/unit/test_workflow_local_world.py new file mode 100644 index 00000000..1ed896fd --- /dev/null +++ b/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)]