feat: Push notification changes

Author: kthota-g

src/a2a/server/request_handlers/default_request_handler.py

@@ -1,6 +1,5 @@
 import asyncio
 import logging
-import uuid
 
 from collections.abc import AsyncGenerator
 from typing import cast
@@ -21,13 +20,13 @@
 )
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.tasks import (
-    PushNotifier,
     ResultAggregator,
+    PushNotificationConfigStore,
+    PushNotificationSender,
     TaskManager,
     TaskStore,
 )
 from a2a.types import (
-    GetTaskPushNotificationConfigParams,
     InternalError,
     Message,
     MessageSendConfiguration,
@@ -39,6 +38,7 @@
     TaskPushNotificationConfig,
     TaskQueryParams,
     UnsupportedOperationError,
+    GetTaskPushNotificationConfigParams,
 )
 from a2a.utils.errors import ServerError
 from a2a.utils.telemetry import SpanKind, trace_class
@@ -63,7 +63,8 @@ def __init__(
         agent_executor: AgentExecutor,
         task_store: TaskStore,
         queue_manager: QueueManager | None = None,
-        push_notifier: PushNotifier | None = None,
+        push_config_store: PushNotificationConfigStore | None = None,
+        push_sender: PushNotificationSender | None = None,
         request_context_builder: RequestContextBuilder | None = None,
     ) -> None:
         """Initializes the DefaultRequestHandler.
@@ -72,14 +73,16 @@ def __init__(
             agent_executor: The `AgentExecutor` instance to run agent logic.
             task_store: The `TaskStore` instance to manage task persistence.
             queue_manager: The `QueueManager` instance to manage event queues. Defaults to `InMemoryQueueManager`.
-            push_notifier: The `PushNotifier` instance for sending push notifications. Defaults to None.
+            push_config_store: The `PushNotificationConfigStore` instance for managing push notification configurations. Defaults to None.
+            push_sender: The `PushNotificationSender` instance for sending push notifications. Defaults to None.
             request_context_builder: The `RequestContextBuilder` instance used
               to build request contexts. Defaults to `SimpleRequestContextBuilder`.
         """
         self.agent_executor = agent_executor
         self.task_store = task_store
         self._queue_manager = queue_manager or InMemoryQueueManager()
