diff --git a/src/a2a/server/tasks/task_updater.py b/src/a2a/server/tasks/task_updater.py index 8298920da..77f7a0009 100644 --- a/src/a2a/server/tasks/task_updater.py +++ b/src/a2a/server/tasks/task_updater.py @@ -84,8 +84,6 @@ async def update_status( raise RuntimeError( f'Task {self.task_id} is already in a terminal state.' ) - if state in self._terminal_states: - self._terminal_state_reached = True # Create proto timestamp from datetime ts = Timestamp() @@ -101,14 +99,15 @@ async def update_status( status.message.CopyFrom(message) status.timestamp.CopyFrom(ts) - await self.event_queue.enqueue_event( - TaskStatusUpdateEvent( - task_id=self.task_id, - context_id=self.context_id, - metadata=metadata, - status=status, - ) + event = TaskStatusUpdateEvent( + task_id=self.task_id, + context_id=self.context_id, + metadata=metadata, + status=status, ) + if state in self._terminal_states: + self._terminal_state_reached = True + await self.event_queue.enqueue_event(event) async def add_artifact( # noqa: PLR0913 self, 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 3ac14862e..a57be1fa0 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -129,6 +129,66 @@ def create_server_call_context() -> ServerCallContext: return ServerCallContext(user=UnauthenticatedUser()) +class _StatusValidationRecoveryAgent(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() + try: + await updater.update_status( + TaskState.TASK_STATE_COMPLETED, + timestamp='not-a-timestamp', + ) + except ValueError: + await updater.failed() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + pass + + +@pytest.mark.asyncio +async def test_status_validation_failure_returns_failed_task() -> None: + store = InMemoryTaskStore() + handler = DefaultRequestHandlerV2( + agent_executor=_StatusValidationRecoveryAgent(), + task_store=store, + agent_card=create_default_agent_card(), + ) + context = create_server_call_context() + try: + result = await asyncio.wait_for( + handler.on_message_send( + SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='validation-recovery', + parts=[Part(text='test status recovery')], + ) + ), + context, + ), + timeout=5, + ) + assert isinstance(result, Task) + assert result.status.state == TaskState.TASK_STATE_FAILED + stored = await store.get(result.id, context) + assert stored is not None + assert stored.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_updater.py b/tests/server/tasks/test_task_updater.py index 49d9dee43..4756ee087 100644 --- a/tests/server/tasks/test_task_updater.py +++ b/tests/server/tasks/test_task_updater.py @@ -115,6 +115,47 @@ async def test_update_status_final( assert event.status.state == TaskState.TASK_STATE_COMPLETED +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'state', + [ + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_CANCELED, + TaskState.TASK_STATE_FAILED, + TaskState.TASK_STATE_REJECTED, + ], +) +async def test_timestamp_validation_failure_allows_terminal_retry( + task_updater: TaskUpdater, event_queue: AsyncMock, state: TaskState +) -> None: + with pytest.raises(ValueError): + await task_updater.update_status(state, timestamp='not-a-timestamp') + + event_queue.enqueue_event.assert_not_awaited() + await task_updater.failed() + + event_queue.enqueue_event.assert_awaited_once() + event = event_queue.enqueue_event.call_args.args[0] + assert event.status.state == TaskState.TASK_STATE_FAILED + + +@pytest.mark.asyncio +async def test_metadata_validation_failure_allows_terminal_retry( + task_updater: TaskUpdater, event_queue: AsyncMock +) -> None: + with pytest.raises(ValueError): + await task_updater.update_status( + TaskState.TASK_STATE_COMPLETED, metadata={'bad': object()} + ) + + event_queue.enqueue_event.assert_not_awaited() + await task_updater.failed() + + event_queue.enqueue_event.assert_awaited_once() + event = event_queue.enqueue_event.call_args.args[0] + assert event.status.state == TaskState.TASK_STATE_FAILED + + @pytest.mark.asyncio async def test_add_artifact_with_custom_id_and_name( task_updater: TaskUpdater, event_queue: AsyncMock, sample_parts: list[Part]