fix(server): send push notifications for terminal states the SDK writes

src/a2a/server/agent_execution/active_task.py

@@ -109,6 +109,29 @@ def __init__(self, request_id: uuid.UUID):
         self.request_id = request_id
 
 
+async def send_push_notification(
+    push_sender: PushNotificationSender | None,
+    task_id: str,
+    event: PushNotificationEvent,
+) -> None:
+    """Sends a push notification for `event`; failures are logged, not raised."""
+    if push_sender is None or not task_id:
+        return
+    logger.debug('Sending push notification for task %s', task_id)
+    try:
+        await push_sender.send_notification(task_id, event)
+    except Exception:
+        # Push delivery is best-effort: a sender failure must not fail the
+        # task or disturb the event stream. This guard is defense-in-depth on
+        # top of the boundary inside BasePushNotificationSender, and also
+        # protects against custom PushNotificationSender implementations.
+        logger.exception(
+            'Push notification sender raised for task %s; ignoring the '
+            'failure.',
+            task_id,
+        )
+
+
 class EventConsumer:
     """Consumes events from the agent and updates system state."""
 
@@ -188,6 +211,9 @@ async def _write_failed_status(self) -> Task | None:
                 task.status.state = TaskState.TASK_STATE_FAILED
                 task.status.timestamp.FromDatetime(datetime.now(timezone.utc))
                 await task_manager.save_task_event(task)
+                await send_push_notification(
+                    self.active_task._push_sender, task.id, task
+                )
             self.active_task._task_created.set()
             return task
         if not task or task.status.state in TERMINAL_TASK_STATES:
@@ -205,11 +231,14 @@ async def _write_failed_status(self) -> Task | None:
             # 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()
             return None
+        await send_push_notification(
+            self.active_task._push_sender, task.id, handled_event
+        )
+        return await task_manager.get_task()
 
     async def _process_event(self, event: Event) -> None:
         if self.fenced:
