From e6474e238b285e316737c3ed780ee70effaff520 Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Wed, 28 Oct 2026 18:27:24 -0700 Subject: [PATCH] fix(server): record FAILED when a task's terminal save fails TaskManager.save_task_event applied status and artifact events to the cached task in place and only then called the task store. When the store raised, the cache was ahead of the store. ActiveTask's failure handler reads that cache to decide whether to record FAILED, saw the unsaved COMPLETED, and wrote nothing, so the task stayed WORKING in the store. Apply the event to a copy instead. _save_task already caches the task only after the store accepts it, so a failed save leaves the cache as last persisted, and a VersionedTaskStore now receives the task as previously read in `prev`, as its contract documents. --- src/a2a/server/agent_execution/active_task.py | 8 +-- src/a2a/server/tasks/task_manager.py | 7 +- .../test_default_request_handler.py | 10 +-- .../test_default_request_handler_v2.py | 72 +++++++++++++++++++ tests/server/tasks/test_task_manager.py | 65 +++++++++++++++++ 5 files changed, 152 insertions(+), 10 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index c5dcbb9dd..743094f0c 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -938,10 +938,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 diff --git a/src/a2a/server/tasks/task_manager.py b/src/a2a/server/tasks/task_manager.py index 420f9f239..78722e570 100644 --- a/src/a2a/server/tasks/task_manager.py +++ b/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( diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 5ed47093d..27bd4f16d 100644 --- a/tests/server/request_handlers/test_default_request_handler.py +++ b/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( diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index 66cf9f4ea..9043df77a 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/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() diff --git a/tests/server/tasks/test_task_manager.py b/tests/server/tasks/test_task_manager.py index 7daa3ab3d..dc8306d06 100644 --- a/tests/server/tasks/test_task_manager.py +++ b/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