feat: add support for notification tokens in PushNotificationSender (#266)

# Description

Thank you for opening a Pull Request!
Before submitting your PR, there are a few things you can do to make
sure it goes smoothly:

- [x] Follow the [`CONTRIBUTING`
Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md).
- [x] Make your Pull Request title in the
<https://www.conventionalcommits.org/> specification.
- Important Prefixes for
[release-please](https://github.com/googleapis/release-please):
- `fix:` which represents bug fixes, and correlates to a
[SemVer](https://semver.org/) patch.
- `feat:` represents a new feature, and correlates to a SemVer minor.
- `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking
change (indicated by the `!`) and will result in a SemVer major.
- [x] Ensure the tests and linter pass (Run `nox -s format` from the
repository root to format)
- [ ] Appropriate docs were updated (if necessary)

Fixes #284

---------

Co-authored-by: kthota-g <kcthota@google.com>

Author: fernando-torres-blip-ai

src/a2a/server/tasks/base_push_notification_sender.py

@@ -52,8 +52,13 @@ async def _dispatch_notification(
     ) -> 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/request_handlers/test_jsonrpc_handler.py

@@ -585,6 +585,7 @@ async def streaming_coro():
                         'kind': 'task',
                         'status': {'state': 'submitted'},
                     },
+                    headers=None
                 ),
                 call(
                     'http://example.com',
@@ -605,6 +606,7 @@ async def streaming_coro():
                         'kind': 'task',
                         'status': {'state': 'submitted'},
                     },
+                    headers=None
                 ),
                 call(
                     'http://example.com',
@@ -625,6 +627,7 @@ async def streaming_coro():
                         'kind': 'task',
                         'status': {'state': 'completed'},
                     },
+                    headers=None
                 ),
             ]
             mock_httpx_client.post.assert_has_calls(calls)

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)

tests/server/tasks/test_push_notification_sender.py

@@ -25,9 +25,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 TestBasePushNotificationSender(unittest.IsolatedAsyncioTestCase):
@@ -61,6 +61,29 @@ async def test_send_notification_success(self):
         self.mock_httpx_client.post.assert_awaited_once_with(
             config.url,
             json=task_data.model_dump(mode='json', exclude_none=True),
+            headers=None
+        )
+        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')
+        self.mock_config_store.get_info.return_value = [config]
+
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        self.mock_httpx_client.post.return_value = mock_response
+
+        await self.sender.send_notification(task_data)
+
+        self.mock_config_store.get_info.assert_awaited_once_with
+
+        # assert httpx_client post method got invoked with right parameters
+        self.mock_httpx_client.post.assert_awaited_once_with(
+            config.url,
+            json=task_data.model_dump(mode='json', exclude_none=True),
+            headers={'X-A2A-Notification-Token': 'unique_token'}
         )
         mock_response.raise_for_status.assert_called_once()
 
@@ -97,6 +120,7 @@ async def test_send_notification_http_status_error(
         self.mock_httpx_client.post.assert_awaited_once_with(
             config.url,
             json=task_data.model_dump(mode='json', exclude_none=True),
+            headers=None
         )
         mock_logger.error.assert_called_once()
 
@@ -124,10 +148,12 @@ async def test_send_notification_multiple_configs(self):
         self.mock_httpx_client.post.assert_any_call(
             config1.url,
             json=task_data.model_dump(mode='json', exclude_none=True),
+            headers=None
         )
         # Check calls for config2
         self.mock_httpx_client.post.assert_any_call(
             config2.url,
             json=task_data.model_dump(mode='json', exclude_none=True),
+            headers=None
         )
         mock_response.raise_for_status.call_count = 2