Fix: idle detection only working for wait_for_event (#359)

.changeset/modern-planets-tan.md

@@ -0,0 +1,5 @@
+---
+"llama-index-workflows": patch
+---
+
+Fix idle detection only working for wait_for_event, not for steps waiting on InputRequiredEvent

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

@@ -36,6 +36,7 @@ from workflows.runtime.types.commands import (
     CommandPublishEvent,
     CommandQueueEvent,
     CommandRunWorker,
+    CommandScheduleIdleCheck,
     WorkflowCommand,
     indicates_exit,
 )
@@ -64,6 +65,7 @@ from workflows.runtime.types.results import (
 from workflows.runtime.types.ticks import (
     TickAddEvent,
     TickCancelRun,
+    TickIdleCheck,
     TickPublishEvent,
     TickStepResult,
     TickTimeout,
@@ -130,6 +132,8 @@ class _ControlLoopRunner:
         self._pull_sequence = 0
         # Map from worker task to (step_name, worker_id) key
         self._task_keys: dict[asyncio.Task[TickStepResult], tuple[str, int]] = {}
+        # Whether a TickIdleCheck is currently in tick_buffer
+        self._idle_check_pending = False
 
     def schedule_tick(self, tick: WorkflowTick, at_time: float) -> None:
         """Schedule a tick to be processed at a specific time."""
@@ -242,6 +246,11 @@ class _ControlLoopRunner:
         elif isinstance(command, CommandFailWorkflow):
             await self.cleanup_tasks()
             raise command.exception
+        elif isinstance(command, CommandScheduleIdleCheck):
+            if not self._idle_check_pending:
+                self.tick_buffer.append(TickIdleCheck())
+                self._idle_check_pending = True
+            return None
         else:
             raise ValueError(f"Unknown command type: {type(command)}")
 
@@ -327,6 +336,8 @@ class _ControlLoopRunner:
                 # Drain and process buffered ticks first (from rehydration, queue_tick, etc.)
                 while self.tick_buffer:
                     tick = self.tick_buffer.pop(0)
+                    if isinstance(tick, TickIdleCheck):
+                        self._idle_check_pending = False
                     result = await self._process_tick(tick)
                     if result is not None:
                         return result
@@ -334,6 +345,7 @@ class _ControlLoopRunner:
                 # optimization
                 if was_buffered:
                     now = await self.adapter.get_now()
+
                 # Calculate timeout for next scheduled wakeup
                 timeout = self.next_wakeup_timeout(now)
 
@@ -491,18 +503,29 @@ def _reduce_tick(
     tick: WorkflowTick, init: BrokerState, now_seconds: float
 ) -> tuple[BrokerState, list[WorkflowCommand]]:
     if isinstance(tick, TickStepResult):
-        return _process_step_result_tick(tick, init, now_seconds)
+        state, commands = _process_step_result_tick(tick, init, now_seconds)
     elif isinstance(tick, TickAddEvent):
-        return _process_add_event_tick(tick, init, now_seconds)
+        state, commands = _process_add_event_tick(tick, init, now_seconds)
     elif isinstance(tick, TickCancelRun):
-        return _process_cancel_run_tick(tick, init)
+        state, commands = _process_cancel_run_tick(tick, init)
     elif isinstance(tick, TickPublishEvent):
-        return _process_publish_event_tick(tick, init)
+        state, commands = _process_publish_event_tick(tick, init)
     elif isinstance(tick, TickTimeout):
-        return _process_timeout_tick(tick, init)
+        state, commands = _process_timeout_tick(tick, init)
+    elif isinstance(tick, TickIdleCheck):
+        # Return early — idle check ticks don't schedule further idle checks
+        if _check_idle_state(init):
+            return init, [CommandPublishEvent(WorkflowIdleEvent())]
+        return init, []
     else:
         raise ValueError(f"Unknown tick type: {type(tick)}")
 
+    # After any non-idle-check tick, schedule an idle check if state is quiescent
+    if _check_idle_state(state):
+        commands.append(CommandScheduleIdleCheck())
+
+    return state, commands
+
 
 def rewind_in_progress(
     state: BrokerState,
@@ -534,13 +557,12 @@ def rewind_in_progress(
 
 
 def _check_idle_state(state: BrokerState) -> bool:
-    """Returns True if workflow is idle (waiting only on external events).
+    """Returns True if workflow is idle (no work can advance internally).
 
     A workflow is idle when:
     1. The workflow is running (hasn't completed/failed/cancelled)
     2. All steps have no pending events in their queues
     3. All steps have no workers currently executing
-    4. At least one step has an active waiter (from ctx.wait_for_event())
     """
     if not state.is_running:
         return False
@@ -549,7 +571,7 @@ def _check_idle_state(state: BrokerState) -> bool:
         if worker_state.queue or worker_state.in_progress:
             return False
 
-    return any(ws.collected_waiters for ws in state.workers.values())
+    return True
 
 
 def _process_step_result_tick(
@@ -752,13 +774,6 @@ def _process_step_result_tick(
             )
             commands.extend(subcommands)
 
-    # Check for idle transition at end of processing
-    was_idle = _check_idle_state(init)
-    now_idle = _check_idle_state(state)
-
-    if now_idle and not was_idle:
-        commands.append(CommandPublishEvent(WorkflowIdleEvent()))
-
     return state, commands
 
 

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

@@ -59,6 +59,19 @@ class CommandPublishEvent:
     event: Event
 
 
+@dataclass(frozen=True)
+class CommandScheduleIdleCheck:
+    """Schedule a deferred idle check via TickIdleCheck.
+
+    Returned by the reducer when state looks quiescent after processing a tick.
+    The runner appends a TickIdleCheck to tick_buffer so that idle is confirmed
+    on the next loop iteration, after an asyncio.sleep(0) yield gives in-flight
+    ctx.send_event() calls a chance to drain.
+    """
+
+    pass
+
+
 WorkflowCommand = Union[
     CommandRunWorker,
     CommandQueueEvent,
@@ -66,6 +79,7 @@ WorkflowCommand = Union[
     CommandCompleteRun,
     CommandFailWorkflow,
     CommandPublishEvent,
+    CommandScheduleIdleCheck,
 ]
 
 

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

@@ -67,8 +67,27 @@ class TickTimeout(BaseModel):
     timeout: float
 
 
+class TickIdleCheck(BaseModel):
+    """Scheduled after state appears idle, to re-check after async events drain.
+
+    Appended to tick_buffer when the reducer sees quiescent state. Processed
+    on the next loop iteration after asyncio.sleep(0), giving in-flight
+    ctx.send_event() calls a chance to deliver via the pull task.
+    """
+
+    model_config = ConfigDict(frozen=True)
+    type: Literal["idle_check"] = "idle_check"
+
+
 WorkflowTick = Annotated[
-    Union[TickStepResult, TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout],
+    Union[
+        TickStepResult,
+        TickAddEvent,
+        TickCancelRun,
+        TickPublishEvent,
+        TickTimeout,
+        TickIdleCheck,
+    ],
     Discriminator("type"),
 ]
 

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

@@ -842,23 +842,24 @@ async def test_control_loop_emits_idle_event_when_waiting(
     test_plugin: MockRunAdapter,
 ) -> None:
     """
-    Test that WorkflowIdleEvent is emitted when workflow becomes idle waiting for event.
+    Test that WorkflowIdleEvent is emitted when workflow becomes idle.
 
-    A workflow is idle when it has active waiters but no pending work. This test
-    validates that the idle event is published to the stream when this state is reached.
+    A workflow is idle when all steps have empty queues and no in-progress
+    workers. This uses a two-step pattern: the first step completes (leaving
+    state idle), and the second step accepts an external event to finish.
     """
 
-    class AwaitedEvent(Event):
+    class ExternalEvent(HumanResponseEvent):
         value: str
 
     class IdleTrackingWorkflow(Workflow):
         @step
-        async def waiter(self, ev: StartEvent, ctx: Context) -> StopEvent:
-            awaited = await ctx.wait_for_event(
-                AwaitedEvent,
-                waiter_event=InputRequiredEvent(),
-            )
-            return StopEvent(result=f"received_{awaited.value}")
+        async def start(self, ev: StartEvent) -> None:
+            pass
+
+        @step
+        async def finish(self, ev: ExternalEvent) -> StopEvent:
+            return StopEvent(result=f"received_{ev.value}")
 
     wf = IdleTrackingWorkflow(timeout=2.0)
     task = asyncio.create_task(
@@ -870,9 +871,62 @@ async def test_control_loop_emits_idle_event_when_waiting(
     )
 
     # Collect events until we see the WorkflowIdleEvent
+    idle_event_found = False
+
+    while True:
+        ev = await test_plugin.get_stream_event(timeout=1.0)
+        if isinstance(ev, WorkflowIdleEvent):
+            idle_event_found = True
+            break
+        if isinstance(ev, StopEvent):
+            break
+
+    assert idle_event_found, (
+        "WorkflowIdleEvent should be emitted when workflow has no pending work"
+    )
+
+    # Now send the external event to complete the workflow
+    await test_plugin.send_event(TickAddEvent(event=ExternalEvent(value="test")))
+
+    result = await asyncio.wait_for(task, timeout=1.0)
+    assert isinstance(result, StopEvent)
+    assert result.result == "received_test"
+
+
+@pytest.mark.asyncio
+async def test_control_loop_emits_idle_event_with_wait_for_event(
+    test_plugin: MockRunAdapter,
+) -> None:
+    """WorkflowIdleEvent fires when a step uses ctx.wait_for_event().
+
+    wait_for_event raises an internal exception that registers a waiter and
+    releases the worker. After that, the state has no queued events and no
+    in-progress workers, so the workflow is idle.
+    """
+
+    class AwaitedEvent(Event):
+        value: str
+
+    class WaitForEventWorkflow(Workflow):
+        @step
+        async def waiter(self, ev: StartEvent, ctx: Context) -> StopEvent:
+            awaited = await ctx.wait_for_event(
+                AwaitedEvent,
+                waiter_event=InputRequiredEvent(),
+            )
+            return StopEvent(result=f"received_{awaited.value}")
+
+    wf = WaitForEventWorkflow(timeout=2.0)
+    task = asyncio.create_task(
+        run_control_loop(
+            workflow=wf,
+            start_event=StartEvent(),
+            test_runtime=test_plugin,
+        )
+    )
+
     idle_event_found = False
     input_required_found = False
-    events_before_idle: list[Event] = []
 
     while True:
         ev = await test_plugin.get_stream_event(timeout=1.0)
@@ -881,17 +935,15 @@ async def test_control_loop_emits_idle_event_when_waiting(
             break
         if isinstance(ev, InputRequiredEvent):
             input_required_found = True
-        events_before_idle.append(ev)
         if isinstance(ev, StopEvent):
             break
 
-    # WorkflowIdleEvent should be emitted after the workflow enters wait state
     assert idle_event_found, (
         "WorkflowIdleEvent should be emitted when workflow is waiting for external event"
     )
     assert input_required_found, "InputRequiredEvent should be emitted before idle"
 
-    # Now send the awaited event to complete the workflow
+    # Send the awaited event to complete the workflow
     await test_plugin.send_event(TickAddEvent(event=AwaitedEvent(value="test")))
 
     result = await asyncio.wait_for(task, timeout=1.0)

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

@@ -715,10 +715,9 @@ def test_check_idle_state_has_in_progress(base_state: BrokerState) -> None:
     assert _check_idle_state(base_state) is False
 
 
-def test_check_idle_state_no_waiters(base_state: BrokerState) -> None:
-    """A workflow with no waiters is not idle (even with empty queues)."""
-    # State is running, no queue, no in_progress, but no waiters either
-    assert _check_idle_state(base_state) is False
+def test_check_idle_state_no_work_is_idle(base_state: BrokerState) -> None:
+    """A running workflow with empty queues and no in-progress work is idle."""
+    assert _check_idle_state(base_state) is True
 
 
 def test_check_idle_state_is_idle_with_waiter(base_state: BrokerState) -> None:
@@ -735,47 +734,31 @@ def test_check_idle_state_is_idle_with_waiter(base_state: BrokerState) -> None:
     assert _check_idle_state(base_state) is True
 
 
-def test_idle_event_emitted_on_transition_to_idle(base_state: BrokerState) -> None:
-    """WorkflowIdleEvent is emitted when workflow transitions to idle."""
+def test_step_result_does_not_emit_idle(base_state: BrokerState) -> None:
+    """Step result tick never emits WorkflowIdleEvent directly.
+
+    Idle detection is handled at the runner level via TickIdleCheck, not in
+    the pure reducer. This test confirms the reducer doesn't emit idle.
+    """
     event = MyTestEvent(value=42)
     add_worker(base_state, event)
 
-    # Add a waiter so the workflow can become idle
-    waiter = StepWorkerWaiter(
-        waiter_id="w1",
-        event=event,
-        waiting_for_event=OtherEvent,
-        requirements={},
-        has_requirements=False,
-        resolved_event=None,
-    )
-    base_state.workers["test_step"].collected_waiters.append(waiter)
-
-    # Process result that completes the worker but leaves waiter active
-    result = AddWaiter(
-        waiter_id="w1",
-        waiter_event=None,
-        requirements={},
-        timeout=None,
-        event_type=OtherEvent,
-    )
-
     tick: TickStepResult = TickStepResult(
         step_name="test_step",
         worker_id=0,
         event=event,
-        result=[cast(StepFunctionResult, result)],
+        result=[StepWorkerResult(result=None)],
     )
 
     new_state, commands = _process_step_result_tick(tick, base_state, now_seconds=110.0)
 
-    # Should have WorkflowIdleEvent as the last command
     idle_commands = [
         c
         for c in commands
         if isinstance(c, CommandPublishEvent) and isinstance(c.event, WorkflowIdleEvent)
     ]
-    assert len(idle_commands) == 1
+    assert len(idle_commands) == 0
+    # State IS idle (no queued work, no in-progress), but emission is the runner's job
     assert _check_idle_state(new_state) is True
 
 

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

@@ -131,7 +131,6 @@ async def test_internal_events_multiple_workers(
         start_event=StartEvent(message="hello"),  # type: ignore
         exclude_events=[StopEvent],
     )
-    assert all(isinstance(ev, StepStateChanged) for ev in result.collected)
     collected = [ev for ev in result.collected if isinstance(ev, StepStateChanged)]
     run_ids = [
         str(r.worker_id) + r.name