diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index ea1955fdd..67d248dfe 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -36,6 +36,7 @@ from __future__ import annotations import asyncio +import inspect import logging import uuid @@ -560,17 +561,26 @@ 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. + # Persist FAILED (store + push) before finally closes the + # queues (issue #1175). Do not emit a FAILED status event: + # blocking on_message_send would treat that Task as success. + # The producer exception is the stream signal. 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() + try: + await self._task_manager.ensure_task_id( + self._task_id, + request_context.context_id or '', + ) + await self._persist_and_publish_terminal( + TaskState.TASK_STATE_FAILED, + publish_to_stream=False, + ) + self._task_created.set() + except Exception: + logger.exception( + 'Producer[%s]: Failed to persist FAILED state', + self._task_id, + ) await self._event_queue_agent.enqueue_event(cast('Event', e)) finally: @@ -730,7 +740,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task: logger.debug( 'Cancel[%s]: Cancelling producer task', self._task_id ) - self._producer_task.cancel() + # Await executor.cancel before cancelling the producer so a + # terminal write can still reach the open subscriber queue + # (#1172 / #1175). Producer cancel stays in finally so a + # BaseException from executor.cancel cannot leak the producer. try: await self._agent_executor.cancel( request_context, self._event_queue_agent @@ -741,6 +754,27 @@ async def cancel(self, call_context: ServerCallContext) -> Task: ) await self._mark_task_as_failed(e) raise + finally: + try: + task = await self._task_manager.get_task() + if ( + task is not None + and task.status.state not in TERMINAL_TASK_STATES + ): + # Cleanup-only executor.cancel() or a parked + # input-required task leaves no terminal event. + # Write CANCELED and publish it while queues + # are still open so a live subscriber does not + # have to poll. + await self._persist_and_publish_terminal( + TaskState.TASK_STATE_CANCELED + ) + except Exception: + logger.exception( + 'Cancel[%s]: Failed to persist CANCELED state', + self._task_id, + ) + self._producer_task.cancel() else: logger.debug( 'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling', @@ -748,6 +782,14 @@ async def cancel(self, call_context: ServerCallContext) -> Task: self._is_finished.is_set(), self._producer_task, ) + task = await self._task_manager.get_task() + if ( + task is not None + and task.status.state not in TERMINAL_TASK_STATES + ): + await self._persist_and_publish_terminal( + TaskState.TASK_STATE_CANCELED + ) await self._is_finished.wait() task = await self._task_manager.get_task() @@ -820,6 +862,54 @@ async def _maybe_cleanup(self) -> None: logger.debug('Cleanup[%s]: Triggering cleanup', self._task_id) self._on_cleanup(self) + async def _persist_and_publish_terminal( + self, state: TaskState, *, publish_to_stream: bool = True + ) -> Task | None: + """Write a terminal state to the store and notify live observers. + + Direct ``save_task_event`` after the subscriber queue is closed is + invisible to ``SubscribeToTask`` / ``message/stream`` and to push + (issue #1175). This helper persists a *copy* of the current task + (the shared ``get_task()`` object must not mutate under a reader) + and, when ``publish_to_stream`` is true, emits a + ``TaskStatusUpdateEvent`` to subscribers *before* teardown. + + Producer-failure keeps ``publish_to_stream=False``: blocking + ``on_message_send`` treats a FAILED ``Task`` as a successful + terminal result, so the crash must still surface as the + producer exception on the stream. Store and push still get + FAILED before the queues close. + """ + task = await self._task_manager.get_task() + if task is None: + return None + + if task.status.state not in TERMINAL_TASK_STATES: + updated = Task() + updated.CopyFrom(task) + updated.status.state = state + await self._task_manager.save_task_event(updated) + task = updated + + event = TaskStatusUpdateEvent( + task_id=task.id, + context_id=task.context_id, + status=TaskStatus(state=task.status.state), + ) + if publish_to_stream: + updated_task_copy = Task() + updated_task_copy.CopyFrom(task) + await self._event_queue_subscribers.enqueue_event( + cast('Any', (event, updated_task_copy)) + ) + if self._push_sender and self._task_id: + notification = self._push_sender.send_notification( + self._task_id, event + ) + if inspect.isawaitable(notification): + await notification + return task + async def _mark_task_as_failed(self, exception: Exception) -> Task | None: logger.debug('Marking task %s as failed: %s', self._task_id, exception) task = None diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 1be233ee1..67226640e 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -129,6 +129,148 @@ async def execute_mock(req, q): agent_executor.cancel.assert_called_once() stop_event.set() + @staticmethod + def _wire_current_task(task_manager: Mock, task: Task) -> None: + """Make get_task / save_task_event share one current Task.""" + + async def get_task() -> Task: + return task_manager._current_task + + async def ensure_task_id(task_id: str, context_id: str) -> Task: + return task_manager._current_task + + async def save_task_event(event: Task) -> None: + if isinstance(event, Task): + task_manager._current_task = event + + task_manager._current_task = task + task_manager.get_task = AsyncMock(side_effect=get_task) + task_manager.save_task_event = AsyncMock(side_effect=save_task_event) + task_manager.ensure_task_id = AsyncMock(side_effect=ensure_task_id) + + @pytest.mark.asyncio + async def test_cancel_publishes_canceled_to_subscriber_stream( + self, + active_task: ActiveTask, + agent_executor: Mock, + request_context: Mock, + task_manager: Mock, + push_sender: Mock, + ) -> None: + """Issue #1175: cancel() must emit CANCELED on a live subscriber stream. + + A cleanup-only executor.cancel() writes nothing. The helper persists a + copy and enqueues TaskStatusUpdateEvent before the producer tears the + subscriber queue down. + """ + stop_event = asyncio.Event() + + async def execute_mock(req, q): + await stop_event.wait() + + shared = Task( + id='test-task-id', + context_id='test-context-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + self._wire_current_task(task_manager, shared) + agent_executor.execute = AsyncMock(side_effect=execute_mock) + agent_executor.cancel = AsyncMock() + + await active_task.enqueue_request(request_context) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + await asyncio.sleep(0.05) + + events: list[object] = [] + + async def collect() -> None: + try: + async for event in active_task.subscribe(): + events.append(event) + except Exception: # noqa: BLE001 + pass + + collector = asyncio.create_task(collect()) + await asyncio.sleep(0.05) + + result = await active_task.cancel(request_context) + stop_event.set() + await asyncio.wait_for(collector, timeout=2) + + assert result.status.state == TaskState.TASK_STATE_CANCELED + assert result is not shared + assert shared.status.state == TaskState.TASK_STATE_WORKING + status_events = [ + e + for e in events + if isinstance(e, TaskStatusUpdateEvent) + and e.status.state == TaskState.TASK_STATE_CANCELED + ] + assert status_events, f'subscriber saw {events!r}, expected CANCELED' + push_sender.send_notification.assert_awaited() + + @pytest.mark.asyncio + async def test_producer_failure_persists_failed_and_notifies_push( + self, + active_task: ActiveTask, + agent_executor: Mock, + request_context: Mock, + task_manager: Mock, + push_sender: Mock, + ) -> None: + """Issue #1175: producer-failure persists FAILED and notifies push. + + The crash still surfaces as ValueError on the stream: blocking + on_message_send would treat a FAILED Task as a successful result. + """ + shared = Task( + id='test-task-id', + context_id='test-context-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + self._wire_current_task(task_manager, shared) + request_context.context_id = 'test-context-id' + crash = asyncio.Event() + + async def execute_mock(req, q): + await crash.wait() + raise ValueError('Producer crashed') + + agent_executor.execute = AsyncMock(side_effect=execute_mock) + + await active_task.enqueue_request(request_context) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + # Let the consumer flush _RequestStarted so this tap sees only the + # terminal publish (same timing as the cancel-stream test). + await asyncio.sleep(0.05) + + collector_error: list[BaseException] = [] + + async def collect() -> None: + try: + async for _event in active_task.subscribe(): + pass + except ValueError as exc: + collector_error.append(exc) + return + + collector = asyncio.create_task(collect()) + await asyncio.sleep(0.05) + crash.set() + await asyncio.wait_for(collector, timeout=2) + + assert collector_error, 'subscriber hung up instead of seeing the crash' + assert ( + task_manager._current_task.status.state + == TaskState.TASK_STATE_FAILED + ) + assert shared.status.state == TaskState.TASK_STATE_WORKING + push_sender.send_notification.assert_awaited() + @pytest.mark.asyncio async def test_active_task_interrupted_auth( self,