11import asyncio
22import logging
3- import uuid
43
54from collections .abc import AsyncGenerator
65from typing import cast
2120)
2221from a2a .server .request_handlers .request_handler import RequestHandler
2322from a2a .server .tasks import (
24- PushNotifier ,
23+ PushNotificationConfigStore ,
24+ PushNotificationSender ,
2525 ResultAggregator ,
2626 TaskManager ,
2727 TaskStore ,
2828)
2929from a2a .types import (
30+ DeleteTaskPushNotificationConfigParams ,
3031 GetTaskPushNotificationConfigParams ,
3132 InternalError ,
3233 InvalidParamsError ,
34+ ListTaskPushNotificationConfigParams ,
3335 Message ,
3436 MessageSendConfiguration ,
3537 MessageSendParams ,
@@ -67,12 +69,13 @@ class DefaultRequestHandler(RequestHandler):
6769
6870 _running_agents : dict [str , asyncio .Task ]
6971
70- def __init__ (
72+ def __init__ ( # noqa: PLR0913
7173 self ,
7274 agent_executor : AgentExecutor ,
7375 task_store : TaskStore ,
7476 queue_manager : QueueManager | None = None ,
75- push_notifier : PushNotifier | None = None ,
77+ push_config_store : PushNotificationConfigStore | None = None ,
78+ push_sender : PushNotificationSender | None = None ,
7679 request_context_builder : RequestContextBuilder | None = None ,
7780 ) -> None :
7881 """Initializes the DefaultRequestHandler.
@@ -81,14 +84,16 @@ def __init__(
8184 agent_executor: The `AgentExecutor` instance to run agent logic.
8285 task_store: The `TaskStore` instance to manage task persistence.
8386 queue_manager: The `QueueManager` instance to manage event queues. Defaults to `InMemoryQueueManager`.
84- push_notifier: The `PushNotifier` instance for sending push notifications. Defaults to None.
87+ push_config_store: The `PushNotificationConfigStore` instance for managing push notification configurations. Defaults to None.
88+ push_sender: The `PushNotificationSender` instance for sending push notifications. Defaults to None.
8589 request_context_builder: The `RequestContextBuilder` instance used
8690 to build request contexts. Defaults to `SimpleRequestContextBuilder`.
8791 """
8892 self .agent_executor = agent_executor
8993 self .task_store = task_store
9094 self ._queue_manager = queue_manager or InMemoryQueueManager ()
91- self ._push_notifier = push_notifier
95+ self ._push_config_store = push_config_store
96+ self ._push_sender = push_sender
9297 self ._request_context_builder = (
9398 request_context_builder
9499 or SimpleRequestContextBuilder (
@@ -198,15 +203,15 @@ async def _setup_message_execution(
198203
199204 task = task_manager .update_with_message (params .message , task )
200205 if self .should_add_push_info (params ):
201- assert isinstance ( self ._push_notifier , PushNotifier )
206+ assert self ._push_config_store is not None
202207 assert isinstance (
203208 params .configuration , MessageSendConfiguration
204209 )
205210 assert isinstance (
206211 params .configuration .pushNotificationConfig ,
207212 PushNotificationConfig ,
208213 )
209- await self ._push_notifier .set_info (
214+ await self ._push_config_store .set_info (
210215 task .id , params .configuration .pushNotificationConfig
211216 )
212217
@@ -247,10 +252,10 @@ async def _send_push_notification_if_needed(
247252 self , task_id : str , result_aggregator : ResultAggregator
248253 ) -> None :
249254 """Sends push notification if configured and task is available."""
250- if self ._push_notifier and task_id :
255+ if self ._push_sender and task_id :
251256 latest_task = await result_aggregator .current_result
252257 if isinstance (latest_task , Task ):
253- await self ._push_notifier .send_notification (latest_task )
258+ await self ._push_sender .send_notification (latest_task )
254259
255260 async def on_message_send (
256261 self ,
@@ -329,11 +334,11 @@ async def on_message_send_stream(
329334 self ._validate_task_id_match (task_id , event .id )
330335
331336 if (
332- self ._push_notifier
337+ self ._push_config_store
333338 and params .configuration
334339 and params .configuration .pushNotificationConfig
335340 ):
336- await self ._push_notifier .set_info (
341+ await self ._push_config_store .set_info (
337342 task_id ,
338343 params .configuration .pushNotificationConfig ,
339344 )
@@ -372,16 +377,14 @@ async def on_set_task_push_notification_config(
372377
373378 Requires a `PushNotifier` to be configured.
374379 """
375- if not self ._push_notifier :
380+ if not self ._push_config_store :
376381 raise ServerError (error = UnsupportedOperationError ())
377382
378383 task : Task | None = await self .task_store .get (params .taskId )
379384 if not task :
380385 raise ServerError (error = TaskNotFoundError ())
381386
382- # Generate a unique id for the notification
383- params .pushNotificationConfig .id = str (uuid .uuid4 ())
384- await self ._push_notifier .set_info (
387+ await self ._push_config_store .set_info (
385388 params .taskId ,
386389 params .pushNotificationConfig ,
387390 )
@@ -395,21 +398,27 @@ async def on_get_task_push_notification_config(
395398 ) -> TaskPushNotificationConfig :
396399 """Default handler for 'tasks/pushNotificationConfig/get'.
397400
398- Requires a `PushNotifier ` to be configured.
401+ Requires a `PushConfigStore ` to be configured.
399402 """
400- if not self ._push_notifier :
403+ if not self ._push_config_store :
401404 raise ServerError (error = UnsupportedOperationError ())
402405
403406 task : Task | None = await self .task_store .get (params .id )
404407 if not task :
405408 raise ServerError (error = TaskNotFoundError ())
406409
407- push_notification_config = await self ._push_notifier .get_info (params .id )
408- if not push_notification_config :
409- raise ServerError (error = InternalError ())
410+ push_notification_config = await self ._push_config_store .get_info (
411+ params .id
412+ )
413+ if not push_notification_config or not push_notification_config [0 ]:
414+ raise ServerError (
415+ error = InternalError (
416+ message = 'Push notification config not found'
417+ )
418+ )
410419
411420 return TaskPushNotificationConfig (
412- taskId = params .id , pushNotificationConfig = push_notification_config
421+ taskId = params .id , pushNotificationConfig = push_notification_config [ 0 ]
413422 )
414423
415424 async def on_resubscribe_to_task (
@@ -450,10 +459,61 @@ async def on_resubscribe_to_task(
450459 async for event in result_aggregator .consume_and_emit (consumer ):
451460 yield event
452461
462+ async def on_list_task_push_notification_config (
463+ self ,
464+ params : ListTaskPushNotificationConfigParams ,
465+ context : ServerCallContext | None = None ,
466+ ) -> list [TaskPushNotificationConfig ]:
467+ """Default handler for 'tasks/pushNotificationConfig/list'.
468+
469+ Requires a `PushConfigStore` to be configured.
470+ """
471+ if not self ._push_config_store :
472+ raise ServerError (error = UnsupportedOperationError ())
473+
474+ task : Task | None = await self .task_store .get (params .id )
475+ if not task :
476+ raise ServerError (error = TaskNotFoundError ())
477+
478+ push_notification_config_list = await self ._push_config_store .get_info (
479+ params .id
480+ )
481+
482+ task_push_notification_config = []
483+ if push_notification_config_list :
484+ for config in push_notification_config_list :
485+ task_push_notification_config .append (
486+ TaskPushNotificationConfig (
487+ taskId = params .id , pushNotificationConfig = config
488+ )
489+ )
490+
491+ return task_push_notification_config
492+
493+ async def on_delete_task_push_notification_config (
494+ self ,
495+ params : DeleteTaskPushNotificationConfigParams ,
496+ context : ServerCallContext | None = None ,
497+ ) -> None :
498+ """Default handler for 'tasks/pushNotificationConfig/delete'.
499+
500+ Requires a `PushConfigStore` to be configured.
501+ """
502+ if not self ._push_config_store :
503+ raise ServerError (error = UnsupportedOperationError ())
504+
505+ task : Task | None = await self .task_store .get (params .id )
506+ if not task :
507+ raise ServerError (error = TaskNotFoundError ())
508+
509+ await self ._push_config_store .delete_info (
510+ params .id , params .pushNotificationConfigId
511+ )
512+
453513 def should_add_push_info (self , params : MessageSendParams ) -> bool :
454514 """Determines if push notification info should be set for a task."""
455515 return bool (
456- self ._push_notifier
516+ self ._push_config_store
457517 and params .configuration
458518 and params .configuration .pushNotificationConfig
459519 )
0 commit comments