Python: wake workflow streaming on iteration completion

python/packages/core/agent_framework/_workflows/_runner.py

@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
-import contextlib
 import logging
 import warnings
 from collections import defaultdict
@@ -125,24 +124,24 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
             logger.info(f"Starting superstep {self._iteration + 1}")
             yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
 
-            # Run iteration concurrently with live event streaming: we poll
-            # for new events while the iteration coroutine progresses.
+            # Wake on either a live event or iteration completion, including silent supersteps.
             iteration_task = asyncio.create_task(self._run_iteration())
+            event_task: asyncio.Task[WorkflowEvent] | None = None
             try:
                 while not iteration_task.done():
-                    try:
-                        # Wait briefly for any new event; timeout allows progress checks
-                        event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
-                        yield event
-                    except asyncio.TimeoutError:
-                        # Periodically continue to let iteration advance
-                        continue
-            except asyncio.CancelledError:
-                # Propagate cancellation to the iteration task to avoid orphaned work
-                iteration_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError):
-                    await iteration_task
-                raise
+                    event_task = asyncio.create_task(self._ctx.next_event())
+                    done, _ = await asyncio.wait((iteration_task, event_task), return_when=asyncio.FIRST_COMPLETED)
+                    if event_task in done:
+                        yield event_task.result()
+            finally:
+                # Cancellation and generator closure must not leave an event waiter or executor running.
+                tasks: list[asyncio.Task[Any]] = (
+                    [iteration_task] if event_task is None else [iteration_task, event_task]
+                )
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
 
             # Propagate errors from iteration, but first surface any pending events
             try:

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

@@ -499,6 +499,100 @@ async def test_runner_iteration_exception_drains_events():
     assert len(events) > 0
 
 
+async def test_runner_completion_does_not_require_a_timer(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A completed superstep must wake an otherwise idle event consumer."""
+    ctx = InProcRunnerContext()
+    runner = Runner([], {}, State(), ctx, "test", graph_signature_hash="test")
+    loop = asyncio.get_running_loop()
+    watchdog = loop.create_future()
+    # Schedule the test watchdog before deferring timers used by the runner.
+    watchdog_handle = loop.call_later(2, watchdog.set_result, None)
+    call_at = loop.call_at
+    monkeypatch.setattr(
+        loop, "call_at", lambda when, callback, *args, **kwargs: call_at(loop.time() + 3600, callback, *args, **kwargs)
+    )
+
+    async def consume() -> list[WorkflowEvent]:
+        return [event async for event in runner.run_until_convergence()]
+
+    task = asyncio.create_task(consume())
+    try:
+        done, _ = await asyncio.wait({task, watchdog}, return_when=asyncio.FIRST_COMPLETED)
+        assert task in done, "Iteration completion was waiting for a polling timer"
+        assert [event.type for event in task.result()] == ["superstep_started", "superstep_completed"]
+    finally:
+        watchdog_handle.cancel()
+        watchdog.cancel()
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+
+
+@pytest.mark.parametrize("fail", [False, True])
+async def test_runner_streams_live_and_tail_events(monkeypatch: pytest.MonkeyPatch, fail: bool) -> None:
+    """Yield live events before completion and retain ordered tail events on success or failure."""
+    ctx = InProcRunnerContext()
+    runner = Runner([], {}, State(), ctx, "test", graph_signature_hash="test")
+    release = asyncio.Event()
+
+    async def iteration() -> None:
+        await ctx.add_event(WorkflowEvent(type="output", data="live", executor_id="test"))
+        await release.wait()
+        for value in ("tail_1", "tail_2"):
+            await ctx.add_event(WorkflowEvent(type="output", data=value, executor_id="test"))
+        if fail:
+            raise RuntimeError("iteration failed")
+
+    monkeypatch.setattr(runner, "_run_iteration", iteration)
+    stream = runner.run_until_convergence()
+    try:
+        assert (await anext(stream)).type == "superstep_started"
+        assert (await anext(stream)).data == "live"
+        release.set()
+        outputs = []
+        try:
+            async for event in stream:
+                if event.type == "output":
+                    outputs.append(event.data)
+        except RuntimeError as exc:
+            assert fail and str(exc) == "iteration failed"
+        else:
+            assert not fail
+        assert outputs == ["tail_1", "tail_2"]
+    finally:
+        await stream.aclose()
+
+
+async def test_runner_stream_close_stops_active_iteration(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Closing after a live event must not leave background executor work running."""
+    ctx = InProcRunnerContext()
+    runner = Runner([], {}, State(), ctx, "test", graph_signature_hash="test")
+    stopped = asyncio.Event()
+    iteration_task: asyncio.Task[Any] | None = None
+
+    async def iteration() -> None:
+        nonlocal iteration_task
+        iteration_task = asyncio.current_task()
+        try:
+            await ctx.add_event(WorkflowEvent(type="output", data="live", executor_id="test"))
+            await asyncio.Event().wait()
+        finally:
+            stopped.set()
+
+    monkeypatch.setattr(runner, "_run_iteration", iteration)
+    stream = runner.run_until_convergence()
+    try:
+        await anext(stream)
+        assert (await anext(stream)).data == "live"
+        await stream.aclose()
+        assert stopped.is_set()
+        assert iteration_task is not None and iteration_task.done()
+    finally:
+        await stream.aclose()
+        if iteration_task is not None:
+            iteration_task.cancel()
+            await asyncio.gather(iteration_task, return_exceptions=True)
+
+
 async def test_runner_reset_iteration_count():
     """Test that reset_iteration_count works correctly."""
     executor_a = MockExecutor(id="executor_a")

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

@@ -834,8 +834,17 @@ async def test_workflow_with_simple_cycle_and_exit_condition():
 
 async def test_workflow_concurrent_execution_prevention():
     """Test that concurrent workflow executions are prevented."""
-    # Create a simple workflow that takes some time to execute
-    executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    class GatedExecutor(Executor):
+        @handler
+        async def handle(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
+            started.set()
+            await release.wait()
+            await ctx.yield_output(message.data)
+
+    executor = GatedExecutor(id="gated_executor")
     workflow = WorkflowBuilder(start_executor=executor).build()
 
     # Create a task that will run the workflow
@@ -845,17 +854,17 @@ async def run_workflow():
     # Start the first workflow execution
     task1 = asyncio.create_task(run_workflow())
 
-    # Give it a moment to start
-    await asyncio.sleep(0.01)
-
-    # Try to start a second concurrent execution - this should fail
-    with pytest.raises(
-        WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
-    ):
-        await workflow.run(NumberMessage(data=0))
-
-    # Wait for the first task to complete
-    result = await task1
+    try:
+        await started.wait()
+        # The first run stays active regardless of how quickly the runner schedules work.
+        with pytest.raises(
+            WorkflowException,
+            match="Workflow is already running; concurrent runs are not allowed on the same instance.",
+        ):
+            await workflow.run(NumberMessage(data=0))
+    finally:
+        release.set()
+        result = await task1
     assert result.get_final_state() == WorkflowRunState.IDLE
 
     # After the first execution completes, we should be able to run again