From 69d9c146ef8d4de8c12cb1cc8b129c040e2d36ee Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Fri, 23 Oct 2026 15:37:21 +0000 Subject: [PATCH] fix(server): send push notifications for terminal states the SDK writes Push notifications were only sent for events emitted by the agent. Three terminal states are written by the SDK itself, and none were pushed: - FAILED after an executor raises (EventConsumer._write_failed_status). - CANCELED when a cancel closes out a task the agent did not finish, e.g. one waiting for input (ActiveTask.cancel). - CANCELED written by a replica that is not running the task (DefaultRequestHandlerV2._write_cancel). A push-only client never learned that such a task had ended. Each of these writes now sends a push through one best-effort helper, which the consumer also uses for agent events. --- src/a2a/server/agent_execution/active_task.py | 58 ++++---- .../default_request_handler_v2.py | 2 + tests/server/cluster/conftest.py | 3 + .../cluster/test_cancel_multireplica.py | 45 ++++++ .../test_default_request_handler_v2.py | 130 ++++++++++++++++++ 5 files changed, 214 insertions(+), 24 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 520faa8b9..62b665d6b 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/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 diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index ee32880e7..c96bf8c9e 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/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: diff --git a/tests/server/cluster/conftest.py b/tests/server/cluster/conftest.py index b2e7ddb50..c26d2a7d8 100644 --- a/tests/server/cluster/conftest.py +++ b/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,6 +408,7 @@ 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( @@ -414,6 +416,7 @@ def make_replica( task_store=store, agent_card=streaming_agent_card(), event_stream=stream, + push_sender=push_sender, ) diff --git a/tests/server/cluster/test_cancel_multireplica.py b/tests/server/cluster/test_cancel_multireplica.py index d85167c5b..e3f90ccf3 100644 --- a/tests/server/cluster/test_cancel_multireplica.py +++ b/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() 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 a57be1fa0..b074aae8f 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/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()