feat: Add support for notification tokens in PushNotificationSender

Author: fernando-torres-blip-ai

src/a2a/server/tasks/base_push_notification_sender.py

@@ -9,17 +9,16 @@
 from a2a.server.tasks.push_notification_sender import PushNotificationSender
 from a2a.types import PushNotificationConfig, Task
 
-
 logger = logging.getLogger(__name__)
 
 
 class BasePushNotificationSender(PushNotificationSender):
     """Base implementation of PushNotificationSender interface."""
 
     def __init__(
-        self,
-        httpx_client: httpx.AsyncClient,
-        config_store: PushNotificationConfigStore,
+            self,
+            httpx_client: httpx.AsyncClient,
+            config_store: PushNotificationConfigStore,
     ) -> None:
         """Initializes the BasePushNotificationSender.
 
@@ -48,12 +47,17 @@ async def send_notification(self, task: Task) -> None:
             )
 
     async def _dispatch_notification(
-        self, task: Task, push_info: PushNotificationConfig
+            self, task: Task, push_info: PushNotificationConfig
     ) -> 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=task.model_dump(mode='json', exclude_none=True)
+                url,
+                json=task.model_dump(mode='json', exclude_none=True),
+                headers=headers
             )
             response.raise_for_status()
             logger.info(

tests/server/tasks/test_inmemory_push_notifications.py

@@ -26,9 +26,9 @@ def create_sample_task(task_id='task123', status_state=TaskState.completed):
 
 
 def create_sample_push_config(
-    url='http://example.com/callback', config_id='cfg1'
+    url='http://example.com/callback', config_id='cfg1', token=None
 ):
-    return PushNotificationConfig(id=config_id, url=url)
+    return PushNotificationConfig(id=config_id, url=url, token=token)
 
 
 class TestInMemoryPushNotifier(unittest.IsolatedAsyncioTestCase):
@@ -158,6 +158,35 @@ async def test_send_notification_success(self):
         )  # auth is not passed by current implementation
         mock_response.raise_for_status.assert_called_once()
 
+    async def test_send_notification_with_token_success(self):
+        task_id = 'task_send_success'
+        task_data = create_sample_task(task_id=task_id)
+        config = create_sample_push_config(url='http://notify.me/here', token='unique_token')
+        await self.config_store.set_info(task_id, config)
+
+        # Mock the post call to simulate success
+        mock_response = AsyncMock(spec=httpx.Response)
+        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
+
+        self.mock_httpx_client.post.assert_awaited_once()
+        called_args, called_kwargs = self.mock_httpx_client.post.call_args
+        self.assertEqual(called_args[0], config.url)
+        self.assertEqual(
+            called_kwargs['json'],
+            task_data.model_dump(mode='json', exclude_none=True),
+        )
+        self.assertEqual(
+            called_kwargs['headers'],
+            {"X-A2A-Notification-Token": "unique_token"},
+        )
+        self.assertNotIn(
+            'auth', called_kwargs
+        )  # auth is not passed by current implementation
+        mock_response.raise_for_status.assert_called_once()
+
     async def test_send_notification_no_config(self):
         task_id = 'task_send_no_config'
         task_data = create_sample_task(task_id=task_id)