feat: Retain is_running=True during cancellation (#235)

Co-authored-by: Cursor Agent <[email protected]>

Author: adrianlyjak

.changeset/wide-islands-start.md

@@ -0,0 +1,5 @@
+---
+"llama-index-workflows": patch
+---
+
+Fix resuming from serialized context for workflows that uses typed events

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

@@ -648,7 +648,9 @@ def _process_cancel_run_tick(
     tick: TickCancelRun, init: BrokerState
 ) -> tuple[BrokerState, list[WorkflowCommand]]:
     state = init.deepcopy()
-    state.is_running = False
+    # retain running state, for resumption.
+    # TODO - when/if we persist stream events, this StopEvent should be reconsidered, as there should only ever be one stop event.
+    # Perhaps on resumption, if the workflow is running, then any existing stop events of a "cancellation" type should be omitted from the stream.
     return state, [
         CommandPublishEvent(event=StopEvent()),
         CommandHalt(exception=WorkflowCancelledByUser()),

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

@@ -166,8 +166,9 @@ def from_serialized(
 
         # Start with a base state from the workflow
         base_state = BrokerState.from_workflow(workflow)
-        # Always set is_running to False on deserialization - the workflow will set it to True when it starts
-        base_state.is_running = False
+        # Unfortunately, important to preserve this state, since the workflow needs to know this to decide
+        # whether to create a start_event from kwargs (it only constructs and passes a start event if not already running)
+        base_state.is_running = serialized.is_running
 
         # Restore worker state (queues, collected events, waiters)
         # We do this regardless of is_running state so workflows can resume from where they left off

packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py

@@ -421,7 +421,10 @@ def test_cancel_run(base_state: BrokerState) -> None:
     tick = TickCancelRun()
     new_state, commands = _process_cancel_run_tick(tick, base_state)
 
-    assert new_state.is_running is False
+    # This is perhaps unintuitive, but it's important to be able to cancel and resume a workflow
+    # based on this state--Workflow uses this as a signal to determine whether to pass or construct
+    # a start event
+    assert new_state.is_running is True
     assert len(commands) == 2
     assert isinstance(commands[0], CommandPublishEvent)
     assert isinstance(commands[1], CommandHalt)

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

@@ -55,6 +55,10 @@ class MyStop(StopEvent):
     outcome: str
 
 
+class ResumeStartEvent(StartEvent):
+    topic: str
+
+
 def test_fn() -> None:
     print("test_fn")
 
@@ -557,6 +561,52 @@ async def test_human_in_the_loop_with_resume() -> None:
     assert step2_runs == 1
 
 
+@pytest.mark.asyncio
+async def test_human_in_the_loop_resume_custom_start_event_inactive_ctx() -> None:
+    class CustomHumanWorkflow(Workflow):
+        @step
+        async def ask(self, ctx: Context, ev: ResumeStartEvent) -> InputRequiredEvent:
+            runs = await ctx.store.get("ask_runs", default=0)
+            await ctx.store.set("ask_runs", runs + 1)
+            return InputRequiredEvent(prefix=ev.topic)  # type: ignore[arg-type]
+
+        @step
+        async def complete(self, ctx: Context, ev: HumanResponseEvent) -> StopEvent:
+            runs = await ctx.store.get("complete_runs", default=0)
+            await ctx.store.set("complete_runs", runs + 1)
+            return StopEvent(result=ev.response)
+
+    workflow = CustomHumanWorkflow()
+    handler: WorkflowHandler = workflow.run(topic="pizza")
+    assert handler.ctx
+
+    async for event in handler.stream_events():
+        if isinstance(event, InputRequiredEvent):
+            break
+
+    await handler.cancel_run()
+    ctx_dict = handler.ctx.to_dict()
+    assert ctx_dict["is_running"]
+
+    resumed_ctx = Context.from_dict(workflow, ctx_dict)
+    resumed_handler = workflow.run(ctx=resumed_ctx)
+    resumed_handler.ctx.send_event(HumanResponseEvent(response="42"))  # type: ignore[arg-type]
+
+    events = []
+    async for event in resumed_handler.stream_events():
+        events.append(event)
+
+    assert events == [StopEvent(result="42")]
+
+    final_result = await resumed_handler
+    assert final_result == "42"
+
+    ask_runs = await resumed_handler.ctx.store.get("ask_runs")  # type: ignore[arg-type]
+    complete_runs = await resumed_handler.ctx.store.get("complete_runs")  # type: ignore[arg-type]
+    assert ask_runs == 1
+    assert complete_runs == 1
+
+
 class DummyWorkflowForConcurrentRunsTest(Workflow):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)