diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5e3dab50f..e02ff3315 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 @@ -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 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 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