[Bug]: Timeout in EventConsumer causes loss of events

[pilartomas]

What happened?

This is a code snippet from the EventConsumer implementation:

event = await asyncio.wait_for(
    self.queue.dequeue_event(), timeout=self._timeout
)

It forcefully times out the dequeue_event() coroutine which might have already picked up the event internally. This will result in a loss of that event.

I'm experiencing this behaviour with Redis implementation of the EventQueue. It does a blocking XREADGROUP followed by XACK. It appears that even XREADGROUP (as implemented in redis-py) in isolation isn't capable of handling the timeout.

Code of Conduct

• I agree to follow this project's Code of Conduct

[rohityan]

Hi @pilartomas , Thank you for reporting this issue. We will get back to you regarding any update.

[rohityan]

Hi @ishymko , PTAL

[mtsmyassin]

Root Cause Analysis + Proposed Fix

I've reviewed the EventConsumer.consume_all() implementation and can confirm the bug.

Root Cause

The problem is in the asyncio.wait_for() call inside consume_all():

event = await asyncio.wait_for(
    self.queue.dequeue_event(), timeout=self._timeout
)

asyncio.wait_for() cancels the inner coroutine when the timeout expires. For the in-memory asyncio.Queue, the get() coroutine atomically removes the item from the queue before returning it to the caller. When wait_for cancels the coroutine at the wrong moment, the item has already been dequeued from the queue but is then lost because the caller receives a TimeoutError instead of the item.

This is a well-known Python asyncio footgun. See the Python docs note that cancellation of wait_for may leave the wrapped coroutine running in an inconsistent state.

The Fix

Instead of using asyncio.wait_for, use a shield + poll pattern, or restructure to avoid timing out the dequeue itself. The safest approach is to decouple the exception-checking from the dequeue:

async def consume_all(self) -> AsyncGenerator[Event]:
    """Consume all the generated streaming events from the agent."""
    logger.debug('Starting to consume all events from the queue.')
    while True:
        if self._exception:
            raise self._exception
        try:
            # Use shield so that cancellation of wait_for does NOT cancel
            # the underlying dequeue_event() coroutine. The event will not
            # be lost if the timeout fires.
            event = await asyncio.wait_for(
                asyncio.shield(self.queue.dequeue_event()),
                timeout=self._timeout,
            )
        except TimeoutError:
            continue
        except asyncio.TimeoutError:  # Python < 3.11 compat
            continue
        except (QueueClosed, asyncio.QueueEmpty):
            if self.queue.is_closed():
                break
            continue
        except ValidationError:
            logger.exception('Invalid event format received')
            continue
        except Exception as e:
            logger.exception('Stopping event consumption due to exception')
            self._exception = e
            continue

        logger.debug('Dequeued event of type: %s in consume_all.', type(event))
        self.queue.task_done()

        is_final_event = (
            (isinstance(event, TaskStatusUpdateEvent) and event.final)
            or isinstance(event, Message)
            or (
                isinstance(event, Task)
                and event.status.state in (
                    TaskState.completed, TaskState.canceled, TaskState.failed,
                    TaskState.rejected, TaskState.unknown, TaskState.input_required,
                )
            )
        )
        if is_final_event:
            logger.debug('Stopping event consumption in consume_all.')
            await self.queue.close(True)
            yield event
            break
        yield event

The key change is wrapping self.queue.dequeue_event() in asyncio.shield(). With shield, when wait_for times out, the cancellation is sent to the shield wrapper, not to the underlying dequeue_event() coroutine. The underlying coroutine keeps running and will complete normally — it just means the event will sit in a temporary holding state until the next await.

Alternative Fix: Separate Polling from Exception Checking

Another approach that avoids wait_for entirely is to use asyncio.wait() with a separate "exception occurred" event:

async def consume_all(self) -> AsyncGenerator[Event]:
    exception_event = asyncio.Event()
    
    # Signal the event when an exception is set
    original_callback = self.agent_task_callback
    
    while True:
        if self._exception:
            raise self._exception
        try:
            # Try non-blocking first
            event = await self.queue.dequeue_event(no_wait=True)
        except asyncio.QueueEmpty:
            # Fall back to async wait with a short sleep to allow
            # exception checking without risking event loss
            await asyncio.sleep(self._timeout)
            continue
        # ... rest of processing

The asyncio.shield() approach is cleaner and the recommended fix.

Impact

This bug affects any EventQueue implementation where dequeue_event() uses non-atomic get operations — notably Redis-backed queues (XREADGROUP + XACK pattern) where the event is "consumed" from the stream before being acknowledged. Losing the event means it can never be re-consumed from the pending entries list without manual intervention.

Happy to submit a PR with the asyncio.shield() fix if this analysis is correct.

[ishymko]

Hi @pilartomas, this infrastructure changed significantly in 1.0 alpha versions (the latest is v1.0.0-alpha.1), could you please check if you still experience the issue there?

Unfortunately fixing this for 0.3.x is not feasible at this point.

[pilartomas]

@ishymko yes, this implementation hasn't changed. The issue is still present.

[ishymko]

@pilartomas yeah the code is there but EventConsumer which has this code is not used from the new request handler.

[pilartomas]

Ah thanks. If there is a new implementation entirely, then we can probably close this issue as it appears no longer relevant.