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

The bug is in a2a-sdk's DefaultRequestHandlerV2 and ActiveTaskRegistry. Start by reading active_task.py and task_manager.py to understand the caching of _current_task. The workaround in the issue suggests modifying get_or_create in ActiveTaskRegistry to clear the snapshot when reusing an idle task. Test with a multi-replica setup and a task that hits input_required twice. Verify that artifacts are not lost by checking the shared TaskStore after each step.

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):

# 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):

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.