vercel/vercel-py#343 repro

pr343_set_iteration_repro.py

"""Reproduce PR #343's timer-error busy loop with real set nondeterminism.

Run from the repository root with:

    uv run python pr343_set_iteration_repro.py

The parent process starts the same workflow in two child processes backed by
the same LocalWorld data directory.  PYTHONHASHSEED=3 selects the hook branch
and records a real hook suspension; PYTHONHASHSEED=1 selects asyncio.sleep()
when that run is replayed.  The resulting hook/wait divergence completes the
timer's workflow Future with NondeterminismError.  PR #343 reports that error
but leaves the asyncio.sleep() task pending, so the child process busy-loops.
The parent kills it after a short timeout instead of letting it spin forever.
"""

from __future__ import annotations

import asyncio
import dataclasses
import importlib
import os
import subprocess
import sys
import tempfile
from pathlib import Path

from vercel.workflow import BaseHook, Workflows, start
from vercel.workflow._internal import runtime, world as w
from vercel.workflow._internal.worlds import local as local_mod

TOKEN = "pr343-set-iteration-repro"
TERMINAL_STATUSES = {"completed", "failed", "cancelled"}
app = Workflows(as_vercel_job=False)


@dataclasses.dataclass
class Signal(BaseHook):
    value: str


@app.workflow
async def nondeterministic_workflow() -> None:
    operation = next(iter({"hook", "sleep"}))
    print(
        f"workflow selected {operation!r} with PYTHONHASHSEED={os.environ['PYTHONHASHSEED']}",
        flush=True,
    )

    if operation == "hook":
        await Signal.wait(token=TOKEN)
    else:
        await asyncio.sleep(3600)


async def open_world(data_dir: str) -> local_mod.LocalWorld:
    os.environ["WORKFLOW_LOCAL_DATA_DIR"] = data_dir
    world = local_mod.LocalWorld()
    w.set_world(world)
    runtime.workflow_entrypoint(app)
    await world._get_queue_client()
    return world


async def record_hook(data_dir: str) -> None:
    world = await open_world(data_dir)
    try:
        run = await start(nondeterministic_workflow)
        for _ in range(200):
            try:
                hook = await world.hooks_get_by_token(TOKEN)
            except w.HookNotFoundError:
                await asyncio.sleep(0.05)
                continue
            print(
                f"recorded {hook.hook_id!r} for run {run.run_id!r}",
                flush=True,
            )
            return
        raise RuntimeError("the workflow never recorded its hook suspension")
    finally:
        await world.aclose()
        w.set_world(None)


async def replay_as_sleep(data_dir: str) -> None:
    world = await open_world(data_dir)
    try:
        hook = await Signal(value="resume").resume(TOKEN)
        print(f"queued replay for run {hook.run_id!r}", flush=True)

        # With the bug fixed, replay reaches a terminal failed state and this
        # loop exits.  With PR #343 as written, the embedded worker blocks in
        # WorkflowLoop.run_forever(), so even this outer poll cannot run again.
        while True:
            run = await world.runs_get(hook.run_id)
            if run.status in TERMINAL_STATUSES:
                print(f"run reached terminal status {run.status!r}", flush=True)
                return
            await asyncio.sleep(0.05)
    finally:
        await world.aclose()
        w.set_world(None)


def child_command(mode: str, data_dir: str, *, hash_seed: int) -> subprocess.CompletedProcess[str]:
    env = os.environ.copy()
    env["PYTHONHASHSEED"] = str(hash_seed)
    return subprocess.run(
        [sys.executable, "-u", str(Path(__file__).resolve()), mode, data_dir],
        env=env,
        text=True,
        capture_output=True,
        timeout=10 if mode == "record" else 3,
        check=False,
    )


def show_output(stdout: str | bytes | None, stderr: str | bytes | None) -> str:
    def as_text(value: str | bytes | None) -> str:
        if value is None:
            return ""
        if isinstance(value, bytes):
            return value.decode(errors="replace")
        return value

    output = as_text(stdout) + as_text(stderr)
    if output:
        print(output, end="" if output.endswith("\n") else "\n")
    return output


def parent() -> int:
    with tempfile.TemporaryDirectory(prefix="pr343-repro-") as data_dir:
        print("[1/2] Recording a real hook suspension with PYTHONHASHSEED=3")
        recorded = child_command("record", data_dir, hash_seed=3)
        show_output(recorded.stdout, recorded.stderr)
        if recorded.returncode != 0:
            print(f"recording failed with exit code {recorded.returncode}")
            return 1

        print("[2/2] Replaying the same run as asyncio.sleep() with PYTHONHASHSEED=1")
        try:
            replayed = child_command("replay", data_dir, hash_seed=1)
        except subprocess.TimeoutExpired as exc:
            output = show_output(exc.stdout, exc.stderr)
            if "NondeterminismError" not in output:
                print("child timed out, but no NondeterminismError was observed")
                return 1
            print(
                "REPRODUCED: NondeterminismError was reported, but replay did not "
                "terminate and had to be killed."
            )
            return 0

        show_output(replayed.stdout, replayed.stderr)
        print(
            "NOT REPRODUCED: replay exited instead of busy-looping "
            f"(exit code {replayed.returncode})."
        )
        return 1


def main() -> int:
    if len(sys.argv) == 1:
        return parent()
    if len(sys.argv) != 3 or sys.argv[1] not in {"record", "replay"}:
        print(f"usage: {sys.argv[0]} [record|replay DATA_DIR]", file=sys.stderr)
        return 2

    operation, data_dir = sys.argv[1:]
    asyncio.run(record_hook(data_dir) if operation == "record" else replay_as_sleep(data_dir))
    return 0


if __name__ == "__main__":
    # Import under a stable module name so the workflow sandbox can re-import
    # its defining module during replay instead of referring to ``__main__``.
    module = importlib.import_module(Path(__file__).stem)
    raise SystemExit(module.main())