Merge branch 'main' into improve-error-handling

Author: holtskinner

src/a2a/server/request_handlers/default_request_handler.py

@@ -33,9 +33,7 @@
     InvalidParamsError,
     ListTaskPushNotificationConfigParams,
     Message,
-    MessageSendConfiguration,
     MessageSendParams,
-    PushNotificationConfig,
     Task,
     TaskIdParams,
     TaskNotFoundError,
@@ -202,18 +200,6 @@ async def _setup_message_execution(
                 )
 
             task = task_manager.update_with_message(params.message, task)
-            if self.should_add_push_info(params):
-                assert self._push_config_store is not None
-                assert isinstance(
-                    params.configuration, MessageSendConfiguration
-                )
-                assert isinstance(
-                    params.configuration.pushNotificationConfig,
-                    PushNotificationConfig,
-                )
-                await self._push_config_store.set_info(
-                    task.id, params.configuration.pushNotificationConfig
-                )
 
         # Build request context
         request_context = await self._request_context_builder.build(
@@ -228,6 +214,16 @@ async def _setup_message_execution(
         # Always assign a task ID. We may not actually upgrade to a task, but
         # dictating the task ID at this layer is useful for tracking running
         # agents.
+
+        if (
+            self._push_config_store
+            and params.configuration
+            and params.configuration.pushNotificationConfig
+        ):
+            await self._push_config_store.set_info(
+                task_id, params.configuration.pushNotificationConfig
+            )
+
         queue = await self._queue_manager.create_or_tap(task_id)
         result_aggregator = ResultAggregator(task_manager)
         # TODO: to manage the non-blocking flows.
@@ -333,16 +329,6 @@ async def on_message_send_stream(
                 if isinstance(event, Task):
                     self._validate_task_id_match(task_id, event.id)
 
-                if (
-                    self._push_config_store
-                    and params.configuration
-                    and params.configuration.pushNotificationConfig
-                ):
-                    await self._push_config_store.set_info(
-                        task_id,
-                        params.configuration.pushNotificationConfig,
-                    )
-
                 await self._send_push_notification_if_needed(
                     task_id, result_aggregator
                 )
@@ -509,11 +495,3 @@ async def on_delete_task_push_notification_config(
         await self._push_config_store.delete_info(
             params.id, params.pushNotificationConfigId
         )
-
-    def should_add_push_info(self, params: MessageSendParams) -> bool:
-        """Determines if push notification info should be set for a task."""
-        return bool(
-            self._push_config_store
-            and params.configuration
-            and params.configuration.pushNotificationConfig
-        )

src/a2a/server/tasks/__init__.py

@@ -1,9 +1,10 @@
 """Components for managing tasks within the A2A server."""
 
+import logging
+
 from a2a.server.tasks.base_push_notification_sender import (
     BasePushNotificationSender,
 )
-from a2a.server.tasks.database_task_store import DatabaseTaskStore
 from a2a.server.tasks.inmemory_push_notification_config_store import (
     InMemoryPushNotificationConfigStore,
 )
@@ -18,6 +19,30 @@
 from a2a.server.tasks.task_updater import TaskUpdater
 
 
+logger = logging.getLogger(__name__)
+
+try:
+    from a2a.server.tasks.database_task_store import (
+        DatabaseTaskStore,  # type: ignore
+    )
+except ImportError as e:
+    _original_error = e
+    # If the database task store is not available, we can still use in-memory stores.
+    logger.debug(
+        'DatabaseTaskStore not loaded. This is expected if database dependencies are not installed. Error: %s',
+        e,
+    )
+
+    class DatabaseTaskStore:  # type: ignore
+        """Placeholder for DatabaseTaskStore when dependencies are not installed."""
+
+        def __init__(self, *args, **kwargs):
+            raise ImportError(
+                'To use DatabaseTaskStore, its dependencies must be installed. '
+                'You can install them with \'pip install "a2a-sdk[sql]"\''
+            ) from _original_error
+
+
 __all__ = [
     'BasePushNotificationSender',
     'DatabaseTaskStore',

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_default_request_handler.py

@@ -401,6 +401,90 @@ async def get_current_result():
     mock_agent_executor.execute.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_on_message_send_with_push_notification_no_existing_Task():
+    """Test on_message_send for new task sets push notification info if provided."""
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_push_notification_store = AsyncMock(spec=PushNotificationConfigStore)
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+    mock_request_context_builder = AsyncMock(spec=RequestContextBuilder)
+
+    task_id = 'push_task_1'
+    context_id = 'push_ctx_1'
+
+    mock_task_store.get.return_value = (
+        None  # Simulate new task scenario for TaskManager
+    )
+
+    # Mock _request_context_builder.build to return a context with the generated/confirmed IDs
+    mock_request_context = MagicMock(spec=RequestContext)
+    mock_request_context.task_id = task_id
+    mock_request_context.context_id = context_id
+    mock_request_context_builder.build.return_value = mock_request_context
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        push_config_store=mock_push_notification_store,
+        request_context_builder=mock_request_context_builder,
+    )
+
+    push_config = PushNotificationConfig(url='http://callback.com/push')
+    message_config = MessageSendConfiguration(
+        pushNotificationConfig=push_config,
+        acceptedOutputModes=['text/plain'],  # Added required field
+    )
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            messageId='msg_push',
+            parts=[],
+            taskId=task_id,
+            contextId=context_id,
+        ),
+        configuration=message_config,
+    )
+
+    # Mock ResultAggregator and its consume_and_break_on_interrupt
+    mock_result_aggregator_instance = AsyncMock(spec=ResultAggregator)
+    final_task_result = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.completed
+    )
+    mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
+        final_task_result,
+        False,
+    )
+
+    # Mock the current_result property to return the final task result
+    async def get_current_result():
+        return final_task_result
+
+    # Configure the 'current_result' property on the type of the mock instance
+    type(mock_result_aggregator_instance).current_result = PropertyMock(
+        return_value=get_current_result()
+    )
+
+    with (
+        patch(
+            'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+            return_value=mock_result_aggregator_instance,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.get_task',
+            return_value=None,
+        ),
+    ):
+        await request_handler.on_message_send(
+            params, create_server_call_context()
+        )
+
+    mock_push_notification_store.set_info.assert_awaited_once_with(
+        task_id, push_config
+    )
+    # Other assertions for full flow if needed (e.g., agent execution)
+    mock_agent_executor.execute.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_on_message_send_no_result_from_aggregator():
     """Test on_message_send when aggregator returns (None, False)."""

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