Undo background task for close.

Author: bartek-w

src/a2a/server/events/event_consumer.py

@@ -32,7 +32,6 @@ def __init__(self, queue: EventQueue):
         self.queue = queue
         self._timeout = 0.5
         self._exception: BaseException | None = None
-        self._close_task: asyncio.Task[None] | None = None
         logger.debug('EventConsumer initialized')
 
     async def consume_one(self) -> Event:
@@ -107,25 +106,13 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                     )
                 )
 
-                # Initiate a graceful close in the background.
-                # We use immediate=False to ensure tapped child queues can finish
-                # draining their events (preventing data loss for slow consumers).
-                # We use clear_parent_events=True to prevent deadlocks since this
-                # consumer is stopping and won't dequeue any remaining trailing events.
-                # We use create_task instead of awaiting it directly because awaiting
-                # immediate=False would block this consumer until all child queues
-                # are also drained.
                 # Make sure the yield is after the close events, otherwise
                 # the caller may end up in a blocked state where this
                 # generator isn't called again to close things out and the
                 # other part is waiting for an event or a closed queue.
                 if is_final_event:
                     logger.debug('Stopping event consumption in consume_all.')
-                    self._close_task = asyncio.create_task(
-                        self.queue.close(
-                            immediate=False, clear_parent_events=True
-                        )
-                    )
+                    await self.queue.close(True)
                     yield event
                     break
                 yield event

src/a2a/server/events/event_queue.py

@@ -172,9 +172,7 @@ def tap(self) -> 'EventQueue':
         self._children.append(queue)
         return queue
 
-    async def close(
-        self, immediate: bool = False, clear_parent_events: bool = False
-    ) -> None:
+    async def close(self, immediate: bool = False) -> None:
         """Closes the queue for future push events and also closes all child queues.
 
         Args:
@@ -183,9 +181,6 @@ async def close(
                 `QueueShutDown`. If False (default), the queue is marked as closed to new
                 events, but existing events can still be dequeued and processed until the
                 queue is fully drained.
-            clear_parent_events: If True, completely clears all pending events from this
-                specific parent queue without processing them. This parameter is ignored if
-                `immediate=True`.
         """
         logger.debug('Closing EventQueue.')
         async with self._lock:
@@ -195,9 +190,6 @@ async def close(
 
         self.queue.shutdown(immediate)
 
-        if clear_parent_events and not immediate:
-            await self.clear_events(clear_child_queues=False)
-
         await asyncio.gather(
             *(child.close(immediate) for child in self._children)
         )

tests/server/events/test_event_consumer.py

@@ -461,6 +461,7 @@ async def test_consume_all_handles_validation_error(
         )
 
 
+@pytest.mark.xfail(reason='https://github.com/a2aproject/a2a-python/issues/869')
 @pytest.mark.asyncio
 async def test_graceful_close_allows_tapped_queues_to_drain() -> None:
 
@@ -515,30 +516,25 @@ async def slow_consume() -> list:
     assert len(slow_events) == 3
 
 
+@pytest.mark.xfail(reason='https://github.com/a2aproject/a2a-python/issues/869', raises=asyncio.TimeoutError)
 @pytest.mark.asyncio
 async def test_background_close_deadlocks_on_trailing_events() -> None:
     queue = EventQueue()
-    consumer = EventConsumer(queue)
 
     # Producer enqueues a final event, but then enqueues another event
     # (e.g., simulating a delayed log message, race condition, or multiple messages).
     await queue.enqueue_event(Message(message_id='final'))
     await queue.enqueue_event(Message(message_id='trailing_log'))
 
-    # Consume events. The consumer will break its internal loop on the 'final' message,
-    # leaving 'trailing_log' in the queue forever.
-    events = [event async for event in consumer.consume_all()]
+    # Consumer dequeues 'final' but stops there (e.g. because it is a final event).
+    event = await queue.dequeue_event()
+    assert isinstance(event, Message) and event.message_id == 'final'
+    queue.task_done()
 
-    assert len(events) == 1
-    assert consumer._close_task is not None
-
-    # The background close task awaits queue.close(immediate=False), which waits for
-    # queue.join(). Since we passed clear_parent_events=True, it clears trailing events
-    # and successfully completes without deadlocking.
-    try:
-        await asyncio.wait_for(consumer._close_task, timeout=0.1)
-    except asyncio.TimeoutError:
-        pytest.fail('Background close task deadlocked on trailing events!')
+    # Now attempt a graceful close. This demonstrates the deadlock that
+    # the previous implementation (with background task and clear_parent_events)
+    # was trying to solve.
+    await asyncio.wait_for(queue.close(immediate=False), timeout=0.1)
 
 
 @pytest.mark.asyncio

tests/server/events/test_event_queue.py

@@ -482,15 +482,14 @@ async def test_close_propagates_to_children(event_queue: EventQueue) -> None:
     assert child_queue2.is_closed()
 
 
+@pytest.mark.xfail(reason='https://github.com/a2aproject/a2a-python/issues/869')
 @pytest.mark.asyncio
 async def test_enqueue_close_race_condition() -> None:
     queue = EventQueue()
     event = create_sample_message()
 
     enqueue_task = asyncio.create_task(queue.enqueue_event(event))
-    close_task = asyncio.create_task(
-        queue.close(immediate=False, clear_parent_events=True)
-    )
+    close_task = asyncio.create_task(queue.close(immediate=False))
 
     try:
         results = await asyncio.wait_for(
@@ -582,19 +581,16 @@ async def getter():
 
 
 @pytest.mark.parametrize(
-    'immediate, clear_parent_events, expected_events, close_blocks',
+    'immediate, expected_events, close_blocks',
     [
-        (False, False, (1, 1), True),
-        (False, True, (0, 1), True),
-        (True, False, (0, 0), False),
-        (True, True, (0, 0), False),
+        (False, (1, 1), True),
+        (True, (0, 0), False),
     ],
 )
 @pytest.mark.asyncio
 async def test_event_queue_close_behaviors(
     event_queue: EventQueue,
     immediate: bool,
-    clear_parent_events: bool,
     expected_events: tuple[int, int],
     close_blocks: bool,
 ) -> None:
@@ -611,11 +607,7 @@ async def test_event_queue_close_behaviors(
     event_queue.queue = QueueJoinWrapper(event_queue.queue, join_reached)
     child_queue.queue = QueueJoinWrapper(child_queue.queue, join_reached)
 
-    close_task = asyncio.create_task(
-        event_queue.close(
-            immediate=immediate, clear_parent_events=clear_parent_events
-        )
-    )
+    close_task = asyncio.create_task(event_queue.close(immediate=immediate))
 
     if close_blocks:
         await join_reached.wait()