From 80fbe5e54fec4ebd5fc27251ad9a06d8fcb1b9dd Mon Sep 17 00:00:00 2001 From: Omnishambles <160447654+HowieYin@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:34:55 +0800 Subject: [PATCH 1/3] fix(server): validate events before marking tasks terminal --- src/a2a/server/tasks/task_updater.py | 17 ++-- .../test_default_request_handler_v2.py | 60 ++++++++++++ tests/server/tasks/test_task_updater.py | 97 +++++++++++++++++++ 3 files changed, 165 insertions(+), 9 deletions(-) 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 b30fdc695..1b9813177 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -128,6 +128,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..630f24161 100644 --- a/tests/server/tasks/test_task_updater.py +++ b/tests/server/tasks/test_task_updater.py @@ -6,6 +6,7 @@ import pytest from a2a.server.events import EventQueue +from a2a.server.events.event_queue import EventQueueLegacy from a2a.server.id_generator import IDGenerator from a2a.server.tasks import TaskUpdater from a2a.types.a2a_pb2 import ( @@ -115,6 +116,102 @@ 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_message_validation_failure_allows_terminal_retry( + task_updater: TaskUpdater, event_queue: AsyncMock +) -> None: + with pytest.raises(TypeError): + await task_updater.update_status( + TaskState.TASK_STATE_COMPLETED, + message=Part(text='not a Message'), + ) + + 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_legacy_partial_publish_cancellation_keeps_terminal_guard() -> ( + None +): + parent = EventQueueLegacy(max_queue_size=1) + child = await parent.tap(max_queue_size=1) + updater = TaskUpdater(parent, 'task-id', 'context-id') + pending: asyncio.Task[None] | None = None + try: + await updater.start_work() + await parent.dequeue_event() + parent.task_done() + # The child remains full while the parent can accept COMPLETED. + pending = asyncio.create_task(updater.complete()) + event = await asyncio.wait_for(parent.dequeue_event(), timeout=1) + parent.task_done() + assert event.status.state == TaskState.TASK_STATE_COMPLETED + assert not pending.done() + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + await child.dequeue_event() + child.task_done() + with pytest.raises(RuntimeError, match='already in a terminal state'): + await updater.failed() + assert parent.queue.empty() + assert child.queue.empty() + finally: + if pending is not None and not pending.done(): + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + await parent.close(immediate=True) + await child.close(immediate=True) + + @pytest.mark.asyncio async def test_add_artifact_with_custom_id_and_name( task_updater: TaskUpdater, event_queue: AsyncMock, sample_parts: list[Part] From 6e8fe1af705dc036c3ded603e708498a818de9b8 Mon Sep 17 00:00:00 2001 From: HowieYin <160447654+HowieYin@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:46:43 +0800 Subject: [PATCH 2/3] test(server): remove redundant task updater regressions --- tests/server/tasks/test_task_updater.py | 56 ------------------------- 1 file changed, 56 deletions(-) diff --git a/tests/server/tasks/test_task_updater.py b/tests/server/tasks/test_task_updater.py index 630f24161..4756ee087 100644 --- a/tests/server/tasks/test_task_updater.py +++ b/tests/server/tasks/test_task_updater.py @@ -6,7 +6,6 @@ import pytest from a2a.server.events import EventQueue -from a2a.server.events.event_queue import EventQueueLegacy from a2a.server.id_generator import IDGenerator from a2a.server.tasks import TaskUpdater from a2a.types.a2a_pb2 import ( @@ -157,61 +156,6 @@ async def test_metadata_validation_failure_allows_terminal_retry( assert event.status.state == TaskState.TASK_STATE_FAILED -@pytest.mark.asyncio -async def test_message_validation_failure_allows_terminal_retry( - task_updater: TaskUpdater, event_queue: AsyncMock -) -> None: - with pytest.raises(TypeError): - await task_updater.update_status( - TaskState.TASK_STATE_COMPLETED, - message=Part(text='not a Message'), - ) - - 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_legacy_partial_publish_cancellation_keeps_terminal_guard() -> ( - None -): - parent = EventQueueLegacy(max_queue_size=1) - child = await parent.tap(max_queue_size=1) - updater = TaskUpdater(parent, 'task-id', 'context-id') - pending: asyncio.Task[None] | None = None - try: - await updater.start_work() - await parent.dequeue_event() - parent.task_done() - # The child remains full while the parent can accept COMPLETED. - pending = asyncio.create_task(updater.complete()) - event = await asyncio.wait_for(parent.dequeue_event(), timeout=1) - parent.task_done() - assert event.status.state == TaskState.TASK_STATE_COMPLETED - assert not pending.done() - - pending.cancel() - with pytest.raises(asyncio.CancelledError): - await pending - - await child.dequeue_event() - child.task_done() - with pytest.raises(RuntimeError, match='already in a terminal state'): - await updater.failed() - assert parent.queue.empty() - assert child.queue.empty() - finally: - if pending is not None and not pending.done(): - pending.cancel() - await asyncio.gather(pending, return_exceptions=True) - await parent.close(immediate=True) - await child.close(immediate=True) - - @pytest.mark.asyncio async def test_add_artifact_with_custom_id_and_name( task_updater: TaskUpdater, event_queue: AsyncMock, sample_parts: list[Part] From 6a318dd7591bd28a3e89418920039ab2b0321605 Mon Sep 17 00:00:00 2001 From: HowieYin <160447654+HowieYin@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:48:25 +0800 Subject: [PATCH 3/3] chore: sign task updater review update