feat/step-hook-timeout-resumption

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"

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",

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:

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,
     )
 

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,
     )

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