a2aproject / a2aproject/a2a-python

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

Abierto
#1,188 2 comentarios 0 reacciones 1 asignado Reclamado por @long2bui-andpad Ver en GitHub
component: server
Lenguaje dominante
Python
Estrellas
2.1k
Forks
496
Merge medio
4 d 17 h
PR fusionados (30 d)
12

Descripción

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

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.