diff --git a/pyworkflow/celery/tasks.py b/pyworkflow/celery/tasks.py index dcc7706d..74045b82 100644 --- a/pyworkflow/celery/tasks.py +++ b/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,6 +664,12 @@ 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 @@ -665,6 +677,11 @@ async def _record_step_suspended( 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 ) diff --git a/pyworkflow/context/base.py b/pyworkflow/context/base.py index 4338a0f7..93355b32 100644 --- a/pyworkflow/context/base.py +++ b/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 # ========================================================================= diff --git a/pyworkflow/core/step.py b/pyworkflow/core/step.py index 5f7f5459..424b26da 100644 --- a/pyworkflow/core/step.py +++ b/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 diff --git a/pyworkflow/engine/executor.py b/pyworkflow/engine/executor.py index d4e93810..29bc200f 100644 --- a/pyworkflow/engine/executor.py +++ b/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, diff --git a/pyworkflow/primitives/step_checkpoint.py b/pyworkflow/primitives/step_checkpoint.py index 8f296758..5c7fe59e 100644 --- a/pyworkflow/primitives/step_checkpoint.py +++ b/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. diff --git a/pyworkflow/primitives/step_hook.py b/pyworkflow/primitives/step_hook.py index 06cda937..87c6d58a 100644 --- a/pyworkflow/primitives/step_hook.py +++ b/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, ) diff --git a/pyworkflow/runtime/local.py b/pyworkflow/runtime/local.py index 66b81d53..d13b09f8 100644 --- a/pyworkflow/runtime/local.py +++ b/pyworkflow/runtime/local.py @@ -246,6 +246,7 @@ async def start_workflow( # Record WORKFLOW_SUSPENDED event from pyworkflow.engine.events import create_workflow_suspended_event + from pyworkflow.engine.executor import record_step_suspended_if_needed step_id = e.data.get("step_id") if e.data else None step_name = e.data.get("step_name") if e.data else None @@ -264,6 +265,19 @@ async def start_workflow( ) await storage.record_event(suspended_event) + # Record STEP_SUSPENDED for inline step_hook suspensions. + await record_step_suspended_if_needed(storage, run_id, e) + + # Fast-answer race: resume_hook() may have delivered the answer + # from within on_created(), before this suspension was recorded. + # The resume it scheduled ran too early to make progress (the + # STEP_SUSPENDED marker did not exist yet). Now that the + # bookkeeping exists, re-drive the workflow so the answer is not + # lost. + if await self._hook_already_answered(e, run_id, hook_id, storage): + await self.resume_workflow(run_id, storage) + return run_id + # Enhanced logging for retry suspensions if e.reason.startswith("retry:"): step_id = e.data.get("step_id") if e.data else "unknown" @@ -436,6 +450,7 @@ async def resume_workflow( # Record WORKFLOW_SUSPENDED event from pyworkflow.engine.events import create_workflow_suspended_event + from pyworkflow.engine.executor import record_step_suspended_if_needed step_id = e.data.get("step_id") if e.data else None step_name = e.data.get("step_name") if e.data else None @@ -454,6 +469,9 @@ async def resume_workflow( ) await storage.record_event(suspended_event) + # Record STEP_SUSPENDED for inline step_hook suspensions. + await record_step_suspended_if_needed(storage, run_id, e) + logger.info( f"Workflow suspended again: {e.reason}", run_id=run_id, @@ -461,6 +479,12 @@ async def resume_workflow( reason=e.reason, ) + # Fast-answer race: if the answer already arrived (e.g. resume_hook() + # fired from within on_created during this resume), re-drive the + # workflow now that the suspension bookkeeping exists. + if await self._hook_already_answered(e, run_id, hook_id, storage): + return await self.resume_workflow(run_id, storage) + return None except ContinueAsNewSignal as e: @@ -506,6 +530,30 @@ async def resume_workflow( raise + async def _hook_already_answered( + self, + signal: SuspensionSignal, + run_id: str, + hook_id: str | None, + storage: "StorageBackend", + ) -> bool: + """ + Report whether a hook/step_hook suspension was already answered. + + Covers the "fast answer" race where resume_hook() records HOOK_RECEIVED + (and, on the local runtime, synchronously drives a resume) from within + on_created(), before the suspension bookkeeping exists. That early resume + cannot make progress, so the caller must re-drive the workflow once the + suspension (including STEP_SUSPENDED) has been recorded. + """ + from pyworkflow.engine.events import EventType + + if not hook_id or not signal.reason: + return False + if not (signal.reason.startswith("hook:") or signal.reason.startswith("step_hook:")): + return False + return await storage.has_event(run_id, EventType.HOOK_RECEIVED.value, hook_id=hook_id) + async def schedule_resume( self, run_id: str, diff --git a/pyworkflow/storage/memory.py b/pyworkflow/storage/memory.py index cfb17e8a..58c39643 100644 --- a/pyworkflow/storage/memory.py +++ b/pyworkflow/storage/memory.py @@ -228,9 +228,12 @@ async def get_events( with self._lock: events = list(self._events.get(run_id, [])) - # Filter by event types + # Filter by event types. event_types is a list of string values + # (e.g. "hook.received"); e.type is an EventType enum, so compare on + # its .value to match — comparing the enum directly never matches a + # string and would silently drop every event. if event_types: - events = [e for e in events if e.type in event_types] + events = [e for e in events if e.type.value in event_types] # Sort by sequence events.sort(key=lambda e: e.sequence or 0) diff --git a/tests/integration/test_step_hook_inline_suspension_e2e.py b/tests/integration/test_step_hook_inline_suspension_e2e.py new file mode 100644 index 00000000..5a7a5f91 --- /dev/null +++ b/tests/integration/test_step_hook_inline_suspension_e2e.py @@ -0,0 +1,285 @@ +""" +End-to-end tests for inline ``step_hook()`` suspension against on-disk storage. + +These complement the single-process, in-memory unit tests in +``tests/unit/test_step_hook.py`` by exercising the fix through a *real* +``FileStorageBackend`` and, crucially, resuming through **freshly constructed +backend instances** (with the config reset in between). That reproduces the +production shape of the bug: the workflow suspends in one process, the event log +is persisted to disk, and a *different* process replays it from that log. + +Before the fix, a step that suspended via ``step_hook()`` on the inline (local +runtime / ``force_local``) path recorded no ``STEP_SUSPENDED`` event, so a fresh +replay saw the step's ``STEP_STARTED`` with no terminal event, treated it as +still in progress, and re-suspended forever -- the run was stranded. +""" + +import pytest + +from pyworkflow import configure, reset_config, resume, start +from pyworkflow.core.step import _generate_step_id, step +from pyworkflow.core.workflow import workflow +from pyworkflow.engine.events import ( + EventType, + create_hook_received_event, +) +from pyworkflow.primitives.resume_hook import create_hook_token, resume_hook +from pyworkflow.primitives.step_hook import step_hook +from pyworkflow.serialization.decoder import deserialize_args +from pyworkflow.serialization.encoder import serialize +from pyworkflow.storage.file import FileStorageBackend +from pyworkflow.storage.schemas import RunStatus + +pytestmark = pytest.mark.integration + + +@pytest.fixture(autouse=True) +def reset_pyworkflow_config(): + """Reset configuration before and after each test.""" + reset_config() + yield + reset_config() + + +def _fresh_backend(storage_path) -> FileStorageBackend: + """A brand-new FileStorageBackend over ``storage_path``. + + Constructing a new instance (rather than reusing the one that started the + run) simulates a separate process: nothing is shared but the on-disk log. + """ + return FileStorageBackend(base_path=str(storage_path)) + + +async def _events(storage_path, run_id): + return await _fresh_backend(storage_path).get_events(run_id) + + +async def _status(storage_path, run_id) -> RunStatus: + run = await _fresh_backend(storage_path).get_run(run_id) + return run.status + + +class TestInlineSuspensionFileStorage: + """step_hook() suspend -> persist -> cross-backend resume, on real disk.""" + + @pytest.mark.asyncio + async def test_suspend_persists_step_suspended_and_resume_completes(self, tmp_path): + """STEP_SUSPENDED is written to disk on suspend; a fresh backend resumes + the run to completion instead of stranding it.""" + storage_path = tmp_path / "storage" + + @step(name="review_step") + async def review_step(): + feedback = await step_hook("review") + return {"feedback": feedback} + + @workflow(name="review_workflow") + async def review_workflow(): + return await review_step() + + # --- "Process 1": start the workflow; it suspends on the hook. --- + storage_start = _fresh_backend(storage_path) + configure(storage=storage_start, default_durable=True) + run_id = await start(review_workflow, durable=True, storage=storage_start) + + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # The STEP_SUSPENDED marker must be durably on disk, readable by any + # process, so replay does not treat the step as still in progress. + events = await _events(storage_path, run_id) + suspended = [e for e in events if e.type == EventType.STEP_SUSPENDED] + assert len(suspended) == 1, "STEP_SUSPENDED must be persisted on the inline path" + assert suspended[0].data.get("step_id") == _generate_step_id("review_step", (), {}) + + # --- "Process 2": fresh config + fresh backend deliver the answer. --- + reset_config() + storage_resume = _fresh_backend(storage_path) + configure(storage=storage_resume, default_durable=True) + + token = create_hook_token(run_id, "step_hook_review_0") + await resume_hook(token, {"approved": True}, storage=storage_resume) + + assert await _status(storage_path, run_id) == RunStatus.COMPLETED + run = await _fresh_backend(storage_path).get_run(run_id) + assert deserialize_args(run.result)[0] == {"feedback": {"approved": True}} + + @pytest.mark.asyncio + async def test_crash_replay_resume_completes_without_duplicate_hook(self, tmp_path): + """Simulated crash between hook delivery and resume: recording + HOOK_RECEIVED and replaying from disk with a fresh executor completes the + run and re-creates neither the hook nor a second STEP_SUSPENDED.""" + storage_path = tmp_path / "storage" + + @step(name="crash_step") + async def crash_step(): + feedback = await step_hook("approval") + return {"feedback": feedback} + + @workflow(name="crash_workflow") + async def crash_workflow(): + return await crash_step() + + storage_start = _fresh_backend(storage_path) + configure(storage=storage_start, default_durable=True) + run_id = await start(crash_workflow, durable=True, storage=storage_start) + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # Crash: the answer was recorded (HOOK_RECEIVED) but the process died + # before resuming. Record it directly, then replay from a fresh backend. + hook_id = "step_hook_approval_0" + reset_config() + storage_resume = _fresh_backend(storage_path) + configure(storage=storage_resume, default_durable=True) + await storage_resume.record_event( + create_hook_received_event( + run_id=run_id, hook_id=hook_id, payload=serialize({"approved": True}) + ) + ) + + result = await resume(run_id, storage=storage_resume) + assert result == {"feedback": {"approved": True}} + assert await _status(storage_path, run_id) == RunStatus.COMPLETED + + # The visitor-facing hook must not be re-created on replay: exactly one + # HOOK_CREATED in the durable log. + events = await _events(storage_path, run_id) + hook_created = [ + e + for e in events + if e.type == EventType.HOOK_CREATED and e.data.get("hook_id") == hook_id + ] + assert len(hook_created) == 1, "hook must not be re-created on replay" + + @pytest.mark.asyncio + async def test_two_sequential_hooks_complete_across_fresh_backends(self, tmp_path): + """A step suspending twice on the SAME step_id (two sequential rounds) + completes when each round is answered through a fresh backend, and the + durable log holds exactly one STEP_SUSPENDED per suspended start.""" + storage_path = tmp_path / "storage" + + @step(name="two_round_step") + async def two_round_step(): + first = await step_hook("round_one") + second = await step_hook("round_two") + return {"first": first, "second": second} + + @workflow(name="two_round_workflow") + async def two_round_workflow(): + return await two_round_step() + + storage_start = _fresh_backend(storage_path) + configure(storage=storage_start, default_durable=True) + run_id = await start(two_round_workflow, durable=True, storage=storage_start) + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # Answer round one through a fresh backend -> re-suspends on round two. + reset_config() + storage_r1 = _fresh_backend(storage_path) + configure(storage=storage_r1, default_durable=True) + await resume_hook( + create_hook_token(run_id, "step_hook_round_one_0"), {"n": 1}, storage=storage_r1 + ) + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # Answer round two through yet another fresh backend -> completes. + reset_config() + storage_r2 = _fresh_backend(storage_path) + configure(storage=storage_r2, default_durable=True) + await resume_hook( + create_hook_token(run_id, "step_hook_round_two_1"), {"n": 2}, storage=storage_r2 + ) + + assert await _status(storage_path, run_id) == RunStatus.COMPLETED + run = await _fresh_backend(storage_path).get_run(run_id) + assert deserialize_args(run.result)[0] == {"first": {"n": 1}, "second": {"n": 2}} + + # Exactly one STEP_SUSPENDED per suspended STEP_STARTED cycle: every start + # for this step_id ended suspended except the final completing one. + step_id = _generate_step_id("two_round_step", (), {}) + events = await _events(storage_path, run_id) + started = [ + e + for e in events + if e.type == EventType.STEP_STARTED and e.data.get("step_id") == step_id + ] + suspended = [ + e + for e in events + if e.type == EventType.STEP_SUSPENDED and e.data.get("step_id") == step_id + ] + completed = [ + e + for e in events + if e.type == EventType.STEP_COMPLETED and e.data.get("step_id") == step_id + ] + assert len(completed) == 1 + assert len(suspended) == 2 + assert len(suspended) == len(started) - 1 + + @pytest.mark.asyncio + async def test_fast_answer_from_on_created_completes(self, tmp_path): + """A hook answered from within on_created (before the suspension is fully + recorded) is not lost against on-disk storage: the run completes.""" + storage_path = tmp_path / "storage" + storage_start = _fresh_backend(storage_path) + + @step(name="fast_step") + async def fast_step(): + async def answer_immediately(token): + await resume_hook(token, {"approved": True, "fast": True}, storage=storage_start) + + feedback = await step_hook("fast_review", on_created=answer_immediately) + return {"feedback": feedback} + + @workflow(name="fast_workflow") + async def fast_workflow(): + return await fast_step() + + configure(storage=storage_start, default_durable=True) + run_id = await start(fast_workflow, durable=True, storage=storage_start) + + assert await _status(storage_path, run_id) == RunStatus.COMPLETED + run = await _fresh_backend(storage_path).get_run(run_id) + assert deserialize_args(run.result)[0] == {"feedback": {"approved": True, "fast": True}} + + @pytest.mark.asyncio + async def test_repeated_replay_without_answer_stays_resumable(self, tmp_path): + """Regression guard for the stranding bug: replaying a suspended run + WITHOUT delivering the answer must keep it cleanly SUSPENDED (each replay + re-suspends), and answering afterwards must still complete it.""" + storage_path = tmp_path / "storage" + + @step(name="patient_step") + async def patient_step(): + feedback = await step_hook("patient") + return {"feedback": feedback} + + @workflow(name="patient_workflow") + async def patient_workflow(): + return await patient_step() + + storage_start = _fresh_backend(storage_path) + configure(storage=storage_start, default_durable=True) + run_id = await start(patient_workflow, durable=True, storage=storage_start) + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # Replay several times with no answer available. Each fresh replay must + # re-suspend cleanly (return None) rather than error or hang. + for _ in range(3): + reset_config() + storage_replay = _fresh_backend(storage_path) + configure(storage=storage_replay, default_durable=True) + result = await resume(run_id, storage=storage_replay) + assert result is None + assert await _status(storage_path, run_id) == RunStatus.SUSPENDED + + # Now deliver the answer -> the run completes. + reset_config() + storage_answer = _fresh_backend(storage_path) + configure(storage=storage_answer, default_durable=True) + await resume_hook( + create_hook_token(run_id, "step_hook_patient_0"), {"ok": True}, storage=storage_answer + ) + assert await _status(storage_path, run_id) == RunStatus.COMPLETED + run = await _fresh_backend(storage_path).get_run(run_id) + assert deserialize_args(run.result)[0] == {"feedback": {"ok": True}} diff --git a/tests/unit/test_step_hook.py b/tests/unit/test_step_hook.py index 60abfb17..1dddc0db 100644 --- a/tests/unit/test_step_hook.py +++ b/tests/unit/test_step_hook.py @@ -202,6 +202,369 @@ async def test_step_hook_deterministic_ids(self): reset_context(ctx_token) +class TestStepHookInlineSuspension: + """End-to-end tests for step_hook() suspension on the inline (local runtime) path. + + Regression coverage for the bug where a step suspending via step_hook() on the + inline path recorded no STEP_SUSPENDED event, so replay saw the step's + STEP_STARTED without a STEP_SUSPENDED, treated it as still in progress, and + re-suspended forever. + """ + + @pytest.mark.asyncio + async def test_local_runtime_completes_after_resume_hook(self): + """A local-runtime step that suspends via step_hook() completes after resume.""" + from pyworkflow import configure, reset_config, start + from pyworkflow.core.step import step + from pyworkflow.core.workflow import workflow + from pyworkflow.primitives.resume_hook import create_hook_token, resume_hook + from pyworkflow.primitives.step_hook import step_hook + from pyworkflow.serialization.decoder import deserialize_args + + reset_config() + + @step(name="inline_review_step") + async def review_step(): + feedback = await step_hook("review") + return {"feedback": feedback} + + @workflow(name="inline_review_workflow") + async def review_workflow(): + return await review_step() + + try: + storage = InMemoryStorageBackend() + configure(storage=storage) + run_id = await start(review_workflow, durable=True, storage=storage) + + # The workflow suspends waiting for the hook. + run = await storage.get_run(run_id) + assert run.status == RunStatus.SUSPENDED + + # A STEP_SUSPENDED event must be recorded on the inline path so that + # replay does not treat the step as still in progress. + from pyworkflow.engine.events import EventType + + events = await storage.get_events(run_id) + suspended = [e for e in events if e.type == EventType.STEP_SUSPENDED] + assert len(suspended) == 1, "STEP_SUSPENDED must be recorded on the inline path" + + # Deliver the external answer. resume_hook() drives the local runtime + # to resume the workflow synchronously. + token = create_hook_token(run_id, "step_hook_review_0") + await resume_hook(token, {"approved": True}, storage=storage) + + # The run must now be COMPLETE and carry the payload returned by step_hook(). + run = await storage.get_run(run_id) + assert run.status == RunStatus.COMPLETED, ( + f"run should be COMPLETED after resume, got {run.status}" + ) + result = deserialize_args(run.result)[0] + assert result == {"feedback": {"approved": True}} + finally: + reset_config() + + @pytest.mark.asyncio + async def test_crash_replay_reexecutes_without_duplicate_hook(self): + """After suspension, a fresh executor replay re-executes the step and + completes once resumed, without re-creating the visitor-facing hook.""" + from pyworkflow import configure, reset_config, resume, start + from pyworkflow.core.step import step + from pyworkflow.core.workflow import workflow + from pyworkflow.engine.events import EventType + from pyworkflow.primitives.step_hook import step_hook + from pyworkflow.serialization.decoder import deserialize_args + + reset_config() + + body_executions = 0 + on_created_calls = 0 + + @step(name="crash_replay_step") + async def review_step(): + nonlocal body_executions + body_executions += 1 + + async def notify(_token): + nonlocal on_created_calls + on_created_calls += 1 + + feedback = await step_hook("approval", on_created=notify) + return {"feedback": feedback, "runs": body_executions} + + @workflow(name="crash_replay_workflow") + async def review_workflow(): + return await review_step() + + try: + storage = InMemoryStorageBackend() + configure(storage=storage) + run_id = await start(review_workflow, durable=True, storage=storage) + + assert body_executions == 1 + assert on_created_calls == 1 + run = await storage.get_run(run_id) + assert run.status == RunStatus.SUSPENDED + + # Simulate delivery of the hook WITHOUT the local runtime auto-resuming + # (i.e. crash between resume_hook recording and resume). Record the + # HOOK_RECEIVED event directly. + from pyworkflow.engine.events import create_hook_received_event + from pyworkflow.serialization.encoder import serialize + + hook_id = "step_hook_approval_0" + await storage.record_event( + create_hook_received_event( + run_id=run_id, + hook_id=hook_id, + payload=serialize({"approved": True}), + ) + ) + + # Fresh executor replays from the event log and resumes. + resume_result = await resume(run_id, storage=storage) + assert resume_result["feedback"] == {"approved": True} + + # The step body re-executed on resume (idempotent re-execution). + assert body_executions == 2 + # The visitor-facing hook must NOT be re-created on replay: on_created + # fires only once, and only one HOOK_CREATED event exists. + assert on_created_calls == 1 + events = await storage.get_events(run_id) + hook_created = [ + e + for e in events + if e.type == EventType.HOOK_CREATED and e.data.get("hook_id") == hook_id + ] + assert len(hook_created) == 1, "hook must not be re-created on replay" + + run = await storage.get_run(run_id) + assert run.status == RunStatus.COMPLETED + payload = deserialize_args(run.result)[0] + assert payload["feedback"] == {"approved": True} + finally: + reset_config() + + @pytest.mark.asyncio + async def test_fast_answer_from_on_created_completes(self): + """A hook answered from within on_created (before the suspension bookkeeping + exists) is not lost: the run completes on the local runtime.""" + from pyworkflow import configure, reset_config, start + from pyworkflow.core.step import step + from pyworkflow.core.workflow import workflow + from pyworkflow.primitives.resume_hook import resume_hook + from pyworkflow.primitives.step_hook import step_hook + from pyworkflow.serialization.decoder import deserialize_args + + reset_config() + + @step(name="fast_answer_step") + async def review_step(): + async def answer_immediately(token): + # Deliver the answer synchronously, before the workflow has + # finished recording its suspension. + await resume_hook(token, {"approved": True, "fast": True}, storage=storage) + + feedback = await step_hook("fast_review", on_created=answer_immediately) + return {"feedback": feedback} + + @workflow(name="fast_answer_workflow") + async def review_workflow(): + return await review_step() + + try: + storage = InMemoryStorageBackend() + configure(storage=storage) + run_id = await start(review_workflow, durable=True, storage=storage) + + run = await storage.get_run(run_id) + assert run.status == RunStatus.COMPLETED, ( + f"fast-answered run should be COMPLETED, got {run.status}" + ) + result = deserialize_args(run.result)[0] + assert result == {"feedback": {"approved": True, "fast": True}} + finally: + reset_config() + + @pytest.mark.asyncio + async def test_two_sequential_step_hooks_in_one_step_completes(self): + """A step that suspends twice on the SAME step_id (two sequential + step_hook rounds) completes — the second suspension must record its own + STEP_SUSPENDED, not be deduped away by the first round's event.""" + from pyworkflow import configure, reset_config, start + from pyworkflow.core.step import _generate_step_id, step + from pyworkflow.core.workflow import workflow + from pyworkflow.engine.events import EventType + from pyworkflow.primitives.resume_hook import create_hook_token, resume_hook + from pyworkflow.primitives.step_hook import step_hook + from pyworkflow.serialization.decoder import deserialize_args + + reset_config() + + @step(name="two_round_step") + async def review_step(): + first = await step_hook("round_one") + second = await step_hook("round_two") + return {"first": first, "second": second} + + @workflow(name="two_round_workflow") + async def review_workflow(): + return await review_step() + + try: + storage = InMemoryStorageBackend() + configure(storage=storage) + run_id = await start(review_workflow, durable=True, storage=storage) + + # Suspended on round one. + assert (await storage.get_run(run_id)).status == RunStatus.SUSPENDED + + # Answer round one -> local runtime resumes -> suspends on round two. + await resume_hook( + create_hook_token(run_id, "step_hook_round_one_0"), + {"n": 1}, + storage=storage, + ) + assert (await storage.get_run(run_id)).status == RunStatus.SUSPENDED + + # Answer round two -> run completes. + await resume_hook( + create_hook_token(run_id, "step_hook_round_two_1"), + {"n": 2}, + storage=storage, + ) + run = await storage.get_run(run_id) + assert run.status == RunStatus.COMPLETED, ( + f"two-round run should be COMPLETED, got {run.status}" + ) + result = deserialize_args(run.result)[0] + assert result == {"first": {"n": 1}, "second": {"n": 2}} + + # One STEP_SUSPENDED per suspended STEP_STARTED cycle for this step_id: + # every start ended in a suspension except the final one that completed. + events = await storage.get_events(run_id) + step_id = _generate_step_id("two_round_step", (), {}) + started = [ + e + for e in events + if e.type == EventType.STEP_STARTED and e.data.get("step_id") == step_id + ] + suspended = [ + e + for e in events + if e.type == EventType.STEP_SUSPENDED and e.data.get("step_id") == step_id + ] + completed = [ + e + for e in events + if e.type == EventType.STEP_COMPLETED and e.data.get("step_id") == step_id + ] + assert len(completed) == 1 + assert len(suspended) == 2 + assert len(suspended) == len(started) - 1 + finally: + reset_config() + + @pytest.mark.asyncio + async def test_crash_replay_between_sequential_step_hooks_completes(self): + """Crash-replay across the SECOND suspension: after the round-one answer + is replayed by a fresh executor and the step re-suspends on round two, + a fresh executor replays again and completes after the round-two answer.""" + from pyworkflow import configure, reset_config, resume, start + from pyworkflow.core.step import step + from pyworkflow.core.workflow import workflow + from pyworkflow.engine.events import create_hook_received_event + from pyworkflow.primitives.step_hook import step_hook + from pyworkflow.serialization.encoder import serialize + + reset_config() + + @step(name="two_round_crash_step") + async def review_step(): + first = await step_hook("r1") + second = await step_hook("r2") + return {"first": first, "second": second} + + @workflow(name="two_round_crash_workflow") + async def review_workflow(): + return await review_step() + + try: + storage = InMemoryStorageBackend() + configure(storage=storage) + run_id = await start(review_workflow, durable=True, storage=storage) + assert (await storage.get_run(run_id)).status == RunStatus.SUSPENDED + + # Crash between rounds: record the round-one answer directly (no + # auto-resume), then replay with a fresh executor. + await storage.record_event( + create_hook_received_event( + run_id=run_id, hook_id="step_hook_r1_0", payload=serialize({"n": 1}) + ) + ) + await resume(run_id, storage=storage) + assert (await storage.get_run(run_id)).status == RunStatus.SUSPENDED + + # Crash again mid-second-suspension: record the round-two answer and + # replay once more; the run must complete. + await storage.record_event( + create_hook_received_event( + run_id=run_id, hook_id="step_hook_r2_1", payload=serialize({"n": 2}) + ) + ) + result = await resume(run_id, storage=storage) + assert result == {"first": {"n": 1}, "second": {"n": 2}} + assert (await storage.get_run(run_id)).status == RunStatus.COMPLETED + finally: + reset_config() + + +class TestStepHookCounterAPI: + """Tests for the public step-hook counter API used by cross-process restore.""" + + @pytest.mark.asyncio + async def test_counter_get_set_round_trip(self): + """get_step_hook_counter/set_step_hook_counter round-trip through the context.""" + storage = InMemoryStorageBackend() + ctx = LocalContext( + run_id="counter_run", workflow_name="test_workflow", storage=storage, durable=True + ) + assert ctx.get_step_hook_counter() == 0 + ctx.set_step_hook_counter(7) + assert ctx.get_step_hook_counter() == 7 + # Backing storage stays the private attribute. + assert ctx._step_hook_counter == 7 + + @pytest.mark.asyncio + async def test_restored_counter_reproduces_deterministic_hook_id(self): + """A restored counter reproduces the same deterministic hook id on re-execution.""" + storage = InMemoryStorageBackend() + run_id = "counter_run_2" + await storage.create_run( + WorkflowRun(run_id=run_id, workflow_name="test_workflow", status=RunStatus.RUNNING) + ) + ctx = LocalContext( + run_id=run_id, workflow_name="test_workflow", storage=storage, durable=True + ) + ctx._is_step_worker = True + + # Simulate a checkpoint/restore across a process boundary: the app poked + # the counter to 3 before re-executing the step. + ctx.set_step_hook_counter(3) + + 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("resume_point") + assert exc_info.value.data["hook_id"] == "step_hook_resume_point_3" + # The counter advanced past the restored value. + assert ctx.get_step_hook_counter() == 4 + finally: + reset_step_execution_context(step_tokens) + reset_context(ctx_token) + + class TestStepHookTimeout: """Tests for step_hook() timeout semantics (on_timeout="return")."""