fix(server): record FAILED when a task's terminal save fails

src/a2a/server/agent_execution/active_task.py

@@ -964,10 +964,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
             # Write a copy rather than mutating the task in place: get_task()
             # returns the shared _task_manager._current_task, which may already
             # have been yielded to a subscriber. Copying keeps THIS terminal
-            # write off the yielded reference. It is not a file-wide guarantee:
-            # ordinary status writes (save_task_event -> ensure_task -> CopyFrom)
-            # still update the shared object in place. Making every write
-            # copy-on-yield is the general property tracked in #1175 and #1191.
+            # write off the yielded reference. Status and artifact events are
+            # applied to a copy too (TaskManager.save_task_event), but a Task
+            # event or update_with_message still writes the object it is given.
+            # Making every write copy-on-yield is tracked in #1175 and #1191.
             updated = Task()
             updated.CopyFrom(task)
             updated.status.state = TaskState.TASK_STATE_CANCELED

src/a2a/server/tasks/task_manager.py

@@ -267,7 +267,12 @@ async def save_task_event(
                 self._pending_event = None
             return event
 
-        task: Task = await self.ensure_task(event)
+        # Apply the event to a copy; _save_task caches it only once the store
+        # accepts it. If the save raises, the cached task still matches the
+        # store, so the failure path does not take an unsaved terminal state
+        # for a saved one, and `prev` is the task as last read.
+        task = Task()
+        task.CopyFrom(await self.ensure_task(event))
 
         if isinstance(event, TaskStatusUpdateEvent):
             logger.debug(

tests/server/request_handlers/test_default_request_handler.py

@@ -1098,11 +1098,10 @@ async def test_on_message_send_non_blocking(agent_card):
 
     assert task is not None
     assert task.status.state == TaskState.TASK_STATE_COMPLETED
-    assert (
-        result.history
-        and task.history
-        and len(result.history) == len(task.history)
-    )
+    # The result is the task as it was returned, with its whole history; the
+    # agent's later status messages are added to the stored task only.
+    assert result.history
+    assert list(task.history[: len(result.history)]) == list(result.history)
 
 
 @pytest.mark.asyncio
@@ -2941,6 +2940,7 @@ async def test_on_message_send_error_does_not_hang(agent_card):
     """Test that if the consumer raises an exception during blocking wait, the producer is cancelled and no deadlock occurs."""
     agent = HelloWorldAgentExecutor()
     task_store = AsyncMock(spec=TaskStore)
+    task_store.get.return_value = None
     task_store.save.side_effect = RuntimeError('This is an Error!')
 
     request_handler = DefaultRequestHandler(

tests/server/request_handlers/test_default_request_handler_v2.py

@@ -190,6 +190,78 @@ async def test_status_validation_failure_returns_failed_task() -> None:
         await handler.aclose()
 
 
+class _CompletingAgent(AgentExecutor):
+    async def execute(
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
+        assert context.message is not None
+        await event_queue.enqueue_event(
+            new_task_from_user_message(context.message)
+        )
+        updater = TaskUpdater(
+            event_queue,
+            str(context.task_id or ''),
+            str(context.context_id or ''),
+        )
+        await updater.start_work()
+        await updater.complete()
+
+    async def cancel(
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
+        pass
+
+
+class _FailCompletedSaveTaskStore(InMemoryTaskStore):
+    """Rejects the first COMPLETED write, as a lost database connection would."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.failed_save = False
+
+    async def save(self, task: Task, context: ServerCallContext) -> None:
+        if (
+            task.status.state == TaskState.TASK_STATE_COMPLETED
+            and not self.failed_save
+        ):
+            self.failed_save = True
+            raise ConnectionError('database connection lost')
+        await super().save(task, context)
+
+
+@pytest.mark.asyncio
+async def test_failed_terminal_save_records_failed_task() -> None:
+    store = _FailCompletedSaveTaskStore()
+    handler = DefaultRequestHandlerV2(
+        agent_executor=_CompletingAgent(),
+        task_store=store,
+        agent_card=create_default_agent_card(),
+    )
+    context = create_server_call_context()
+    try:
+        with pytest.raises(ConnectionError, match='database connection lost'):
+            await asyncio.wait_for(
+                handler.on_message_send(
+                    SendMessageRequest(
+                        message=Message(
+                            role=Role.ROLE_USER,
+                            message_id='failed-terminal-save',
+                            parts=[Part(text='test failed save')],
+                        )
+                    ),
+                    context,
+                ),
+                timeout=5,
+            )
+        assert store.failed_save
+        listed = await store.list(ListTasksRequest(), context)
+        assert len(listed.tasks) == 1
+        # The COMPLETED write never landed, so the task must not stay WORKING.
+        assert listed.tasks[0].status.state == TaskState.TASK_STATE_FAILED
+    finally:
+        await handler.aclose()
+
+
 def test_init_default_dependencies():
     """Test that default dependencies are created if not provided."""
     agent_executor = MockAgentExecutor()

tests/server/tasks/test_task_manager.py

@@ -4,6 +4,7 @@
 import pytest
 
 from a2a.auth.user import User
+from a2a.server.cluster import VersionedTaskStore
 from a2a.server.context import ServerCallContext
 from a2a.server.tasks import TaskManager
 from a2a.server.tasks.task_manager import append_artifact_to_task
@@ -196,6 +197,70 @@ async def test_save_task_event_metadata_update(
     assert updated_task.metadata == new_metadata
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'event',
+    [
+        TaskStatusUpdateEvent(
+            task_id=MINIMAL_TASK_ID,
+            context_id=MINIMAL_CONTEXT_ID,
+            status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
+        ),
+        TaskArtifactUpdateEvent(
+            task_id=MINIMAL_TASK_ID,
+            context_id=MINIMAL_CONTEXT_ID,
+            artifact=Artifact(
+                artifact_id='artifact-id', parts=[Part(text='chunk')]
+            ),
+        ),
+    ],
+    ids=['status', 'artifact'],
+)
+async def test_save_task_event_failed_save_keeps_cached_task(
+    task_manager: TaskManager,
+    mock_task_store: AsyncMock,
+    event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
+) -> None:
+    """A rejected write leaves the cached task as it was last persisted."""
+    mock_task_store.get.return_value = create_minimal_task()
+    await task_manager.get_task()
+    mock_task_store.save.side_effect = ConnectionError('connection lost')
+
+    with pytest.raises(ConnectionError):
+        await task_manager.save_task_event(event)
+
+    task = await task_manager.get_task()
+    assert task is not None
+    assert task.status.state == TaskState.TASK_STATE_SUBMITTED
+    assert not task.artifacts
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_passes_previous_task_as_prev() -> None:
+    """A versioned store receives the task as it was before the event."""
+    store = AsyncMock(spec=VersionedTaskStore)
+    task_manager = TaskManager(
+        task_id=MINIMAL_TASK_ID,
+        context_id=MINIMAL_CONTEXT_ID,
+        task_store=store,
+        initial_message=None,
+        context=TEST_CONTEXT,
+    )
+    await task_manager.save_task_event(create_minimal_task())
+
+    await task_manager.save_task_event(
+        TaskStatusUpdateEvent(
+            task_id=MINIMAL_TASK_ID,
+            context_id=MINIMAL_CONTEXT_ID,
+            status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
+        )
+    )
+
+    call = store.save.call_args
+    assert call.args[0].status.state == TaskState.TASK_STATE_WORKING
+    assert call.kwargs['prev'].status.state == TaskState.TASK_STATE_SUBMITTED
+
+
 @pytest.mark.asyncio
 async def test_ensure_task_existing(
     task_manager: TaskManager, mock_task_store: AsyncMock