Merge branch 'main' into fix/logging-and-string-formatting

Author: holtskinner

src/a2a/server/events/event_consumer.py

@@ -125,7 +125,7 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 # other part is waiting for an event or a closed queue.
                 if is_final_event:
                     logger.debug('Stopping event consumption in consume_all.')
-                    await self.queue.close()
+                    await self.queue.close(True)
                     yield event
                     break
                 yield event
@@ -135,7 +135,7 @@ async def consume_all(self) -> AsyncGenerator[Event]:
             except asyncio.TimeoutError:  # pyright: ignore [reportUnusedExcept]
                 # This class was made an alias of build-in TimeoutError after 3.11
                 continue
-            except QueueClosed:
+            except (QueueClosed, asyncio.QueueEmpty):
                 # Confirm that the queue is closed, e.g. we aren't on
                 # python 3.12 and get a queue empty error on an open queue
                 if self.queue.is_closed():

src/a2a/server/events/event_queue.py

@@ -90,7 +90,12 @@ async def dequeue_event(self, no_wait: bool = False) -> Event:
             asyncio.QueueShutDown: If the queue has been closed and is empty.
         """
         async with self._lock:
-            if self._is_closed and self.queue.empty():
+            if (
+                sys.version_info < (3, 13)
+                and self._is_closed
+                and self.queue.empty()
+            ):
+                # On 3.13+, skip early raise; await self.queue.get() will raise QueueShutDown after shutdown()
                 logger.warning('Queue is closed. Event will not be dequeued.')
                 raise asyncio.QueueEmpty('Queue is closed.')
 
@@ -127,25 +132,38 @@ def tap(self) -> 'EventQueue':
         self._children.append(queue)
         return queue
 
-    async def close(self) -> None:
-        """Closes the queue for future push events.
+    async def close(self, immediate: bool = False) -> None:
+        """Closes the queue for future push events and also closes all child queues.
+
+        Once closed, no new events can be enqueued. For Python 3.13+, this will trigger
+        `asyncio.QueueShutDown` when the queue is empty and a consumer tries to dequeue.
+        For lower versions, the queue will be marked as closed and optionally cleared.
+
+        Args:
+            immediate (bool):
+                - True: Immediately closes the queue and clears all unprocessed events without waiting for them to be consumed. This is suitable for scenarios where you need to forcefully interrupt and quickly release resources.
+                - False (default): Gracefully closes the queue, waiting for all queued events to be processed (i.e., the queue is drained) before closing. This is suitable when you want to ensure all events are handled.
 
-        Once closed, `dequeue_event` will eventually raise `asyncio.QueueShutDown`
-        when the queue is empty. Also closes all child queues.
         """
         logger.debug('Closing EventQueue.')
         async with self._lock:
             # If already closed, just return.
-            if self._is_closed:
+            if self._is_closed and not immediate:
                 return
-            self._is_closed = True
+            if not self._is_closed:
+                self._is_closed = True
         # If using python 3.13 or higher, use the shutdown method
         if sys.version_info >= (3, 13):
-            self.queue.shutdown()
+            self.queue.shutdown(immediate)
             for child in self._children:
-                await child.close()
+                await child.close(immediate)
         # Otherwise, join the queue
         else:
+            if immediate:
+                await self.clear_events(True)
+                for child in self._children:
+                    await child.close(immediate)
+                return
             tasks = [asyncio.create_task(self.queue.join())]
             tasks.extend(
                 asyncio.create_task(child.close()) for child in self._children
@@ -155,3 +173,53 @@ async def close(self) -> 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(
+                        f'Discarding unprocessed event of type: {type(event)}, content: {event}'
+                    )
+                    self.queue.task_done()
+                    cleared_count += 1
+            except asyncio.QueueEmpty:
+                pass
+            except Exception as e:
+                # Handle Python 3.13+ QueueShutDown
+                if (
+                    sys.version_info >= (3, 13)
+                    and type(e).__name__ == 'QueueShutDown'
+                ):
+                    pass
+                else:
+                    raise
+
+            if cleared_count > 0:
+                logger.debug(
+                    f'Cleared {cleared_count} unprocessed events from EventQueue.'
+                )
+
+        # 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
+            ]
+
+            if child_tasks:
+                await asyncio.gather(*child_tasks, return_exceptions=True)

src/a2a/server/request_handlers/default_request_handler.py

