diff --git a/pyproject.toml b/pyproject.toml index dd13f113..5bfdcbcf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}] [project] name = "pyworkflow-engine" -version = "0.3.6" +version = "0.3.7" description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow" readme = "README.md" requires-python = ">=3.11" diff --git a/pyworkflow/__init__.py b/pyworkflow/__init__.py index 62d4fceb..1af083bb 100644 --- a/pyworkflow/__init__.py +++ b/pyworkflow/__init__.py @@ -29,7 +29,7 @@ >>> run_id = await start(my_workflow, "Alice") """ -__version__ = "0.3.6" +__version__ = "0.3.7" # Configuration from pyworkflow.config import ( @@ -140,7 +140,11 @@ load_step_checkpoint, save_step_checkpoint, ) -from pyworkflow.primitives.step_hook import step_hook +from pyworkflow.primitives.step_hook import ( + STEP_HOOK_TIMEOUT, + StepHookTimeout, + step_hook, +) # Runtime from pyworkflow.runtime import LocalRuntime, Runtime, get_runtime, register_runtime @@ -290,6 +294,8 @@ "load_step_checkpoint", "delete_step_checkpoint", "step_hook", + "StepHookTimeout", + "STEP_HOOK_TIMEOUT", # Streams "stream_workflow", "stream_step", diff --git a/pyworkflow/celery/singleton.py b/pyworkflow/celery/singleton.py index 8eb3e940..8aefd891 100644 --- a/pyworkflow/celery/singleton.py +++ b/pyworkflow/celery/singleton.py @@ -323,7 +323,11 @@ def generate_lock( # Bind arguments to function signature sig = inspect.signature(self.run) - bound = sig.bind(*task_args, **task_kwargs).arguments + bound_sig = sig.bind(*task_args, **task_kwargs) + # Include declared defaults so unique_on may reference parameters + # the caller did not pass explicitly (e.g. dedup_scope). + bound_sig.apply_defaults() + bound = bound_sig.arguments unique_args: list[Any] = [] for key in unique_on: diff --git a/pyworkflow/celery/tasks.py b/pyworkflow/celery/tasks.py index 5271edab..dcc7706d 100644 --- a/pyworkflow/celery/tasks.py +++ b/pyworkflow/celery/tasks.py @@ -371,6 +371,19 @@ def execute_step_task( hook_id=hook_id or "", ) ) + # step_hook(on_timeout="return") carries the hook deadline: schedule + # a resume at expiry so the workflow wakes even if nobody ever calls + # resume_hook (the step re-executes and receives STEP_HOOK_TIMEOUT). + # A racing resume_hook is harmless — duplicate resumes are dropped + # by try_claim_run. + resume_at = e.data.get("resume_at") + if resume_at is not None: + schedule_workflow_resumption( + run_id, + resume_at, + storage_config, + triggered_by="step_hook_timeout", + ) return None # Other SuspensionSignals are not supported from steps @@ -2025,12 +2038,13 @@ async def _start_workflow_on_worker( name="pyworkflow.resume_workflow", base=SingletonWorkflowTask, queue="pyworkflow.schedules", - unique_on=["run_id"], + unique_on=["run_id", "dedup_scope"], ) def resume_workflow_task( run_id: str, storage_config: dict[str, Any] | None = None, triggered_by_hook_id: str | None = None, + dedup_scope: str = "immediate", ) -> Any | None: """ Resume a suspended workflow. @@ -2703,10 +2717,15 @@ def schedule_workflow_resumption( triggered_by=triggered_by, ) - # Schedule the resume task + # Schedule the resume task. Delayed (deadline) resumes hold their + # singleton lock for the whole countdown; scoping them separately keeps + # them from swallowing immediate resumes (e.g. resume_hook wake-ups) + # enqueued while the countdown is pending. Extra resume executions are + # harmless: try_claim_run drops all but the first. + dedup_scope = "deadline" if delay_seconds > 0 else "immediate" resume_workflow_task.apply_async( args=[run_id], - kwargs={"storage_config": storage_config}, + kwargs={"storage_config": storage_config, "dedup_scope": dedup_scope}, countdown=delay_seconds, ) diff --git a/pyworkflow/primitives/step_hook.py b/pyworkflow/primitives/step_hook.py index 3459fdf2..06cda937 100644 --- a/pyworkflow/primitives/step_hook.py +++ b/pyworkflow/primitives/step_hook.py @@ -25,7 +25,8 @@ async def agent_step(): """ from collections.abc import Awaitable, Callable -from typing import Any +from datetime import UTC, datetime, timedelta +from typing import Any, Literal from loguru import logger from pydantic import BaseModel @@ -35,12 +36,29 @@ async def agent_step(): from pyworkflow.primitives.step_checkpoint import get_step_run_id +class StepHookTimeout: + """Sentinel returned by ``step_hook(on_timeout="return")`` when the hook expires. + + Check with ``isinstance(result, StepHookTimeout)`` or compare against the + ``STEP_HOOK_TIMEOUT`` singleton. + """ + + __slots__ = () + + def __repr__(self) -> str: + return "STEP_HOOK_TIMEOUT" + + +STEP_HOOK_TIMEOUT = StepHookTimeout() + + async def step_hook( name: str, *, timeout: str | int | None = None, on_created: Callable[[str], Awaitable[None]] | None = None, payload_schema: type[BaseModel] | None = None, + on_timeout: Literal["suspend", "return"] = "suspend", ) -> Any: """ Wait for an external event from within a @step function. @@ -58,9 +76,14 @@ async def step_hook( timeout: Optional max wait time (str duration or seconds) on_created: Optional async callback with the hook token payload_schema: Optional Pydantic model for payload validation + on_timeout: What to do when ``timeout`` elapses without a resume. + "suspend" (default): keep waiting for resume_hook() — legacy behavior. + "return": the runtime schedules a resume at the deadline and this + call returns the ``STEP_HOOK_TIMEOUT`` sentinel on re-execution. Returns: - Payload from resume_hook() + Payload from resume_hook(), or ``STEP_HOOK_TIMEOUT`` when the hook + expired and ``on_timeout="return"`` Raises: RuntimeError: If called outside a step context @@ -107,11 +130,11 @@ async def notify(token): events = await storage.get_events(ctx.run_id) hook_received = None - hook_created = False + hook_created_event = None for event in events: if event.type == EventType.HOOK_CREATED and event.data.get("hook_id") == hook_id: - hook_created = True + hook_created_event = event elif event.type == EventType.HOOK_RECEIVED and event.data.get("hook_id") == hook_id: hook_received = event @@ -127,17 +150,53 @@ async def notify(token): ) return payload - # If hook was already created but not received, re-suspend - if hook_created: + # If hook was already created but not received, check expiry, else re-suspend + if hook_created_event is not None: + expires_at: datetime | None = None + expires_at_raw = hook_created_event.data.get("expires_at") + if expires_at_raw: + expires_at = datetime.fromisoformat(expires_at_raw) + + if on_timeout == "return" and expires_at is not None and datetime.now(UTC) >= expires_at: + from pyworkflow.engine.events import create_hook_expired_event + from pyworkflow.storage.schemas import HookStatus + + await storage.record_event( + create_hook_expired_event(run_id=ctx.run_id, hook_id=hook_id) + ) + try: + await storage.update_hook_status(hook_id, HookStatus.EXPIRED) + except Exception: + # Best-effort: the HOOK_EXPIRED event is authoritative for replay; + # a failed status update must not fail the step. + logger.warning( + f"Step hook '{name}' expired but status update failed", + run_id=ctx.run_id, + hook_id=hook_id, + ) + logger.info( + f"Step hook '{name}' expired, returning timeout sentinel", + run_id=ctx.run_id, + hook_id=hook_id, + ) + return STEP_HOOK_TIMEOUT + logger.debug( f"Step hook '{name}' already created, re-suspending", run_id=ctx.run_id, hook_id=hook_id, ) + # Re-arm the deadline resume only for on_timeout="return" with a future + # deadline: the runtime schedules a resume at resume_at, and scheduling + # one in the past would busy-loop resume → re-suspend. + resume_data: dict[str, Any] = {} + if on_timeout == "return" and expires_at is not None and datetime.now(UTC) < expires_at: + resume_data["resume_at"] = expires_at raise SuspensionSignal( reason=f"step_hook:{hook_id}", hook_id=hook_id, step_id=step_run_id, + **resume_data, ) # Parse timeout @@ -179,8 +238,6 @@ async def notify(token): payload_schema=payload_schema.__name__ if payload_schema else None, ) if timeout_seconds: - from datetime import UTC, datetime, timedelta - hook_record.expires_at = datetime.now(UTC) + timedelta(seconds=timeout_seconds) await storage.create_hook(hook_record) @@ -197,9 +254,15 @@ async def notify(token): if on_created: await on_created(token) - # Raise SuspensionSignal to suspend the step + # Raise SuspensionSignal to suspend the step. For on_timeout="return", + # carry the deadline so the runtime schedules a resume at expiry — without + # it, an expired hook nobody resumes would suspend the workflow forever. + suspend_data: dict[str, Any] = {} + if on_timeout == "return" and hook_record.expires_at is not None: + suspend_data["resume_at"] = hook_record.expires_at raise SuspensionSignal( reason=f"step_hook:{hook_id}", hook_id=hook_id, step_id=step_run_id, + **suspend_data, ) diff --git a/tests/unit/test_singleton.py b/tests/unit/test_singleton.py index c19fc348..3f0f629d 100644 --- a/tests/unit/test_singleton.py +++ b/tests/unit/test_singleton.py @@ -1007,3 +1007,41 @@ def run(self, run_id, step_id, args_json=None, storage_config=None): key_prefix="pyworkflow:lock:", ) backend.unlock.assert_called_once_with(expected_key) + + +class TestResumeDedupScope: + """The resume task's singleton key must separate immediate vs deadline resumes. + + A deadline (countdown) resume holds its singleton lock for the whole + countdown; without scope separation it swallows every immediate resume + (e.g. resume_hook wake-ups) enqueued meanwhile. + """ + + def test_lock_key_differs_by_dedup_scope(self): + from pyworkflow.celery.tasks import resume_workflow_task + + immediate = resume_workflow_task.generate_lock( + resume_workflow_task.name, + ["run_x"], + {"storage_config": None, "dedup_scope": "immediate"}, + ) + deadline = resume_workflow_task.generate_lock( + resume_workflow_task.name, + ["run_x"], + {"storage_config": None, "dedup_scope": "deadline"}, + ) + assert immediate != deadline + + def test_lock_key_defaults_to_immediate_scope(self): + """Callers not passing dedup_scope must match explicit 'immediate'.""" + from pyworkflow.celery.tasks import resume_workflow_task + + default_key = resume_workflow_task.generate_lock( + resume_workflow_task.name, ["run_x"], {"storage_config": None} + ) + explicit_key = resume_workflow_task.generate_lock( + resume_workflow_task.name, + ["run_x"], + {"storage_config": None, "dedup_scope": "immediate"}, + ) + assert default_key == explicit_key diff --git a/tests/unit/test_step_hook.py b/tests/unit/test_step_hook.py index 9bdd7c51..60abfb17 100644 --- a/tests/unit/test_step_hook.py +++ b/tests/unit/test_step_hook.py @@ -200,3 +200,201 @@ async def test_step_hook_deterministic_ids(self): finally: reset_step_execution_context(step_tokens) reset_context(ctx_token) + + +class TestStepHookTimeout: + """Tests for step_hook() timeout semantics (on_timeout="return").""" + + def _make_ctx(self, storage, run_id): + ctx = LocalContext( + run_id=run_id, workflow_name="test_workflow", storage=storage, durable=True + ) + ctx._is_step_worker = True + return ctx + + @pytest.mark.asyncio + async def test_first_call_with_timeout_return_carries_resume_at(self): + """First call with on_timeout="return" suspends with the deadline in signal data.""" + from datetime import UTC, datetime + + storage = InMemoryStorageBackend() + run_id = "test_run_to_1" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + before = datetime.now(UTC) + with pytest.raises(SuspensionSignal) as exc_info: + await step_hook("tick", timeout=60, on_timeout="return") + resume_at = exc_info.value.data.get("resume_at") + assert resume_at is not None + delta = (resume_at - before).total_seconds() + assert 55 <= delta <= 65 + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + @pytest.mark.asyncio + async def test_first_call_default_mode_has_no_resume_at(self): + """Default on_timeout="suspend" keeps legacy behavior: no deadline resume.""" + storage = InMemoryStorageBackend() + run_id = "test_run_to_2" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + with pytest.raises(SuspensionSignal) as exc_info: + await step_hook("tick", timeout=60) + assert exc_info.value.data.get("resume_at") is None + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + @pytest.mark.asyncio + async def test_expired_hook_returns_sentinel(self): + """Re-execution after the deadline returns STEP_HOOK_TIMEOUT (on_timeout="return").""" + from datetime import UTC, datetime, timedelta + + from pyworkflow.engine.events import create_hook_created_event + from pyworkflow.primitives.step_hook import STEP_HOOK_TIMEOUT, StepHookTimeout + + storage = InMemoryStorageBackend() + run_id = "test_run_to_3" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + # Simulate a prior execution that created the hook with a now-past deadline + await storage.record_event( + create_hook_created_event( + run_id=run_id, + hook_id="step_hook_tick_0", + token="tok", + expires_at=datetime.now(UTC) - timedelta(seconds=5), + name="tick", + ) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + result = await step_hook("tick", timeout=60, on_timeout="return") + assert result is STEP_HOOK_TIMEOUT + assert isinstance(result, StepHookTimeout) + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + @pytest.mark.asyncio + async def test_expired_hook_default_mode_resuspends(self): + """Legacy mode keeps waiting: expired hook re-suspends without resume_at.""" + from datetime import UTC, datetime, timedelta + + from pyworkflow.engine.events import create_hook_created_event + + storage = InMemoryStorageBackend() + run_id = "test_run_to_4" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + await storage.record_event( + create_hook_created_event( + run_id=run_id, + hook_id="step_hook_tick_0", + token="tok", + expires_at=datetime.now(UTC) - timedelta(seconds=5), + name="tick", + ) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + with pytest.raises(SuspensionSignal) as exc_info: + await step_hook("tick", timeout=60) + assert exc_info.value.data.get("resume_at") is None + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + @pytest.mark.asyncio + async def test_received_payload_wins_over_expiry(self): + """A payload received before re-execution is returned even past the deadline.""" + from datetime import UTC, datetime, timedelta + + from pyworkflow.engine.events import ( + create_hook_created_event, + create_hook_received_event, + ) + from pyworkflow.serialization.encoder import serialize + + storage = InMemoryStorageBackend() + run_id = "test_run_to_5" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + await storage.record_event( + create_hook_created_event( + run_id=run_id, + hook_id="step_hook_tick_0", + token="tok", + expires_at=datetime.now(UTC) - timedelta(seconds=5), + name="tick", + ) + ) + await storage.record_event( + create_hook_received_event( + run_id=run_id, + hook_id="step_hook_tick_0", + payload=serialize({"index": 3}), + ) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + result = await step_hook("tick", timeout=60, on_timeout="return") + assert result == {"index": 3} + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + @pytest.mark.asyncio + async def test_unexpired_hook_rearms_resume_at(self): + """Re-execution before the deadline re-suspends carrying the original deadline.""" + from datetime import UTC, datetime, timedelta + + from pyworkflow.engine.events import create_hook_created_event + + storage = InMemoryStorageBackend() + run_id = "test_run_to_6" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + deadline = datetime.now(UTC) + timedelta(seconds=120) + await storage.record_event( + create_hook_created_event( + run_id=run_id, + hook_id="step_hook_tick_0", + token="tok", + expires_at=deadline, + name="tick", + ) + ) + ctx = self._make_ctx(storage, run_id) + ctx_token = set_context(ctx) + step_tokens = set_step_execution_context(f"{run_id}:step_test_abc123", storage) + try: + with pytest.raises(SuspensionSignal) as exc_info: + await step_hook("tick", timeout=60, on_timeout="return") + resume_at = exc_info.value.data.get("resume_at") + assert resume_at is not None + assert abs((resume_at - deadline).total_seconds()) < 1 + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token)