@@ -401,30 +430,10 @@ async def _update_task_state(
         if is_terminal:
             await self._handle_terminal_state(updated_task)
 
-        if (
-            self.active_task._push_sender
-            and self.active_task._task_id
-            and isinstance(event, PushNotificationEvent)
-        ):
-            logger.debug(
-                'Consumer[%s]: Sending push notification',
-                self.active_task._task_id,
+        if isinstance(event, PushNotificationEvent):
+            await send_push_notification(
+                self.active_task._push_sender, self.active_task._task_id, event
             )
-            try:
-                await self.active_task._push_sender.send_notification(
-                    self.active_task._task_id, event
-                )
-            except Exception:
-                # Push delivery is best-effort: a sender failure must not
-                # fail the task or disturb the event stream. This guard is
-                # defense-in-depth on top of the boundary inside
-                # BasePushNotificationSender, and also protects against
-                # custom PushNotificationSender implementations.
-                logger.exception(
-                    'Consumer[%s]: Push notification sender raised; '
-                    'ignoring the failure.',
-                    self.active_task._task_id,
-                )
 
     async def _handle_terminal_state(self, updated_task: Task) -> None:
         logger.debug(
@@ -970,6 +979,7 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
             updated.CopyFrom(task)
             updated.status.state = TaskState.TASK_STATE_CANCELED
             await self._task_manager.save_task_event(updated)
+            await send_push_notification(self._push_sender, updated.id, updated)
             task = updated
         return task
 

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -16,6 +16,7 @@
 from a2a.server.agent_execution.active_task import (
     INTERRUPTED_TASK_STATES,
     TERMINAL_TASK_STATES,
+    send_push_notification,
 )
 from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry
 from a2a.server.cluster.event_stream import VersionedEvent
@@ -333,6 +334,7 @@ async def _write_cancel(
             await self._event_stream.publish(
                 task_id, VersionedEvent(event=event, version=new_version)
             )
+        await send_push_notification(self._push_sender, task_id, event)
         return cancelled
 
     def _validate_task_id_match(self, task_id: str, event_task_id: str) -> None:

tests/server/cluster/conftest.py

@@ -24,6 +24,7 @@
     DefaultRequestHandlerV2,
 )
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.server.tasks.push_notification_sender import PushNotificationSender
 from a2a.server.tasks.task_updater import TaskUpdater
 from a2a.types.a2a_pb2 import (
     AgentCapabilities,
@@ -407,13 +408,15 @@ def make_replica(
     store: VersionedTaskStore,
     stream: TaskEventStream,
     agent: AgentExecutor,
+    push_sender: PushNotificationSender | None = None,
 ) -> DefaultRequestHandlerV2:
     """Builds a handler ('replica') wired to a shared store and stream."""
     return DefaultRequestHandlerV2(
         agent_executor=agent,
         task_store=store,
         agent_card=streaming_agent_card(),
         event_stream=stream,
+        push_sender=push_sender,
     )
 
 

tests/server/cluster/test_cancel_multireplica.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from a2a.server.tasks.push_notification_sender import PushNotificationSender
 from a2a.types.a2a_pb2 import (
     CancelTaskRequest,
     Task,
@@ -288,3 +289,47 @@ async def collect() -> None:
     finally:
         await replica_a.aclose()
         await replica_b.aclose()
+
+
+class _RecordingPushSender(PushNotificationSender):
+    def __init__(self) -> None:
+        self.sent: list[tuple[str, int]] = []
+
+    async def send_notification(self, task_id, event) -> None:  # noqa: ANN001
+        self.sent.append((task_id, event.status.state))
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(15)
+async def test_remote_cancel_sends_canceled_push(
+    shared_store, shared_stream, context
+) -> None:
+    """The replica running the agent never pushes a cancel written by another
+    replica, so the cancelling replica does."""
+    agent = LongRunningAgent()
+    replica_a = make_replica(shared_store, shared_stream, agent)
+    push_sender = _RecordingPushSender()
+    replica_b = make_replica(shared_store, shared_stream, agent, push_sender)
+    try:
+        a_task = asyncio.create_task(
+            drain(
+                replica_a.on_message_send_stream(
+                    build_send_request('go', message_id='m1'), context
+                )
+            )
+        )
+        await asyncio.wait_for(agent.working.wait(), timeout=5)
+        task_id = next(
+            iter(replica_a._active_task_registry._active_tasks)  # noqa: SLF001
+        )
+        await wait_for_state(
+            shared_store, task_id, TaskState.TASK_STATE_WORKING, context
+        )
+
+        await replica_b.on_cancel_task(CancelTaskRequest(id=task_id), context)
+
+        assert push_sender.sent == [(task_id, TaskState.TASK_STATE_CANCELED)]
+        await asyncio.wait_for(a_task, timeout=8)
+    finally:
+        await replica_a.aclose()
+        await replica_b.aclose()

tests/server/request_handlers/test_default_request_handler_v2.py

@@ -3304,3 +3304,133 @@ async def test_failed_cancel_restarts_idle_release() -> None:
 
     assert await _wait_for_release(handler, task.id)
     await handler.aclose()
+
+
+class _RecordingPushSender(PushNotificationSender):
+    def __init__(self) -> None:
+        self.sent: list[tuple[str, int]] = []
+
+    async def send_notification(self, task_id, event) -> None:
+        self.sent.append((task_id, event.status.state))
+
+
+class _FailingAgent(AgentExecutor):
+    """Raises before or after creating its task."""
+
+    def __init__(self, create_task_first: bool) -> None:
+        self.create_task_first = create_task_first
+
+    async def execute(self, context: RequestContext, event_queue: EventQueue):
+        if self.create_task_first:
+            assert context.message is not None
+            await event_queue.enqueue_event(
+                new_task_from_user_message(context.message)
+            )
+            updater = TaskUpdater(
+                event_queue, str(context.task_id), str(context.context_id)
+            )
+            await updater.start_work()
+        raise RuntimeError('agent failed')
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue):
+        pass
+
+
+class _InputRequiredAgent(AgentExecutor):
+    async def execute(self, context: RequestContext, event_queue: EventQueue):
+        assert context.message is not None
+        await event_queue.enqueue_event(
+            new_task_from_user_message(context.message)
+        )
+        updater = TaskUpdater(
+            event_queue, str(context.task_id), str(context.context_id)
+        )
+        await updater.requires_input()
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue):
+        pass
+
+
+def _send_request(message_id: str) -> SendMessageRequest:
+    return SendMessageRequest(
+        message=Message(
+            role=Role.ROLE_USER, message_id=message_id, parts=[Part(text='hi')]
+        )
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('create_task_first', [True, False])
+async def test_executor_failure_sends_failed_push(create_task_first: bool):
+    """The FAILED state written after an executor raises is pushed too."""
+    push_sender = _RecordingPushSender()
+    task_store = InMemoryTaskStore()
+    handler = DefaultRequestHandlerV2(
+        agent_executor=_FailingAgent(create_task_first),
+        task_store=task_store,
+        push_sender=push_sender,
+        agent_card=create_default_agent_card(),
+    )
+    context = create_server_call_context()
+
+    with pytest.raises(RuntimeError, match='agent failed'):
+        await handler.on_message_send(_send_request('m1'), context)
+
+    task_id, last_state = push_sender.sent[-1]
+    assert last_state == TaskState.TASK_STATE_FAILED
+    stored = await task_store.get(task_id, context)
+    assert stored is not None
+    assert stored.status.state == TaskState.TASK_STATE_FAILED
+    await handler.aclose()
+
+
+@pytest.mark.asyncio
+async def test_cancel_of_waiting_task_sends_canceled_push():
+    """Cancelling a task that waits for input closes it out as CANCELED
+    without the agent writing anything; that write is pushed too."""
+    push_sender = _RecordingPushSender()
+    handler = DefaultRequestHandlerV2(
+        agent_executor=_InputRequiredAgent(),
+        task_store=InMemoryTaskStore(),
+        push_sender=push_sender,
+        agent_card=create_default_agent_card(),
+    )
+    context = create_server_call_context()
+    task = await handler.on_message_send(_send_request('m1'), context)
+    assert task.status.state == TaskState.TASK_STATE_INPUT_REQUIRED
+
+    cancelled = await handler.on_cancel_task(
+        CancelTaskRequest(id=task.id), context
+    )
+
+    assert cancelled.status.state == TaskState.TASK_STATE_CANCELED
+    assert push_sender.sent[-1] == (task.id, TaskState.TASK_STATE_CANCELED)
+    await handler.aclose()
+
+
+@pytest.mark.asyncio
+async def test_raising_push_sender_does_not_break_sdk_terminal_writes():
+    task_store = InMemoryTaskStore()
+    context = create_server_call_context()
+    failing = DefaultRequestHandlerV2(
+        agent_executor=_FailingAgent(create_task_first=True),
+        task_store=task_store,
+        push_sender=RaisingPushSender(),
+        agent_card=create_default_agent_card(),
+    )
+    with pytest.raises(RuntimeError, match='agent failed'):
+        await failing.on_message_send(_send_request('m1'), context)
+    await failing.aclose()
+
+    waiting = DefaultRequestHandlerV2(
+        agent_executor=_InputRequiredAgent(),
+        task_store=task_store,
+        push_sender=RaisingPushSender(),
+        agent_card=create_default_agent_card(),
+    )
+    task = await waiting.on_message_send(_send_request('m2'), context)
+    cancelled = await waiting.on_cancel_task(
+        CancelTaskRequest(id=task.id), context
+    )
+    assert cancelled.status.state == TaskState.TASK_STATE_CANCELED
+    await waiting.aclose()