From febc1b92a21c69bf088f76a6699c828380658f2e Mon Sep 17 00:00:00 2001 From: Sankalp Gilda Date: Sat, 13 Jun 2026 16:46:42 -0400 Subject: [PATCH 1/3] feat(server): add aclose() to drain ActiveTask background tasks (#1101) At shutdown the ActiveTask producer can stay pending and surface as "Task was destroyed but it is pending!". The producer's finally closes the subscriber queue with immediate=False, which joins every subscriber sink; an abandoned subscriber leaves an undrained sink so the join never returns and the producer hangs. No public API drains these background tasks today. Add aclose() to ActiveTask, ActiveTaskRegistry, and DefaultRequestHandlerV2. It force-closes the event queues (immediate=True), which releases the wedged producer, then cancels and awaits the producer and consumer tasks. ActiveTask sets _is_finished under _lock so it is mutually exclusive with start(); the registry marks itself closed so get_or_create refuses new work during teardown, closing the orphan-task race. Fixes #1101 --- src/a2a/server/agent_execution/active_task.py | 37 ++++++ .../agent_execution/active_task_registry.py | 39 +++++++ .../default_request_handler_v2.py | 9 ++ .../agent_execution/test_active_task.py | 88 ++++++++++++++ .../test_active_task_registry.py | 107 ++++++++++++++++++ .../test_default_request_handler_v2.py | 34 ++++++ 6 files changed, 314 insertions(+) create mode 100644 tests/server/agent_execution/test_active_task_registry.py diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5bb27c0c0..29a9393a0 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -755,6 +755,43 @@ async def cancel(self, call_context: ServerCallContext) -> Task: raise RuntimeError('Task should have been created') return task + async def aclose(self) -> None: + """Force-closes the task's queues and drains its background tasks. + + Provides a bounded, public teardown for the producer and consumer + ``asyncio.Task``s spawned in ``start()``. Without it, a producer + wedged in its ``finally`` closing an abandoned subscriber sink can + survive until event-loop shutdown and surface as + ``Task was destroyed but it is pending!``. + + Always forces: the queues are closed with ``immediate=True`` and the + background tasks are cancelled, so teardown is bounded even when a + subscriber sink was never drained. It is safe to call multiple times. + """ + await self._event_queue_agent.close(immediate=True) + await self._event_queue_subscribers.close(immediate=True) + # Set `_is_finished` and collect the background tasks under `_lock` so + # this is mutually exclusive with `start()`, which refuses to spawn + # once `_is_finished` is set. The lock is released before awaiting the + # tasks, because their teardown re-acquires it. + async with self._lock: + self._is_finished.set() + background_tasks = [ + task + for task in (self._producer_task, self._consumer_task) + if task is not None + ] + for task in background_tasks: + task.cancel() + if background_tasks: + results = await asyncio.gather( + *background_tasks, return_exceptions=True + ) + for result in results: + # CancelledError is a BaseException, so it is excluded here. + if isinstance(result, Exception): + logger.error('Error during aclose', exc_info=result) + async def _maybe_cleanup(self) -> None: """Triggers cleanup if task is finished and has no subscribers. diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index 4f5822f68..ab7d6a11c 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -35,6 +35,7 @@ def __init__( self._active_tasks: dict[str, ActiveTask] = {} self._lock = asyncio.Lock() self._cleanup_tasks: set[asyncio.Task[None]] = set() + self._closed = False async def get_or_create( self, @@ -46,6 +47,8 @@ async def get_or_create( ) -> ActiveTask: """Retrieves an existing ActiveTask or creates a new one.""" async with self._lock: + if self._closed: + raise RuntimeError('ActiveTaskRegistry is closed') if task_id in self._active_tasks: return self._active_tasks[task_id] @@ -88,3 +91,39 @@ async def get(self, task_id: str) -> ActiveTask | None: """Retrieves an existing task.""" async with self._lock: return self._active_tasks.get(task_id) + + async def aclose(self) -> None: + """Closes the registry and drains all active tasks. + + Marks the registry closed so ``get_or_create`` refuses new work, then + force-closes every registered ``ActiveTask`` and awaits the in-flight + ``_remove_task`` cleanup tasks they schedule, so no SDK-owned + ``asyncio.Task`` is left pending at event-loop shutdown. Safe to call + multiple times. + + The close flag is set and the active-task snapshot is taken under + ``_lock``, and the lock is then released before awaiting, because + ``_remove_task`` re-acquires ``_lock``; holding it while draining + would deadlock. Marking closed under the same lock prevents a + concurrent ``get_or_create`` from registering a task that the drain + would miss. + """ + async with self._lock: + self._closed = True + active_tasks = list(self._active_tasks.values()) + + if active_tasks: + results = await asyncio.gather( + *(task.aclose() for task in active_tasks), + return_exceptions=True, + ) + for result in results: + if isinstance(result, Exception): + logger.error('Error draining active task', exc_info=result) + + cleanup_tasks = list(self._cleanup_tasks) + if cleanup_tasks: + await asyncio.gather(*cleanup_tasks, return_exceptions=True) + + async with self._lock: + self._active_tasks.clear() diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 599e38ce5..872a3bfa2 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -112,6 +112,15 @@ def __init__( # noqa: PLR0913 ) self._background_tasks = set() + async def aclose(self) -> None: + """Shuts down the handler, draining all active tasks. + + Drains the ``ActiveTaskRegistry`` so a server shutdown leaves no + pending ``asyncio.Task``. Intended to be wired into an ASGI + ``lifespan`` / ``on_shutdown`` hook. Safe to call multiple times. + """ + await self._active_task_registry.aclose() + @validate_request_params async def on_get_task( # noqa: D102 self, diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 47597c6a0..fea2fc4b4 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -20,6 +20,7 @@ TaskStatus, TaskStatusUpdateEvent, ) +from a2a.utils._async_queue_compat import QueueShutDown, create_async_queue from a2a.utils.errors import InvalidParamsError @@ -979,3 +980,90 @@ async def execute_mock(req, q): assert cleanup_called, ( 'on_cleanup should be called after producer is drained' ) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_active_task_aclose_reaps_background_tasks(): + """aclose() drains a live producer and consumer.""" + agent_executor = Mock() + task_manager = Mock() + request_context = Mock(spec=RequestContext) + + active_task = ActiveTask( + agent_executor=agent_executor, + task_id='test-task-id', + task_manager=task_manager, + push_sender=Mock(), + ) + + async def slow_execute(req, q): + await asyncio.sleep(10) + + agent_executor.execute = AsyncMock(side_effect=slow_execute) + task_manager.get_task = AsyncMock( + return_value=Task( + id='test-task-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + + await active_task.enqueue_request(request_context) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + + await active_task.aclose() + + assert active_task._producer_task is not None + assert active_task._producer_task.done() + assert active_task._consumer_task is not None + assert active_task._consumer_task.done() + assert active_task._is_finished.is_set() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_active_task_aclose_force_closes_undrained_subscriber(): + """aclose() unblocks past an undrained subscriber sink. + + Reproduces issue #1101: a graceful close(immediate=False) would block + forever on the leaked sink's join(). + """ + agent_executor = Mock() + task_manager = Mock() + request_context = Mock(spec=RequestContext) + + active_task = ActiveTask( + agent_executor=agent_executor, + task_id='test-task-id', + task_manager=task_manager, + push_sender=Mock(), + ) + + async def slow_execute(req, q): + await asyncio.sleep(10) + + agent_executor.execute = AsyncMock(side_effect=slow_execute) + task_manager.get_task = AsyncMock( + return_value=Task( + id='test-task-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + + await active_task.enqueue_request(request_context) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + + # Leak a subscriber sink and push an event into it without draining it. + leaked = await active_task._event_queue_subscribers.tap() + await active_task._event_queue_subscribers.enqueue_event(Message()) + await asyncio.sleep(0.05) + + await active_task.aclose() + + assert active_task._producer_task is not None + assert active_task._producer_task.done() + assert leaked.is_closed() diff --git a/tests/server/agent_execution/test_active_task_registry.py b/tests/server/agent_execution/test_active_task_registry.py new file mode 100644 index 000000000..16d9c8797 --- /dev/null +++ b/tests/server/agent_execution/test_active_task_registry.py @@ -0,0 +1,107 @@ +import asyncio +import logging + +from unittest.mock import AsyncMock + +import pytest + +from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry +from a2a.server.agent_execution.agent_executor import AgentExecutor +from a2a.server.agent_execution.context import RequestContext +from a2a.server.context import ServerCallContext +from a2a.server.events.event_queue_v2 import EventQueue +from a2a.server.tasks import InMemoryTaskStore + + +class _SlowExecutor(AgentExecutor): + """An executor whose execute() blocks until cancelled.""" + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + await asyncio.sleep(10) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + return None + + +def _make_registry() -> ActiveTaskRegistry: + return ActiveTaskRegistry( + agent_executor=_SlowExecutor(), + task_store=InMemoryTaskStore(), + ) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_reaps_active_tasks_and_empties_registry(): + """aclose() reaps background tasks and removes them.""" + registry = _make_registry() + active = await registry.get_or_create( + 'task-1', + call_context=ServerCallContext(), + create_task_if_missing=True, + ) + + await registry.aclose() + + assert active._producer_task is not None + assert active._producer_task.done() + assert active._consumer_task is not None + assert active._consumer_task.done() + assert await registry.get('task-1') is None + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_is_idempotent(): + """Calling aclose() repeatedly is a safe no-op.""" + registry = _make_registry() + await registry.get_or_create( + 'task-1', + call_context=ServerCallContext(), + create_task_if_missing=True, + ) + + await registry.aclose() + await registry.aclose() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_on_empty_registry(): + """aclose() with no active tasks returns immediately.""" + registry = _make_registry() + await registry.aclose() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_get_or_create_rejected_after_aclose(): + """A closed registry refuses to create new tasks (no orphan race).""" + registry = _make_registry() + await registry.aclose() + + with pytest.raises(RuntimeError): + await registry.get_or_create( + 'task-1', + call_context=ServerCallContext(), + create_task_if_missing=True, + ) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_logs_and_swallows_task_errors(caplog): + """A failing ActiveTask.aclose is logged, not propagated.""" + registry = _make_registry() + failing = AsyncMock() + failing.aclose = AsyncMock(side_effect=ValueError('boom')) + registry._active_tasks['bad'] = failing + + with caplog.at_level(logging.ERROR): + await registry.aclose() + + assert 'Error draining active task' in caplog.text diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index da8b2c561..b276fb77a 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -1712,3 +1712,37 @@ async def test_on_get_task_push_notification_config_is_owner_scoped(): ), _ctx('bob'), ) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_drains_registry(): + """aclose() drains the active-task registry on shutdown.""" + handler = DefaultRequestHandlerV2( + agent_executor=MockAgentExecutor(), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + await handler._active_task_registry.get_or_create( + 'task-1', + call_context=ServerCallContext(user=UnauthenticatedUser()), + create_task_if_missing=True, + ) + + await handler.aclose() + + assert await handler._active_task_registry.get('task-1') is None + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_aclose_is_idempotent_and_handles_empty(): + """aclose() is safe with no active tasks and when called twice.""" + handler = DefaultRequestHandlerV2( + agent_executor=MockAgentExecutor(), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + + await handler.aclose() + await handler.aclose() From 70f129c7241d40ff0c49448736baf0068458c650 Mon Sep 17 00:00:00 2001 From: Sankalp Gilda Date: Thu, 9 Jul 2026 15:58:51 -0400 Subject: [PATCH 2/3] fix(server): shut down request queue in aclose() for never-started tasks If start() was never called, the producer and consumer finally blocks that normally shut down _request_queue never run, so a caller parked in enqueue_request() would never be released. Shut the queue down first in aclose(), mirroring the producer's finally ordering. --- src/a2a/server/agent_execution/active_task.py | 5 +++ .../agent_execution/test_active_task.py | 35 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 29a9393a0..ea1955fdd 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -768,6 +768,11 @@ async def aclose(self) -> None: background tasks are cancelled, so teardown is bounded even when a subscriber sink was never drained. It is safe to call multiple times. """ + # Shut down the request queue first, mirroring the producer's + # ``finally``. If `start()` was never called, the producer/consumer + # ``finally`` blocks never run, and without this a caller parked in + # `enqueue_request()` would wait forever. + self._request_queue.shutdown(immediate=True) await self._event_queue_agent.close(immediate=True) await self._event_queue_subscribers.close(immediate=True) # Set `_is_finished` and collect the background tasks under `_lock` so diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index fea2fc4b4..1be233ee1 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -1067,3 +1067,38 @@ async def slow_execute(req, q): assert active_task._producer_task is not None assert active_task._producer_task.done() assert leaked.is_closed() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_active_task_aclose_never_started_shuts_down_request_queue(): + """aclose() on a never-started task shuts down the request queue. + + If start() was never called, the producer and consumer `finally` blocks + that normally shut down `_request_queue` never run, so without this a + caller parked in `enqueue_request()` would never be released. + """ + active_task = ActiveTask( + agent_executor=Mock(), + task_id='test-task-id', + task_manager=Mock(), + push_sender=Mock(), + ) + + # Swap in a bounded queue and fill it so the next enqueue_request parks. + active_task._request_queue = create_async_queue(maxsize=1) + await active_task.enqueue_request(Mock(spec=RequestContext)) + waiter = asyncio.create_task( + active_task.enqueue_request(Mock(spec=RequestContext)) + ) + await asyncio.sleep(0.05) + assert not waiter.done() + + await active_task.aclose() + + # The parked waiter is released with the shutdown error, and any + # subsequent enqueue fails fast instead of hanging. + with pytest.raises(QueueShutDown): + await waiter + with pytest.raises(QueueShutDown): + await active_task.enqueue_request(Mock(spec=RequestContext)) From e135ce9d556d5bddefc614cf9fe284209785fa3e Mon Sep 17 00:00:00 2001 From: Sankalp Gilda Date: Fri, 17 Jul 2026 16:04:44 -0400 Subject: [PATCH 3/3] fix(server): await cancelled dispatcher in EventQueueSource.close() EventQueueSource.close() cancelled its dispatcher task but only awaited the sinks, so a source with no sinks (create_default_sink=False and no taps) returned with the cancelled dispatcher still pending, producing "Task was destroyed but it is pending!". aclose() calls close() on both sources, so it hit this for the subscriber-less _event_queue_subscribers. Await the dispatcher after cancelling it, suppressing both the expected CancelledError and any error it raised while shutting down: close() is teardown (reachable from __aexit__) and must not raise, and _dispatch_loop already logs its own exceptions, so a dispatcher failure stays observable without propagating out of close(). --- src/a2a/server/events/event_queue_v2.py | 15 +++++++ tests/server/events/test_event_queue_v2.py | 48 ++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/src/a2a/server/events/event_queue_v2.py b/src/a2a/server/events/event_queue_v2.py index 3386a732e..7f9a53e87 100644 --- a/src/a2a/server/events/event_queue_v2.py +++ b/src/a2a/server/events/event_queue_v2.py @@ -239,6 +239,21 @@ async def close(self, immediate: bool = False) -> None: *(sink.close(immediate=False) for sink in sinks_to_close) ) + # Both branches above cancel the dispatcher task, but the sink gather + # only awaits the sinks. Await the cancelled dispatcher here so it is + # never left pending when close() returns -- e.g. a source created with + # create_default_sink=False and no taps has no sinks for the gather to + # await, which would otherwise leave the cancelled dispatcher pending + # and trigger "Task was destroyed but it is pending!". + # + # Suppress the expected CancelledError and any error the dispatcher + # raised while shutting down: close() is teardown (reachable from + # __aexit__) and must not raise, and _dispatch_loop already logs its own + # exceptions, so a dispatcher failure stays observable without + # propagating out of close(). + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._dispatcher_task + def is_closed(self) -> bool: """[DEPRECATED] Checks if the queue is closed. diff --git a/tests/server/events/test_event_queue_v2.py b/tests/server/events/test_event_queue_v2.py index 27bceea4c..5e93ab870 100644 --- a/tests/server/events/test_event_queue_v2.py +++ b/tests/server/events/test_event_queue_v2.py @@ -592,6 +592,54 @@ async def test_dispatch_task_failed(event_queue: EventQueueSource) -> None: await asyncio.wait_for(event_queue.close(immediate=False), timeout=0.1) +@pytest.mark.asyncio +async def test_close_immediate_awaits_dispatcher_without_sinks() -> None: + """A subscriber-less source must not leave its dispatcher pending on close. + + Regression: a source created with create_default_sink=False and no taps has + no sinks for close()'s gather to await. close(immediate=True) cancels the + dispatcher task, so with nothing else to await it used to return while the + cancelled dispatcher was still pending, producing "Task was destroyed but it + is pending!". close() must await its own cancelled dispatcher task. + """ + source = EventQueueSource(create_default_sink=False) + assert not source._dispatcher_task.done() + + await source.close(immediate=True) + + assert source._dispatcher_task.done() + + +@pytest.mark.asyncio +async def test_close_does_not_propagate_dispatcher_crash() -> None: + """close() awaits its dispatcher without surfacing a non-cancel crash. + + close() is teardown (reachable from __aexit__) and must not raise, and + _dispatch_loop already logs its own exceptions. If the dispatcher died with a + non-cancel error, awaiting it in close() must not re-raise that error. + """ + source = EventQueueSource(create_default_sink=False) + + # Retire the real dispatcher, then stand in a task that has already failed + # with a non-cancel error, simulating a _dispatch_loop crash. + real = source._dispatcher_task + real.cancel() + try: + await real + except asyncio.CancelledError: + pass + + async def crashed() -> None: + raise RuntimeError('dispatcher boom') + + source._dispatcher_task = asyncio.ensure_future(crashed()) + await asyncio.sleep(0.01) # let it fail so the exception is stored + + # Must return cleanly, not re-raise the dispatcher's RuntimeError. + await source.close(immediate=True) + assert source._dispatcher_task.done() + + @pytest.mark.asyncio async def test_concurrent_close_immediate_false() -> None: """Test that concurrent close(immediate=False) calls both wait for join() deterministically."""