Fix concurrent step cancellation regression in 2.14.0 (#340)

* Add xfailing test for concurrent step cancellation regression in 2.14.0

This test exposes a regression introduced in 2.14.0:
- In 2.13.1: A single asyncio.sleep(0) yield after stop_step returns is
  sufficient for cancellation to propagate, blocking other_step's write.
- In 2.14.0: The yield is not enough - other_step's write still goes through.

Test sequence:
1. stop_step returns StopEvent
2. Single yield point (asyncio.sleep(0))
3. other_step writes to stream

In 2.13.1: write is blocked (PASS)
In 2.14.0: write goes through (FAIL)

The architectural change from asyncio.Queue to tick_buffer + wait_for_next_task
in commit 45e7614 ("Remove asyncio queue from control_loop #315") changed the
timing characteristics, requiring longer delays for cancellation to take effect.

In 2.13.1, workers immediately put ticks in the queue via queue_tick(), and
the continuous _pull() task meant quick processing. In 2.14.0, workers return
results collected by wait_for_next_task, adding overhead that delays cancellation.

Related: https://llama-index.slack.com/archives/C0A3TFERFHV/p1770391664693399

https://claude.ai/code/session_01VWUvDsJ88uaVn4jT68Ny8X

* Fix concurrent step cancellation regression in 2.14.0

When a worker returns a StopEvent, immediately cancel other running
workers before adding the tick to the buffer. This restores the 2.13.1
behavior where a single asyncio.sleep(0) yield after stop_step returns
is sufficient for cancellation to propagate.

The regression was introduced in commit 45e7614 ("Remove asyncio queue
from control_loop #315") which changed the architecture from asyncio.Queue
to tick_buffer + wait_for_next_task. The new architecture had more overhead,
causing cancellation to take longer.

The fix checks if the completed worker's result contains a StopEvent, and
if so, calls cleanup_tasks() immediately to cancel other workers before
they can write to the event stream.

Also removes internal implementation references from comments.

Fixes: https://llama-index.slack.com/archives/C0A3TFERFHV/p1770391664693399

https://claude.ai/code/session_01VWUvDsJ88uaVn4jT68Ny8X

* Add patch changeset for concurrent step cancellation fix

* fix lint

* Improve StopEvent cancellation speed for sibling steps

---------

Co-authored-by: Claude <noreply@anthropic.com>

.changeset/fix-concurrent-step-cancellation.md

@@ -0,0 +1,5 @@
+---
+"llama-index-workflows": patch
+---
+
+Fix concurrent step cancellation regression where StopEvent no longer cancelled as quickly as previously

packages/llama-index-workflows/src/workflows/runtime/control_loop.py

@@ -101,7 +101,7 @@ class _ControlLoopRunner:
     This control loop uses a sequential, deterministic design:
     - Scheduled wakeups are tracked in a heap (for timeouts/delays)
     - External events come via wait_receive
-    - No concurrent timeout tasks, ensuring deterministic DBOS function_id ordering
+    - No concurrent timeout tasks, ensuring deterministic ordering for replay
     """
 
     def __init__(
@@ -237,6 +237,7 @@ class _ControlLoopRunner:
             if command.exception is not None:
                 raise command.exception
         elif isinstance(command, CommandCompleteRun):
+            await self.cleanup_tasks()
             return command.result
         elif isinstance(command, CommandPublishEvent):
             await self.adapter.write_to_event_stream(command.event)
@@ -249,7 +250,7 @@ class _ControlLoopRunner:
 
     async def cleanup_tasks(self) -> None:
         """Cancel and cleanup all running worker tasks."""
-        # Signal adapter to stop waiting (wakes blocked DBOS.recv)
+        # Signal adapter to stop waiting
         try:
             await self.adapter.close()
         except Exception:
@@ -279,7 +280,7 @@ class _ControlLoopRunner:
 
         This uses a sequential, deterministic design that combines timeout
         handling with event waiting in a single operation, ensuring
-        deterministic DBOS function_id ordering for replay.
+        deterministic ordering for replay.
 
         Args:
             start_event: Optional initial event to process
@@ -396,6 +397,15 @@ class _ControlLoopRunner:
                             "Worker task failed unexpectedly", exc_info=True
                         )
                     else:
+                        # Check if this worker returned a StopEvent - if so,
+                        # cancel other workers immediately to prevent them from
+                        # writing to the event stream after workflow completion
+                        for res in tick_result.result:
+                            if isinstance(res, StepWorkerResult) and isinstance(
+                                res.result, StopEvent
+                            ):
+                                await self.cleanup_tasks()
+                                break
                         self.tick_buffer.append(tick_result)
 
         finally:

packages/llama-index-workflows/src/workflows/runtime/types/plugin.py

@@ -157,7 +157,7 @@ class InternalRunAdapter(ABC):
         Signal shutdown to wake any blocked wait operations.
 
         Called during cleanup to allow the adapter to exit gracefully.
-        Default is no-op. DBOS adapter sends a shutdown signal to wake blocked recv.
+        Default is no-op. Custom adapters may send a shutdown signal to wake blocked recv.
         """
         pass
 
@@ -200,8 +200,8 @@ class InternalRunAdapter(ABC):
 
         Default implementation uses asyncio.wait(FIRST_COMPLETED) and returns
         the highest-priority completed task (workers before pull).
-        DBOS overrides to coordinate based on journal for deterministic replay,
-        using the stable keys from NamedTask to identify tasks.
+        Custom adapters may override to coordinate based on journal for
+        deterministic replay, using the stable keys from NamedTask to identify tasks.
         """
         tasks = NamedTask.all_tasks(task_set)
         if not tasks:

packages/llama-index-workflows/tests/test_workflow.py

@@ -838,3 +838,81 @@ async def test_workflow_parallel_resume() -> None:
     resume_event.set()
     result = await new_handler
     assert result == "Done"
+
+
+class OtherEvent(Event):
+    """Event written to stream by concurrent step."""
+
+    pass
+
+
+@pytest.mark.asyncio
+async def test_stop_event_cancels_concurrent_step_stream_write() -> None:
+    """Test that StopEvent cancels concurrent step before it writes to the event stream.
+
+    This test exposes a regression introduced in 2.14.0:
+    - In 2.13.1: A single asyncio.sleep(0) yield after stop_step returns is
+      sufficient for cancellation to propagate, blocking other_step's write.
+    - In 2.14.0: The yield is not enough - other_step's write still goes through.
+
+    The architectural change from asyncio.Queue to tick_buffer + wait_for_next_task
+    changed the timing characteristics, requiring longer delays for cancellation
+    to take effect.
+
+    To reproduce:
+    1. stop_step returns StopEvent
+    2. Single yield point (asyncio.sleep(0))
+    3. other_step writes to stream
+
+    In 2.13.1: write is blocked (PASS)
+    In 2.14.0: write goes through (FAIL)
+    """
+    stop_started = asyncio.Event()
+    other_started = asyncio.Event()
+    stop_proceed = asyncio.Event()
+    write_proceed = asyncio.Event()
+    return_proceed = asyncio.Event()
+
+    class ConcurrentStreamWriteWorkflow(Workflow):
+        @step
+        async def stop_step(self, ev: StartEvent) -> StopEvent:
+            stop_started.set()
+            await stop_proceed.wait()
+            return StopEvent(result="done")
+
+        @step
+        async def other_step(self, ctx: Context, ev: StartEvent) -> None:
+            other_started.set()
+            await write_proceed.wait()
+            ctx.write_event_to_stream(OtherEvent())
+            await return_proceed.wait()
+
+    wf = ConcurrentStreamWriteWorkflow(timeout=5)
+    handler = wf.run()
+
+    # Wait for both steps to start
+    await asyncio.wait_for(stop_started.wait(), timeout=2)
+    await asyncio.wait_for(other_started.wait(), timeout=2)
+
+    # Sequence: stop returns, yield, then write attempts
+    stop_proceed.set()
+    await asyncio.sleep(0)  # Single yield - enough in 2.13.1, not in 2.14.0
+    write_proceed.set()
+    return_proceed.set()
+
+    # Collect events from stream
+    events: list[Event] = []
+    async for event in handler.stream_events():
+        events.append(event)
+        if isinstance(event, StopEvent):
+            break
+
+    await handler
+
+    # In 2.13.1: other_step is cancelled before write, only StopEvent in stream
+    # In 2.14.0: write goes through, OtherEvent appears before StopEvent
+    other_events = [e for e in events if isinstance(e, OtherEvent)]
+    assert len(other_events) == 0, (
+        f"OtherEvent should not appear in stream - other_step should have been "
+        f"cancelled after stop_step returned. Got events: {events}"
+    )