From 9113929ca968194e21c1b02aa8ff039aff530c35 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Fri, 23 Oct 2026 15:29:29 +0000 Subject: [PATCH] fix(server): keep terminal task states final across replicas When another replica cancels a task (or otherwise finishes it) while the executing replica's agent is still running, the executing replica's next save conflicts and it reloads the terminal state. Events its agent had already enqueued were still applied afterwards: the reload refreshed the version, so the next save passed the compare-and-swap and overwrote the terminal state (e.g. CANCELED became COMPLETED). - VersionedDatabaseTaskStore refuses any write on top of a terminal task, not only CANCELED writes. The VersionedTaskStore contract says so. - The consumer stops applying a stale executor's events once a reload shows another writer finished the task. The in-memory test store now copies tasks in and out like a database and follows the same rule; it can also model a store without the rule, to check the consumer side on its own. --- src/a2a/server/agent_execution/active_task.py | 17 ++++ src/a2a/server/cluster/database_task_store.py | 9 ++ src/a2a/server/cluster/task_store.py | 4 +- tests/server/cluster/conftest.py | 27 +++++- .../cluster/test_cancel_multireplica.py | 93 ++++++++++++++++++- .../test_database_versioned_task_store.py | 38 ++++++++ .../server/cluster/test_send_multireplica.py | 5 +- 7 files changed, 185 insertions(+), 8 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 0b0be4e85..1d5f34c76 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/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: ( diff --git a/src/a2a/server/cluster/database_task_store.py b/src/a2a/server/cluster/database_task_store.py index fb9c8ffa0..a55717405 100644 --- a/src/a2a/server/cluster/database_task_store.py +++ b/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) diff --git a/src/a2a/server/cluster/task_store.py b/src/a2a/server/cluster/task_store.py index aed8777fe..f9d8854d6 100644 --- a/src/a2a/server/cluster/task_store.py +++ b/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: diff --git a/tests/server/cluster/conftest.py b/tests/server/cluster/conftest.py index 97da1eaf6..b2e7ddb50 100644 --- a/tests/server/cluster/conftest.py +++ b/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,14 +93,15 @@ 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: @@ -102,8 +109,16 @@ async def save( 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, diff --git a/tests/server/cluster/test_cancel_multireplica.py b/tests/server/cluster/test_cancel_multireplica.py index bc270d50f..d85167c5b 100644 --- a/tests/server/cluster/test_cancel_multireplica.py +++ b/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() diff --git a/tests/server/cluster/test_database_versioned_task_store.py b/tests/server/cluster/test_database_versioned_task_store.py index 624dee448..0e2d1c028 100644 --- a/tests/server/cluster/test_database_versioned_task_store.py +++ b/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, diff --git a/tests/server/cluster/test_send_multireplica.py b/tests/server/cluster/test_send_multireplica.py index fabd02dca..8480b8939 100644 --- a/tests/server/cluster/test_send_multireplica.py +++ b/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