From f42a68e6d8f74bfb695b3b55cb2e799591f6ad43 Mon Sep 17 00:00:00 2001 From: fernandotorres Date: Tue, 1 Jul 2025 22:22:22 -0300 Subject: [PATCH 1/3] feat: Add support for notification tokens in PushNotificationSender --- .../tasks/base_push_notification_sender.py | 16 +++++---- .../tasks/test_inmemory_push_notifications.py | 33 +++++++++++++++++-- 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 308ed9780..9b8d10171 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -9,7 +9,6 @@ from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.types import PushNotificationConfig, Task - logger = logging.getLogger(__name__) @@ -17,9 +16,9 @@ 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( diff --git a/tests/server/tasks/test_inmemory_push_notifications.py b/tests/server/tasks/test_inmemory_push_notifications.py index 586c78632..9fe4eee79 100644 --- a/tests/server/tasks/test_inmemory_push_notifications.py +++ b/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) From 978adcbe0688e5136e13faa2bafa826756b34d5d Mon Sep 17 00:00:00 2001 From: fernandotorres Date: Tue, 1 Jul 2025 22:30:22 -0300 Subject: [PATCH 2/3] chore: adjust tests and lint --- src/a2a/server/tasks/base_push_notification_sender.py | 9 +++++---- tests/server/request_handlers/test_jsonrpc_handler.py | 3 +++ 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 9b8d10171..5f9efba2e 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -9,6 +9,7 @@ from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.types import PushNotificationConfig, Task + logger = logging.getLogger(__name__) @@ -16,9 +17,9 @@ 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. @@ -47,7 +48,7 @@ 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: diff --git a/tests/server/request_handlers/test_jsonrpc_handler.py b/tests/server/request_handlers/test_jsonrpc_handler.py index de6c1453d..a4400ead2 100644 --- a/tests/server/request_handlers/test_jsonrpc_handler.py +++ b/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) From 020c7bdbc850f0e8203d8cad584159a5f0716d69 Mon Sep 17 00:00:00 2001 From: fernandotorres Date: Tue, 1 Jul 2025 22:41:20 -0300 Subject: [PATCH 3/3] Ensure unit testing and lint --- .../tasks/test_push_notification_sender.py | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/server/tasks/test_push_notification_sender.py b/tests/server/tasks/test_push_notification_sender.py index 4a11680d6..50cfed688 100644 --- a/tests/server/tasks/test_push_notification_sender.py +++ b/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