@@ -286,11 +286,19 @@ async def on_message_send(
 
         interrupted_or_non_blocking = False
         try:
+            # Create async callback for push notifications
+            async def push_notification_callback() -> None:
+                await self._send_push_notification_if_needed(
+                    task_id, result_aggregator
+                )
+
             (
                 result,
                 interrupted_or_non_blocking,
             ) = await result_aggregator.consume_and_break_on_interrupt(
-                consumer, blocking=blocking
+                consumer,
+                blocking=blocking,
+                event_callback=push_notification_callback,
             )
             if not result:
                 raise ServerError(error=InternalError())  # noqa: TRY301

src/a2a/server/tasks/result_aggregator.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 
-from collections.abc import AsyncGenerator, AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
 
 from a2a.server.events import Event, EventConsumer
 from a2a.server.tasks.task_manager import TaskManager
@@ -24,7 +24,10 @@ class ResultAggregator:
        Task object and emit that Task object.
     """
 
-    def __init__(self, task_manager: TaskManager):
+    def __init__(
+        self,
+        task_manager: TaskManager,
+    ) -> None:
         """Initializes the ResultAggregator.
 
         Args:
@@ -92,7 +95,10 @@ async def consume_all(
         return await self.task_manager.get_task()
 
     async def consume_and_break_on_interrupt(
-        self, consumer: EventConsumer, blocking: bool = True
+        self,
+        consumer: EventConsumer,
+        blocking: bool = True,
+        event_callback: Callable[[], Awaitable[None]] | None = None,
     ) -> tuple[Task | Message | None, bool]:
         """Processes the event stream until completion or an interruptable state is encountered.
 
@@ -105,6 +111,9 @@ async def consume_and_break_on_interrupt(
             consumer: The `EventConsumer` to read events from.
             blocking: If `False`, the method returns as soon as a task/message
                       is available. If `True`, it waits for a terminal state.
+            event_callback: Optional async callback function to be called after each event
+                           is processed in the background continuation.
+                           Mainly used for push notifications currently.
 
         Returns:
             A tuple containing:
@@ -150,13 +159,17 @@ async def consume_and_break_on_interrupt(
             if should_interrupt:
                 # Continue consuming the rest of the events in the background.
                 # TODO: We should track all outstanding tasks to ensure they eventually complete.
-                asyncio.create_task(self._continue_consuming(event_stream))  # noqa: RUF006
+                asyncio.create_task(  # noqa: RUF006
+                    self._continue_consuming(event_stream, event_callback)
+                )
                 interrupted = True
                 break
         return await self.task_manager.get_task(), interrupted
 
     async def _continue_consuming(
-        self, event_stream: AsyncIterator[Event]
+        self,
+        event_stream: AsyncIterator[Event],
+        event_callback: Callable[[], Awaitable[None]] | None = None,
     ) -> None:
         """Continues processing an event stream in a background task.
 
@@ -165,6 +178,9 @@ async def _continue_consuming(
 
         Args:
             event_stream: The remaining `AsyncIterator` of events from the consumer.
+            event_callback: Optional async callback function to be called after each event is processed.
         """
         async for event in event_stream:
             await self.task_manager.process(event)
+            if event_callback:
+                await event_callback()

tests/server/events/test_event_consumer.py

@@ -324,6 +324,59 @@ async def test_consume_all_continues_on_queue_empty_if_not_really_closed(
     assert mock_event_queue.is_closed.call_count == 1
 
 
+@pytest.mark.asyncio
+async def test_consume_all_handles_queue_empty_when_closed_python_version_agnostic(
+    event_consumer: EventConsumer, mock_event_queue: AsyncMock, monkeypatch
+):
+    """Ensure consume_all stops with no events when queue is closed and dequeue_event raises asyncio.QueueEmpty (Python version-agnostic)."""
+    # Make QueueClosed a distinct exception (not QueueEmpty) to emulate py3.13 semantics
+    from a2a.server.events import event_consumer as ec
+
+    class QueueShutDown(Exception):
+        pass
+
+    monkeypatch.setattr(ec, 'QueueClosed', QueueShutDown, raising=True)
+
+    # Simulate queue reporting closed while dequeue raises QueueEmpty
+    mock_event_queue.dequeue_event.side_effect = asyncio.QueueEmpty(
+        'closed/empty'
+    )
+    mock_event_queue.is_closed.return_value = True
+
+    consumed_events = []
+    async for event in event_consumer.consume_all():
+        consumed_events.append(event)
+
+    assert consumed_events == []
+    mock_event_queue.dequeue_event.assert_called_once()
+    mock_event_queue.is_closed.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_consume_all_continues_on_queue_empty_when_not_closed(
+    event_consumer: EventConsumer, mock_event_queue: AsyncMock, monkeypatch
+):
+    """Ensure consume_all continues after asyncio.QueueEmpty when queue is open, yielding the next (final) event."""
+    # First dequeue raises QueueEmpty (transient empty), then a final Message arrives
+    final = Message(role='agent', parts=[{'text': 'done'}], message_id='final')
+    mock_event_queue.dequeue_event.side_effect = [
+        asyncio.QueueEmpty('temporarily empty'),
+        final,
+    ]
+    mock_event_queue.is_closed.return_value = False
+
+    # Make the polling responsive in tests
+    event_consumer._timeout = 0.001
+
+    consumed = []
+    async for ev in event_consumer.consume_all():
+        consumed.append(ev)
+
+    assert consumed == [final]
+    assert mock_event_queue.dequeue_event.call_count == 2
+    mock_event_queue.is_closed.assert_called_once()
+
+
 def test_agent_task_callback_sets_exception(event_consumer: EventConsumer):
     """Test that agent_task_callback sets _exception if the task had one."""
     mock_task = MagicMock(spec=asyncio.Task)