Add push notification support

Author: swapydapy

examples/google_adk/birthday_planner/__main__.py

@@ -5,6 +5,7 @@
 
 import click
 import uvicorn
+
 from adk_agent_executor import ADKAgentExecutor
 from dotenv import load_dotenv
 
@@ -18,6 +19,7 @@
     AgentSkill,
 )
 
+
 load_dotenv()
 
 logging.basicConfig()

examples/google_adk/birthday_planner/adk_agent_executor.py

@@ -1,10 +1,12 @@
 import asyncio
 import logging
-from collections.abc import AsyncGenerator
-from typing import Any, AsyncIterable
+
+from collections.abc import AsyncGenerator, AsyncIterable
+from typing import Any
 from uuid import uuid4
 
 import httpx
+
 from google.adk import Runner
 from google.adk.agents import LlmAgent, RunConfig
 from google.adk.artifacts import InMemoryArtifactService
@@ -42,6 +44,7 @@
 from a2a.utils import get_text_parts
 from a2a.utils.errors import ServerError
 
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -66,7 +69,7 @@ def __init__(self, calendar_agent_url):
             name='birthday_planner_agent',
             description='An agent that helps manage birthday parties.',
             after_tool_callback=self._handle_auth_required_task,
-            instruction=f"""
+            instruction="""
     You are an agent that helps plan birthday parties.
 
     Your job as a party planner is to act as a sounding board and idea generator for
