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()