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