Fix workflow cancellation not propagating to active executors (#3663)

python/packages/core/tests/workflow/test_runner.py

@@ -191,3 +191,73 @@ async def test_runner_emits_runner_completion_for_agent_response_without_targets
     # The runner should complete without errors when handling AgentExecutorResponse without targets
     # No specific events are expected since there are no executors to process the message
     assert isinstance(events, list)  # Just verify the runner completed without errors
+
+
+class SlowExecutor(Executor):
+    """An executor that takes time to process, used for cancellation testing."""
+
+    def __init__(self, id: str, work_duration: float = 0.5):
+        super().__init__(id=id)
+        self.started_count = 0
+        self.completed_count = 0
+        self.work_duration = work_duration
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        self.started_count += 1
+        await asyncio.sleep(self.work_duration)
+        self.completed_count += 1
+        if message.data < 2:
+            await ctx.send_message(MockMessage(data=message.data + 1))
+        else:
+            await ctx.yield_output(message.data)
+
+
+async def test_runner_cancellation_stops_active_executor():
+    """Test that cancelling a workflow properly cancels the active executor."""
+    executor_a = SlowExecutor(id="executor_a", work_duration=0.3)
+    executor_b = SlowExecutor(id="executor_b", work_duration=1.0)
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    shared_state = SharedState()
+    ctx = InProcRunnerContext()
+
+    runner = Runner(edges, executors, shared_state, ctx)
+
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        shared_state,
+        ctx,
+    )
+
+    async def run_workflow():
+        async for _ in runner.run_until_convergence():
+            pass
+
+    task = asyncio.create_task(run_workflow())
+
+    # Wait for executor_a to complete (0.3s) and executor_b to start but not finish
+    await asyncio.sleep(0.5)
+
+    # Cancel while executor_b is mid-execution (it takes 1.0s)
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Give time for any leaked tasks to complete (if cancellation didn't work)
+    await asyncio.sleep(1.5)
+
+    # executor_a should have completed once, executor_b should have started but not completed
+    assert executor_a.completed_count == 1
+    assert executor_b.started_count == 1
+    assert executor_b.completed_count == 0  # Should NOT have completed due to cancellation