fix(server): keep terminal task states final across replicas

src/a2a/server/agent_execution/active_task.py

@@ -119,6 +119,9 @@ def __init__(self, active_task: ActiveTask) -> None:
         # Set when the producer failed, so a request that failed before any
         # task was created still leaves a FAILED task behind.
         self.producer_failed = False
+        # Set when another writer finished the task while this executor was
+        # still running; the events it enqueues afterwards are dropped.
+        self.fenced = False
 
     async def run(self) -> None:
         """Consumes events from the agent and updates system state."""
@@ -209,6 +212,9 @@ async def _write_failed_status(self) -> Task | None:
             return None
 
     async def _process_event(self, event: Event) -> None:
+        if self.fenced:
+            self._drop_fenced_event(event)
+            return
         try:
             await self._process_event_inner(event)
         except ConcurrentTaskModificationError:
@@ -220,6 +226,7 @@ async def _process_event(self, event: Event) -> None:
             self.active_task._task_manager.invalidate()
             task = await self.active_task._task_manager.get_task()
             if task is not None and task.status.state in TERMINAL_TASK_STATES:
+                self.fenced = True
                 await self._enqueue_to_subscribers(task, task)
                 producer = self.active_task._producer_task
                 if producer is not None and not producer.done():
@@ -231,6 +238,16 @@ async def _process_event(self, event: Event) -> None:
                 return
             await self._process_event_inner(event)
 
