fix(server): drop stale ActiveTask snapshot when reusing an idle task - #1227
long2bui-andpad wants to merge 2 commits into
Conversation
A reused ActiveTask keeps its TaskManager._current_task snapshot across non-terminal interrupts (input_required/auth_required) because cleanup only runs on terminal states. The per-request get_task() in _run_producer intends to refresh the task each request, but get_task() short-circuits on the cached snapshot, so the refresh is a no-op. With more than one replica sharing a TaskStore, a resume routed back to the pod holding the stale snapshot silently overwrites artifacts, history and status that another replica persisted while the task was parked. Drop the cached snapshot when the registry reuses an idle ActiveTask (reference_count <= 1, no subscriber stream in flight), forcing the next request to re-read the store at the request boundary. A still-streaming task keeps its snapshot, so an open artifact is not lost mid-stream.
🧪 Code Coverage (vs
|
| Base | PR | Delta | |
|---|---|---|---|
| src/a2a/server/agent_execution/active_task.py | 95.02% | 95.07% | 🟢 +0.05% |
| src/a2a/server/agent_execution/active_task_registry.py | 97.37% | 97.50% | 🟢 +0.13% |
| src/a2a/server/tasks/task_manager.py | 98.86% | 98.88% | 🟢 +0.01% |
| Total | 92.95% | 92.96% | 🟢 +0.01% |
Generated by coverage-comment.yml
|
@long2bui-andpad thank you for working on a fix! Multi-replica cluster support has now landed on main in #1281 (a2a.server.cluster), which introduces VersionedTaskStore (so replicas never silently overwrite each other's updates) and TaskEventStream (cross-replica streaming, resubscribe, and cancellation). As part of #1281, this stale-snapshot issue was resolved by adding TaskManager.invalidate() and calling it in ActiveTask._run_producer right after acquiring _request_lock. Because _request_lock is held until EventConsumer finishes persisting all events for the current request (_RequestCompleted), a reused ActiveTask now always re-reads the TaskStore at request boundaries after input_required / auth_required while keeping in-flight streams and append=True artifact chunks intact (covered by test_multiturn_input_required_across_replicas). |
2 participants
Description
Fixes a silent data-loss bug in
DefaultRequestHandlerV2(ActiveTaskRegistry) for multi-replica deployments.A reused
ActiveTaskkeeps itsTaskManager._current_tasksnapshot across non-terminal interrupts (input_required/auth_required) — cleanup only runs on terminal states. The per-requestget_task()in_run_producerintends to refresh the task each request (see the existing# TODO: Should we create task manager every time?), butget_task()short-circuits on the cached snapshot, so the refresh is a no-op.With more than one replica sharing a
TaskStoreand no task affinity, a resume routed back to the pod holding the stale snapshot silently overwrites artifacts, history and status that another replica persisted while the task was parked. There is no error — the save succeeds and the client simply sees the other replica's artifacts disappear.Fix
When the registry reuses an idle
ActiveTask(_reference_count <= 1, i.e. no subscriber stream in flight), drop the cached snapshot so the next request re-reads the store at the request boundary. This matches the intent of the existing TODO while preserving V2's per-task serialization and streaming semantics.A still-streaming task (
_reference_count > 1) keeps its snapshot on purpose: re-reading the store mid-stream would lose the open artifact and make the nextappend=Truechunk fail withInvalidAgentResponseError.This follows the approach suggested by the maintainer in the issue (folding the downstream
FreshTaskRegistryworkaround into the SDK).Changes:
TaskManager.invalidate_cached_task()— drops the cached snapshot so the nextget_task()re-reads the store.ActiveTask.refresh_on_reuse()— invalidates the snapshot only when the task is idle.ActiveTaskRegistry.get_or_create()— callsrefresh_on_reuse()on the cache-hit path.Tests
test_task_manager.py: invalidation forces a fresh store read.test_active_task.py:refresh_on_reusedrops the snapshot when idle, keeps it while streaming.test_active_task_registry.py: a reused idle task drops its stale snapshot; a reused streaming task keeps it.Out of scope
Concurrent writes from two replicas to the same task still race (last-writer-wins at the
TaskStore). Addressing that needs a store-level version / compare-and-swap contract and is tracked separately in the issue.CONTRIBUTINGGuide.bash scripts/format.shfrom the repository root to format)Fixes #1188 🦕