a2aproject / a2aproject/a2a-python

[Bug] DefaultRequestHandlerV2: ActiveTaskRegistry resumes from a stale task snapshot after input_required, silently overwriting store state (multi-replica deployments)

Aperta
#1,188 2 commenti 0 reazioni 1 assegnatario Rivendicata da @long2bui-andpad Vedi su GitHub
component: server
Lingua principale
Python
Stelle
2.1k
Fork
496
Merge medio
4g 17h
PR unite (30g)
12

Descrizione

## 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.

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.