fix: return background task from consume_and_break_on_interrupt to prevent GC (#775)

# Description

ResultAggregator.consume_and_break_on_interrupt creates a background
asyncio.Task to continue consuming events after an interruption
(non-blocking or auth_required), but discards the task reference. On
Python 3.12+ the event loop only holds weak references to tasks, so the
garbage collector can silently collect the task before it completes —
dropping remaining events (completed/failed status) and push
notification callbacks.

Return the background task as a third tuple element so callers can hold
a strong reference. DefaultRequestHandler.on_message_send now tracks it
via _track_background_task(), the same mechanism already used for other
background work.

- [x] Follow the [`CONTRIBUTING`
Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md).
- [x] Make your Pull Request title in the
<https://www.conventionalcommits.org/> specification.
- Important Prefixes for
[release-please](https://github.com/googleapis/release-please):
- `fix:` which represents bug fixes, and correlates to a
[SemVer](https://semver.org/) patch.
- `feat:` represents a new feature, and correlates to a SemVer minor.
- `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking
change (indicated by the `!`) and will result in a SemVer major.
- [x] Ensure the tests and linter pass (Run `bash scripts/format.sh`
from the repository root to format)

Fixes #774 🦕

---------

Co-authored-by: Ivan Shymko <ishymko@google.com>

Author: kevinlu310

src/a2a/server/request_handlers/default_request_handler.py

@@ -322,12 +322,17 @@ async def push_notification_callback() -> None:
             (
                 result,
                 interrupted_or_non_blocking,
+                bg_consume_task,
             ) = await result_aggregator.consume_and_break_on_interrupt(
                 consumer,
                 blocking=blocking,
                 event_callback=push_notification_callback,
             )
 
+            if bg_consume_task is not None:
+                bg_consume_task.set_name(f'continue_consuming:{task_id}')
+                self._track_background_task(bg_consume_task)
+
         except Exception:
             logger.exception('Agent execution failed')
             producer_task.cancel()

src/a2a/server/tasks/result_aggregator.py

@@ -99,7 +99,7 @@ async def consume_and_break_on_interrupt(
         consumer: EventConsumer,
         blocking: bool = True,
         event_callback: Callable[[], Awaitable[None]] | None = None,
-    ) -> tuple[Task | Message | None, bool]:
+    ) -> tuple[Task | Message | None, bool, asyncio.Task | None]:
         """Processes the event stream until completion or an interruptable state is encountered.
 
         If `blocking` is False, it returns after the first event that creates a Task or Message.
@@ -119,16 +119,23 @@ async def consume_and_break_on_interrupt(
             A tuple containing:
             - The current aggregated result (`Task` or `Message`) at the point of completion or interruption.
             - A boolean indicating whether the consumption was interrupted (`True`) or completed naturally (`False`).
+            - The background ``asyncio.Task`` that continues consuming events
+              after an interruption, or ``None`` when no background work was
+              spawned.  **Callers must hold a strong reference** to this task
+              (e.g. in a ``set``) to prevent the garbage collector from
+              collecting it before it finishes — the event loop only keeps
+              weak references to tasks.
 
         Raises:
             BaseException: If the `EventConsumer` raises an exception during consumption.
         """
         event_stream = consumer.consume_all()
         interrupted = False
+        bg_task: asyncio.Task | None = None
         async for event in event_stream:
             if isinstance(event, Message):
                 self._message = event
-                return event, False
+                return event, False, None
             await self.task_manager.process(event)
 
             should_interrupt = False
@@ -158,13 +165,13 @@ 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(  # noqa: RUF006
+                # The caller is responsible for tracking this task to prevent GC.
+                bg_task = asyncio.create_task(
                     self._continue_consuming(event_stream, event_callback)
                 )
                 interrupted = True
                 break
-        return await self.task_manager.get_task(), interrupted
+        return await self.task_manager.get_task(), interrupted, bg_task
 
     async def _continue_consuming(
         self,

tests/server/request_handlers/test_default_request_handler.py

@@ -421,6 +421,7 @@ async def test_on_message_send_with_push_notification():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         final_task_result,
         False,
+        None,
     )
 
     # Mock the current_result property to return the final task result
@@ -520,6 +521,7 @@ async def test_on_message_send_with_push_notification_in_non_blocking_request():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         initial_task,
         True,  # interrupted = True for non-blocking
+        MagicMock(spec=asyncio.Task),  # background task
     )
 
     # Mock the current_result property to return the final task
@@ -540,7 +542,11 @@ async def mock_consume_and_break_on_interrupt(
         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
+        return (
+            initial_task,
+            True,
+            MagicMock(spec=asyncio.Task),
+        )  # interrupted = True for non-blocking
 
     mock_result_aggregator_instance.consume_and_break_on_interrupt = (
         mock_consume_and_break_on_interrupt
@@ -631,6 +637,7 @@ async def test_on_message_send_with_push_notification_no_existing_Task():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         final_task_result,
         False,
+        None,
     )
 
     # Mock the current_result property to return the final task result
@@ -689,6 +696,7 @@ async def test_on_message_send_no_result_from_aggregator():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         None,
         False,
+        None,
     )
 
     from a2a.utils.errors import ServerError  # Local import
@@ -740,6 +748,7 @@ async def test_on_message_send_task_id_mismatch():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         mismatched_task,
         False,
+        None,
     )
 
     from a2a.utils.errors import ServerError  # Local import
@@ -950,6 +959,7 @@ async def test_on_message_send_interrupted_flow():
     mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
         interrupt_task_result,
         True,
+        MagicMock(spec=asyncio.Task),  # background task
     )  # Interrupted = True
 
     # Patch asyncio.create_task to verify _cleanup_producer is scheduled

tests/server/tasks/test_result_aggregator.py

@@ -228,12 +228,14 @@ async def mock_consume_generator():
         (
             result,
             interrupted,
+            bg_task,
         ) = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer
         )
 
         self.assertEqual(result, sample_message)
         self.assertFalse(interrupted)
+        self.assertIsNone(bg_task)
         self.mock_task_manager.process.assert_not_called()  # Process is not called for the Message if returned directly
         # _continue_consuming should not be called if it's a message interrupt
         # and no auth_required state.
@@ -265,12 +267,14 @@ async def mock_consume_generator():
         (
             result,
             interrupted,
+            bg_task,
         ) = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer
         )
 
         self.assertEqual(result, auth_task)
         self.assertTrue(interrupted)
+        self.assertIsNotNone(bg_task)
         self.mock_task_manager.process.assert_called_once_with(auth_task)
         mock_create_task.assert_called_once()  # Check that create_task was called
         # self.aggregator._continue_consuming is an AsyncMock.
@@ -317,12 +321,14 @@ async def mock_consume_generator():
         (
             result,
             interrupted,
+            bg_task,
         ) = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer
         )
 
         self.assertEqual(result, current_task_state_after_update)
         self.assertTrue(interrupted)
+        self.assertIsNotNone(bg_task)
         self.mock_task_manager.process.assert_called_once_with(
             auth_status_update
         )
@@ -353,13 +359,15 @@ async def mock_consume_generator():
         (
             result,
             interrupted,
+            bg_task,
         ) = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer
         )
 
         # If the first event is a Message, it's returned directly.
         self.assertEqual(result, event1)
         self.assertFalse(interrupted)
+        self.assertIsNone(bg_task)
         # process() is NOT called for the Message if it's the one causing the return
         self.mock_task_manager.process.assert_not_called()
         self.mock_task_manager.get_task.assert_not_called()
@@ -415,12 +423,14 @@ async def mock_consume_generator():
         (
             result,
             interrupted,
+            bg_task,
         ) = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer, blocking=False
         )
 
         self.assertEqual(result, first_event)
         self.assertTrue(interrupted)
+        self.assertIsNotNone(bg_task)
         self.mock_task_manager.process.assert_called_once_with(first_event)
         mock_create_task.assert_called_once()
         # The background task should be created with the remaining stream
@@ -468,7 +478,7 @@ async def initial_consume_generator():
         mock_create_task.side_effect = lambda coro: asyncio.ensure_future(coro)
 
         # Call the main method that triggers _continue_consuming via create_task
-        _, _ = await self.aggregator.consume_and_break_on_interrupt(
+        _, _, _ = await self.aggregator.consume_and_break_on_interrupt(
             self.mock_event_consumer
         )