pyworkflow-engine 0.1.17__tar.gz → 0.1.19__tar.gz

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/CLAUDE.md

@@ -360,14 +360,14 @@ Cancellation is checked at these points:
 **Important Limitation:**
 Cancellation does NOT interrupt a step that is already executing. If a step function takes a long time (e.g., a 10-minute API call), cancellation will only be detected after the step completes. This is by design to avoid leaving operations in an inconsistent state.
 
-For long-running operations that need to be cancellable mid-execution, the step should periodically call `ctx.check_cancellation()` to cooperatively check for cancellation:
+For long-running operations that need to be cancellable mid-execution, the step should periodically call `await ctx.check_cancellation()` to cooperatively check for cancellation. This async method checks both the in-memory flag and the storage backend's cancellation flag (in durable mode), enabling detection of external cancellation requests:
 
 ```python
 @step()
 async def long_running_step():
     ctx = get_context()
     for chunk in process_large_dataset():
-        ctx.check_cancellation()  # Cooperative cancellation check
+        await ctx.check_cancellation()  # Cooperative cancellation check
         await process_chunk(chunk)
     return result
 ```

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/PKG-INFO

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

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/docs/concepts/cancellation.mdx

@@ -119,7 +119,7 @@ async def process_large_dataset(dataset_id: str):
     results = []
     for chunk in dataset.chunks():
         # Check for cancellation periodically
-        ctx.check_cancellation()
+        await ctx.check_cancellation()
 
         result = await process_chunk(chunk)
         results.append(result)
@@ -127,6 +127,8 @@ async def process_large_dataset(dataset_id: str):
     return results
 ```
 
+The `check_cancellation()` method is async because in durable mode it queries the storage backend's cancellation flag, enabling detection of external cancellation requests (e.g., from `cancel_workflow()`). It checks the in-memory flag first as a fast path, then falls back to storage if needed.
+
 This allows the step to respond to cancellation requests between chunks rather than waiting until the entire dataset is processed.
 
 ## Handling Cancellation
@@ -167,7 +169,7 @@ async with shield():
 ```
 
 While inside a `shield()` block:
-- `ctx.check_cancellation()` will not raise `CancellationError`
+- `await ctx.check_cancellation()` will not raise `CancellationError`
 - The cancellation request is preserved
 - Cancellation will take effect after exiting the shield
 
@@ -286,7 +288,7 @@ $ pyworkflow runs logs run_abc123 --filter cancellation
         ctx = get_context()
         for i, item in enumerate(items):
             if i % 100 == 0:  # Check every 100 items