-        self._push_notifier = push_notifier
+        self._push_config_store = push_config_store
+        self._push_sender = push_sender
         self._request_context_builder = (
             request_context_builder
             or SimpleRequestContextBuilder(
@@ -180,15 +183,15 @@ async def on_message_send(
         if task:
             task = task_manager.update_with_message(params.message, task)
             if self.should_add_push_info(params):
-                assert isinstance(self._push_notifier, PushNotifier)
+                assert self._push_config_store is not None
                 assert isinstance(
                     params.configuration, MessageSendConfiguration
                 )
                 assert isinstance(
                     params.configuration.pushNotificationConfig,
                     PushNotificationConfig,
                 )
-                await self._push_notifier.set_info(
+                await self._push_config_store.set_info(
                     task.id, params.configuration.pushNotificationConfig
                 )
         request_context = await self._request_context_builder.build(
@@ -199,7 +202,7 @@ async def on_message_send(
             context=context,
         )
 
-        task_id = cast('str', request_context.task_id)
+        task_id = cast(str, request_context.task_id)
         # 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.
@@ -237,7 +240,7 @@ async def on_message_send(
         finally:
             if interrupted:
                 # TODO: Track this disconnected cleanup task.
-                asyncio.create_task(  # noqa: RUF006
+                asyncio.create_task(
                     self._cleanup_producer(producer_task, task_id)
                 )
             else:
@@ -267,15 +270,15 @@ async def on_message_send_stream(
             task = task_manager.update_with_message(params.message, task)
 
             if self.should_add_push_info(params):
-                assert isinstance(self._push_notifier, PushNotifier)
+                assert self._push_config_store is not None
                 assert isinstance(
                     params.configuration, MessageSendConfiguration
                 )
                 assert isinstance(
                     params.configuration.pushNotificationConfig,
                     PushNotificationConfig,
                 )
-                await self._push_notifier.set_info(
+                await self._push_config_store.set_info(
                     task.id, params.configuration.pushNotificationConfig
                 )
         else:
@@ -289,7 +292,7 @@ async def on_message_send_stream(
             context=context,
         )
 
-        task_id = cast('str', request_context.task_id)
+        task_id = cast(str, request_context.task_id)
         queue = await self._queue_manager.create_or_tap(task_id)
         producer_task = asyncio.create_task(
             self._run_event_stream(
@@ -315,19 +318,19 @@ async def on_message_send_stream(
                         )
 
                     if (
-                        self._push_notifier
+                        self._push_config_store # Check if store is available for config
                         and params.configuration
                         and params.configuration.pushNotificationConfig
                     ):
-                        await self._push_notifier.set_info(
+                        await self._push_config_store.set_info(
                             task_id,
                             params.configuration.pushNotificationConfig,
                         )
 
-                if self._push_notifier and task_id:
+                if self._push_sender and task_id: # Check if sender is available
                     latest_task = await result_aggregator.current_result
                     if isinstance(latest_task, Task):
-                        await self._push_notifier.send_notification(latest_task)
+                        await self._push_sender.send_notification(latest_task)
                 yield event
         finally:
             await self._cleanup_producer(producer_task, task_id)
@@ -359,16 +362,14 @@ async def on_set_task_push_notification_config(
 
         Requires a `PushNotifier` to be configured.
         """
-        if not self._push_notifier:
+        if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
         task: Task | None = await self.task_store.get(params.taskId)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
-        # Generate a unique id for the notification
-        params.pushNotificationConfig.id = str(uuid.uuid4())
-        await self._push_notifier.set_info(
+        await self._push_config_store.set_info(
             params.taskId,
             params.pushNotificationConfig,
         )
@@ -384,14 +385,14 @@ async def on_get_task_push_notification_config(
 
         Requires a `PushNotifier` to be configured.
         """
-        if not self._push_notifier:
+        if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
         task: Task | None = await self.task_store.get(params.id)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
-        push_notification_config = await self._push_notifier.get_info(params.id)
+        push_notification_config = await self._push_config_store.get_info(params.id)
         if not push_notification_config:
             raise ServerError(error=InternalError())
 
@@ -431,9 +432,8 @@ async def on_resubscribe_to_task(
             yield event
 
     def should_add_push_info(self, params: MessageSendParams) -> bool:
-        """Determines if push notification info should be set for a task."""
         return bool(
-            self._push_notifier
+            self._push_config_store
             and params.configuration
             and params.configuration.pushNotificationConfig
         )

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -34,6 +34,7 @@
     TaskPushNotificationConfig,
     TaskResubscriptionRequest,
     TaskStatusUpdateEvent,
+    GetTaskPushNotificationConfigParams
 )
 from a2a.utils.errors import ServerError
 from a2a.utils.helpers import validate

src/a2a/server/tasks/__init__.py

@@ -1,18 +1,21 @@
 """Components for managing tasks within the A2A server."""
 
-from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
+from a2a.server.tasks.base_push_notification_sender import BasePushNotificationSender
+from a2a.server.tasks.inmemory_push_notification_config_store import InMemoryPushNotificationConfigStore
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
-from a2a.server.tasks.push_notifier import PushNotifier
+from a2a.server.tasks.push_notification_config_store import PushNotificationConfigStore
+from a2a.server.tasks.push_notification_sender import PushNotificationSender
 from a2a.server.tasks.result_aggregator import ResultAggregator
 from a2a.server.tasks.task_manager import TaskManager
 from a2a.server.tasks.task_store import TaskStore
 from a2a.server.tasks.task_updater import TaskUpdater
 
-
 __all__ = [
-    'InMemoryPushNotifier',
+    'BasePushNotificationSender',
+    'InMemoryPushNotificationConfigStore',
     'InMemoryTaskStore',
-    'PushNotifier',
+    'PushNotificationConfigStore',
+    'PushNotificationSender',
     'ResultAggregator',
     'TaskManager',
     'TaskStore',

src/a2a/server/tasks/base_push_notification_sender.py

@@ -0,0 +1,40 @@
+import logging
+
+import httpx
+
+from a2a.server.tasks.push_notification_config_store import PushNotificationConfigStore
+from a2a.server.tasks.push_notification_sender import PushNotificationSender
+from a2a.types import Task
+
+logger = logging.getLogger(__name__)
+
+
+class BasePushNotificationSender(PushNotificationSender):
+    """Base implementation of PushNotificationSender interface.
+    """
+
+    def __init__(self, httpx_client: httpx.AsyncClient, config_store: PushNotificationConfigStore) -> None:
+        """Initializes the BasePushNotificationSender.
+
+        Args:
+            httpx_client: An async HTTP client instance to send notifications.
+            config_store: A PushNotificationConfigStore instance to retrieve configurations.
+        """
+        self._client = httpx_client
+        self._config_store = config_store
+
+    async def send_notification(self, task: Task) -> None:
+        """Sends a push notification for a task if configuration exists."""
+        push_info = await self._config_store.get_info(task.id)
+        if not push_info:
+            return
+        url = push_info.url
+
+        try:
+            response = await self._client.post(
+                url, json=task.model_dump(mode='json', exclude_none=True)
+            )
+            response.raise_for_status()
+            logger.info(f'Push-notification sent for URL: {url}')
+        except Exception as e:
+            logger.error(f'Error sending push-notification: {e}')

src/a2a/server/tasks/inmemory_push_notification_config_store.py

@@ -0,0 +1,43 @@
+import asyncio
+import logging
+
+from a2a.server.tasks.push_notification_config_store import PushNotificationConfigStore
+from a2a.types import PushNotificationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryPushNotificationConfigStore(PushNotificationConfigStore):
+    """In-memory implementation of PushNotificationConfigStore interface.
+
+    Stores push notification configurations in memory and uses an httpx client
+    to send notifications.
+    """
+    def __init__(self) -> None:
+        """Initializes the InMemoryPushNotifier.
+
+        Args:
+            httpx_client: An async HTTP client instance to send notifications.
+        """
+        self.lock = asyncio.Lock()
+        self._push_notification_infos: dict[str, PushNotificationConfig] = {}
+
+    async def set_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ):
+        """Sets or updates the push notification configuration for a task in memory."""
+        async with self.lock:
+            self._push_notification_infos[task_id] = notification_config
+
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        """Retrieves the push notification configuration for a task from memory."""
+        async with self.lock:
+            return self._push_notification_infos.get(task_id)
+
+
+
+    async def delete_info(self, task_id: str):
+        """Deletes the push notification configuration for a task from memory."""
+        async with self.lock:
+            if task_id in self._push_notification_infos:
+                del self._push_notification_infos[task_id]

src/a2a/server/tasks/push_notification_config_store.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from a2a.types import PushNotificationConfig
+
+
+class PushNotificationConfigStore(ABC):
+    """Interface for storing and retrieving push notification configurations for tasks."""
+
+    @abstractmethod
+    async def set_info(self, task_id: str, notification_config: PushNotificationConfig):
+        """Sets or updates the push notification configuration for a task."""
+
+    @abstractmethod
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        """Retrieves the push notification configuration for a task."""
+
+    @abstractmethod
+    async def delete_info(self, task_id: str):
+        """Deletes the push notification configuration for a task."""

src/a2a/server/tasks/push_notification_sender.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+from a2a.types import Task
+
+
+class PushNotificationSender(ABC):
+    """Interface for sending push notifications for tasks."""
+
+    @abstractmethod
+    async def send_notification(self, task: Task) -> None:
+        """Sends a push notification containing the latest task state."""