From 5ff4dc95f80d1488ac98d8fdee1ed47962f310cd Mon Sep 17 00:00:00 2001 From: Yatsuiii Date: Thu, 22 Oct 2026 12:46:29 +0530 Subject: [PATCH 1/2] fix(server): record executor failures through the event consumer When an executor enqueued its Task and raised before that Task was saved, the producer's failure write and the consumer's save both inserted the same row. With DatabaseTaskStore one insert failed: the request returned the SUBMITTED task as a success, or raised IntegrityError in place of the executor's error. The producer now enqueues a _RequestFailed event and the consumer writes FAILED when it dequeues it, after the events the executor enqueued first. The fallback task is still created when the executor never made one, FAILED is stamped with the failure time, and the write still lands when the agent answered with a Message. Fixes #1351 --- src/a2a/server/agent_execution/active_task.py | 66 ++++-- .../test_default_request_handler_v2.py | 190 ++++++++++++++++++ 2 files changed, 238 insertions(+), 18 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5e3dab50f..2d752e0a5 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -40,6 +40,7 @@ import logging import uuid +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, cast from a2a.server.agent_execution.context import RequestContext @@ -108,6 +109,18 @@ def __init__(self, request_id: uuid.UUID): self.request_id = request_id +class _RequestFailed: + """The executor raised while handling a request. + + The consumer records the failure when it dequeues this, so FAILED is + written in order, after every event the executor enqueued before raising. + """ + + def __init__(self, error: Exception, context_id: str): + self.error = error + self.context_id = context_id + + class EventConsumer: """Consumes events from the agent and updates system state.""" @@ -115,6 +128,9 @@ def __init__(self, active_task: ActiveTask) -> None: self.active_task = active_task self.task_mode: bool | None = None self.message_to_save: Message | None = None + # Set when the executor raised, so a failure before any task was + # created still leaves a FAILED task behind. + self.failed_request: _RequestFailed | None = None async def run(self) -> None: """Consumes events from the agent and updates system state.""" @@ -160,16 +176,34 @@ async def run(self) -> None: ) async def _write_failed_status(self) -> Task | None: - task = await self.active_task._task_manager.get_task() + task_manager = self.active_task._task_manager + task = await task_manager.get_task() + if task is None and self.failed_request is not None: + task = await task_manager.ensure_task_id( + self.active_task._task_id, self.failed_request.context_id + ) + if task.status.state not in TERMINAL_TASK_STATES: + task.status.state = TaskState.TASK_STATE_FAILED + task.status.timestamp.FromDatetime(datetime.now(timezone.utc)) + await task_manager.save_task_event(task) + self.active_task._task_created.set() + return None if not task or task.status.state in TERMINAL_TASK_STATES: return None + # Stamped with the failure time, which DatabaseTaskStore also uses for + # last_updated. + status = TaskStatus(state=TaskState.TASK_STATE_FAILED) + status.timestamp.FromDatetime(datetime.now(timezone.utc)) handled_event = TaskStatusUpdateEvent( task_id=task.id, context_id=task.context_id, - status=TaskStatus(state=TaskState.TASK_STATE_FAILED), + status=status, ) try: - return await self._handle_task_event(handled_event) + # Not _handle_task_event: that validates what the agent sent, and + # this write must land even after the agent replied with a Message. + await self._handle_task_modification_event(handled_event) + return await task_manager.get_task() except ConcurrentTaskModificationError: # Another writer advanced the task; keep its state. self.active_task._task_manager.invalidate() @@ -218,6 +252,9 @@ async def _process_event_inner(self, event: Event) -> None: 'Consumer[%s]: Request started', self.active_task._task_id ) self.message_to_save = event.request_context.message + elif isinstance(event, _RequestFailed): + self.failed_request = event + raise event.error elif isinstance(event, BaseException): raise event elif isinstance(event, Message): @@ -667,18 +704,12 @@ async def _run_producer(self) -> None: 'Producer[%s]: Execution failed', self._task_id, ) - # Persist the failure directly instead of relying on the closing - # event queue to carry a final status update. + # The consumer writes FAILED: writing here would race its save of + # an event the executor enqueued just before raising. + failure: Exception | _RequestFailed = e if request_context: - task = await self._task_manager.ensure_task_id( - self._task_id, - request_context.context_id or '', - ) - 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)) + failure = _RequestFailed(e, request_context.context_id or '') + await self._event_queue_agent.enqueue_event(cast('Event', failure)) finally: self._request_queue.shutdown(immediate=True) @@ -907,10 +938,9 @@ async def cancel(self, call_context: ServerCallContext) -> 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, - # and the producer-failure write above) still update the shared - # object in place. Making every write copy-on-yield is the general - # property tracked in #1175 and #1191. + # 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. updated = Task() updated.CopyFrom(task) updated.status.state = TaskState.TASK_STATE_CANCELED 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..2eea656ae 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -5,6 +5,7 @@ import uuid import warnings +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -390,6 +391,69 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue): pass +class FailsAfterTaskAgentExecutor(AgentExecutor): + """Enqueues a Task with its own metadata, then raises when allowed to.""" + + def __init__(self) -> None: + self.may_raise = asyncio.Event() + self.task_id: str | None = None + + async def execute(self, context: RequestContext, event_queue: EventQueue): + assert context.message is not None + task = new_task_from_user_message(context.message) + task.metadata.update({'agent_key': 'kept'}) + self.task_id = task.id + await event_queue.enqueue_event(task) + # Bounded, so the test fails rather than hangs if the hook never fires. + await asyncio.wait_for(self.may_raise.wait(), timeout=5) + raise RuntimeError('model call failed') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +class WorksThenFailsAgentExecutor(AgentExecutor): + """Reports WORKING with a status message, then raises.""" + + def __init__(self) -> None: + self.task_id: str | None = None + self.raised_at: datetime | None = None + + async def execute(self, context: RequestContext, event_queue: EventQueue): + assert context.message is not None + task = new_task_from_user_message(context.message) + self.task_id = task.id + await event_queue.enqueue_event(task) + task_updater = TaskUpdater(event_queue, task.id, task.context_id) + await task_updater.update_status( + TaskState.TASK_STATE_WORKING, + message=task_updater.new_agent_message([Part(text='halfway')]), + ) + # Fail only after the WORKING status, with its timestamp, is stored. + await asyncio.sleep(0.05) + self.raised_at = datetime.now(timezone.utc) + raise RuntimeError('model call failed') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +class RepliesThenFailsAgentExecutor(AgentExecutor): + """Answers a follow-up with a Message, then raises.""" + + async def execute(self, context: RequestContext, event_queue: EventQueue): + task_updater = TaskUpdater( + event_queue, context.task_id, context.context_id + ) + await event_queue.enqueue_event( + task_updater.new_agent_message([Part(text='on it')]) + ) + raise RuntimeError('model call failed') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + class FailedStatusAgentExecutor(AgentExecutor): async def execute( self, context: RequestContext, event_queue: EventQueue @@ -1599,6 +1663,132 @@ async def save_task_and_signal_terminal_state(self, task): assert stored_task.status.state == terminal_state +@pytest.mark.asyncio +async def test_on_message_send_failure_right_after_task_is_recorded_on_that_task( + tmp_path, +): + """The failure lands on the agent's Task, even while that Task is saving. + + The executor raises while the store's save of its Task has looked the + task up but not yet inserted it: the window a database round trip + leaves open. + """ + pytest.importorskip('sqlalchemy') + pytest.importorskip('aiosqlite') + from a2a.server.tasks.database_task_store import DatabaseTaskStore + from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine + + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "tasks.db"}') + task_store = DatabaseTaskStore(engine) + agent_executor = FailsAfterTaskAgentExecutor() + 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_failure_right_after_task', + parts=[Part(text='Hi')], + ) + ) + context = create_server_call_context() + original_merge = AsyncSession.merge + held = False + + async def hold_agent_task_insert(self, instance, **kwargs): + nonlocal held + merged = await original_merge(self, instance, **kwargs) + if not held and 'agent_key' in ( + getattr(instance, 'task_metadata', None) or {} + ): + held = True + agent_executor.may_raise.set() + # The failure path runs before this insert is flushed. + await asyncio.sleep(0.2) + return merged + + try: + with patch.object(AsyncSession, 'merge', hold_agent_task_insert): + with pytest.raises(RuntimeError, match='model call failed'): + await request_handler.on_message_send(params, context) + assert agent_executor.task_id is not None + stored_task = await task_store.get(agent_executor.task_id, context) + finally: + await engine.dispose() + + assert held + assert stored_task is not None + assert stored_task.status.state == TaskState.TASK_STATE_FAILED + assert stored_task.metadata['agent_key'] == 'kept' + + +@pytest.mark.asyncio +async def test_on_message_send_failure_after_work_stamps_failure_time(): + """A task failing after it reported progress is stamped when it failed.""" + task_store = InMemoryTaskStore() + agent_executor = WorksThenFailsAgentExecutor() + 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_failure_after_work', + parts=[Part(text='Hi')], + ) + ) + context = create_server_call_context() + + with pytest.raises(RuntimeError, match='model call failed'): + await request_handler.on_message_send(params, context) + + assert agent_executor.task_id is not None + stored_task = await task_store.get(agent_executor.task_id, context) + assert stored_task is not None + assert stored_task.status.state == TaskState.TASK_STATE_FAILED + assert agent_executor.raised_at is not None + failed_at = stored_task.status.timestamp.ToDatetime(tzinfo=timezone.utc) + assert failed_at >= agent_executor.raised_at + + +@pytest.mark.asyncio +async def test_on_message_send_failure_after_message_reply_marks_task_failed(): + """A follow-up answered with a Message still fails the task on a raise.""" + task_store = InMemoryTaskStore() + context = create_server_call_context() + task = Task( + id='task_awaiting_input', + context_id='ctx_awaiting_input', + status=TaskStatus(state=TaskState.TASK_STATE_INPUT_REQUIRED), + ) + await task_store.save(task, context) + request_handler = DefaultRequestHandlerV2( + agent_executor=RepliesThenFailsAgentExecutor(), + task_store=task_store, + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_follow_up', + parts=[Part(text='large')], + task_id=task.id, + context_id=task.context_id, + ) + ) + + with pytest.raises(RuntimeError, match='model call failed'): + await request_handler.on_message_send(params, context) + + stored_task = await task_store.get(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_failed_task_does_not_hide_producer_exception() -> ( None From 202094804c5b058e6be3fbd0c7a60b2e95cfecaa Mon Sep 17 00:00:00 2001 From: Yatsuiii Date: Thu, 22 Oct 2026 18:09:07 +0530 Subject: [PATCH 2/2] refactor(server): drop _RequestFailed and return the fallback task The consumer now notes a producer failure when it dequeues the exception, since only the producer enqueues exceptions, and takes the context id for the fallback task from the task manager. The fallback path returns the task it created, and _write_failed_status documents its three cases. --- src/a2a/server/agent_execution/active_task.py | 46 ++++++++----------- .../agent_execution/test_active_task.py | 1 + 2 files changed, 21 insertions(+), 26 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 2d752e0a5..e02ff3315 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -109,18 +109,6 @@ def __init__(self, request_id: uuid.UUID): self.request_id = request_id -class _RequestFailed: - """The executor raised while handling a request. - - The consumer records the failure when it dequeues this, so FAILED is - written in order, after every event the executor enqueued before raising. - """ - - def __init__(self, error: Exception, context_id: str): - self.error = error - self.context_id = context_id - - class EventConsumer: """Consumes events from the agent and updates system state.""" @@ -128,9 +116,9 @@ def __init__(self, active_task: ActiveTask) -> None: self.active_task = active_task self.task_mode: bool | None = None self.message_to_save: Message | None = None - # Set when the executor raised, so a failure before any task was - # created still leaves a FAILED task behind. - self.failed_request: _RequestFailed | None = None + # Set when the producer failed, so a request that failed before any + # task was created still leaves a FAILED task behind. + self.producer_failed = False async def run(self) -> None: """Consumes events from the agent and updates system state.""" @@ -176,18 +164,29 @@ async def run(self) -> None: ) async def _write_failed_status(self) -> Task | None: + """Records FAILED once a failure ends the event stream. + + After a producer failure this runs once every event the executor + enqueued before raising has been processed, so the write lands in + order. There are three cases: + + - The request failed before a task was created: create one as FAILED. + - No task exists and the failure is the consumer's own, or the task + already reached a terminal state: nothing is written. + - Otherwise: the existing task is marked FAILED. + """ task_manager = self.active_task._task_manager task = await task_manager.get_task() - if task is None and self.failed_request is not None: + if task is None and self.producer_failed: task = await task_manager.ensure_task_id( - self.active_task._task_id, self.failed_request.context_id + self.active_task._task_id, task_manager.context_id or '' ) if task.status.state not in TERMINAL_TASK_STATES: task.status.state = TaskState.TASK_STATE_FAILED task.status.timestamp.FromDatetime(datetime.now(timezone.utc)) await task_manager.save_task_event(task) self.active_task._task_created.set() - return None + return task if not task or task.status.state in TERMINAL_TASK_STATES: return None # Stamped with the failure time, which DatabaseTaskStore also uses for @@ -252,10 +251,9 @@ async def _process_event_inner(self, event: Event) -> None: 'Consumer[%s]: Request started', self.active_task._task_id ) self.message_to_save = event.request_context.message - elif isinstance(event, _RequestFailed): - self.failed_request = event - raise event.error elif isinstance(event, BaseException): + # Only the producer enqueues exceptions, when a request fails. + self.producer_failed = True raise event elif isinstance(event, Message): self._handle_message_event(event) @@ -636,7 +634,6 @@ async def _run_producer(self) -> None: Runs as a detached asyncio.Task. Safe to cancel. """ logger.debug('Producer[%s]: Started', self._task_id) - request_context = None try: while True: ( @@ -706,10 +703,7 @@ async def _run_producer(self) -> None: ) # The consumer writes FAILED: writing here would race its save of # an event the executor enqueued just before raising. - failure: Exception | _RequestFailed = e - if request_context: - failure = _RequestFailed(e, request_context.context_id or '') - await self._event_queue_agent.enqueue_event(cast('Event', failure)) + await self._event_queue_agent.enqueue_event(cast('Event', e)) finally: self._request_queue.shutdown(immediate=True) diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 7b0a2a012..254423e18 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -42,6 +42,7 @@ def task_manager(self) -> Mock: tm.get_task = AsyncMock(return_value=None) tm.context_id = 'test-context-id' tm._init_task_obj = Mock(return_value=Task(id='test-task-id')) + tm.ensure_task_id = AsyncMock(return_value=Task(id='test-task-id')) tm.save_task_event = AsyncMock() return tm