fix flakes in tests

Author: adrianlyjak

packages/llama-index-workflows/src/workflows/server/server.py

@@ -1826,6 +1826,7 @@ async def _iter_events(self, timeout: float = 1) -> AsyncGenerator[Event, None]:
         Converts the queue to an async generator while the workflow is still running, and there are still events.
         For better or worse, multiple consumers will compete for events
         """
+        queue_get_task: asyncio.Task[Event] | None = None
 
         try:
             while not self.queue.empty() or (
@@ -1836,9 +1837,7 @@ async def _iter_events(self, timeout: float = 1) -> AsyncGenerator[Event, None]:
                     available_events.append(self.queue.get_nowait())
                 for event in available_events:
                     yield event
-                queue_get_task: asyncio.Task[Event] = asyncio.create_task(
-                    self.queue.get()
-                )
+                queue_get_task = asyncio.create_task(self.queue.get())
                 task_waitable = self.task
                 done, pending = await asyncio.wait(
                     {queue_get_task, task_waitable}
@@ -1848,10 +1847,31 @@ async def _iter_events(self, timeout: float = 1) -> AsyncGenerator[Event, None]:
                 )
                 if queue_get_task in done:
                     yield await queue_get_task
+                    queue_get_task = None
                 else:  # otherwise task completed, so nothing else will be published to the queue
                     queue_get_task.cancel()
+                    queue_get_task = None
                     break
         finally:
+            # Cancel any pending queue.get() task to prevent orphaned tasks from
+            # consuming events after the consumer disconnects. If the task already
+            # completed, put the event back in the queue so the next consumer sees it.
+            if queue_get_task is not None:
+                if queue_get_task.done():
+                    try:
+                        event = queue_get_task.result()
+                        # Put the event back at the front of the queue
+                        # We use a workaround since asyncio.Queue has no put_front
+                        self.queue._queue.appendleft(event)  # type: ignore[attr-defined]
+                    except (asyncio.CancelledError, Exception):
+                        pass
+                else:
+                    queue_get_task.cancel()
+                    try:
+                        await queue_get_task
+                    except asyncio.CancelledError:
+                        pass
+
             if self._on_finish is not None and self.run_handler.done():
                 # clean up the resources if the stream has been consumed
                 await self._on_finish()

packages/llama-index-workflows/tests/plugins/test_runtime_matrix.py

@@ -386,10 +386,13 @@ async def original_step(
             self, ctx: Context, ev: StartEvent
         ) -> Union[OneTestEvent, LastEvent]:
             await ctx.store.set("num_to_collect", 3)
+            # Send test4 first to ensure it's pulled from receive_queue
+            # before test_step workers complete. Events are pulled one per
+            # iteration, so ordering in receive_queue determines delivery order.
+            ctx.send_event(AnotherTestEvent(another_test_param="test4"))
             ctx.send_event(OneTestEvent(test_param="test1"))
             ctx.send_event(OneTestEvent(test_param="test2"))
             ctx.send_event(OneTestEvent(test_param="test3"))
-            ctx.send_event(AnotherTestEvent(another_test_param="test4"))
             return LastEvent()
 
         @step(num_workers=3)