fix: publish terminal task state to subscriber streams (#1175)

Cancel and producer-failure wrote CANCELED/FAILED to the store after
queues closed, so a live stream saw WORKING then hang-up. Persist a
copy and notify observers before teardown; cancel also emits the
status event so subscribers do not have to poll.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: anthonyrs06

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,13 +754,42 @@ 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',
                     self._task_id,
                     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

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,