EventQueue: Create interface and make tap() async.

Author: bartek-w

src/a2a/server/events/__init__.py

@@ -1,7 +1,7 @@
 """Event handling components for the A2A server."""
 
 from a2a.server.events.event_consumer import EventConsumer
-from a2a.server.events.event_queue import Event, EventQueue
+from a2a.server.events.event_queue import Event, EventQueue, EventQueueLegacy
 from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager
 from a2a.server.events.queue_manager import (
     NoTaskQueue,
@@ -14,6 +14,7 @@
     'Event',
     'EventConsumer',
     'EventQueue',
+    'EventQueueLegacy',
     'InMemoryQueueManager',
     'NoTaskQueue',
     'QueueManager',

src/a2a/server/events/event_consumer.py

@@ -12,7 +12,6 @@
     TaskState,
     TaskStatusUpdateEvent,
 )
-from a2a.utils.errors import InternalError
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -34,31 +33,6 @@ def __init__(self, queue: EventQueue):
         self._exception: BaseException | None = None
         logger.debug('EventConsumer initialized')
 
-    async def consume_one(self) -> Event:
-        """Consume one event from the agent event queue non-blocking.
-
-        Returns:
-            The next event from the queue.
-
-        Raises:
-            InternalError: If the queue is empty when attempting to dequeue
-                immediately.
-        """
-        logger.debug('Attempting to consume one event.')
-        try:
-            event = await self.queue.dequeue_event(no_wait=True)
-        except asyncio.QueueEmpty as e:
-            logger.warning('Event queue was empty in consume_one.')
-            raise InternalError(
-                message='Agent did not return any response'
-            ) from e
-
-        logger.debug('Dequeued event of type: %s in consume_one.', type(event))
-
-        self.queue.task_done()
-
-        return event
-
     async def consume_all(self) -> AsyncGenerator[Event]:
         """Consume all the generated streaming events from the agent.
 

src/a2a/server/events/event_queue.py

@@ -2,8 +2,9 @@
 import logging
 import sys
 
+from abc import ABC, abstractmethod
 from types import TracebackType
-from typing import Any
+from typing import Any, cast
 
 from typing_extensions import Self
 
@@ -46,8 +47,129 @@ def _create_async_queue(maxsize: int = 0) -> AsyncQueue[Any]:
 DEFAULT_MAX_QUEUE_SIZE = 1024
 
 
+class EventQueue(ABC):
+    """Base class and factory for EventQueueSource.
+
+    EventQueue provides an abstraction for a queue of events that can be tapped
+    by multiple consumers.
+    EventQueue maintain main queue and source and maintain child queues in sync.
+    GUARANTEE: All sinks (including the default one) will receive events in the exact same order.
+
+    WARNING (Concurrency): All events from all sinks (both the default queue and any
+    tapped child queues) must be regularly consumed and marked as done. If any single
+    consumer stops processing and its queue reaches capacity, it can block the event
+    dispatcher and stall the entire system, causing a widespread deadlock.
+
+    WARNING (Memory Leak): Event queues spawn background tasks. To prevent memory
+    and task leaks, all queue objects (both source and sinks) MUST be explicitly
+    closed via `await queue.close()` or by using the async context manager (`async with queue:`).
+    Child queues are automatically closed when parent queue is closed, but you
+    should still close them explicitly to prevent queues from reaching capacity by
+    unconsumed events.
+
+    Typical usage:
+    queue = EventQueue()
+    child_queue1 = await queue.tap()
+    child_queue2 = await queue.tap()
+
+    async for event in child_queue1:
+        do_some_work(event)
+        child_queue1.task_done()
+    """
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self:
+        """Redirects instantiation to EventQueueLegacy for backwards compatibility."""
+        if cls is EventQueue:
+            instance = EventQueueLegacy.__new__(EventQueueLegacy)
+            EventQueueLegacy.__init__(instance, *args, **kwargs)
+            return cast('Self', instance)
+        return super().__new__(cls)
+
+    @abstractmethod
+    async def enqueue_event(self, event: Event) -> None:
+        """Pushes an event into the queue.
+
+        Only main queue can enqueue events. Child queues can only dequeue events.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def dequeue_event(self) -> Event:
+        """Pulls an event from the queue."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def task_done(self) -> None:
+        """Signals that a work on dequeued event is complete."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def tap(
+        self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE
+    ) -> 'EventQueue':
+        """Creates a child queue that receives future events.
+
+        Note: The tapped queue may receive some old events if the incoming event
+        queue is lagging behind and hasn't dispatched them yet.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def close(self, immediate: bool = False) -> None:
+        """Closes the queue.
+
+        For parent queue: it closes the main queue and all its child queues.
+        For child queue: it closes only child queue.
+
+        It is safe to call it multiple times.
+        If immediate is True, the queue will be closed without waiting for all events to be processed.
+        If immediate is False, the queue will be closed after all events are processed (and confirmed with task_done() calls).
+
+        WARNING: Closing the parent queue with immediate=False is a deadlock risk if there are unconsumed events
+        in any of the child sinks and the consumer has crashed without draining its queue.
+        It is highly recommended to wrap graceful shutdowns with a timeout, e.g.,
+        `asyncio.wait_for(queue.close(immediate=False), timeout=...)`.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def is_closed(self) -> bool:
+        """[DEPRECATED] Checks if the queue is closed.
+
+        NOTE: Relying on this for enqueue logic introduces race conditions.
+        It is maintained primarily for backwards compatibility, workarounds for
+        Python 3.10/3.12 async queues in consumers, and for the test suite.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def __aenter__(self) -> Self:
+        """Enters the async context manager, returning the queue itself.
+
+        WARNING: See `__aexit__` for important deadlock risks associated with
+        exiting this context manager if unconsumed events remain.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exits the async context manager, ensuring close() is called.
+
+        WARNING: The context manager calls `close(immediate=False)` by default.
+        If a consumer exits the `async with` block early (e.g., due to an exception
+        or an explicit `break`) while unconsumed events remain in the queue,
+        `__aexit__` will deadlock waiting for `task_done()` to be called on those events.
+        """
+        raise NotImplementedError
+
+
 @trace_class(kind=SpanKind.SERVER)
-class EventQueue:
+class EventQueueLegacy(EventQueue):
     """Event queue for A2A responses from agent.
 
     Acts as a buffer between the agent's asynchronous execution and the