+    def _drop_fenced_event(self, event: Event) -> None:
+        logger.debug(
+            'Consumer[%s]: dropping %s from a fenced executor',
+            self.active_task._task_id,
+            type(event).__name__,
+        )
+        if isinstance(event, _RequestCompleted):
+            self.active_task._request_lock.release()
+        self.active_task._event_queue_agent.task_done()
+
     async def _process_event_inner(self, event: Event) -> None:
         updated_task = None
         handled_event: (

src/a2a/server/cluster/database_task_store.py

@@ -232,6 +232,15 @@ async def _save_once(
                 )
                 if result.rowcount == 0:  # ty:ignore[unresolved-attribute]
                     raise ConcurrentTaskModificationError(task.id)
+                # A terminal state is final. A writer holding the latest
+                # version can still be stale: it reloaded after losing a race
+                # to another writer while its executor kept producing events.
+                current = await self._current_task(session, task.id, owner)
+                if (
+                    current is not None
+                    and current.status.state in _TERMINAL_STATES
+                ):
+                    raise ConcurrentTaskModificationError(task.id)
 
             await session.merge(model)
 

src/a2a/server/cluster/task_store.py

@@ -52,7 +52,9 @@ async def save(
                 write. A write moving `task` to CANCELED overwrites a
                 non-terminal stored task without a version check, and raises
                 `ConcurrentTaskModificationError` if the stored task is already
-                terminal or absent.
+                terminal or absent. Any other write on top of a terminal
+                stored task MUST raise `ConcurrentTaskModificationError`, even
+                when `prev_version` matches: a terminal state is final.
             context: The server call context (used to resolve the owner).
 
         Returns:

tests/server/cluster/conftest.py

@@ -61,9 +61,15 @@ class VersionedInMemoryTaskStore(VersionedTaskStore):
     """
 
     def __init__(
-        self, owner_resolver: OwnerResolver = resolve_user_scope
+        self,
+        owner_resolver: OwnerResolver = resolve_user_scope,
+        *,
+        enforce_terminal_state: bool = True,
     ) -> None:
         self._owner_resolver = owner_resolver
+        # False models a store written against the contract before terminal
+        # states were required to be final.
+        self._enforce_terminal_state = enforce_terminal_state
         self._store = InMemoryTaskStore()
         # Maps owner to a mapping of task_id to its current integer version.
         self._versions: dict[str, dict[str, int]] = {}
@@ -87,23 +93,32 @@ async def save(
         match the currently stored version. `TaskVersion.MISSING` skips the
         check; a write moving `task` to CANCELED overwrites a non-terminal
         stored task without a version check, and raises when it is terminal
-        or absent.
+        or absent. Any other write on top of a terminal task raises too.
+        Tasks are copied in and out, like a database row.
         """
         del event, prev  # in-memory store keeps no event log
         owner = self._owner_resolver(context)
         with self._lock:
             stored = self._current_version_locked(owner, task_id=task.id)
+            current = await self._store.get(task.id, context)
             if task.status.state == TaskState.TASK_STATE_CANCELED:
-                current = await self._store.get(task.id, context)
                 if current is None or current.status.state in _TERMINAL_STATES:
                     raise ConcurrentTaskModificationError(task.id)
             elif not prev_version.is_missing:
                 if stored == 0:
                     raise ConcurrentTaskModificationError(task.id)
                 if TaskVersion(stored) != prev_version:
                     raise ConcurrentTaskModificationError(task.id)
+                if (
+                    self._enforce_terminal_state
+                    and current is not None
+                    and current.status.state in _TERMINAL_STATES
+                ):
+                    raise ConcurrentTaskModificationError(task.id)
             new_version = stored + 1
-            await self._store.save(task, context)
+            copy = Task()
+            copy.CopyFrom(task)
+            await self._store.save(copy, context)
             self._versions.setdefault(owner, {})[task.id] = new_version
             return TaskVersion(new_version)
 
@@ -116,8 +131,10 @@ async def get(
             task = await self._store.get(task_id, context)
             if task is None:
                 return None
+            copy = Task()
+            copy.CopyFrom(task)
             stored = self._current_version_locked(owner, task_id)
-            return StoredTask(task, TaskVersion(stored))
+            return StoredTask(copy, TaskVersion(stored))
 
     async def list(
         self,

tests/server/cluster/test_cancel_multireplica.py

@@ -3,7 +3,13 @@
 
 import pytest
 
-from a2a.types.a2a_pb2 import CancelTaskRequest, TaskState
+from a2a.types.a2a_pb2 import (
+    CancelTaskRequest,
+    Task,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
 from a2a.utils.errors import (
     TaskNotCancelableError,
     TaskNotFoundError,
@@ -12,8 +18,10 @@
 
 from .conftest import (
     CompletingAgent,
+    ControlledAgent,
     InputRequiredThenCompleteAgent,
     LongRunningAgent,
+    VersionedInMemoryTaskStore,
     build_send_request,
     drain,
     make_context,
@@ -197,3 +205,86 @@ async def test_follow_up_after_remote_cancel_of_waiting_task_is_rejected(
     finally:
         await replica_a.aclose()
         await replica_b.aclose()
+
+
+async def _fail_directly(store, task_id: str, context) -> None:  # noqa: ANN001
+    """Writes FAILED the way a writer outside the running replica would."""
+    stored = await store.get(task_id, context)
+    failed = Task()
+    failed.CopyFrom(stored.task)
+    failed.status.CopyFrom(TaskStatus(state=TaskState.TASK_STATE_FAILED))
+    await store.save(
+        failed,
+        event=TaskStatusUpdateEvent(
+            task_id=task_id, context_id=failed.context_id, status=failed.status
+        ),
+        prev=stored.task,
+        prev_version=stored.version,
+        context=context,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(15)
+@pytest.mark.parametrize('writer', ['remote_cancel', 'direct_failed'])
+@pytest.mark.parametrize(
+    'store_enforces_terminal_state',
+    [True, False],
+    ids=['store_guard', 'consumer_fence_only'],
+)
+async def test_terminal_state_survives_events_queued_by_stale_executor(
+    shared_stream, context, writer: str, store_enforces_terminal_state: bool
+) -> None:
+    """Another writer finishes the task while A's agent is paused. When the
+    agent resumes it emits an artifact and COMPLETED back to back: the first
+    save conflicts, and the queued COMPLETED must not overwrite the result,
+    even with a store that does not itself refuse writes on a terminal task."""
+    shared_store = VersionedInMemoryTaskStore(
+        enforce_terminal_state=store_enforces_terminal_state
+    )
+    agent = ControlledAgent()
+    replica_a = make_replica(shared_store, shared_stream, agent)
+    replica_b = make_replica(shared_store, shared_stream, agent)
+    events: list = []
+
+    async def collect() -> None:
+        async for event in replica_a.on_message_send_stream(
+            build_send_request('go', message_id='m1'), context
+        ):
+            events.append(event)
+
+    try:
+        a_task = asyncio.create_task(collect())
+        await asyncio.wait_for(agent.working.wait(), timeout=5)
+        task_id = next(
+            iter(replica_a._active_task_registry._active_tasks)  # noqa: SLF001
+        )
+        await wait_for_state(
+            shared_store, task_id, TaskState.TASK_STATE_WORKING, context
+        )
+
+        if writer == 'remote_cancel':
+            await replica_b.on_cancel_task(
+                CancelTaskRequest(id=task_id), context
+            )
+            expected = TaskState.TASK_STATE_CANCELED
+        else:
+            await _fail_directly(shared_store, task_id, context)
+            expected = TaskState.TASK_STATE_FAILED
+
+        agent.release.set()
+        await asyncio.wait_for(a_task, timeout=8)
+
+        final = await shared_store.get(task_id, context)
+        assert final is not None
+        assert final.task.status.state == expected
+        assert isinstance(events[-1], Task)
+        assert events[-1].status.state == expected
+        assert not any(
+            getattr(e, 'status', None) is not None
+            and e.status.state == TaskState.TASK_STATE_COMPLETED
+            for e in events
+        )
+    finally:
+        await replica_a.aclose()
+        await replica_b.aclose()

tests/server/cluster/test_database_versioned_task_store.py

@@ -260,6 +260,44 @@ async def test_stale_update_raises_conflict(
     assert stored.task.status.state == TaskState.TASK_STATE_WORKING
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'next_state',
+    [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_COMPLETED],
+    ids=['non_terminal', 'same_terminal'],
+)
+async def test_write_on_top_of_terminal_state_raises_conflict(
+    versioned_store: VersionedDatabaseTaskStore, next_state: TaskState
+) -> None:
+    v1 = await versioned_store.save(
+        create_task(state=TaskState.TASK_STATE_WORKING),
+        event=None,
+        prev=None,
+        prev_version=TaskVersion.MISSING,
+        context=TEST_CONTEXT,
+    )
+    v2 = await versioned_store.save(
+        create_task(state=TaskState.TASK_STATE_COMPLETED),
+        event=None,
+        prev=None,
+        prev_version=v1,
+        context=TEST_CONTEXT,
+    )
+    # The writer holds the current version, but the task is already final.
+    with pytest.raises(ConcurrentTaskModificationError):
+        await versioned_store.save(
+            create_task(state=next_state),
+            event=None,
+            prev=None,
+            prev_version=v2,
+            context=TEST_CONTEXT,
+        )
+    stored = await versioned_store.get('task-abc', TEST_CONTEXT)
+    assert stored is not None
+    assert stored.task.status.state == TaskState.TASK_STATE_COMPLETED
+    assert stored.version == v2
+
+
 @pytest.mark.asyncio
 async def test_concurrent_first_insert_one_wins(
     versioned_store: VersionedDatabaseTaskStore,

tests/server/cluster/test_send_multireplica.py

@@ -93,7 +93,10 @@ async def test_stale_version_save_raises_conflict_directly(
     shared_store, shared_stream, context
 ) -> None:
     """The store-level CAS that underpins the handler fix."""
-    handler = make_replica(shared_store, shared_stream, CompletingAgent())
+    # A non-terminal task, so only the version check can reject the save.
+    handler = make_replica(
+        shared_store, shared_stream, InputRequiredThenCompleteAgent()
+    )
     try:
         result = await handler.on_message_send(
             build_send_request('hi'), context