@@ -165,7 +168,7 @@ async def _process_request(
                 task_updater.add_artifact(response)
                 task_updater.complete()
                 break
-            elif calls := event.get_function_calls():
+            if calls := event.get_function_calls():
                 for call in calls:
                     # Provide an update on what we're doing.
                     if call.name == 'message_calendar_agent':
@@ -314,23 +317,21 @@ def convert_a2a_part_to_genai(part: Part) -> types.Part:
     part = part.root
     if isinstance(part, TextPart):
         return types.Part(text=part.text)
-    elif isinstance(part, FilePart):
+    if isinstance(part, FilePart):
         if isinstance(part.file, FileWithUri):
             return types.Part(
                 file_data=types.FileData(
                     file_uri=part.file.uri, mime_type=part.file.mime_type
                 )
             )
-        elif isinstance(part.file, FileWithBytes):
+        if isinstance(part.file, FileWithBytes):
             return types.Part(
                 inline_data=types.Blob(
                     data=part.file.bytes, mime_type=part.file.mime_type
                 )
             )
-        else:
-            raise ValueError(f'Unsupported file type: {type(part.file)}')
-    else:
-        raise ValueError(f'Unsupported part type: {type(part)}')
+        raise ValueError(f'Unsupported file type: {type(part.file)}')
+    raise ValueError(f'Unsupported part type: {type(part)}')
 
 
 def convert_genai_parts_to_a2a(parts: list[types.Part]) -> list[Part]:
@@ -346,14 +347,14 @@ def convert_genai_part_to_a2a(part: types.Part) -> Part:
     """Convert a single Google GenAI Part type into an A2A Part type."""
     if part.text:
         return TextPart(text=part.text)
-    elif part.file_data:
+    if part.file_data:
         return FilePart(
             file=FileWithUri(
                 uri=part.file_data.file_uri,
                 mime_type=part.file_data.mime_type,
             )
         )
-    elif part.inline_data:
+    if part.inline_data:
         return Part(
             root=FilePart(
                 file=FileWithBytes(
@@ -362,5 +363,4 @@ def convert_genai_part_to_a2a(part: types.Part) -> Part:
                 )
             )
         )
-    else:
-        raise ValueError(f'Unsupported part type: {part}')
+    raise ValueError(f'Unsupported part type: {part}')

examples/langgraph/__main__.py

@@ -2,14 +2,15 @@
 import sys
 
 import click
+import httpx
 
 from agent import CurrencyAgent
 from agent_executor import CurrencyAgentExecutor
 from dotenv import load_dotenv
 
 from a2a.server.apps import A2AStarletteApplication
 from a2a.server.request_handlers import DefaultRequestHandler
-from a2a.server.tasks import InMemoryTaskStore
+from a2a.server.tasks import InMemoryPushNotifier, InMemoryTaskStore
 from a2a.types import (
     AgentAuthentication,
     AgentCapabilities,
@@ -29,9 +30,11 @@ def main(host: str, port: int):
         print('GOOGLE_API_KEY environment variable not set.')
         sys.exit(1)
 
+    client = httpx.AsyncClient()
     request_handler = DefaultRequestHandler(
         agent_executor=CurrencyAgentExecutor(),
         task_store=InMemoryTaskStore(),
+        push_notifier=InMemoryPushNotifier(client),
     )
 
     server = A2AStarletteApplication(

src/a2a/server/agent_execution/context.py

@@ -19,7 +19,7 @@ def __init__(
         task_id: str | None = None,
         context_id: str | None = None,
         task: Task | None = None,
-        related_tasks: list[Task] = None,
+        related_tasks: list[Task] | None = None,
     ):
         if related_tasks is None:
             related_tasks = []

src/a2a/server/events/in_memory_queue_manager.py

@@ -18,7 +18,7 @@ class InMemoryQueueManager(QueueManager):
     true scalable deployment.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._task_queue: dict[str, EventQueue] = {}
         self._lock = asyncio.Lock()
 

src/a2a/server/request_handlers/default_request_handler.py

@@ -1,6 +1,6 @@
 import asyncio
-import contextlib
 import logging
+
 from collections.abc import AsyncGenerator
 from typing import cast
 
@@ -10,12 +10,16 @@
     EventConsumer,
     EventQueue,
     InMemoryQueueManager,
-    NoTaskQueue,
     QueueManager,
     TaskQueueExists,
 )
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.server.tasks import ResultAggregator, TaskManager, TaskStore
+from a2a.server.tasks import (
+    PushNotifier,
+    ResultAggregator,
+    TaskManager,
+    TaskStore,
+)
 from a2a.types import (
     InternalError,
     Message,
@@ -29,6 +33,7 @@
 )
 from a2a.utils.errors import ServerError
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -42,10 +47,12 @@ def __init__(
         agent_executor: AgentExecutor,
         task_store: TaskStore,
         queue_manager: QueueManager | None = None,
+        push_notifier: PushNotifier | None = None,
     ) -> None:
         self.agent_executor = agent_executor
         self.task_store = task_store
         self._queue_manager = queue_manager or InMemoryQueueManager()
+        self._push_notifier = push_notifier
         # TODO: Likely want an interface for managing this, like AgentExecutionManager.
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
@@ -116,6 +123,15 @@ async def on_message_send(
         task: Task | None = await task_manager.get_task()
         if task:
             task = task_manager.update_with_message(params.message, task)
+            if (
+                self._push_notifier
+                and params.configuration
+                and params.configuration.pushNotificationConfig
+                and not params.configuration.blocking
+            ):
+                await self._push_notifier.set_info(
+                    task.id, params.configuration.pushNotificationConfig
+                )
         request_context = RequestContext(
             params,
             task.id if task else None,
@@ -173,6 +189,16 @@ async def on_message_send_stream(
         if task:
             task = task_manager.update_with_message(params.message, task)
 
+            if (
+                self._push_notifier
+                and params.configuration
+                and params.configuration.pushNotificationConfig
+            ):
+                await self._push_notifier.set_info(
+                    task.id, params.configuration.pushNotificationConfig
+                )
+        else:
+            queue = EventQueue()
         result_aggregator = ResultAggregator(task_manager)
         request_context = RequestContext(
             params,
@@ -196,12 +222,26 @@ async def on_message_send_stream(
                 # Now we know we have a Task, register the queue
                 if isinstance(event, Task):
                     try:
-                        await self._queue_manager.add(event.id, queue)
-                        task_id = event.id
+                        created_task: Task = event
+                        await self._queue_manager.add(created_task.id, queue)
+                        task_id = created_task.id
                     except TaskQueueExists:
                         logging.info(
                             'Multiple Task objects created in event stream.'
                         )
+                    if (
+                        self._push_notifier
+                        and params.configuration
+                        and params.configuration.pushNotificationConfig
+                    ):
+                        await self._push_notifier.set_info(
+                            created_task.id,
+                            params.configuration.pushNotificationConfig,
+                        )
+                if self._push_notifier and task_id:
+                    latest_task = await result_aggregator.current_result
+                    if isinstance(latest_task, Task):
+                        await self._push_notifier.send_notification(latest_task)
                 yield event
         finally:
             await self._cleanup_producer(producer_task, task_id)
@@ -220,13 +260,38 @@ async def on_set_task_push_notification_config(
         self, params: TaskPushNotificationConfig
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/set'."""
-        raise ServerError(error=UnsupportedOperationError())
+        if not self._push_notifier:
+            raise ServerError(error=UnsupportedOperationError())
+
+        task: Task | None = await self.task_store.get(params.taskId)
+        if not task:
+            raise ServerError(error=TaskNotFoundError())
+
+        await self._push_notifier.set_info(
+            params.taskId,
+            params.pushNotificationConfig,
+        )
+
+        return params
 
     async def on_get_task_push_notification_config(
         self, params: TaskIdParams
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/get'."""
-        raise ServerError(error=UnsupportedOperationError())
+        if not self._push_notifier:
+            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)
+        if not push_notification_config:
+            raise ServerError(error=TaskNotFoundError())
+
+        return TaskPushNotificationConfig(
+            taskId=params.id, pushNotificationConfig=push_notification_config
+        )
 
     async def on_resubscribe_to_task(
         self, params: TaskIdParams

src/a2a/server/tasks/__init__.py

@@ -1,12 +1,16 @@
+from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.server.tasks.push_notifier import PushNotifier
 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',
     'InMemoryTaskStore',
+    'PushNotifier',
     'ResultAggregator',
     'TaskManager',
     'TaskStore',

src/a2a/server/tasks/inmemory_push_notifier.py

@@ -0,0 +1,49 @@
+import asyncio
+import logging
+
+import httpx
+
+from a2a.server.tasks.push_notifier import PushNotifier
+from a2a.types import PushNotificationConfig, Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryPushNotifier(PushNotifier):
+    """In-memory implementation of PushNotifier interface."""
+
+    def __init__(self, httpx_client: httpx.AsyncClient) -> None:
+        self._client = httpx_client
+        self.lock = asyncio.Lock()
+        self._push_notification_infos: dict[str, PushNotificationConfig] = {}
+
+    async def set_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ):
+        async with self.lock:
+            self._push_notification_infos[task_id] = notification_config
+
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        async with self.lock:
+            return self._push_notification_infos.get(task_id)
+
+    async def delete_info(self, task_id: str):
+        async with self.lock:
+            if task_id in self._push_notification_infos:
+                del self._push_notification_infos[task_id]
+
+    async def send_notification(self, task: Task):
+        push_info = await self.get_info(task.id)
+        if not push_info:
+            return
+        url = push_info.url
+
+        try:
+            response = await self._client.post(
+                url, json=task.model_dump(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/push_notifier.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+from a2a.types import PushNotificationConfig, Task
+
+
+class PushNotifier(ABC):
+    """PushNotifier interface to store, retrieve push notification for tasks and send push notifications."""
+
+    @abstractmethod
+    async def set_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ):
+        pass
+
+    @abstractmethod
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        pass
+
+    @abstractmethod
+    async def delete_info(self, task_id: str):
+        pass
+
+    @abstractmethod
+    async def send_notification(self, task: Task):
+        pass