feat: use StreamResponse as push notifications payload

Fixes #678

Author: ishymko

src/a2a/server/request_handlers/default_request_handler.py

@@ -21,6 +21,7 @@
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.tasks import (
     PushNotificationConfigStore,
+    PushNotificationEvent,
     PushNotificationSender,
     ResultAggregator,
     TaskManager,
@@ -310,13 +311,15 @@ def _validate_task_id_match(self, task_id: str, event_task_id: str) -> None:
             )
 
     async def _send_push_notification_if_needed(
-        self, task_id: str, result_aggregator: ResultAggregator
+        self, task_id: str, event: Event
     ) -> None:
-        """Sends push notification if configured and task is available."""
-        if self._push_sender and task_id:
-            latest_task = await result_aggregator.current_result
-            if isinstance(latest_task, Task):
-                await self._push_sender.send_notification(latest_task)
+        """Sends push notification if configured."""
+        if (
+            self._push_sender
+            and task_id
+            and isinstance(event, PushNotificationEvent)
+        ):
+            await self._push_sender.send_notification(task_id, event)
 
     async def on_message_send(
         self,
@@ -346,10 +349,8 @@ async def on_message_send(
         interrupted_or_non_blocking = False
         try:
             # Create async callback for push notifications
-            async def push_notification_callback() -> None:
-                await self._send_push_notification_if_needed(
-                    task_id, result_aggregator
-                )
+            async def push_notification_callback(event: Event) -> None:
+                await self._send_push_notification_if_needed(task_id, event)
 
             (
                 result,
@@ -384,8 +385,6 @@ async def push_notification_callback() -> None:
                     result, params.configuration.history_length
                 )
 
-        await self._send_push_notification_if_needed(task_id, result_aggregator)
-
         return result
 
     async def on_message_send_stream(
@@ -413,9 +412,7 @@ async def on_message_send_stream(
                 if isinstance(event, Task):
                     self._validate_task_id_match(task_id, event.id)
 
-                await self._send_push_notification_if_needed(
-                    task_id, result_aggregator
-                )
+                await self._send_push_notification_if_needed(task_id, event)
                 yield event
         except (asyncio.CancelledError, GeneratorExit):
             # Client disconnected: continue consuming and persisting events in the background

src/a2a/server/tasks/__init__.py

@@ -12,7 +12,10 @@
 from a2a.server.tasks.push_notification_config_store import (
     PushNotificationConfigStore,
 )
-from a2a.server.tasks.push_notification_sender import PushNotificationSender
+from a2a.server.tasks.push_notification_sender import (
+    PushNotificationEvent,
+    PushNotificationSender,
+)
 from a2a.server.tasks.result_aggregator import ResultAggregator
 from a2a.server.tasks.task_manager import TaskManager
 from a2a.server.tasks.task_store import TaskStore

src/a2a/server/tasks/base_push_notification_sender.py

@@ -8,8 +8,12 @@
 from a2a.server.tasks.push_notification_config_store import (
     PushNotificationConfigStore,
 )
-from a2a.server.tasks.push_notification_sender import PushNotificationSender
-from a2a.types.a2a_pb2 import PushNotificationConfig, StreamResponse, Task
+from a2a.server.tasks.push_notification_sender import (
+    PushNotificationEvent,
+    PushNotificationSender,
+)
+from a2a.types.a2a_pb2 import PushNotificationConfig
+from a2a.utils.proto_utils import to_stream_response
 
 
 logger = logging.getLogger(__name__)
@@ -32,44 +36,50 @@ def __init__(
         self._client = httpx_client
         self._config_store = config_store
 
-    async def send_notification(self, task: Task) -> None:
-        """Sends a push notification for a task if configuration exists."""
-        push_configs = await self._config_store.get_info(task.id)
+    async def send_notification(
+        self, task_id: str, event: PushNotificationEvent
+    ) -> None:
+        """Sends a push notification for an event if configuration exists."""
+        push_configs = await self._config_store.get_info(task_id)
         if not push_configs:
             return
 
         awaitables = [
-            self._dispatch_notification(task, push_info)
+            self._dispatch_notification(event, push_info, task_id)
             for push_info in push_configs
         ]
         results = await asyncio.gather(*awaitables)
 
         if not all(results):
             logger.warning(
-                'Some push notifications failed to send for task_id=%s', task.id
+                'Some push notifications failed to send for task_id=%s', task_id
             )
 
     async def _dispatch_notification(
-        self, task: Task, push_info: PushNotificationConfig
+        self,
+        event: PushNotificationEvent,
+        push_info: PushNotificationConfig,
+        task_id: str,
     ) -> bool:
         url = push_info.url
         try:
             headers = None
             if push_info.token:
                 headers = {'X-A2A-Notification-Token': push_info.token}
+
             response = await self._client.post(
                 url,
-                json=MessageToDict(StreamResponse(task=task)),
+                json=MessageToDict(to_stream_response(event)),
                 headers=headers,
             )
             response.raise_for_status()
             logger.info(
-                'Push-notification sent for task_id=%s to URL: %s', task.id, url
+                'Push-notification sent for task_id=%s to URL: %s', task_id, url
             )
         except Exception:
             logger.exception(
                 'Error sending push-notification for task_id=%s to URL: %s.',
-                task.id,
+                task_id,
                 url,
             )
             return False

src/a2a/server/tasks/push_notification_sender.py

@@ -1,11 +1,20 @@
 from abc import ABC, abstractmethod
 
-from a2a.types.a2a_pb2 import Task
+from a2a.types.a2a_pb2 import (
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskStatusUpdateEvent,
+)
+
+
+PushNotificationEvent = Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
 
 
 class PushNotificationSender(ABC):
     """Interface for sending push notifications for tasks."""
 
     @abstractmethod
-    async def send_notification(self, task: Task) -> None:
+    async def send_notification(
+        self, task_id: str, event: PushNotificationEvent
+    ) -> None:
         """Sends a push notification containing the latest task state."""

src/a2a/server/tasks/result_aggregator.py

@@ -98,7 +98,7 @@ async def consume_and_break_on_interrupt(
         self,
         consumer: EventConsumer,
         blocking: bool = True,
-        event_callback: Callable[[], Awaitable[None]] | None = None,
+        event_callback: Callable[[Event], Awaitable[None]] | None = None,
     ) -> tuple[Task | Message | None, bool]:
         """Processes the event stream until completion or an interruptible state is encountered.
 
@@ -131,6 +131,9 @@ async def consume_and_break_on_interrupt(
                 return event, False
             await self.task_manager.process(event)
 
+            if event_callback:
+                await event_callback(event)
+
             should_interrupt = False
             is_auth_required = (
                 isinstance(event, Task | TaskStatusUpdateEvent)
@@ -169,7 +172,7 @@ async def consume_and_break_on_interrupt(
     async def _continue_consuming(
         self,
         event_stream: AsyncIterator[Event],
-        event_callback: Callable[[], Awaitable[None]] | None = None,
+        event_callback: Callable[[Event], Awaitable[None]] | None = None,
     ) -> None:
         """Continues processing an event stream in a background task.
 
@@ -183,4 +186,4 @@ async def _continue_consuming(
         async for event in event_stream:
             await self.task_manager.process(event)
             if event_callback:
-                await event_callback()
+                await event_callback(event)

tests/e2e/push_notifications/notifications_app.py

@@ -12,7 +12,7 @@
 class Notification(BaseModel):
     """Encapsulates default push notification data."""
 
-    task: dict[str, Any]
+    event: dict[str, Any]
     token: str
 
 
@@ -36,20 +36,33 @@ async def add_notification(request: Request):
         try:
             json_data = await request.json()
             stream_response = ParseDict(json_data, StreamResponse())
-            if not stream_response.HasField('task'):
-                raise HTTPException(
-                    status_code=400, detail='Missing task in StreamResponse'
-                )
-            task = stream_response.task
+
+            task_id = None
+            if stream_response.HasField('task'):
+                task_id = stream_response.task.id
+            elif stream_response.HasField('status_update'):
+                task_id = stream_response.status_update.task_id
+            elif stream_response.HasField('artifact_update'):
+                task_id = stream_response.artifact_update.task_id
+
+            if not task_id:
+                # Ignore events without task_id (e.g. Message) for now, or log them?
+                # For tests, we just want to ensure we don't 400 on valid updates.
+                # If we return 200 but don't store, the test waiting for n=2 might timeout if n=2 expected stored.
+                # But Message usually accompanies a Task update.
+                return {'status': 'ignored_no_task_id'}
+
         except Exception as e:
             raise HTTPException(status_code=400, detail=str(e))
 
         async with store_lock:
-            if task.id not in store:
-                store[task.id] = []
-            store[task.id].append(
+            if task_id not in store:
+                store[task_id] = []
+            store[task_id].append(
                 Notification(
-                    task=MessageToDict(task, preserving_proto_field_name=True),
+                    event=MessageToDict(
+                        stream_response, preserving_proto_field_name=True
+                    ),
                     token=token,
                 )
             )

tests/e2e/push_notifications/test_default_push_notification_support.py

@@ -139,12 +139,22 @@ async def test_notification_triggering_with_in_message_config_e2e(
     notifications = await wait_for_n_notifications(
         http_client,
         f'{notifications_server}/{task.id}/notifications',
-        n=1,
+        n=2,
     )
     assert notifications[0].token == token
-    # Notification.task is a dict from proto serialization
-    assert notifications[0].task['id'] == task.id
-    assert notifications[0].task['status']['state'] == 'TASK_STATE_COMPLETED'
+
+    # Verify exactly two consecutive events: SUBMITTED -> COMPLETED
+    assert len(notifications) == 2
+
+    # 1. First event: SUBMITTED (Task)
+    event0 = notifications[0].event
+    state0 = event0['task'].get('status', {}).get('state')
+    assert state0 == 'TASK_STATE_SUBMITTED'
+
+    # 2. Second event: COMPLETED (TaskStatusUpdateEvent)
+    event1 = notifications[1].event
+    state1 = event1['status_update'].get('status', {}).get('state')
+    assert state1 == 'TASK_STATE_COMPLETED'
 
 
 @pytest.mark.asyncio
@@ -220,10 +230,12 @@ async def test_notification_triggering_after_config_change_e2e(
         f'{notifications_server}/{task.id}/notifications',
         n=1,
     )
-    # Notification.task is a dict from proto serialization
-    assert notifications[0].task['id'] == task.id
-    assert notifications[0].task['status']['state'] == 'TASK_STATE_COMPLETED'
     assert notifications[0].token == token
+    # It should be a status update or task update
+    event = notifications[0].event
+    state = event['status_update'].get('status', {}).get('state', '')
+
+    assert state == 'TASK_STATE_COMPLETED'
 
 
 async def wait_for_n_notifications(

tests/server/request_handlers/test_default_request_handler.py

@@ -1,6 +1,7 @@
 import asyncio
 import contextlib
 import logging
+import uuid
 import time
 
 from unittest.mock import (
@@ -695,8 +696,11 @@ async def mock_consume_and_break_on_interrupt(
         'event_callback should not be None'
     )
 
+    # Simulate background processing invoking the callback
+    await event_callback_received(final_task)
+
     # Verify that the push notification was sent with the final task
-    mock_push_sender.send_notification.assert_called_with(final_task)
+    mock_push_sender.send_notification.assert_called_with(task_id, final_task)
 
     # Verify that the push notification config was stored
     mock_push_notification_store.set_info.assert_awaited_once_with(
@@ -1412,8 +1416,12 @@ def sync_get_event_stream_gen_for_prop_test(*args, **kwargs):
 
     # 2. send_notification called for each task event yielded by aggregator
     assert mock_push_sender.send_notification.await_count == 2
-    mock_push_sender.send_notification.assert_any_await(event1_task_update)
-    mock_push_sender.send_notification.assert_any_await(event2_final_task)
+    mock_push_sender.send_notification.assert_any_await(
+        task_id, event1_task_update
+    )
+    mock_push_sender.send_notification.assert_any_await(
+        task_id, event2_final_task
+    )
 
     mock_agent_executor.execute.assert_awaited_once()
 

tests/server/tasks/test_inmemory_push_notifications.py

@@ -156,7 +156,7 @@ async def test_send_notification_success(self) -> None:
         mock_response.status_code = 200
         self.mock_httpx_client.post.return_value = mock_response
 
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_awaited_once()
         called_args, called_kwargs = self.mock_httpx_client.post.call_args
@@ -183,7 +183,7 @@ async def test_send_notification_with_token_success(self) -> None:
         mock_response.status_code = 200
         self.mock_httpx_client.post.return_value = mock_response
 
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_awaited_once()
         called_args, called_kwargs = self.mock_httpx_client.post.call_args
@@ -205,7 +205,7 @@ async def test_send_notification_no_config(self) -> None:
         task_id = 'task_send_no_config'
         task_data = create_sample_task(task_id=task_id)
 
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_not_called()
 
@@ -229,7 +229,7 @@ async def test_send_notification_http_status_error(
         self.mock_httpx_client.post.side_effect = http_error
 
         # The method should catch the error and log it, not re-raise
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_awaited_once()
         mock_logger.exception.assert_called_once()
@@ -251,7 +251,7 @@ async def test_send_notification_request_error(
         request_error = httpx.RequestError('Network issue', request=MagicMock())
         self.mock_httpx_client.post.side_effect = request_error
 
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_awaited_once()
         mock_logger.exception.assert_called_once()
@@ -281,7 +281,7 @@ async def test_send_notification_with_auth(
         mock_response.status_code = 200
         self.mock_httpx_client.post.return_value = mock_response
 
-        await self.notifier.send_notification(task_data)  # Pass only task_data
+        await self.notifier.send_notification(task_id, task_data)
 
         self.mock_httpx_client.post.assert_awaited_once()
         called_args, called_kwargs = self.mock_httpx_client.post.call_args