fix: record STEP_SUSPENDED for inline step_hook suspensions (stranded resume)

pyworkflow/celery/tasks.py

@@ -45,6 +45,10 @@
     create_workflow_started_event,
     create_workflow_suspended_event,
 )
+from pyworkflow.engine.executor import (
+    record_step_suspended_if_needed,
+    step_suspended_already_recorded,
+)
 from pyworkflow.serialization.decoder import deserialize_args, deserialize_kwargs
 from pyworkflow.serialization.encoder import serialize_args, serialize_kwargs
 from pyworkflow.storage.base import StorageBackend
@@ -267,7 +271,9 @@ def execute_step_task(
             set_step_execution_context,
         )
 
-        step_exec_tokens = set_step_execution_context(step_exec_key, storage)
+        step_exec_tokens = set_step_execution_context(
+            step_exec_key, storage, step_id=step_id, step_name=step_name
+        )
     except Exception as e:
         logger.warning(f"Failed to set up step execution context: {e}")
 
@@ -658,13 +664,24 @@ async def _record_step_suspended(
 
     Does NOT schedule workflow resumption — that happens when resume_hook()
     is called externally.
+
+    Deduplicated per suspension round: if a STEP_SUSPENDED already exists for
+    this step *since its most recent STEP_STARTED* (e.g. the engine recorded it
+    inline on a force_local step, or a prior worker attempt recorded it), this is
+    a no-op. A step that suspends again on the same step_id in a later round
+    still records its own STEP_SUSPENDED.
     """
     from pyworkflow.engine.events import create_step_suspended_event
 
     storage = _get_storage_backend(storage_config)
     if hasattr(storage, "connect"):
         await storage.connect()
 
+    # Short-circuit only if already recorded for the current start (engine inline
+    # path or prior attempt) — not for an earlier round on the same step_id.
+    if await step_suspended_already_recorded(storage, run_id, step_id):
+        return
+
     # Wait for WORKFLOW_SUSPENDED event to avoid sequence number race
     max_wait_attempts = 50
     wait_interval = 0.01
@@ -1132,6 +1149,9 @@ async def _execute_child_workflow_on_worker(
         )
         await storage.record_event(suspended_event)
 
+        # Record STEP_SUSPENDED for an inline (force_local) step_hook suspension.
+        await record_step_suspended_if_needed(storage, child_run_id, e)
+
         logger.debug(
             f"Child workflow suspended: {workflow_name}",
             parent_run_id=parent_run_id,
@@ -1532,6 +1552,9 @@ async def _recover_workflow_on_worker(
         )
         await storage.record_event(suspended_event)
 
+        # Record STEP_SUSPENDED for an inline (force_local) step_hook suspension.
+        await record_step_suspended_if_needed(storage, run_id, e)
+
         logger.info(
             f"Recovered workflow suspended: {e.reason}",
             run_id=run_id,
@@ -1925,6 +1948,9 @@ async def _start_workflow_on_worker(
         )
         await storage.record_event(suspended_event)
 
+        # Record STEP_SUSPENDED for an inline (force_local) step_hook suspension.
+        await record_step_suspended_if_needed(storage, run_id, e)
+
         logger.info(
             f"Workflow suspended on worker: {e.reason}",
             run_id=run_id,
@@ -1959,7 +1985,8 @@ async def _start_workflow_on_worker(
         # resume_hook() may have recorded HOOK_RECEIVED and scheduled a resume task
         # while the workflow was still running (before status was set to SUSPENDED).
         # That resume task would have failed try_claim_run and been discarded.
-        if hook_id and e.reason.startswith("hook:"):
+        # Covers both workflow-level hook() ("hook:") and step_hook() suspensions.
+        if hook_id and (e.reason.startswith("hook:") or e.reason.startswith("step_hook:")):
             hook_received = await storage.has_event(
                 run_id, EventType.HOOK_RECEIVED.value, hook_id=hook_id
             )
@@ -2544,6 +2571,9 @@ async def _resume_workflow_on_worker(
         )
         await storage.record_event(suspended_event)
 
+        # Record STEP_SUSPENDED for an inline (force_local) step_hook suspension.
+        await record_step_suspended_if_needed(storage, run_id, e)
+
         logger.info(
             f"Workflow suspended again on worker: {e.reason}",
             run_id=run_id,
@@ -2576,8 +2606,9 @@ async def _resume_workflow_on_worker(
                 )
                 return None
 
-        # For hook suspensions, check if hook was already received (race condition)
-        if hook_id and e.reason.startswith("hook:"):
+        # For hook suspensions, check if hook was already received (race condition).
+        # Covers both workflow-level hook() ("hook:") and step_hook() suspensions.
+        if hook_id and (e.reason.startswith("hook:") or e.reason.startswith("step_hook:")):
             hook_received = await storage.has_event(
                 run_id, EventType.HOOK_RECEIVED.value, hook_id=hook_id
             )

pyworkflow/context/base.py

@@ -136,6 +136,8 @@ def __init__(
         self._tracing: dict[str, Any] | None = None
         self._tracing_provider: Any = None  # TracingProvider instance
         self._trace_id: str | None = None  # Stable trace ID
+        # Monotonic counter backing deterministic step_hook() ids within this run.
+        self._step_hook_counter: int = 0
 
     @property
     def run_id(self) -> str:
@@ -161,6 +163,39 @@ def tracing_provider(self) -> Any:
         """Get the active TracingProvider instance, or None."""
         return self._tracing_provider
 
+    # =========================================================================
+    # Step-hook counter (public API for cross-process checkpoint/restore)
+    # =========================================================================
+
+    def get_step_hook_counter(self) -> int:
+        """
+        Get the current step_hook() counter for this run.
+
+        step_hook() derives deterministic hook ids as
+        ``f"step_hook_{name}_{counter}"`` from a monotonic per-run counter.
+        Apps that checkpoint a step and restore it in a fresh process (so the
+        step body re-executes from the top) must restore this counter before
+        re-execution so the same step_hook() call reproduces the same hook id
+        and matches the recorded HOOK_CREATED/HOOK_RECEIVED events.
+
+        Returns:
+            The number of step_hook() calls made so far in this run.
+        """
+        return getattr(self, "_step_hook_counter", 0)
+
+    def set_step_hook_counter(self, value: int) -> None:
+        """
+        Set the step_hook() counter for this run.
+
+        Use together with :meth:`get_step_hook_counter` to restore the counter
+        when re-executing a checkpointed step across a process boundary. See
+        :meth:`get_step_hook_counter` for the full rationale.
+
+        Args:
+            value: The counter value to restore.
+        """
+        self._step_hook_counter = value
+
     # =========================================================================
     # Abstract methods - must be implemented by subclasses
     # =========================================================================

pyworkflow/core/step.py

@@ -274,7 +274,9 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
 
             # Set up step execution context for checkpoint/hook primitives
             step_exec_key = f"{ctx.run_id}:{step_id}"
-            step_exec_tokens = set_step_execution_context(step_exec_key, ctx.storage)
+            step_exec_tokens = set_step_execution_context(
+                step_exec_key, ctx.storage, step_id=step_id, step_name=step_name
+            )
 
             try:
                 # Execute step function

pyworkflow/engine/executor.py

@@ -25,7 +25,9 @@
 from pyworkflow.core.registry import get_workflow_by_func
 from pyworkflow.core.workflow import execute_workflow_with_context
 from pyworkflow.engine.events import (
+    EventType,
     create_cancellation_requested_event,
+    create_step_suspended_event,
     create_workflow_cancelled_event,
     create_workflow_continued_as_new_event,
 )
@@ -40,6 +42,90 @@ class ConfigurationError(Exception):
     pass
 
 
+async def step_suspended_already_recorded(
+    storage: StorageBackend,
+    run_id: str,
+    step_id: str,
+) -> bool:
+    """
+    Report whether a STEP_SUSPENDED is already recorded for the current
+    suspension of ``step_id``.
+
+    A step can suspend multiple times on the same deterministic step_id —
+    sequential step_hook() rounds within one step, and crash-replay
+    re-suspension. Each round records a fresh STEP_STARTED, so dedup must be
+    scoped to the *current* start: a STEP_SUSPENDED counts only if it appears
+    after the most recent STEP_STARTED for this step_id. Deduping on "a
+    STEP_SUSPENDED ever exists" would drop every suspension after the first,
+    leaving a STEP_STARTED with no matching STEP_SUSPENDED so replay keeps the
+    step in progress and re-suspends forever.
+
+    One type-filtered ``get_events`` scan (ordered by sequence) keeps the cost
+    bounded on every backend.
+    """
+    events = await storage.get_events(
+        run_id,
+        event_types=[EventType.STEP_STARTED.value, EventType.STEP_SUSPENDED.value],
+    )
+    suspended_since_last_start = False
+    for event in events:
+        if event.data.get("step_id") != step_id:
+            continue
+        if event.type == EventType.STEP_STARTED:
+            suspended_since_last_start = False
+        elif event.type == EventType.STEP_SUSPENDED:
+            suspended_since_last_start = True
+    return suspended_since_last_start
+
+
+async def record_step_suspended_if_needed(
+    storage: StorageBackend,
+    run_id: str,
+    signal: SuspensionSignal,
+) -> None:
+    """
+    Record a STEP_SUSPENDED event when a step suspended via ``step_hook()``.
+
+    Call this immediately after recording WORKFLOW_SUSPENDED on any *inline*
+    execution path (the local runtime, or a ``force_local`` step running inside
+    a Celery orchestration task). On the inline path the step function raises the
+    SuspensionSignal in-process, so no step worker is involved to record the
+    event. Without STEP_SUSPENDED, replay sees the step's STEP_STARTED with no
+    terminal event, treats it as still in progress, and re-suspends forever.
+
+    Being the single writer immediately after WORKFLOW_SUSPENDED gives correct
+    ordering by construction (no polling for the suspend event). The write is
+    deduplicated per-start (see :func:`step_suspended_already_recorded`) so it is
+    safe to call unconditionally and alongside the Celery step worker's own
+    recorder, while still recording one event per suspension round.
+
+    No-op unless ``signal`` is a step_hook suspension (reason prefix
+    ``step_hook:``).
+    """
+    reason = signal.reason
+    if not reason or not reason.startswith("step_hook:"):
+        return
+
+    data = signal.data or {}
+    step_id = data.get("step_id")
+    if not step_id:
+        return
+    hook_id = data.get("hook_id") or ""
+    step_name = data.get("step_name") or ""
+
+    if await step_suspended_already_recorded(storage, run_id, step_id):
+        return
+
+    await storage.record_event(
+        create_step_suspended_event(
+            run_id=run_id,
+            step_id=step_id,
+            step_name=step_name,
+            hook_id=hook_id,
+        )
+    )
+
+
 async def start(
     workflow_func: Callable,
     *args: Any,
@@ -310,6 +396,11 @@ async def _execute_workflow_local(
         )
         await storage.record_event(suspended_event)
 
+        # Record STEP_SUSPENDED for inline step_hook suspensions (no step worker
+        # records it on this path). Immediately after WORKFLOW_SUSPENDED so the
+        # ordering is correct by construction.
+        await record_step_suspended_if_needed(storage, run_id, e)
+
         logger.info(
             f"Workflow suspended: {e.reason}",
             run_id=run_id,

pyworkflow/primitives/step_checkpoint.py

@@ -26,11 +26,18 @@ async def my_agent_step():
 # Step execution context for checkpoint operations
 _step_run_id: ContextVar[str | None] = ContextVar("_step_run_id", default=None)
 _step_storage: ContextVar[Any] = ContextVar("_step_storage", default=None)
+# Deterministic step id / name of the currently executing step. These identify
+# the step in the event log (STEP_STARTED / STEP_SUSPENDED / STEP_COMPLETED),
+# whereas _step_run_id is the composite "run_id:step_id" checkpoint key.
+_step_id: ContextVar[str | None] = ContextVar("_step_id", default=None)
+_step_name: ContextVar[str | None] = ContextVar("_step_name", default=None)
 
 
 def set_step_execution_context(
     step_run_id: str,
     storage: Any,
+    step_id: str | None = None,
+    step_name: str | None = None,
 ) -> tuple:
     """
     Set the step execution context for checkpoint and hook operations.
@@ -40,20 +47,28 @@ def set_step_execution_context(
     Args:
         step_run_id: Unique identifier for this step execution (run_id:step_id)
         storage: Storage backend instance
+        step_id: Deterministic step id, as used in the event log. Lets
+            step_hook() attach the correct id to its SuspensionSignal so the
+            engine can record a matching STEP_SUSPENDED event.
+        step_name: Human-readable step name, recorded alongside STEP_SUSPENDED.
 
     Returns:
         Tokens for resetting context
     """
     t1 = _step_run_id.set(step_run_id)
     t2 = _step_storage.set(storage)
-    return (t1, t2)
+    t3 = _step_id.set(step_id)
+    t4 = _step_name.set(step_name)
+    return (t1, t2, t3, t4)
 
 
 def reset_step_execution_context(tokens: tuple) -> None:
     """Reset the step execution context."""
-    t1, t2 = tokens
+    t1, t2, t3, t4 = tokens
     _step_run_id.reset(t1)
     _step_storage.reset(t2)
+    _step_id.reset(t3)
+    _step_name.reset(t4)
 
 
 def get_step_run_id() -> str | None:
@@ -66,6 +81,16 @@ def get_step_storage() -> Any:
     return _step_storage.get()
 
 
+def get_step_id() -> str | None:
+    """Get the deterministic step id of the currently executing step, if known."""
+    return _step_id.get()
+
+
+def get_step_name() -> str | None:
+    """Get the name of the currently executing step, if known."""
+    return _step_name.get()
+
+
 async def save_step_checkpoint(data: dict) -> None:
     """
     Save checkpoint data for the current step.

pyworkflow/primitives/step_hook.py

@@ -33,7 +33,7 @@ async def agent_step():
 
 from pyworkflow.context import get_context, has_context
 from pyworkflow.core.exceptions import SuspensionSignal
-from pyworkflow.primitives.step_checkpoint import get_step_run_id
+from pyworkflow.primitives.step_checkpoint import get_step_id, get_step_name, get_step_run_id
 
 
 class StepHookTimeout:
@@ -118,6 +118,20 @@ async def notify(token):
             "Use hook() for workflow-level hooks."
         )
 
+    # Identify the step in the event log. The engine records a STEP_SUSPENDED
+    # event from the SuspensionSignal below (so replay does not treat the step
+    # as still in progress), and that event must carry the *deterministic*
+    # step id — not the composite "run_id:step_id" checkpoint key. Prefer the
+    # id/name plumbed through the step execution context; fall back to stripping
+    # the run_id prefix off step_run_id for callers that did not provide them.
+    signal_step_id = get_step_id()
+    if signal_step_id is None:
+        prefix = f"{ctx.run_id}:"
+        signal_step_id = (
+            step_run_id[len(prefix) :] if step_run_id.startswith(prefix) else step_run_id
+        )
+    signal_step_name = get_step_name()
+
     # Generate deterministic hook_id based on step_run_id and hook name
     # This ensures the same hook call gets the same ID on re-execution
     hook_counter = getattr(ctx, "_step_hook_counter", 0)
@@ -195,7 +209,8 @@ async def notify(token):
         raise SuspensionSignal(
             reason=f"step_hook:{hook_id}",
             hook_id=hook_id,
-            step_id=step_run_id,
+            step_id=signal_step_id,
+            step_name=signal_step_name,
             **resume_data,
         )
 
@@ -263,6 +278,7 @@ async def notify(token):
     raise SuspensionSignal(
         reason=f"step_hook:{hook_id}",
         hook_id=hook_id,
-        step_id=step_run_id,
+        step_id=signal_step_id,
+        step_name=signal_step_name,
         **suspend_data,
     )