a2aproject / a2aproject/a2a-python
[Bug] DefaultRequestHandlerV2: ActiveTaskRegistry resumes from a stale task snapshot after input_required, silently overwriting store state (multi-replica deployments)
- 主要言語
- Python
- スター
- 2.1k
- フォーク
- 496
- 平均マージ
- 4日 17時間
- マージ済み PR(30日)
- 12
説明
## Environment
- a2a-sdk 1.1.2 (also verified in 1.0.3 — same code)
- Python 3.12, `DefaultRequestHandlerV2` (the `DefaultRequestHandler` alias), custom `TaskStore` backed by a shared database
- Server runs as 2 replicas behind a round-robin load balancer, no task affinity
## Summary
`ActiveTaskRegistry` keeps one `ActiveTask` per `task_id` for the lifetime of the process. The `ActiveTask` (and its `TaskManager._current_task` snapshot) survives HITL `input_required` interrupts — cleanup only happens on terminal states — and on reuse `ActiveTask.start()` early-returns without re-reading the store.
The producer loop clearly *intends* to refresh the task per request (`active_task.py`, `_run_producer`):
```python
# TODO: Should we create task manager every time?
self._task_manager._call_context = request_context.call_context
request_context.current_task = (
await self._task_manager.get_task()
)
```
but `TaskManager.get_task()` short-circuits on the cached `_current_task`, so this refresh is a no-op for a reused `ActiveTask`.
In a single-process deployment this is harmless. With more than one replica sharing a `TaskStore`, it loses data:
1. `message/send` lands on **pod A** → task reaches `input_required` 1. Pod A's registry keeps the snapshot.
2. The resume lands on **pod B** → fresh `ActiveTask`, reads the store, produces new artifacts + `input_required` 2, persists everything.
3. The next resume lands on **pod A** → the registry reuses the stale `ActiveTask`; `EventConsumer._handle_task_modification_event` calls `update_with_message(resume_message, stale_snapshot)` and `save_task_event(...)` persists the **interrupt-1 snapshot + new message**, silently deleting pod B's artifacts, history entries, and status.
There is no error anywhere — the save succeeds and the client simply sees the artifacts from step 2 disappear.
## Reproduction
Any agent with two sequential `input_required` interrupts, two server replicas over one `TaskStore`, round-robin routing:
1. send → interrupt 1 (pod A)
2. resume → agent produces artifacts, interrupt 2 (pod B)
3. resume → routed back to pod A
Observed at step 3: the first `save` from pod A carries interrupt 1's `status.timestamp`, a history truncated to the interrupt-1 snapshot plus the new user message, and an artifact list missing everything pod B wrote.
## Expected behavior
A reused `ActiveTask` should not trust its in-memory snapshot across interrupts: the per-request `get_task()` in `_run_producer` should re-read the `TaskStore` (matching the intent of the existing TODO), or the registry should evict/revalidate the `ActiveTask` when it is reused after an interrupt.
## Notes on a possible fix
Simply making `get_task()` always re-read the store is NOT safe: `get_task()` is also called mid-stream (e.g. `_handle_task_modification_event` calls it for every `TaskStatusUpdateEvent` until `_task_created` is set), and a `TaskStore` implementation may legitimately defer writes while an artifact is streaming — re-reading there loses the open artifact and the next `append=True` chunk fails with `InvalidAgentResponseError`. We hit exactly this while testing.
What worked for us as a downstream workaround: drop the snapshot only when the registry reuses an **idle** `ActiveTask` (no subscriber streams in flight, i.e. `_reference_count <= 1`, meaning every event of the previous request has been persisted):
```python
class FreshTaskRegistry(ActiveTaskRegistry):
async def get_or_create(self, task_id, call_context, context_id=None,
create_task_if_missing=False, initial_message=None):
async with self._lock:
existing = self._active_tasks.get(task_id)
if existing is not None and existing._reference_count <= 1:
existing._task_manager._current_task = None
return await super().get_or_create(
task_id, call_context, context_id=context_id,
create_task_if_missing=create_task_if_missing,
initial_message=initial_message,
)
```
This keeps V2's per-task serialization and streaming semantics intact while making a reused task re-read the store at the request boundary. Happy to turn this into a PR if the approach sounds right; an alternative direction would be evicting the `ActiveTask` from the registry once a task enters an interrupted state with no subscribers.
Related: even with this fix, two replicas that write the same task *concurrently* still race (last-writer-wins at the `TaskStore`); a store-level version/compare-and-swap contract might be worth considering separately, or documenting that V2 currently assumes task-to-process affinity.
コントリビューションガイド
評価
この issue はまだ評価されていません。