diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 7a365ec122..dd4c48b622 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/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: diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index c5219a6744..d2dfbe7e79 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/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") diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 014694b9b8..557e644b52 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/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