feat(server): add aclose() to drain ActiveTask background tasks (#1101)#1105
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces aclose() methods to ActiveTask, ActiveTaskRegistry, and DefaultRequestHandlerV2 to facilitate a bounded, force-closed teardown of background tasks and queues during server shutdown, preventing pending task warnings. It also adds comprehensive unit tests to verify the teardown behavior, idempotency, and error handling. The review feedback suggests explicitly shutting down self._request_queue in ActiveTask.aclose() to ensure proper cleanup of all queues if start() was never called or failed early.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
🧪 Code Coverage (vs
|
| Base | PR | Delta | |
|---|---|---|---|
| src/a2a/server/agent_execution/active_task.py | 95.32% | 95.09% | 🔴 -0.24% |
| src/a2a/server/agent_execution/active_task_registry.py | 93.75% | 96.61% | 🟢 +2.86% |
| src/a2a/server/events/event_queue_v2.py | 91.19% | 91.28% | 🟢 +0.09% |
| src/a2a/server/request_handlers/default_request_handler_v2.py | 94.12% | 94.17% | 🟢 +0.05% |
| Total | 92.95% | 92.97% | 🟢 +0.02% |
Generated by coverage-comment.yml
403fe48 to
7e42e06
Compare
|
Addressed the review suggestion: aclose() now shuts down _request_queue (immediate=True) first, covering the never-started case, with a test. CI is green on the new push. Ready for review whenever you have a moment @rohityan. |
|
Rebased onto current main to clear the "behind" state after #1053 and the recent refactors. No content changes, and the full suite is green on the rebased branch. Quick recap of where things stand: the approach was agreed in #1101 (thanks @rohityan), the one review finding (never-started request queue) is addressed with a test, and CI has been green since the July push. Ready for review whenever you have a moment. |
5347d4b to
da13286
Compare
0728696 to
cd52b14
Compare
|
@mykytanetipa Your finding is addressed in 0728696: the drain now happens at the root in EventQueueSource.close(), so every caller gets the guarantee, and there's a regression test for the subscriber-less immediate-close case. Pushed on the 17th, the day after your changes-requested review; details in the thread reply above. The branch has also been rebased onto current main (the 1.1.1 release commit) with no content changes to the three PR commits, and the touched-area suite and ruff are clean. Requesting a re-review, since the change this review asked for was pushed on the 17th and is in the current head. |
|
@astrogilda thank you for the quick fix, I see that tests/server/agent_execution/test_active_task.py is conflicting with main. Can you please resolve it? |
…roject#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 a2aproject#1101
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.
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().
cd52b14 to
e135ce9
Compare
|
@mykytanetipa Done. Rebased onto current main and fixed the overlap in that test file. Your teardown-warning test from #1122 and our three drain tests cover different paths, so I kept all four. Green here: agent_execution plus request_handlers, 131 passed, ruff clean. @rohityan if you get a minute: this is the third rebase forced by unrelated edits to the same test file (#1053, then #1122). We settled the design in #1101 and the one review note is handled, so a merge would spare us the next round. Happy to turn around anything you flag. |
2 participants
Summary
Adds a public
aclose()teardown toActiveTask,ActiveTaskRegistry, andDefaultRequestHandlerV2that force-drains the SDK-owned producer, consumer,and dispatcher
asyncio.Tasks so none are left pending at event-loop shutdown.ActiveTask.aclose()force-closes both event queues and cancels the producerand consumer tasks, then awaits them. It sets
_is_finishedunder_lock, soit is mutually exclusive with
start()(which refuses to spawn once_is_finishedis set).ActiveTaskRegistry.aclose()marks the registry closed soget_or_createrefuses new work, drains every active task, then awaits the in-flight
_remove_taskcleanup tasks. The lock is released before awaiting because_remove_taskre-acquires it.DefaultRequestHandlerV2.aclose()delegates to the registry drain, for wiringinto an ASGI
lifespan/on_shutdownhook.Why
Fixes #1101. At shutdown the
ActiveTaskproducer can stay pending and surfaceas
Task was destroyed but it is pending!. The producer'sfinallycalls_event_queue_subscribers.close(immediate=False), which awaitsjoin()onevery subscriber sink; an abandoned subscriber leaves an undrained sink, so the
join()never returns and the producer hangs. There is no public way to drainthese background tasks today.
aclose()closes the queues withimmediate=True,which releases the wedged producer, and reaps the tasks.
The teardown always forces rather than exposing a graceful
immediate=Falseoption, because that path inherits the documented
close(immediate=False)deadlock and a shutdown hook must be bounded.
Test plan
uv run pytest tests/server/agent_execution/ tests/server/request_handlers/ tests/server/events/— passuv run pytest --cov=a2a --cov-fail-under=88— pass./scripts/lint.sh— ruff, ruff-format, and ty cleanNew tests cover the registry drain, idempotency, empty registry, new-work
rejected after close, and an errored task being logged rather than propagated;
ActiveTaskreaping a running producer and force-closing past an undrainedsubscriber (the #1101 repro); and the handler drain.
Fixes #1101 🦕