fix(server): save a resumed turn's message with its first event

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,

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

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