fix(server): wrap task in StreamResponse for push notifications

Signed-off-by: Luca Muscariello <muscariello@ieee.org>

Author: muscariello

src/a2a/server/tasks/base_push_notification_sender.py

@@ -9,7 +9,7 @@
     PushNotificationConfigStore,
 )
 from a2a.server.tasks.push_notification_sender import PushNotificationSender
-from a2a.types.a2a_pb2 import PushNotificationConfig, Task
+from a2a.types.a2a_pb2 import PushNotificationConfig, StreamResponse, Task
 
 
 logger = logging.getLogger(__name__)
@@ -59,7 +59,7 @@ async def _dispatch_notification(
                 headers = {'X-A2A-Notification-Token': push_info.token}
             response = await self._client.post(
                 url,
-                json=MessageToDict(task),
+                json=MessageToDict(StreamResponse(task=task)),
                 headers=headers,
             )
             response.raise_for_status()

tests/e2e/push_notifications/notifications_app.py

@@ -5,7 +5,7 @@
 from fastapi import FastAPI, HTTPException, Path, Request
 from pydantic import BaseModel, ConfigDict, ValidationError
 
-from a2a.types.a2a_pb2 import Task
+from a2a.types.a2a_pb2 import StreamResponse, Task
 from google.protobuf.json_format import ParseDict, MessageToDict
 
 
@@ -35,7 +35,12 @@ async def add_notification(request: Request):
             )
         try:
             json_data = await request.json()
-            task = ParseDict(json_data, Task())
+            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
         except Exception as e:
             raise HTTPException(status_code=400, detail=str(e))
 

tests/server/tasks/test_inmemory_push_notifications.py

@@ -13,6 +13,7 @@
 )
 from a2a.types.a2a_pb2 import (
     PushNotificationConfig,
+    StreamResponse,
     Task,
     TaskState,
     TaskStatus,
@@ -162,7 +163,7 @@ async def test_send_notification_success(self) -> None:
         self.assertEqual(called_args[0], config.url)
         self.assertEqual(
             called_kwargs['json'],
-            MessageToDict(task_data),
+            MessageToDict(StreamResponse(task=task_data)),
         )
         self.assertNotIn(
             'auth', called_kwargs
@@ -189,7 +190,7 @@ async def test_send_notification_with_token_success(self) -> None:
         self.assertEqual(called_args[0], config.url)
         self.assertEqual(
             called_kwargs['json'],
-            MessageToDict(task_data),
+            MessageToDict(StreamResponse(task=task_data)),
         )
         self.assertEqual(
             called_kwargs['headers'],
@@ -287,7 +288,7 @@ async def test_send_notification_with_auth(
         self.assertEqual(called_args[0], config.url)
         self.assertEqual(
             called_kwargs['json'],
-            MessageToDict(task_data),
+            MessageToDict(StreamResponse(task=task_data)),
         )
         self.assertNotIn(
             'auth', called_kwargs

tests/server/tasks/test_push_notification_sender.py

@@ -10,6 +10,7 @@
 )
 from a2a.types.a2a_pb2 import (
     PushNotificationConfig,
+    StreamResponse,
     Task,
     TaskState,
     TaskStatus,
@@ -65,7 +66,7 @@ async def test_send_notification_success(self) -> None:
         # assert httpx_client post method got invoked with right parameters
         self.mock_httpx_client.post.assert_awaited_once_with(
             config.url,
-            json=MessageToDict(task_data),
+            json=MessageToDict(StreamResponse(task=task_data)),
             headers=None,
         )
         mock_response.raise_for_status.assert_called_once()
@@ -89,7 +90,7 @@ async def test_send_notification_with_token_success(self) -> None:
         # assert httpx_client post method got invoked with right parameters
         self.mock_httpx_client.post.assert_awaited_once_with(
             config.url,
-            json=MessageToDict(task_data),
+            json=MessageToDict(StreamResponse(task=task_data)),
             headers={'X-A2A-Notification-Token': 'unique_token'},
         )
         mock_response.raise_for_status.assert_called_once()
@@ -126,7 +127,7 @@ async def test_send_notification_http_status_error(
         self.mock_config_store.get_info.assert_awaited_once_with(task_id)
         self.mock_httpx_client.post.assert_awaited_once_with(
             config.url,
-            json=MessageToDict(task_data),
+            json=MessageToDict(StreamResponse(task=task_data)),
             headers=None,
         )
         mock_logger.exception.assert_called_once()
@@ -154,13 +155,13 @@ async def test_send_notification_multiple_configs(self) -> None:
         # Check calls for config1
         self.mock_httpx_client.post.assert_any_call(
             config1.url,
-            json=MessageToDict(task_data),
+            json=MessageToDict(StreamResponse(task=task_data)),
             headers=None,
         )
         # Check calls for config2
         self.mock_httpx_client.post.assert_any_call(
             config2.url,
-            json=MessageToDict(task_data),
+            json=MessageToDict(StreamResponse(task=task_data)),
             headers=None,
         )
         mock_response.raise_for_status.call_count = 2