fix: make event_consumer tolerant to closed queues on py3.13 (#407)

# Issue

On Python 3.13, the closed-queue signal is `asyncio.QueueShutDown`.
`event_consumer.py` aliased `asyncio.QueueShutDown` to `ClosedQueue`,
but `asyncio.QueueEmpty` can still arise on py3.13 and it's not handled
properly in `consume_all()`'s except blocks.

# How it's reproduced

In any `AgentExecutor` class:
```
# (...)
async def execute(
        self,
        request: RequestContext,
        event_queue: EventQueue,
    ) -> None:
        await event_queue.enqueue_event(
            new_task(request.message)
        )
```

In python < 3.13:

- Sending a `message/send` request works, a task is added to the
`InMemoryTaskStore`.

In python >= 3.13:

- Sending a `message/send` request crashes with exception:
```
File "venv/lib/python3.13/site-packages/a2a/server/apps/jsonrpc/jsonrpc_app.py", line 219, in _handle_requests
    return await self._process_non_streaming_request(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        request_id, a2a_request, call_context
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "venv/lib/python3.13/site-packages/a2a/server/apps/jsonrpc/jsonrpc_app.py", line 306, in _process_non_streaming_request
    handler_result = await self.handler.on_message_send(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        request_obj, context
        ^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "venv/lib/python3.13/site-packages/a2a/utils/telemetry.py", line 162, in async_wrapper
    result = await func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.13/site-packages/a2a/server/request_handlers/jsonrpc_handler.py", line 87, in on_message_send
    task_or_message = await self.request_handler.on_message_send(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        request.params, context
        ^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "venv/lib/python3.13/site-packages/a2a/utils/telemetry.py", line 162, in async_wrapper
    result = await func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.13/site-packages/a2a/server/request_handlers/default_request_handler.py", line 282, in on_message_send
    ) = await result_aggregator.consume_and_break_on_interrupt(consumer)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.13/site-packages/a2a/server/tasks/result_aggregator.py", line 115, in consume_and_break_on_interrupt
    async for event in event_stream:
    ...<20 lines>...
            break
  File "venv/lib/python3.13/site-packages/a2a/server/events/event_consumer.py", line 87, in consume_all
    raise self._exception
  File "venv/lib/python3.13/site-packages/a2a/server/events/event_consumer.py", line 94, in consume_all
    event = await asyncio.wait_for(
            ^^^^^^^^^^^^^^^^^^^^^^^
        self.queue.dequeue_event(), timeout=self._timeout
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/opt/homebrew/Cellar/[email protected]/3.13.5/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/tasks.py", line 507, in wait_for
    return await fut
           ^^^^^^^^^
  File "venv/lib/python3.13/site-packages/a2a/utils/telemetry.py", line 162, in async_wrapper
    result = await func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.13/site-packages/a2a/server/events/event_queue.py", line 95, in dequeue_event
    raise asyncio.QueueEmpty('Queue is closed.')
asyncio.queues.QueueEmpty: Queue is closed.
```
 
 # Fix

## Code
- `event_consumer.consume_all`:
- Catch `(QueueClosed, asyncio.QueueEmpty)` and break only when
`queue.is_closed()` is True; otherwise continue polling.

- `event_queue.dequeue_event`:
- Version-guard the early-raise: on <3.13 keep raising `QueueEmpty` when
closed+empty; on ≥3.13 skip the early raise and rely on
`queue.shutdown()` to surface `QueueShutDown` exceptions.

## Tests
Added 2 tests which fail on current implementation, but pass after the
fix.

---------

Co-authored-by: taralesc <[email protected]>

Author: ovidiutaralesca

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

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()`