fix(server): record executor failures through the event consumer

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
@@ -115,6 +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 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."""
@@ -160,16 +164,45 @@ async def run(self) -> None:
                 )
 
     async def _write_failed_status(self) -> Task | None:
-        task = await self.active_task._task_manager.get_task()
+        """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.producer_failed:
+            task = await task_manager.ensure_task_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 task
         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()
@@ -219,6 +252,8 @@ async def _process_event_inner(self, event: Event) -> None:
             )
             self.message_to_save = event.request_context.message
         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)
@@ -599,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:
                 (
@@ -667,17 +701,8 @@ 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.
-            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()
+            # The consumer writes FAILED: writing here would race its save of
+            # an event the executor enqueued just before raising.
             await self._event_queue_agent.enqueue_event(cast('Event', e))
 
         finally:
@@ -907,10 +932,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

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
 

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