-                ctx.check_cancellation()
+                await ctx.check_cancellation()
             await process_item(item)
     ```
   </Accordion>

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/examples/celery/durable/workflows/cancellation.py

@@ -189,7 +189,7 @@ async def cancel_demo_long_running_step(items: list) -> list:
     Example of cooperative cancellation within a long-running step.
 
     Since cancellation doesn't interrupt steps mid-execution, use
-    ctx.check_cancellation() for responsive cancellation in long loops.
+    await ctx.check_cancellation() for responsive cancellation in long loops.
     """
     ctx = get_context()
     results = []
@@ -197,7 +197,7 @@ async def cancel_demo_long_running_step(items: list) -> list:
     for i, item in enumerate(items):
         # Check for cancellation periodically
         if i % 10 == 0:
-            ctx.check_cancellation()  # Raises CancellationError if cancelled
+            await ctx.check_cancellation()  # Raises CancellationError if cancelled
 
         # Process item
         await asyncio.sleep(0.1)

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/examples/celery/durable/workflows/continue_as_new.py

@@ -101,7 +101,7 @@ async def batch_processor(offset: int = 0, batch_size: int = 10) -> str:
     print(f"  [Batch] Processed {len(items)} items")
 
     # Continue with next batch
-    continue_as_new(offset=offset + batch_size, batch_size=batch_size)
+    await continue_as_new(offset=offset + batch_size, batch_size=batch_size)
 
 
 # --- Message Consumer Workflow ---
@@ -133,7 +133,7 @@ async def message_consumer(cursor: str | None = None, messages_processed: int =
     print(f"  [Consumer] Handled {count} messages (total: {total})")
 
     # Continue with new cursor
-    continue_as_new(cursor=new_cursor, messages_processed=total)
+    await continue_as_new(cursor=new_cursor, messages_processed=total)
 
 
 # --- Recurring Task Workflow ---
@@ -159,7 +159,7 @@ async def recurring_report(iteration: int = 1, max_iterations: int = 3) -> str:
         return f"All {max_iterations} reports generated!"
 
     # Continue with next iteration
-    continue_as_new(iteration=iteration + 1, max_iterations=max_iterations)
+    await continue_as_new(iteration=iteration + 1, max_iterations=max_iterations)
 
 
 async def run_examples():

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/examples/local/durable/11_continue_as_new.py

@@ -86,7 +86,7 @@ async def batch_processor(offset: int = 0, batch_size: int = 10) -> str:
     print("  [Workflow] Batch complete. Continuing with next batch...")
 
     # Continue with next batch - fresh event history!
-    continue_as_new(offset=offset + batch_size, batch_size=batch_size)
+    await continue_as_new(offset=offset + batch_size, batch_size=batch_size)
 
 
 # --- Example 2: Polling Workflow ---
@@ -116,7 +116,7 @@ async def polling_workflow(cursor: str | None = None, poll_count: int = 0) -> st
 
     # Continue polling with new cursor
     print(f"  [Workflow] Continuing with new cursor: {new_cursor}")
-    continue_as_new(cursor=new_cursor, poll_count=poll_count + 1)
+    await continue_as_new(cursor=new_cursor, poll_count=poll_count + 1)
 
 
 # --- Example 3: Counter Workflow (Simple Demo) ---
@@ -132,7 +132,7 @@ async def countdown_workflow(count: int) -> str:
         return "Countdown complete!"
 
     # Continue with decremented count
-    continue_as_new(count=count - 1)
+    await continue_as_new(count=count - 1)
 
 
 async def example_batch_processing(storage):

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.1.17"
+version = "0.1.19"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.17"
+__version__ = "0.1.19"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/aws/context.py

@@ -209,7 +209,7 @@ class AWSWorkflowContext(WorkflowContext):
         """Request cancellation (AWS manages this internally)."""
         logger.warning("Cancellation not supported in AWS context")
 
-    def check_cancellation(self) -> None:
+    async def check_cancellation(self) -> None:
         """Check cancellation (AWS manages this internally)."""
         pass  # AWS handles this
 

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/celery/tasks.py

@@ -207,6 +207,9 @@ def execute_step_task(
             context_class = _resolve_context_class(context_class_name)
             if context_class is not None:
                 step_ctx = context_class.from_dict(context_data)
+                # Inject cancellation metadata so check_cancellation() works on workers
+                object.__setattr__(step_ctx, "_cancellation_run_id", run_id)
+                object.__setattr__(step_ctx, "_cancellation_storage", storage)
                 step_context_token = _set_step_context_internal(step_ctx)
                 # Set readonly mode to prevent mutation in steps
                 readonly_token = _set_step_context_readonly(True)

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/context/base.py

@@ -239,12 +239,15 @@ class WorkflowContext(ABC):
         ...
 
     @abstractmethod
-    def check_cancellation(self) -> None:
+    async def check_cancellation(self) -> None:
         """
         Check for cancellation and raise if requested.
 
         This should be called at interruptible points (before steps,
-        during sleeps, etc.) to allow graceful cancellation.
+        during sleeps, etc.) to allow graceful cancellation. In durable
+        mode, this also checks the storage backend's cancellation flag
+        to detect external cancellation requests (e.g., from
+        ``cancel_workflow()``).
 
         Raises:
             CancellationError: If cancellation was requested and not blocked

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/context/local.py

@@ -744,7 +744,7 @@ class LocalContext(WorkflowContext):
             return
 
         # Check for cancellation before sleeping
-        self.check_cancellation()
+        await self.check_cancellation()
 
         # Durable mode - suspend workflow
         sleep_id = self._generate_sleep_id(duration_seconds)
@@ -913,7 +913,7 @@ class LocalContext(WorkflowContext):
             )
 
         # Check for cancellation before waiting for hook
-        self.check_cancellation()
+        await self.check_cancellation()
 
         # Generate deterministic hook_id
         self._step_counter += 1
@@ -1038,19 +1038,22 @@ class LocalContext(WorkflowContext):
             reason=reason,
         )
 
-    def check_cancellation(self) -> None:
+    async def check_cancellation(self) -> None:
         """
         Check for cancellation and raise if requested.
 
         This should be called at interruptible points (before steps,
-        during sleeps, etc.) to allow graceful cancellation.
+        during sleeps, etc.) to allow graceful cancellation. In durable
+        mode, this also queries the storage backend to detect external
+        cancellation requests (e.g., from ``cancel_workflow()``).
 
         Raises:
             CancellationError: If cancellation was requested and not blocked
         """
-        if self._cancellation_requested and not self._cancellation_blocked:
-            from pyworkflow.core.exceptions import CancellationError
+        from pyworkflow.core.exceptions import CancellationError
 
+        # Fast path: in-memory flag
+        if self._cancellation_requested and not self._cancellation_blocked:
             logger.info(
                 "Cancellation check triggered - raising CancellationError",
                 run_id=self._run_id,
@@ -1061,6 +1064,27 @@ class LocalContext(WorkflowContext):
                 reason=self._cancellation_reason,
             )
 
+        # Storage check: detect external cancellation (durable mode only)
+        if not self._cancellation_blocked and self._durable and self._storage is not None:
+            try:
+                if await self._storage.check_cancellation_flag(self._run_id):
+                    self._cancellation_requested = True
+                    logger.info(
+                        "Cancellation detected via storage flag - raising CancellationError",
+                        run_id=self._run_id,
+                    )
+                    raise CancellationError(
+                        message="Workflow was cancelled: detected via storage flag",
+                        reason=self._cancellation_reason,
+                    )
+            except CancellationError:
+                raise
+            except Exception as e:
+                logger.warning(
+                    f"Failed to check cancellation flag in storage: {e}",
+                    run_id=self._run_id,
+                )
+
     @property
     def cancellation_blocked(self) -> bool:
         """

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/context/mock.py

@@ -322,7 +322,7 @@ class MockContext(WorkflowContext):
         self._cancellation_requested = True
         self._cancellation_reason = reason
 
-    def check_cancellation(self) -> None:
+    async def check_cancellation(self) -> None:
         """Check if cancellation was requested and raise if not blocked."""
         from pyworkflow.core.exceptions import CancellationError
 

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/context/step_context.py

@@ -47,11 +47,17 @@ Usage:
         return {"valid": True}
 """
 
+from __future__ import annotations
+
 from contextvars import ContextVar, Token
-from typing import Any, Self
+from typing import TYPE_CHECKING, Any, Self
 
+from loguru import logger
 from pydantic import BaseModel, ConfigDict
 
+if TYPE_CHECKING:
+    from pyworkflow.storage.base import StorageBackend
+
 
 class StepContext(BaseModel):
     """
@@ -80,6 +86,61 @@ class StepContext(BaseModel):
 
     model_config = ConfigDict(frozen=True, extra="forbid")
 
+    # Private attributes injected by the framework (not serialized).
+    # These enable check_cancellation() to work even when WorkflowContext
+    # is not available (e.g., on Celery workers, inside LangGraph tools).
+    _cancellation_run_id: str | None = None
+    _cancellation_storage: StorageBackend | None = None
+
+    async def check_cancellation(self) -> None:
+        """
+        Check for cancellation and raise CancellationError if requested.
+
+        This works in all execution contexts:
+        - If a WorkflowContext is available, delegates to it (checks both
+          in-memory flag and storage).
+        - Otherwise, checks the storage cancellation flag directly using
+          the run_id and storage injected by the framework.
+
+        This is especially useful for long-running operations inside steps
+        or tool adapters where WorkflowContext may not be available (e.g.,
+        on Celery workers, inside LangGraph tool execution).
+
+        Raises:
+            CancellationError: If cancellation was requested
+
+        Example:
+            @step()
+            async def long_running_step():
+                ctx = get_step_context()
+                for chunk in chunks:
+                    await ctx.check_cancellation()
+                    await process(chunk)
+        """
+        from pyworkflow.context import get_context, has_context
+        from pyworkflow.core.exceptions import CancellationError
+
+        # Fast path: delegate to WorkflowContext if available
+        if has_context():
+            await get_context().check_cancellation()
+            return
+
+        # Fallback: check storage flag directly
+        if self._cancellation_run_id is not None and self._cancellation_storage is not None:
+            try:
+                if await self._cancellation_storage.check_cancellation_flag(
+                    self._cancellation_run_id
+                ):
+                    raise CancellationError(
+                        message="Workflow was cancelled: detected via storage flag",
+                    )
+            except CancellationError:
+                raise
+            except Exception as e:
+                logger.warning(
+                    f"Failed to check cancellation flag in storage: {e}",
+                )
+
     def with_updates(self: Self, **kwargs: Any) -> Self:
         """
         Create a new context with updated values.
@@ -203,12 +264,21 @@ async def set_step_context(ctx: StepContext) -> None:
     if not isinstance(ctx, StepContext):
         raise TypeError(f"Expected StepContext instance, got {type(ctx).__name__}")
 
+    # Inject cancellation metadata from WorkflowContext if available.
+    # This enables check_cancellation() to work even when WorkflowContext
+    # is not accessible (e.g., on Celery workers, inside LangGraph tools).
+    from pyworkflow.context import get_context, has_context
+
+    if has_context():
+        workflow_ctx = get_context()
+        object.__setattr__(ctx, "_cancellation_run_id", workflow_ctx.run_id)
+        if workflow_ctx.is_durable and workflow_ctx.storage is not None:
+            object.__setattr__(ctx, "_cancellation_storage", workflow_ctx.storage)
+
     # Set the context in the contextvar
     _step_context.set(ctx)
 
     # Persist to storage if we're in a durable workflow
-    from pyworkflow.context import get_context, has_context
-
     if has_context():
         workflow_ctx = get_context()
         if workflow_ctx.is_durable and workflow_ctx.storage is not None:

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/core/exceptions.py

@@ -47,7 +47,7 @@ class CancellationError(WorkflowError):
     Note:
         CancellationError is raised at checkpoint boundaries (before steps,
         sleeps, hooks), not during step execution. Long-running steps can
-        call ``ctx.check_cancellation()`` for cooperative cancellation.
+        call ``await ctx.check_cancellation()`` for cooperative cancellation.
 
     Example:
         @workflow

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/core/step.py

@@ -110,7 +110,7 @@ def step(
             ctx = get_context()
 
             # Check for cancellation before executing step
-            ctx.check_cancellation()
+            await ctx.check_cancellation()
 
             # Transient mode: execute directly without event sourcing
             # Retries are still supported via direct execution
@@ -243,7 +243,7 @@ def step(
             )
 
             # Check for cancellation before executing step
-            ctx.check_cancellation()
+            await ctx.check_cancellation()
 
             # Validate parameters before execution
             validate_step_parameters(func, args, kwargs, step_name)

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/core/workflow.py

@@ -233,6 +233,9 @@ async def execute_workflow_with_context(
             context_data = await storage.get_run_context(run_id)
             if context_data:
                 step_ctx = context_class.from_dict(context_data)
+                # Inject cancellation metadata so check_cancellation() works
+                object.__setattr__(step_ctx, "_cancellation_run_id", run_id)
+                object.__setattr__(step_ctx, "_cancellation_storage", storage)
                 step_context_token = _set_step_context_internal(step_ctx)
 
     try:

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/engine/executor.py

@@ -566,7 +566,7 @@ async def cancel_workflow(
         Cancellation does NOT interrupt a step that is already executing.
         If a step takes a long time, cancellation will only be detected after
         the step completes. For long-running steps that need mid-execution
-        cancellation, call ``ctx.check_cancellation()`` periodically within
+        cancellation, call ``await ctx.check_cancellation()`` periodically within
         the step function.
 
     Args:
@@ -649,9 +649,14 @@ async def cancel_workflow(
         current_status=run.status.value,
     )
 
+    # Always set the cancellation flag in storage so that distributed
+    # components (Celery workers, LangGraph tools, StepContext.check_cancellation())
+    # can detect the cancellation regardless of workflow status.
+    await storage.set_cancellation_flag(run_id)
+
     # Handle based on current status
     if run.status == RunStatus.SUSPENDED:
-        # For suspended workflows, update status to CANCELLED immediately
+        # For suspended workflows, also update status to CANCELLED immediately
         # The workflow will see cancellation when it tries to resume
         cancelled_event = create_workflow_cancelled_event(
             run_id=run_id,
@@ -667,10 +672,6 @@ async def cancel_workflow(
         )
 
     elif run.status in {RunStatus.RUNNING, RunStatus.PENDING}:
-        # For running/pending workflows, set cancellation flag
-        # The workflow will detect this at the next check point
-        await storage.set_cancellation_flag(run_id)
-
         logger.info(
             "Cancellation flag set for running workflow",
             run_id=run_id,

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/primitives/child_workflow.py

@@ -98,7 +98,7 @@ async def start_child_workflow(
         )
 
     # Check for cancellation before starting child
-    ctx.check_cancellation()
+    await ctx.check_cancellation()
 
     # Get workflow metadata
     workflow_meta = get_workflow_by_func(workflow_func)

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/primitives/continue_as_new.py

@@ -14,7 +14,7 @@ from pyworkflow.context import get_context, has_context
 from pyworkflow.core.exceptions import ContinueAsNewSignal
 
 
-def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
+async def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
     """
     Complete current workflow and start a new execution with fresh event history.
 
@@ -48,7 +48,7 @@ def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
 
             # Continue with new cursor if more items
             if new_cursor:
-                continue_as_new(cursor=new_cursor)
+                await continue_as_new(cursor=new_cursor)
 
             return "done"
 
@@ -59,7 +59,7 @@ def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
 
             # Continue with next day
             next_date = get_next_date(date)
-            continue_as_new(date=next_date)
+            await continue_as_new(date=next_date)
 
         @workflow
         async def batch_processor(offset: int = 0, batch_size: int = 100):
@@ -69,7 +69,7 @@ def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
                 for item in items:
                     await process_item(item)
                 # Continue with next batch
-                continue_as_new(offset=offset + batch_size, batch_size=batch_size)
+                await continue_as_new(offset=offset + batch_size, batch_size=batch_size)
 
             return f"Processed {offset} items total"
     """
@@ -88,7 +88,7 @@ def continue_as_new(*args: Any, **kwargs: Any) -> NoReturn:
     ctx = get_context()
 
     # Check for cancellation - don't continue if cancelled
-    ctx.check_cancellation()
+    await ctx.check_cancellation()
 
     logger.info(
         "Workflow continuing as new execution",

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/pyworkflow/primitives/shield.py

@@ -15,7 +15,7 @@ Note:
     Cancellation does NOT interrupt a step mid-execution. If a step takes
     a long time, cancellation will only be detected after it completes.
     For cooperative cancellation within long-running steps, call
-    ``ctx.check_cancellation()`` periodically.
+    ``await ctx.check_cancellation()`` periodically.
 
 Example:
     @workflow

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/tests/integration/test_cancellation.py

@@ -319,7 +319,7 @@ class TestCancellationErrorHandling:
                     durable=False,
                 )
                 ctx.request_cancellation()
-                ctx.check_cancellation()
+                await ctx.check_cancellation()
             except CancellationError:
                 cleanup_called = True
                 raise
@@ -328,3 +328,87 @@ class TestCancellationErrorHandling:
             await workflow_with_cleanup()
 
         assert cleanup_called is True
+
+
+class TestCooperativeCancellation:
+    """Test cooperative cancellation for long-running steps via storage flag."""
+
+    @pytest.mark.asyncio
+    async def test_long_running_step_detects_storage_cancellation(self):
+        """Test that a step using await ctx.check_cancellation() detects external cancellation."""
+        from pyworkflow.context import LocalContext, reset_context, set_context
+
+        storage = InMemoryStorageBackend()
+
+        # Create a workflow run
+        run = WorkflowRun(
+            run_id="coop_cancel_run",
+            workflow_name="test_workflow",
+            status=RunStatus.RUNNING,
+        )
+        await storage.create_run(run)
+
+        ctx = LocalContext(
+            run_id="coop_cancel_run",
+            workflow_name="test_workflow",
+            storage=storage,
+            durable=True,
+        )
+        token = set_context(ctx)
+
+        items_processed = 0
+
+        try:
+            # Set cancellation flag externally (simulating cancel_workflow())
+            await storage.set_cancellation_flag("coop_cancel_run")
+
+            # Simulate a long-running step with cooperative cancellation
+            with pytest.raises(CancellationError):
+                for _i in range(100):
+                    await ctx.check_cancellation()
+                    items_processed += 1
+                    await asyncio.sleep(0.001)
+
+            # Should have been cancelled on the first iteration
+            assert items_processed == 0
+        finally:
+            reset_context(token)
+
+    @pytest.mark.asyncio
+    async def test_cooperative_cancellation_mid_loop(self):
+        """Test cancellation detected mid-loop after external flag is set."""
+        from pyworkflow.context import LocalContext, reset_context, set_context
+
+        storage = InMemoryStorageBackend()
+
+        run = WorkflowRun(
+            run_id="coop_mid_run",
+            workflow_name="test_workflow",
+            status=RunStatus.RUNNING,
+        )
+        await storage.create_run(run)
+
+        ctx = LocalContext(
+            run_id="coop_mid_run",
+            workflow_name="test_workflow",
+            storage=storage,
+            durable=True,
+        )
+        token = set_context(ctx)
+
+        items_processed = 0
+
+        try:
+            with pytest.raises(CancellationError):
+                for i in range(100):
+                    # Set cancellation after processing 5 items
+                    if i == 5:
+                        await storage.set_cancellation_flag("coop_mid_run")
+
+                    await ctx.check_cancellation()
+                    items_processed += 1
+
+            # Should have processed exactly 5 items before cancellation
+            assert items_processed == 5
+        finally:
+            reset_context(token)

{pyworkflow_engine-0.1.17 → pyworkflow_engine-0.1.19}/tests/integration/test_child_workflows.py

@@ -290,10 +290,13 @@ class TestBasicChildWorkflow:
         # Wait for child to complete
         await asyncio.sleep(0.5)
 
-        # Check children
+        # Check children - the child gets cancelled because the parent completed
+        # and _handle_parent_completion_local cancels non-terminal children.
+        # With cooperative cancellation, the child detects the storage flag
+        # and raises CancellationError before completing.
         children = await storage.get_children(run_id)
         assert len(children) == 1
-        assert children[0].status == RunStatus.COMPLETED
+        assert children[0].status in (RunStatus.CANCELLED, RunStatus.FAILED)
 
     @pytest.mark.asyncio
     async def test_child_workflow_handle_result(self, setup_storage):