From e70fdf3532dcb36b09fd66c70f5bffcdd80b48fd Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Fri, 23 Oct 2026 15:33:15 +0000 Subject: [PATCH] fix(server): save a resumed turn's message with its first event When a task waiting for input is resumed, the consumer added the user's message to the history with a separate save before applying the agent's first event. That save wrote the task with its previous state, and in cluster mode it also wrote a Task(INPUT_REQUIRED) row to the shared event log. A SubscribeToTask on another replica treated that row as the end of the stream and missed the rest of the turn. The message is now added in memory and persisted by the save of the event itself, so a resumed turn writes no stale state. If that save conflicts, the message is added again when the event is retried on the reloaded task. --- src/a2a/server/agent_execution/active_task.py | 17 ++- .../server/cluster/test_send_multireplica.py | 120 +++++++++++++++- .../cluster/test_subscribe_multireplica.py | 128 +++++++++++++++++- 3 files changed, 252 insertions(+), 13 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 1d5f34c76..520faa8b9 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -375,19 +375,18 @@ async def _handle_task_modification_event( for message in updated_task.history ) if not message_already_saved: - updated_task = ( - self.active_task._task_manager.update_with_message( - self.message_to_save, - updated_task, - ) + # Saved together with the event below. A separate save here + # would first re-save the previous, interrupted state. + self.active_task._task_manager.update_with_message( + self.message_to_save, + updated_task, ) - await self.active_task._task_manager.save_task_event( - updated_task - ) - self.message_to_save = None self.active_task._task_manager.context_id = event.context_id await self.active_task._task_manager.process(event) + # Kept until the save lands: after a conflict the event is applied + # again on the reloaded task, which needs the message added again. + self.message_to_save = None async def _update_task_state( self, diff --git a/tests/server/cluster/test_send_multireplica.py b/tests/server/cluster/test_send_multireplica.py index 8480b8939..6200220cd 100644 --- a/tests/server/cluster/test_send_multireplica.py +++ b/tests/server/cluster/test_send_multireplica.py @@ -1,9 +1,13 @@ import asyncio +import contextlib import pytest +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster import ConcurrentTaskModificationError, TaskVersion -from a2a.types.a2a_pb2 import Role, Task, TaskState +from a2a.server.tasks.task_updater import TaskUpdater +from a2a.types.a2a_pb2 import Part, Role, Task, TaskState from a2a.utils.errors import InternalError from .conftest import ( @@ -118,19 +122,58 @@ async def test_stale_version_save_raises_conflict_directly( class _ContendedStore(VersionedInMemoryTaskStore): - """Rejects every save while `contended` is set, as if other replicas kept - writing the task.""" + """Rejects every save while `contended` is set, and the next `conflicts` + saves, as if other replicas kept writing the task.""" def __init__(self) -> None: super().__init__() self.contended = False + self.conflicts = 0 async def save(self, task: Task, **kwargs) -> TaskVersion: # noqa: ANN003 if self.contended: raise ConcurrentTaskModificationError(task.id) + if self.conflicts > 0: + self.conflicts -= 1 + raise ConcurrentTaskModificationError(task.id) return await super().save(task, **kwargs) +@pytest.mark.asyncio +@pytest.mark.timeout(10) +async def test_follow_up_message_survives_a_conflict_on_its_first_save( + shared_stream, context +) -> None: + """The follow-up's message is saved with the turn's first event. When + that save conflicts and is retried on the reloaded task, the message + must still be recorded.""" + store = _ContendedStore() + handler = make_replica( + store, shared_stream, InputRequiredThenCompleteAgent() + ) + try: + r1 = await handler.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + store.conflicts = 1 + await handler.on_message_send( + build_send_request( + 'a', task_id=r1.id, context_id=r1.context_id, message_id='m2' + ), + context, + ) + + stored = await store.get(r1.id, context) + assert stored is not None + assert [ + m.message_id + for m in stored.task.history + if m.role == Role.ROLE_USER + ] == ['m1', 'm2'] + finally: + await handler.aclose() + + @pytest.mark.asyncio @pytest.mark.timeout(10) async def test_follow_up_that_keeps_conflicting_fails_with_a2a_error( @@ -170,3 +213,74 @@ async def test_follow_up_that_keeps_conflicting_fails_with_a2a_error( finally: store.contended = False await handler.aclose() + + +class _SecondTurnAgent(AgentExecutor): + """Turn 1 asks for input; turn 2 starts with the configured first event.""" + + def __init__(self, second_turn: str) -> None: + self.second_turn = second_turn + + async def execute(self, context, event_queue) -> None: # noqa: ANN001 + updater = TaskUpdater( + event_queue, str(context.task_id), str(context.context_id) + ) + if context.current_task is None: + await event_queue.enqueue_event( + new_task_from_user_message(context.message) + ) + await updater.requires_input( + message=updater.new_agent_message([Part(text='need input')]) + ) + return + if self.second_turn == 'raises': + raise RuntimeError('agent failed') + if self.second_turn == 'artifact_first': + await updater.add_artifact([Part(text='result')]) + await updater.complete() + + async def cancel(self, context, event_queue) -> None: # noqa: ANN001 + pass + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize( + ('second_turn', 'final_state'), + [ + ('status_first', TaskState.TASK_STATE_COMPLETED), + ('artifact_first', TaskState.TASK_STATE_COMPLETED), + ('raises', TaskState.TASK_STATE_FAILED), + ], +) +async def test_resumed_turn_records_follow_up_message_once( + shared_store, shared_stream, context, second_turn: str, final_state +) -> None: + agent = _SecondTurnAgent(second_turn) + replica_a = make_replica(shared_store, shared_stream, agent) + replica_b = make_replica(shared_store, shared_stream, agent) + try: + first = await replica_a.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + with contextlib.suppress(RuntimeError): + await replica_b.on_message_send( + build_send_request( + 'a1', + task_id=first.id, + context_id=first.context_id, + message_id='m2', + ), + context, + ) + stored = await shared_store.get(first.id, context) + assert stored is not None + assert stored.task.status.state == final_state + assert [ + m.message_id + for m in stored.task.history + if m.role == Role.ROLE_USER + ] == ['m1', 'm2'] + finally: + await replica_a.aclose() + await replica_b.aclose() diff --git a/tests/server/cluster/test_subscribe_multireplica.py b/tests/server/cluster/test_subscribe_multireplica.py index 4cc04c8bb..80f89a02d 100644 --- a/tests/server/cluster/test_subscribe_multireplica.py +++ b/tests/server/cluster/test_subscribe_multireplica.py @@ -1,9 +1,19 @@ import asyncio +import uuid + +from collections.abc import AsyncGenerator import pytest +import pytest_asyncio -from a2a.types.a2a_pb2 import SubscribeToTaskRequest, TaskState +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.server.agent_execution.agent_executor import AgentExecutor +from a2a.server.cluster import DatabaseTaskEventStream +from a2a.server.cluster.database_task_store import VersionedDatabaseTaskStore +from a2a.server.tasks.task_updater import TaskUpdater +from a2a.types.a2a_pb2 import Part, Role, SubscribeToTaskRequest, TaskState from a2a.utils.errors import TaskNotFoundError, UnsupportedOperationError +from sqlalchemy.ext.asyncio import AsyncConnection, create_async_engine from .conftest import ( ControlledAgent, @@ -137,3 +147,119 @@ async def test_resubscribe_non_owner_rejected( pass finally: await handler.aclose() + + +class _PauseThenResumeAgent(AgentExecutor): + """Turn 1 asks for input. Turn 2 goes WORKING, waits, then completes.""" + + def __init__(self) -> None: + self.resumed = asyncio.Event() + self.release = asyncio.Event() + + async def execute(self, context, event_queue) -> None: # noqa: ANN001 + updater = TaskUpdater( + event_queue, str(context.task_id), str(context.context_id) + ) + if context.current_task is None: + await event_queue.enqueue_event( + new_task_from_user_message(context.message) + ) + await updater.requires_input( + message=updater.new_agent_message([Part(text='need input')]) + ) + return + await updater.start_work() + self.resumed.set() + await self.release.wait() + await updater.complete() + + async def cancel(self, context, event_queue) -> None: # noqa: ANN001 + pass + + +@pytest_asyncio.fixture +async def sqlite_dsn() -> AsyncGenerator[str, None]: + """A shared in-memory SQLite database, kept alive for the test.""" + dsn = ( + f'sqlite+aiosqlite:///file:resume_{uuid.uuid4().hex}' + '?mode=memory&cache=shared&uri=true' + ) + engine = create_async_engine(dsn) + keeper: AsyncConnection = await engine.connect() + await VersionedDatabaseTaskStore(engine).initialize() + await DatabaseTaskEventStream(engine).initialize() + yield dsn + await keeper.close() + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_remote_subscriber_follows_resumed_turn_to_completion( + sqlite_dsn: str, context +) -> None: + """A task waits for input after a turn on A. C subscribes, and B runs the + next turn. C must see that turn through to COMPLETED; the user message + that resumed the task must not first re-save the waiting state.""" + agent = _PauseThenResumeAgent() + engines = [create_async_engine(sqlite_dsn) for _ in range(3)] + replica_a, replica_b, replica_c = ( + make_replica( + VersionedDatabaseTaskStore(engine), + DatabaseTaskEventStream(engine, poll_interval_s=0.01), + agent, + ) + for engine in engines + ) + try: + first = await replica_a.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert first.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + + c_states: list = [] + + async def follow_on_c() -> None: + async for event in replica_c.on_subscribe_to_task( + SubscribeToTaskRequest(id=first.id), context + ): + c_states.append(event.status.state) + + c_task = asyncio.create_task(follow_on_c()) + await asyncio.sleep(0.1) + + b_task = asyncio.create_task( + replica_b.on_message_send( + build_send_request( + 'q2', + task_id=first.id, + context_id=first.context_id, + message_id='m2', + ), + context, + ) + ) + await asyncio.wait_for(agent.resumed.wait(), timeout=5) + agent.release.set() + await asyncio.wait_for(b_task, timeout=8) + await asyncio.wait_for(c_task, timeout=8) + + assert c_states == [ + TaskState.TASK_STATE_INPUT_REQUIRED, + TaskState.TASK_STATE_WORKING, + TaskState.TASK_STATE_COMPLETED, + ] + stored = await VersionedDatabaseTaskStore(engines[0]).get( + first.id, context + ) + assert stored is not None + assert [ + m.message_id + for m in stored.task.history + if m.role == Role.ROLE_USER + ] == ['m1', 'm2'] + finally: + for replica in (replica_a, replica_b, replica_c): + await replica.aclose() + for engine in engines: + await engine.dispose()