Merge branch 'main' into chore/improve-coverage-grpc-client

Author: holtskinner

src/a2a/server/events/event_consumer.py

@@ -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.')
 

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)

tests/server/events/test_event_queue.py

@@ -169,7 +169,9 @@ async def test_enqueue_event_propagates_to_children(
 
 
 @pytest.mark.asyncio
-async def test_enqueue_event_when_closed(event_queue: EventQueue) -> None:
+async def test_enqueue_event_when_closed(
+    event_queue: EventQueue, expected_queue_closed_exception
+) -> None:
     """Test that no event is enqueued if the parent queue is closed."""
     await event_queue.close()  # Close the queue first
 
@@ -178,7 +180,7 @@ async def test_enqueue_event_when_closed(event_queue: EventQueue) -> None:
     await event_queue.enqueue_event(event)
 
     # Verify the queue is still empty
-    with pytest.raises(asyncio.QueueEmpty):
+    with pytest.raises(expected_queue_closed_exception):
         await event_queue.dequeue_event(no_wait=True)
 
     # Also verify child queues are not affected directly by parent's enqueue attempt when closed
@@ -192,7 +194,7 @@ async def test_enqueue_event_when_closed(event_queue: EventQueue) -> None:
     await (
         child_queue.close()
     )  # ensure child is also seen as closed for this test's purpose
-    with pytest.raises(asyncio.QueueEmpty):
+    with pytest.raises(expected_queue_closed_exception):
         await child_queue.dequeue_event(no_wait=True)
 
 
@@ -214,7 +216,7 @@ async def test_dequeue_event_closed_and_empty_no_wait(
     with pytest.raises(expected_queue_closed_exception):
         event_queue.queue.get_nowait()
 
-    with pytest.raises(asyncio.QueueEmpty, match='Queue is closed.'):
+    with pytest.raises(expected_queue_closed_exception):
         await event_queue.dequeue_event(no_wait=True)
 
 
@@ -230,7 +232,8 @@ async def test_dequeue_event_closed_and_empty_waits_then_raises(
 
     # This test is tricky because await event_queue.dequeue_event() would hang if not for the close check.
     # The current implementation's dequeue_event checks `is_closed` first.
-    # If closed and empty, it raises QueueEmpty immediately.
+    # If closed and empty, it raises QueueEmpty immediately (on Python <= 3.12).
+    # On Python 3.13+, this check is skipped and asyncio.Queue.get() raises QueueShutDown instead.
     # The "waits_then_raises" scenario described in the subtask implies the `get()` might wait.
     # However, the current code:
     # async with self._lock:
@@ -240,7 +243,7 @@ async def test_dequeue_event_closed_and_empty_waits_then_raises(
     # event = await self.queue.get() -> this line is not reached if closed and empty.
 
     # So, for the current implementation, it will raise QueueEmpty immediately.
-    with pytest.raises(asyncio.QueueEmpty, match='Queue is closed.'):
+    with pytest.raises(expected_queue_closed_exception):
         await event_queue.dequeue_event(no_wait=False)
 
     # If the implementation were to change to allow `await self.queue.get()`

tests/server/request_handlers/test_default_request_handler.py

@@ -405,6 +405,134 @@ async def get_current_result():
     mock_agent_executor.execute.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_on_message_send_with_push_notification_in_non_blocking_request():
+    """Test that push notification callback is called during background event processing for non-blocking requests."""
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_push_notification_store = AsyncMock(spec=PushNotificationConfigStore)
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+    mock_request_context_builder = AsyncMock(spec=RequestContextBuilder)
+    mock_push_sender = AsyncMock()
+
+    task_id = 'non_blocking_task_1'
+    context_id = 'non_blocking_ctx_1'
+
+    # Create a task that will be returned after the first event
+    initial_task = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.working
+    )
+
+    # Create a final task that will be available during background processing
+    final_task = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.completed
+    )
+
+    mock_task_store.get.return_value = None
+
+    # Mock request context
+    mock_request_context = MagicMock(spec=RequestContext)
+    mock_request_context.task_id = task_id
+    mock_request_context.context_id = context_id
+    mock_request_context_builder.build.return_value = mock_request_context
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        push_config_store=mock_push_notification_store,
+        request_context_builder=mock_request_context_builder,
+        push_sender=mock_push_sender,
+    )
+
+    # Configure push notification
+    push_config = PushNotificationConfig(url='http://callback.com/push')
+    message_config = MessageSendConfiguration(
+        push_notification_config=push_config,
+        accepted_output_modes=['text/plain'],
+        blocking=False,  # Non-blocking request
+    )
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            message_id='msg_non_blocking',
+            parts=[],
+            task_id=task_id,
+            context_id=context_id,
+        ),
+        configuration=message_config,
+    )
+
+    # Mock ResultAggregator with custom behavior
+    mock_result_aggregator_instance = AsyncMock(spec=ResultAggregator)
+
+    # First call returns the initial task and indicates interruption (non-blocking)
+    mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
+        initial_task,
+        True,  # interrupted = True for non-blocking
+    )
+
+    # Mock the current_result property to return the final task
+    async def get_current_result():
+        return final_task
+
+    type(mock_result_aggregator_instance).current_result = PropertyMock(
+        return_value=get_current_result()
+    )
+
+    # Track if the event_callback was passed to consume_and_break_on_interrupt
+    event_callback_passed = False
+    event_callback_received = None
+
+    async def mock_consume_and_break_on_interrupt(
+        consumer, blocking=True, event_callback=None
+    ):
+        nonlocal event_callback_passed, event_callback_received
+        event_callback_passed = event_callback is not None
+        event_callback_received = event_callback
+        return initial_task, True  # interrupted = True for non-blocking
+
+    mock_result_aggregator_instance.consume_and_break_on_interrupt = (
+        mock_consume_and_break_on_interrupt
+    )
+
+    with (
+        patch(
+            'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+            return_value=mock_result_aggregator_instance,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.get_task',
+            return_value=initial_task,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.update_with_message',
+            return_value=initial_task,
+        ),
+    ):
+        # Execute the non-blocking request
+        result = await request_handler.on_message_send(
+            params, create_server_call_context()
+        )
+
+    # Verify the result is the initial task (non-blocking behavior)
+    assert result == initial_task
+
+    # Verify that the event_callback was passed to consume_and_break_on_interrupt
+    assert event_callback_passed, (
+        'event_callback should have been passed to consume_and_break_on_interrupt'
+    )
+    assert event_callback_received is not None, (
+        'event_callback should not be None'
+    )
+
+    # Verify that the push notification was sent with the final task
+    mock_push_sender.send_notification.assert_called_with(final_task)
+
+    # Verify that the push notification config was stored
+    mock_push_notification_store.set_info.assert_awaited_once_with(
+        task_id, push_config
+    )
+
+
 @pytest.mark.asyncio
 async def test_on_message_send_with_push_notification_no_existing_Task():
     """Test on_message_send for new task sets push notification info if provided."""