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