diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 51d29e520..121900fa2 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/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 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