fix: Prevent client disconnect from stopping task execution (#440)

# Issue * Client disconnect triggered synchronous cleanup. * That awaited the producer task, effectively tying producer lifetime to the client connection. * Reconnecting with `tasks/resubscribe` would not receive further events because the producer had already been forced to finish. This behaviour no longer raises a `asyncio.exceptions.CancelledError` like claimed in #296 due to this fix: #383, but `tasks/resubscribe` still didn't behave as expected. # How it's reproduced In any streaming agent: Simply sending a (longer-running) `message/stream`, disconnecting, and then reconnecting to the task using `tasks/resubscribe` will no longer yield events, even though the task should have been still running. # Fix ## Code The fix is an one-liner. Now: * Client disconnect schedules cleanup in the background and returns immediately. * Producer continues; resubscribe taps the existing queue and receives subsequent events. * Cleanup still runs once the producer completes. ## Tests **Existing tests:** * Changed existing tests that asserted on `AgentExecutor.execute` by adding an `asyncio.Event` latch to wait until the background producer hits `execute`. **New tests:** * `test_stream_disconnect_then_resubscribe_receives_future_events` -- start streaming, disconnect, resubscribe, and confirm future events are received. * `test_on_message_send_stream_client_disconnect_triggers_background_cleanup_and_producer_continues` -- to validate that disconnecting is non-blocking, producer continues, and cleanup completes afterward. Fixes #296

Author: ovidiutaralesca

src/a2a/server/request_handlers/default_request_handler.py

@@ -67,6 +67,7 @@ class DefaultRequestHandler(RequestHandler):
  """
 
  _running_agents: dict[str, asyncio.Task]
+ _background_tasks: set[asyncio.Task]
 
  def __init__( # noqa: PLR0913
  self,
@@ -102,6 +103,9 @@ def __init__( # noqa: PLR0913
  # TODO: Likely want an interface for managing this, like AgentExecutionManager.
  self._running_agents = {}
  self._running_agents_lock = asyncio.Lock()
+ # Tracks background tasks (e.g., deferred cleanups) to avoid orphaning
+ # asyncio tasks and to surface unexpected exceptions.
+ self._background_tasks = set()
 
  async def on_get_task(
  self,
@@ -355,10 +359,11 @@ async def push_notification_callback() -> None:
  raise
  finally:
  if interrupted_or_non_blocking:
- # TODO: Track this disconnected cleanup task.
- asyncio.create_task( # noqa: RUF006
+ cleanup_task = asyncio.create_task(
  self._cleanup_producer(producer_task, task_id)
  )
+ cleanup_task.set_name(f'cleanup_producer:{task_id}')
+ self._track_background_task(cleanup_task)
  else:
  await self._cleanup_producer(producer_task, task_id)
 
@@ -394,7 +399,11 @@ async def on_message_send_stream(
  )
  yield event
  finally:
- await self._cleanup_producer(producer_task, task_id)
+ cleanup_task = asyncio.create_task(
+ self._cleanup_producer(producer_task, task_id)
+ )
+ cleanup_task.set_name(f'cleanup_producer:{task_id}')
+ self._track_background_task(cleanup_task)
 
  async def _register_producer(
  self, task_id: str, producer_task: asyncio.Task
@@ -403,6 +412,29 @@ async def _register_producer(
  async with self._running_agents_lock:
  self._running_agents[task_id] = producer_task
 
+ def _track_background_task(self, task: asyncio.Task) -> None:
+ """Tracks a background task and logs exceptions on completion.
+
+ This avoids unreferenced tasks (and associated lint warnings) while
+ ensuring any exceptions are surfaced in logs.
+ """
+ self._background_tasks.add(task)
+
+ def _on_done(completed: asyncio.Task) -> None:
+ try:
+ # Retrieve result to raise exceptions, if any
+ completed.result()
+ except asyncio.CancelledError:
+ name = completed.get_name()
+ logger.debug('Background task %s cancelled', name)
+ except Exception:
+ name = completed.get_name()
+ logger.exception('Background task %s failed', name)
+ finally:
+ self._background_tasks.discard(completed)
+
+ task.add_done_callback(_on_done)
+
  async def _cleanup_producer(
  self,
  producer_task: asyncio.Task,