1
1
import asyncio
2
2
import logging
3
- import uuid
4
3
5
4
from collections .abc import AsyncGenerator
6
5
from typing import cast
21
20
)
22
21
from a2a .server .request_handlers .request_handler import RequestHandler
23
22
from a2a .server .tasks import (
24
- PushNotifier ,
23
+ PushNotificationConfigStore ,
24
+ PushNotificationSender ,
25
25
ResultAggregator ,
26
26
TaskManager ,
27
27
TaskStore ,
28
28
)
29
29
from a2a .types import (
30
+ DeleteTaskPushNotificationConfigParams ,
30
31
GetTaskPushNotificationConfigParams ,
31
32
InternalError ,
32
33
InvalidParamsError ,
34
+ ListTaskPushNotificationConfigParams ,
33
35
Message ,
34
36
MessageSendConfiguration ,
35
37
MessageSendParams ,
@@ -67,12 +69,13 @@ class DefaultRequestHandler(RequestHandler):
67
69
68
70
_running_agents : dict [str , asyncio .Task ]
69
71
70
- def __init__ (
72
+ def __init__ ( # noqa: PLR0913
71
73
self ,
72
74
agent_executor : AgentExecutor ,
73
75
task_store : TaskStore ,
74
76
queue_manager : QueueManager | None = None ,
75
- push_notifier : PushNotifier | None = None ,
77
+ push_config_store : PushNotificationConfigStore | None = None ,
78
+ push_sender : PushNotificationSender | None = None ,
76
79
request_context_builder : RequestContextBuilder | None = None ,
77
80
) -> None :
78
81
"""Initializes the DefaultRequestHandler.
@@ -81,14 +84,16 @@ def __init__(
81
84
agent_executor: The `AgentExecutor` instance to run agent logic.
82
85
task_store: The `TaskStore` instance to manage task persistence.
83
86
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.
85
89
request_context_builder: The `RequestContextBuilder` instance used
86
90
to build request contexts. Defaults to `SimpleRequestContextBuilder`.
87
91
"""
88
92
self .agent_executor = agent_executor
89
93
self .task_store = task_store
90
94
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
92
97
self ._request_context_builder = (
93
98
request_context_builder
94
99
or SimpleRequestContextBuilder (
@@ -198,15 +203,15 @@ async def _setup_message_execution(
198
203
199
204
task = task_manager .update_with_message (params .message , task )
200
205
if self .should_add_push_info (params ):
201
- assert isinstance ( self ._push_notifier , PushNotifier )
206
+ assert self ._push_config_store is not None
202
207
assert isinstance (
203
208
params .configuration , MessageSendConfiguration
204
209
)
205
210
assert isinstance (
206
211
params .configuration .pushNotificationConfig ,
207
212
PushNotificationConfig ,
208
213
)
209
- await self ._push_notifier .set_info (
214
+ await self ._push_config_store .set_info (
210
215
task .id , params .configuration .pushNotificationConfig
211
216
)
212
217
@@ -247,10 +252,10 @@ async def _send_push_notification_if_needed(
247
252
self , task_id : str , result_aggregator : ResultAggregator
248
253
) -> None :
249
254
"""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 :
251
256
latest_task = await result_aggregator .current_result
252
257
if isinstance (latest_task , Task ):
253
- await self ._push_notifier .send_notification (latest_task )
258
+ await self ._push_sender .send_notification (latest_task )
254
259
255
260
async def on_message_send (
256
261
self ,
@@ -329,11 +334,11 @@ async def on_message_send_stream(
329
334
self ._validate_task_id_match (task_id , event .id )
330
335
331
336
if (
332
- self ._push_notifier
337
+ self ._push_config_store
333
338
and params .configuration
334
339
and params .configuration .pushNotificationConfig
335
340
):
336
- await self ._push_notifier .set_info (
341
+ await self ._push_config_store .set_info (
337
342
task_id ,
338
343
params .configuration .pushNotificationConfig ,
339
344
)
@@ -372,16 +377,14 @@ async def on_set_task_push_notification_config(
372
377
373
378
Requires a `PushNotifier` to be configured.
374
379
"""
375
- if not self ._push_notifier :
380
+ if not self ._push_config_store :
376
381
raise ServerError (error = UnsupportedOperationError ())
377
382
378
383
task : Task | None = await self .task_store .get (params .taskId )
379
384
if not task :
380
385
raise ServerError (error = TaskNotFoundError ())
381
386
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 (
385
388
params .taskId ,
386
389
params .pushNotificationConfig ,
387
390
)
@@ -395,21 +398,27 @@ async def on_get_task_push_notification_config(
395
398
) -> TaskPushNotificationConfig :
396
399
"""Default handler for 'tasks/pushNotificationConfig/get'.
397
400
398
- Requires a `PushNotifier ` to be configured.
401
+ Requires a `PushConfigStore ` to be configured.
399
402
"""
400
- if not self ._push_notifier :
403
+ if not self ._push_config_store :
401
404
raise ServerError (error = UnsupportedOperationError ())
402
405
403
406
task : Task | None = await self .task_store .get (params .id )
404
407
if not task :
405
408
raise ServerError (error = TaskNotFoundError ())
406
409
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
+ )
410
419
411
420
return TaskPushNotificationConfig (
412
- taskId = params .id , pushNotificationConfig = push_notification_config
421
+ taskId = params .id , pushNotificationConfig = push_notification_config [ 0 ]
413
422
)
414
423
415
424
async def on_resubscribe_to_task (
@@ -450,10 +459,61 @@ async def on_resubscribe_to_task(
450
459
async for event in result_aggregator .consume_and_emit (consumer ):
451
460
yield event
452
461
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
+
453
513
def should_add_push_info (self , params : MessageSendParams ) -> bool :
454
514
"""Determines if push notification info should be set for a task."""
455
515
return bool (
456
- self ._push_notifier
516
+ self ._push_config_store
457
517
and params .configuration
458
518
and params .configuration .pushNotificationConfig
459
519
)
0 commit comments