pyworkflow-engine 0.3.4__tar.gz → 0.3.7__tar.gz

{pyworkflow_engine-0.3.4 → pyworkflow_engine-0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.3.4
+Version: 0.3.7
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.3.4 → pyworkflow_engine-0.3.7}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.3.4"
+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_engine-0.3.4 → pyworkflow_engine-0.3.7}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.3.4"
+__version__ = "0.3.7"
 
 # Configuration
 from pyworkflow.config import (
@@ -140,7 +140,11 @@ from pyworkflow.primitives.step_checkpoint import (
     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 @@ __all__ = [
     "load_step_checkpoint",
     "delete_step_checkpoint",
     "step_hook",
+    "StepHookTimeout",
+    "STEP_HOOK_TIMEOUT",
     # Streams
     "stream_workflow",
     "stream_step",

{pyworkflow_engine-0.3.4 → pyworkflow_engine-0.3.7}/pyworkflow/celery/app.py

@@ -12,12 +12,19 @@ warnings in Celery logs. This is a known compatibility issue between Python 3.13
 garbage collector and Celery's saferepr module. It does not affect functionality.
 """
 
+import contextlib
 import os
 from datetime import timedelta
 from typing import Any
 
 from celery import Celery
-from celery.signals import worker_init, worker_process_init, worker_shutdown
+from celery.signals import (
+    task_postrun,
+    task_prerun,
+    worker_init,
+    worker_process_init,
+    worker_shutdown,
+)
 from kombu import Exchange, Queue
 
 from pyworkflow.observability.logging import configure_logging
@@ -281,6 +288,9 @@ def create_celery_app(
         PYWORKFLOW_CELERY_SENTINEL_MASTER: Sentinel master name (used if sentinel_master_name param not provided)
         PYWORKFLOW_WORKER_MAX_MEMORY: Max memory per worker child (KB) (used if worker_max_memory_per_child param not provided)
         PYWORKFLOW_WORKER_MAX_TASKS: Max tasks per worker child (used if worker_max_tasks_per_child param not provided)
+        PYWORKFLOW_RESCHEDULE_ON_SIGTERM: Re-enqueue in-flight tasks immediately when a worker is shut down
+            (SIGTERM / spot reclaim), instead of waiting out the broker visibility_timeout. Default on; set to
+            "0"/"false"/"no" to disable. Implemented as a consumer bootstep; see pyworkflow/celery/reschedule.py.
 
     Examples:
         # Default configuration (uses env vars if set, otherwise localhost Redis)
@@ -471,6 +481,14 @@ def create_celery_app(
     # to ensure proper initialization AFTER process forking.
     # See on_worker_init() and on_worker_process_init() below.
 
+    # Re-enqueue in-flight tasks on worker shutdown (SIGTERM / spot reclaim) so a
+    # live worker resumes them within seconds instead of waiting out the broker
+    # visibility_timeout. Hooked as a consumer bootstep so it fires exactly as the
+    # consumer is torn down (event-based). See pyworkflow/celery/reschedule.py.
+    from pyworkflow.celery.reschedule import RescheduleConsumerStep
+
+    app.steps["consumer"].add(RescheduleConsumerStep)
+
     # Auto-discover workflows from environment variable or configured modules
     discover_workflows()
 
@@ -555,6 +573,29 @@ def on_worker_shutdown(**kwargs):
         close_worker_loop()
 
 
+@task_prerun.connect
+def on_task_prerun(task_id=None, task=None, args=None, kwargs=None, **_):
+    """
+    Track the task now executing in this child process so it can be re-enqueued
+    if the worker is shut down (SIGTERM) before the task finishes.
+
+    Wrapped defensively: tracking must never interfere with task execution.
+    """
+    from pyworkflow.celery.reschedule import track_task_start
+
+    with contextlib.suppress(Exception):
+        track_task_start(task, task_id, args, kwargs)
+
+
+@task_postrun.connect
+def on_task_postrun(task_id=None, **_):
+    """Clear the in-flight task once it has finished (clean completion)."""
+    from pyworkflow.celery.reschedule import track_task_end
+
+    with contextlib.suppress(Exception):
+        track_task_end(task_id)
+
+
 def get_celery_app() -> Celery:
     """
     Get the global Celery application instance.

pyworkflow_engine-0.3.7/pyworkflow/celery/reschedule.py

@@ -0,0 +1,339 @@
+"""
+Reschedule in-flight tasks when a worker is shut down (SIGTERM).
+
+When a Celery worker pod is reclaimed (e.g. an AWS spot node drain), it receives
+SIGTERM and warm-shuts-down. Any step task that was executing is *not* promptly
+redelivered: with ``task_acks_late=True`` the broker only re-queues it after the
+Redis ``visibility_timeout`` (an hour by default), and the parent workflow stays
+SUSPENDED that whole time -- so the user-facing flow silently hangs with no output.
+
+This module makes the worker re-enqueue its in-flight task(s) *immediately* on
+shutdown so a live worker resumes them within seconds instead of ~1h.
+
+Why a shared Redis registry (not process-local state)
+-----------------------------------------------------
+Tasks execute in forked **child** processes, but the shutdown hook runs in the
+**main** process. The main process can't see a child's in-memory task args, so
+children publish what they are running to a per-pod Redis hash
+(``pyworkflow:inflight:<pod>``) on ``task_prerun`` and remove it on
+``task_postrun``. On shutdown the main process reads that hash and re-enqueues
+every entry.
+
+Why a consumer bootstep (event-based, no sleep)
+-----------------------------------------------
+The re-enqueue must happen *after* this worker's broker consumer is cancelled --
+otherwise the worker's own still-active consumer grabs the message straight back
+into ``unacked`` (stranded until the visibility_timeout). There is no Celery
+*signal* for "consumer cancelled" (``worker_shutting_down`` fires before it;
+``worker_shutdown`` fires after the pool has drained the long task -- far too
+late). So we hook the consumer blueprint directly with a bootstep that
+``requires`` the ``Tasks`` step: its ``stop()`` runs as the consumer is torn
+down (early in shutdown, before the pool drain). There we cancel the task
+consumer ourselves -- guaranteeing this worker stops fetching -- and only then
+re-enqueue, so the message lands in the *ready* queue and another worker's
+blocking BRPOP picks it up immediately. Event-based, deterministic, no sleep.
+
+Safety against double execution
+-------------------------------
+``execute_step_task`` short-circuits on an existing ``STEP_COMPLETED`` event (and
+on a terminal run status), and ``_record_step_completion_and_resume`` re-checks
+before recording, so at most one completion is ever recorded no matter how many
+copies run. Only ``SingletonWorkflowTask`` tasks that declare ``unique_on`` are
+rescheduled -- those are idempotent by run_id/step_id. Tasks without a uniqueness
+key are left alone.
+
+Gate: ``PYWORKFLOW_RESCHEDULE_ON_SIGTERM`` (default on; set to ``0``/``false``/``no``
+to disable and fall back to the old visibility_timeout redelivery behaviour).
+"""
+
+import json
+import os
+import socket
+import threading
+from typing import Any
+
+from celery import bootsteps
+from loguru import logger
+
+# Registry hash self-expires so a hard-killed pod (no postrun) cannot leak entries.
+_REGISTRY_TTL_SECONDS = 6 * 3600
+_REGISTRY_KEY_PREFIX = "pyworkflow:inflight:"
+
+_backend_lock = threading.Lock()
+_backend: Any = None
+_backend_resolved = False
+
+
+def _enabled() -> bool:
+    return os.getenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "1").strip().lower() not in (
+        "0",
+        "false",
+        "no",
+    )
+
+
+def _pod_id() -> str:
+    """Stable per-pod id shared by the main process and its forked children."""
+    return os.environ.get("HOSTNAME") or socket.gethostname()
+
+
+def _registry_key() -> str:
+    return f"{_REGISTRY_KEY_PREFIX}{_pod_id()}"
+
+
+def _get_backend() -> Any:
+    """
+    Cached sentinel-aware Redis backend, reusing the singleton lock configuration
+    (same broker/sentinel the task locks already use). Returns None when no Redis
+    backend is configured, in which case rescheduling is a no-op.
+    """
+    global _backend, _backend_resolved
+    if _backend_resolved:
+        return _backend
+    with _backend_lock:
+        if _backend_resolved:
+            return _backend
+        try:
+            from pyworkflow.celery.app import celery_app
+            from pyworkflow.celery.singleton import RedisLockBackend, SingletonConfig
+
+            cfg = SingletonConfig(celery_app)
+            url = cfg.backend_url
+            if url:
+                _backend = RedisLockBackend(
+                    url,
+                    is_sentinel=cfg.is_sentinel,
+                    sentinel_master=cfg.sentinel_master,
+                )
+        except Exception as exc:
+            logger.warning(f"reschedule: could not initialise redis backend: {exc}")
+            _backend = None
+        _backend_resolved = True
+    return _backend
+
+
+def _resolve_task(name: str) -> Any:
+    """Look up a registered task instance by name (used to re-enqueue)."""
+    from pyworkflow.celery.app import celery_app
+
+    return celery_app.tasks.get(name)
+
+
+def _resolve_queue(task: Any) -> str | None:
+    """Best-effort queue the task was delivered on, so the re-enqueue lands there."""
+    try:
+        delivery_info = getattr(task.request, "delivery_info", None) or {}
+        routing_key = delivery_info.get("routing_key")
+        if routing_key:
+            return routing_key
+    except Exception:
+        pass
+    return getattr(task, "queue", None)
+
+
+def track_task_start(
+    task: Any,
+    task_id: str | None,
+    args: Any,
+    kwargs: Any,
+) -> None:
+    """Publish a task as in-flight (per pod) so it can be re-enqueued on SIGTERM."""
+    if not _enabled() or task is None or task_id is None:
+        return
+
+    # Import here to avoid import cycles at module load.
+    from pyworkflow.celery.singleton import SingletonWorkflowTask
+
+    # Only re-enqueue tasks keyed by a uniqueness field (run_id/step_id). These are
+    # idempotent; blindly re-enqueuing a non-unique task could start duplicate runs.
+    if not isinstance(task, SingletonWorkflowTask) or not task.unique_on:
+        return
+
+    backend = _get_backend()
+    if backend is None:
+        return
+
+    descriptor = json.dumps(
+        {
+            "name": task.name,
+            "args": list(args or ()),
+            "kwargs": dict(kwargs or {}),
+            "queue": _resolve_queue(task),
+        },
+        default=str,
+    )
+    key = _registry_key()
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.hset(key, task_id, descriptor))
+        backend._execute_with_refresh(lambda: backend.redis.expire(key, _REGISTRY_TTL_SECONDS))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to record in-flight task {task_id}: {exc}")
+
+
+def track_task_end(task_id: str | None) -> None:
+    """Remove a task from the in-flight registry once it has finished."""
+    if not _enabled() or task_id is None:
+        return
+    backend = _get_backend()
+    if backend is None:
+        return
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.hdel(_registry_key(), task_id))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to clear in-flight task {task_id}: {exc}")
+
+
+def reschedule_inflight_on_shutdown() -> None:
+    """
+    Re-enqueue every task still recorded as in-flight for this pod. Assumes the
+    caller has already cancelled this worker's task consumer (see
+    :class:`RescheduleConsumerStep`) so the fresh copy is not re-grabbed. Never
+    raises -- shutdown must not be blocked.
+    """
+    if not _enabled():
+        return
+    backend = _get_backend()
+    if backend is None:
+        return
+
+    key = _registry_key()
+    try:
+        entries = backend._execute_with_refresh(lambda: backend.redis.hgetall(key)) or {}
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to read in-flight registry: {exc}")
+        return
+
+    if not entries:
+        return
+
+    for task_id, raw in entries.items():
+        try:
+            descriptor = json.loads(raw)
+            task = _resolve_task(descriptor["name"])
+            if task is None:
+                logger.warning(
+                    f"reschedule: unknown task {descriptor.get('name')!r}, "
+                    f"cannot re-enqueue {task_id}"
+                )
+                continue
+
+            args = descriptor.get("args") or []
+            kwargs = descriptor.get("kwargs") or {}
+            queue = descriptor.get("queue")
+
+            # Release the singleton lock held by the dying task so the re-enqueued
+            # copy is not rejected as a duplicate. release_lock() regenerates the
+            # exact lock key from the same args/kwargs the running task holds.
+            try:
+                task.release_lock(task_args=list(args), task_kwargs=kwargs)
+            except Exception as exc:
+                logger.warning(
+                    f"reschedule: failed to release lock for {task.name} "
+                    f"(old_task_id={task_id}): {exc}"
+                )
+
+            options = {"queue": queue} if queue else {}
+            task.apply_async(args=list(args), kwargs=kwargs, **options)
+            logger.warning(
+                f"reschedule: re-enqueued in-flight task on shutdown "
+                f"task={task.name} old_task_id={task_id} queue={queue}"
+            )
+        except Exception as exc:
+            logger.opt(exception=True).error(
+                f"reschedule: failed to re-enqueue task {task_id}: {exc}"
+            )
+
+    # Drop the registry for this pod; a fresh worker starts with a clean slate.
+    try:
+        backend._execute_with_refresh(lambda: backend.redis.delete(key))
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to clear in-flight registry: {exc}")
+
+
+def _is_worker_stopping() -> bool:
+    """
+    True only when the worker is actually shutting down (SIGTERM warm or cold),
+    not on a transient consumer restart (e.g. broker reconnect), which also tears
+    the consumer down. Celery's SIGTERM handler sets these flags before the
+    consumer is stopped.
+    """
+    try:
+        from celery.worker import state
+
+        # Identity checks (matching celery.worker.state): the flags are set to the
+        # exit code, which can be 0 (EX_OK) -- `0 == False`, so membership/equality
+        # against False would wrongly treat a clean exitcode-0 shutdown as "not
+        # stopping" and skip the reschedule.
+        def _set(v: Any) -> bool:
+            return v is not None and v is not False
+
+        return _set(state.should_stop) or _set(state.should_terminate)
+    except Exception:
+        return False
+
+
+def on_consumer_stop(consumer: Any) -> None:
+    """
+    Invoked from the consumer bootstep as the consumer is torn down. Cancels this
+    worker's task consumer (so it stops fetching) and re-enqueues in-flight work.
+    """
+    if not _enabled() or not _is_worker_stopping():
+        return
+
+    # Cancel our own task consumer first: once we stop BRPOPing, the re-enqueued
+    # message stays in the ready list for another worker instead of being grabbed
+    # back here into `unacked`. Tasks.stop() will cancel again later (idempotent).
+    task_consumer = getattr(consumer, "task_consumer", None)
+    if task_consumer is not None:
+        try:
+            task_consumer.cancel()
+        except Exception as exc:
+            logger.warning(f"reschedule: failed to cancel task consumer: {exc}")
+
+    # Cancelling is not enough on the Redis transport: a BRPOP this worker already
+    # issued is accepted by the server and will still pop the *next* message pushed
+    # to those queues -- including the copy we are about to re-enqueue. The dying
+    # worker then exits before recording it in `unacked`, so the message is lost
+    # entirely (neither in the ready list nor restorable via visibility_timeout).
+    # Drop the consumer's broker connection so any in-flight BRPOP is aborted
+    # server-side *before* we re-publish; the re-enqueue itself uses the producer
+    # pool (a separate connection), so it is unaffected.
+    _abort_consumer_fetching(consumer)
+
+    reschedule_inflight_on_shutdown()
+
+
+def _abort_consumer_fetching(consumer: Any) -> None:
+    """Close the consumer's broker connection to abort any in-flight BRPOP."""
+    conn = getattr(consumer, "connection", None)
+    if conn is None:
+        return
+    try:
+        # collect() tears down channels + the underlying socket, aborting the
+        # pending BRPOP. Never raises into shutdown.
+        conn.collect()
+    except Exception as exc:
+        logger.warning(f"reschedule: failed to drop consumer connection: {exc}")
+
+
+class RescheduleConsumerStep(bootsteps.StartStopStep):
+    """
+    Consumer bootstep that re-enqueues this worker's in-flight tasks the moment
+    the consumer is torn down on shutdown -- event-based, no polling/sleep.
+
+    ``requires`` the task consumer (``Tasks``) so that (a) ``c.task_consumer``
+    exists when our ``stop()`` runs and (b) we run *before* ``Tasks.stop()``,
+    letting us cancel the consumer ourselves and re-enqueue while it's guaranteed
+    not fetching.
+    """
+
+    requires = ("celery.worker.consumer:Tasks",)
+
+    def start(self, c: Any) -> None:  # noqa: D401 - nothing to start
+        pass
+
+    def stop(self, c: Any) -> None:
+        try:
+            on_consumer_stop(c)
+        except Exception as exc:  # never block consumer teardown
+            logger.warning(f"reschedule: consumer-stop hook failed: {exc}")

{pyworkflow_engine-0.3.4 → pyworkflow_engine-0.3.7}/pyworkflow/celery/singleton.py

@@ -323,7 +323,11 @@ class SingletonWorkflowTask(Task):
 
             # 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_engine-0.3.4 → pyworkflow_engine-0.3.7}/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_engine-0.3.4 → pyworkflow_engine-0.3.7}/pyworkflow/primitives/step_hook.py

@@ -25,7 +25,8 @@ Usage:
 """
 
 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 @@ from pyworkflow.core.exceptions import SuspensionSignal
 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 step_hook(
 
     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 step_hook(
         )
         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 step_hook(
         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 step_hook(
     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,
     )