@@ -63,14 +185,19 @@ def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
         if max_queue_size <= 0:
             raise ValueError('max_queue_size must be greater than 0')
 
-        self.queue: AsyncQueue[Event] = _create_async_queue(
+        self._queue: AsyncQueue[Event] = _create_async_queue(
             maxsize=max_queue_size
         )
         self._children: list[EventQueue] = []
         self._is_closed = False
         self._lock = asyncio.Lock()
         logger.debug('EventQueue initialized.')
 
+    @property
+    def queue(self) -> AsyncQueue[Event]:
+        """[DEPRECATED] Returns the underlying asyncio.Queue."""
+        return self._queue
+
     async def __aenter__(self) -> Self:
         """Enters the async context manager, returning the queue itself."""
         return self
@@ -106,7 +233,7 @@ async def enqueue_event(self, event: Event) -> None:
         for child in self._children:
             await child.enqueue_event(event)
 
-    async def dequeue_event(self, no_wait: bool = False) -> Event:
+    async def dequeue_event(self) -> Event:
         """Dequeues an event from the queue.
 
         This implementation expects that dequeue to raise an exception when
@@ -115,38 +242,23 @@ async def dequeue_event(self, no_wait: bool = False) -> Event:
         the user is awaiting the queue.get method. Python<=3.12 this needs to
         manage this lifecycle itself. The current implementation can lead to
         blocking if the dequeue_event is called before the EventQueue has been
-        closed but when there are no events on the queue. Two ways to avoid this
-        are to call this with no_wait = True which won't block, but is the
-        callers responsibility to retry as appropriate. Alternatively, one can
-        use an async Task management solution to cancel the get task if the queue
+        closed but when there are no events on the queue. One way to avoid this
+        is to use an async Task management solution to cancel the get task if the queue
         has closed or some other condition is met. The implementation of the
         EventConsumer uses an async.wait with a timeout to abort the
         dequeue_event call and retry, when it will return with a closed error.
 
-        Args:
-            no_wait: If True, retrieve an event immediately or raise `asyncio.QueueEmpty`.
-                     If False (default), wait until an event is available.
-
         Returns:
             The next event from the queue.
 
         Raises:
-            asyncio.QueueEmpty: If `no_wait` is True and the queue is empty.
             asyncio.QueueShutDown: If the queue has been closed and is empty.
         """
         async with self._lock:
             if self._is_closed and self.queue.empty():
                 logger.warning('Queue is closed. Event will not be dequeued.')
                 raise QueueShutDown('Queue is closed.')
 
-        if no_wait:
-            logger.debug('Attempting to dequeue event (no_wait=True).')
-            event = self.queue.get_nowait()
-            logger.debug(
-                'Dequeued event (no_wait=True) of type: %s', type(event)
-            )
-            return event
-
         logger.debug('Attempting to dequeue event (waiting).')
         event = await self.queue.get()
         logger.debug('Dequeued event (waited) of type: %s', type(event))
@@ -160,15 +272,17 @@ def task_done(self) -> None:
         logger.debug('Marking task as done in EventQueue.')
         self.queue.task_done()
 
-    def tap(self) -> 'EventQueue':
-        """Taps the event queue to create a new child queue that receives all future events.
+    async def tap(
+        self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE
+    ) -> 'EventQueueLegacy':
+        """Taps the event queue to create a new child queue that receives future events.
 
         Returns:
             A new `EventQueue` instance that will receive all events enqueued
             to this parent queue from this point forward.
         """
         logger.debug('Tapping EventQueue to create a child queue.')
-        queue = EventQueue()
+        queue = EventQueueLegacy(max_queue_size=max_queue_size)
         self._children.append(queue)
         return queue
 
@@ -199,48 +313,3 @@ async def close(self, immediate: bool = False) -> None:
     def is_closed(self) -> bool:
         """Checks if the queue is closed."""
         return self._is_closed
-
-    async def clear_events(self, clear_child_queues: bool = True) -> None:
-        """Clears all events from the current queue and optionally all child queues.
-
-        This method removes all pending events from the queue without processing them.
-        Child queues can be optionally cleared based on the clear_child_queues parameter.
-
-        Args:
-            clear_child_queues: If True (default), clear all child queues as well.
-                              If False, only clear the current queue, leaving child queues untouched.
-        """
-        logger.debug('Clearing all events from EventQueue and child queues.')
-
-        # Clear all events from the queue, even if closed
-        cleared_count = 0
-        async with self._lock:
-            try:
-                while True:
-                    event = self.queue.get_nowait()
-                    logger.debug(
-                        'Discarding unprocessed event of type: %s, content: %s',
-                        type(event),
-                        event,
-                    )
-                    self.queue.task_done()
-                    cleared_count += 1
-            except asyncio.QueueEmpty:
-                pass
-            except QueueShutDown:
-                pass
-
-            if cleared_count > 0:
-                logger.debug(
-                    'Cleared %d unprocessed events from EventQueue.',
-                    cleared_count,
-                )
-
-        # Clear all child queues (lock released before awaiting child tasks)
-        if clear_child_queues and self._children:
-            child_tasks = [
-                asyncio.create_task(child.clear_events())
-                for child in self._children
-            ]
-
-            await asyncio.gather(*child_tasks, return_exceptions=True)