From 84a18ddec50f7f0e081a8b0e59b0d0a6d12e7d0b Mon Sep 17 00:00:00 2001 From: Aleksandr Date: Tue, 16 Jun 2026 08:59:34 +0000 Subject: [PATCH 1/4] Fix early producer failure task persistence --- src/a2a/server/agent_execution/active_task.py | 26 ++++-- .../agent_execution/active_task_registry.py | 4 +- .../default_request_handler_v2.py | 1 + .../test_default_request_handler_v2.py | 81 +++++++++++++++++++ 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index b0154c8d6..7b3cfc486 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -260,11 +260,20 @@ async def _handle_task_modification_event( ) if self.message_to_save is not None: - updated_task = self.active_task._task_manager.update_with_message( - self.message_to_save, - updated_task, + message_already_saved = any( + message == self.message_to_save + for message in updated_task.history ) - await self.active_task._task_manager.save_task_event(updated_task) + if not message_already_saved: + updated_task = ( + self.active_task._task_manager.update_with_message( + self.message_to_save, + updated_task, + ) + ) + await self.active_task._task_manager.save_task_event( + updated_task + ) self.message_to_save = None self.active_task._task_manager.context_id = event.context_id @@ -551,12 +560,17 @@ async def _run_producer(self) -> None: 'Producer[%s]: Execution failed', self._task_id, ) - # Create task and mark as failed. + # Persist the failure directly instead of relying on the closing + # event queue to carry a final status update. if request_context: - await self._task_manager.ensure_task_id( + task = await self._task_manager.ensure_task_id( self._task_id, request_context.context_id or '', ) + task.status.CopyFrom( + TaskStatus(state=TaskState.TASK_STATE_FAILED) + ) + await self._task_manager.save_task_event(task) self._task_created.set() await self._event_queue_agent.enqueue_event(cast('Event', e)) diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index 9c1299ab3..4f5822f68 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -11,6 +11,7 @@ from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.server.tasks.task_store import TaskStore + from a2a.types.a2a_pb2 import Message from a2a.server.agent_execution.active_task import ActiveTask from a2a.server.tasks.task_manager import TaskManager @@ -41,6 +42,7 @@ async def get_or_create( call_context: ServerCallContext, context_id: str | None = None, create_task_if_missing: bool = False, + initial_message: Message | None = None, ) -> ActiveTask: """Retrieves an existing ActiveTask or creates a new one.""" async with self._lock: @@ -51,7 +53,7 @@ async def get_or_create( task_id=task_id, context_id=context_id, task_store=self._task_store, - initial_message=None, + initial_message=initial_message, context=call_context, ) diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 30304609a..599e38ce5 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -222,6 +222,7 @@ async def _setup_active_task( context_id=context_id, call_context=call_context, create_task_if_missing=True, + initial_message=request_context.message, ) return active_task, request_context 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 caaa4f88e..a27d56ce2 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -21,6 +21,7 @@ from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry from a2a.server.context import ServerCallContext from a2a.server.events import EventQueue +from a2a.server.events.event_queue_v2 import EventQueueSource from a2a.server.request_handlers import DefaultRequestHandlerV2 from a2a.server.tasks import ( InMemoryPushNotificationConfigStore, @@ -306,6 +307,26 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue): pass +class EarlyFailingAgentExecutor(AgentExecutor): + async def execute(self, context: RequestContext, event_queue: EventQueue): + raise RuntimeError('early producer failure') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +async def send_message_with_early_failure( + request_handler: DefaultRequestHandlerV2, + params: SendMessageRequest, + context: ServerCallContext, +) -> Message | Task | None: + try: + return await request_handler.on_message_send(params, context) + except RuntimeError as e: + assert str(e) == 'early producer failure' + return None + + @pytest.mark.asyncio async def test_on_get_task_limit_history(): task_store = InMemoryTaskStore() @@ -1126,6 +1147,66 @@ async def test_on_message_send_limit_history(): assert task.history is not None and len(task.history) > 1 +@pytest.mark.asyncio +async def test_on_message_send_early_producer_exception_marks_task_failed(): + task_store = InMemoryTaskStore() + request_handler = DefaultRequestHandlerV2( + agent_executor=HelloAgentExecutor(), + task_store=task_store, + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_early_failure_state', + parts=[Part(text='Hi')], + ) + ) + context = create_server_call_context() + original_enqueue_event = EventQueueSource.enqueue_event + + async def fail_before_request_started(self, event): + if type(event).__name__ == '_RequestStarted': + raise RuntimeError('early producer failure') + return await original_enqueue_event(self, event) + + with patch.object( + EventQueueSource, 'enqueue_event', fail_before_request_started + ): + await send_message_with_early_failure(request_handler, params, context) + + stored_task = await task_store.get(params.message.task_id, context) + assert stored_task is not None + assert stored_task.status.state == TaskState.TASK_STATE_FAILED + + +@pytest.mark.asyncio +async def test_on_message_send_early_producer_exception_preserves_originating_message(): + task_store = InMemoryTaskStore() + request_handler = DefaultRequestHandlerV2( + agent_executor=EarlyFailingAgentExecutor(), + task_store=task_store, + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_early_failure_history', + parts=[Part(text='Hi')], + ) + ) + context = create_server_call_context() + + await send_message_with_early_failure(request_handler, params, context) + + stored_task = await task_store.get(params.message.task_id, context) + assert stored_task is not None + assert stored_task.history is not None + assert len(stored_task.history) == 1 + assert stored_task.history[0].message_id == 'msg_early_failure_history' + assert stored_task.history[0].parts[0].text == 'Hi' + + @pytest.mark.asyncio async def test_on_message_send_stream_task_id_mismatch(): mock_task_store = AsyncMock(spec=TaskStore) From 42c44b39451b83eb514cd567eeb754d7c790a508 Mon Sep 17 00:00:00 2001 From: Aleksandr Golovin Date: Tue, 16 Jun 2026 11:20:43 +0200 Subject: [PATCH 2/4] test: reproduce #1067 early producer failure on main --- .../request_handlers/test_default_request_handler_v2.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 a27d56ce2..e49ed5de4 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -1148,7 +1148,7 @@ async def test_on_message_send_limit_history(): @pytest.mark.asyncio -async def test_on_message_send_early_producer_exception_marks_task_failed(): +async def test_on_message_send_early_producer_exception_marks_task_failed_and_preserves_originating_message(): task_store = InMemoryTaskStore() request_handler = DefaultRequestHandlerV2( agent_executor=HelloAgentExecutor(), @@ -1178,6 +1178,9 @@ async def fail_before_request_started(self, event): stored_task = await task_store.get(params.message.task_id, context) assert stored_task is not None assert stored_task.status.state == TaskState.TASK_STATE_FAILED + assert len(stored_task.history) == 1 + assert stored_task.history[0].message_id == 'msg_early_failure_state' + assert stored_task.history[0].parts[0].text == 'Hi' @pytest.mark.asyncio From 5d71f74a4ae68a9d76ea8b1047eea3c8254de427 Mon Sep 17 00:00:00 2001 From: Aleksandr Date: Tue, 16 Jun 2026 09:51:53 +0000 Subject: [PATCH 3/4] refactor: address active task review feedback --- src/a2a/server/agent_execution/active_task.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 7b3cfc486..5f965560d 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -261,7 +261,7 @@ async def _handle_task_modification_event( if self.message_to_save is not None: message_already_saved = any( - message == self.message_to_save + message.message_id == self.message_to_save.message_id for message in updated_task.history ) if not message_already_saved: @@ -567,9 +567,7 @@ async def _run_producer(self) -> None: self._task_id, request_context.context_id or '', ) - task.status.CopyFrom( - TaskStatus(state=TaskState.TASK_STATE_FAILED) - ) + task.status.state = TaskState.TASK_STATE_FAILED await self._task_manager.save_task_event(task) self._task_created.set() await self._event_queue_agent.enqueue_event(cast('Event', e)) From d32abad56c7edaa0fed90cce31b749c29cd9337f Mon Sep 17 00:00:00 2001 From: Aleksandr Date: Thu, 16 Jul 2026 07:46:03 +0000 Subject: [PATCH 4/4] fix: preserve terminal task state after late producer failure --- src/a2a/server/agent_execution/active_task.py | 7 +- .../test_default_request_handler_v2.py | 71 +++++++++++++++++++ 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5f965560d..1420c7a2a 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -141,7 +141,7 @@ async def run(self) -> None: updated_task = None task = await self.active_task._task_manager.get_task() - if task: + if task and task.status.state not in TERMINAL_TASK_STATES: handled_event = TaskStatusUpdateEvent( task_id=task.id, context_id=task.context_id, @@ -567,8 +567,9 @@ async def _run_producer(self) -> None: self._task_id, request_context.context_id or '', ) - task.status.state = TaskState.TASK_STATE_FAILED - await self._task_manager.save_task_event(task) + if task.status.state not in TERMINAL_TASK_STATES: + task.status.state = TaskState.TASK_STATE_FAILED + await self._task_manager.save_task_event(task) self._task_created.set() await self._event_queue_agent.enqueue_event(cast('Event', e)) 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 e49ed5de4..da8b2c561 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -31,6 +31,7 @@ TaskStore, TaskUpdater, ) +from a2a.server.tasks.task_manager import TaskManager from a2a.types import ( InternalError, InvalidAgentResponseError, @@ -315,6 +316,28 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue): pass +class LateFailingTerminalAgentExecutor(AgentExecutor): + def __init__( + self, terminal_state: TaskState, terminal_state_persisted: asyncio.Event + ) -> None: + self.terminal_state = terminal_state + self.terminal_state_persisted = terminal_state_persisted + self.raised = asyncio.Event() + + async def execute(self, context: RequestContext, event_queue: EventQueue): + assert context.message is not None + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + task_updater = TaskUpdater(event_queue, task.id, task.context_id) + await task_updater.update_status(self.terminal_state) + await self.terminal_state_persisted.wait() + self.raised.set() + raise RuntimeError('late producer failure') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + async def send_message_with_early_failure( request_handler: DefaultRequestHandlerV2, params: SendMessageRequest, @@ -1183,6 +1206,54 @@ async def fail_before_request_started(self, event): assert stored_task.history[0].parts[0].text == 'Hi' +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'terminal_state', + [ + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_CANCELED, + ], +) +async def test_on_message_send_late_producer_exception_preserves_persisted_terminal_state( + terminal_state: TaskState, +): + task_store = InMemoryTaskStore() + terminal_state_persisted = asyncio.Event() + agent_executor = LateFailingTerminalAgentExecutor( + terminal_state, terminal_state_persisted + ) + request_handler = DefaultRequestHandlerV2( + agent_executor=agent_executor, + task_store=task_store, + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_late_failure_terminal_state', + parts=[Part(text='Hi')], + ) + ) + context = create_server_call_context() + original_save_task = TaskManager._save_task + + async def save_task_and_signal_terminal_state(self, task): + await original_save_task(self, task) + if task.status.state == terminal_state: + terminal_state_persisted.set() + + with patch.object( + TaskManager, '_save_task', save_task_and_signal_terminal_state + ): + await request_handler.on_message_send(params, context) + await agent_executor.raised.wait() + await asyncio.sleep(0) + + stored_task = await task_store.get(params.message.task_id, context) + assert stored_task is not None + assert stored_task.status.state == terminal_state + + @pytest.mark.asyncio async def test_on_message_send_early_producer_exception_preserves_originating_message(): task_store = InMemoryTaskStore()