1use a2a::*;
4use async_trait::async_trait;
5use futures::{StreamExt, stream::BoxStream};
6use std::{collections::HashMap, sync::Arc};
7use tokio::sync::{RwLock, broadcast};
8
9use crate::middleware::ServiceParams;
10use crate::pagination::resolve_page_size;
11use crate::task_store::apply_history_length;
12
13const EXECUTION_BUFFER_CAPACITY: usize = 32;
14
15#[derive(Clone)]
16struct ExecutionEvent {
17 sequence: u64,
18 result: Result<StreamResponse, A2AError>,
19}
20
21#[derive(Default)]
22struct ActiveExecutionState {
23 sequence: u64,
24 snapshot_task: Option<Task>,
25}
26
27struct ActiveExecution {
28 sender: broadcast::Sender<ExecutionEvent>,
29 state: RwLock<ActiveExecutionState>,
30}
31
32impl ActiveExecution {
33 fn new(task: Task) -> Arc<Self> {
34 let (sender, _) = broadcast::channel(EXECUTION_BUFFER_CAPACITY);
35 Arc::new(Self {
36 sender,
37 state: RwLock::new(ActiveExecutionState {
38 sequence: 0,
39 snapshot_task: Some(task),
40 }),
41 })
42 }
43
44 fn subscribe(&self) -> broadcast::Receiver<ExecutionEvent> {
45 self.sender.subscribe()
46 }
47
48 async fn resubscribe(&self) -> (broadcast::Receiver<ExecutionEvent>, Option<Task>, u64) {
49 let receiver = self.sender.subscribe();
50 let state = self.state.read().await;
51 (receiver, state.snapshot_task.clone(), state.sequence)
52 }
53
54 async fn publish(
55 &self,
56 result: Result<&StreamResponse, &A2AError>,
57 snapshot_task: Option<Task>,
58 ) {
59 let sequence = {
60 let mut state = self.state.write().await;
61 if let Some(task) = snapshot_task {
62 state.snapshot_task = Some(task);
63 }
64 state.sequence += 1;
65 state.sequence
66 };
67
68 if self.sender.receiver_count() == 0 {
71 return;
72 }
73
74 let event = ExecutionEvent {
75 sequence,
76 result: result.cloned().map_err(Clone::clone),
77 };
78 let _ = self.sender.send(event);
79 }
80
81 async fn snapshot_is_terminal(&self) -> bool {
82 let state = self.state.read().await;
83 state
84 .snapshot_task
85 .as_ref()
86 .is_some_and(|task| task.status.state.is_terminal())
87 }
88}
89
90#[derive(Default)]
91struct ExecutionManager {
92 executions: RwLock<HashMap<TaskId, Arc<ActiveExecution>>>,
93}
94
95impl ExecutionManager {
96 async fn start(&self, task: Task) -> Result<Arc<ActiveExecution>, A2AError> {
97 let mut executions = self.executions.write().await;
98 if executions.contains_key(&task.id) {
99 return Err(A2AError::invalid_request(format!(
100 "task execution is already in progress: {}",
101 task.id
102 )));
103 }
104
105 let active = ActiveExecution::new(task.clone());
106 executions.insert(task.id.clone(), active.clone());
107 Ok(active)
108 }
109
110 async fn get(&self, task_id: &str) -> Option<Arc<ActiveExecution>> {
111 let executions = self.executions.read().await;
112 executions.get(task_id).cloned()
113 }
114
115 async fn resubscribe(
116 &self,
117 task_id: &str,
118 ) -> Option<(broadcast::Receiver<ExecutionEvent>, Option<Task>, u64)> {
119 let active = self.get(task_id).await?;
120 let (receiver, snapshot_task, sequence) = active.resubscribe().await;
121 Some((receiver, snapshot_task, sequence))
122 }
123
124 async fn finish(&self, task_id: &str, active: &Arc<ActiveExecution>) {
125 let mut executions = self.executions.write().await;
126 let should_remove = executions
127 .get(task_id)
128 .is_some_and(|current| Arc::ptr_eq(current, active));
129 if should_remove {
130 executions.remove(task_id);
131 }
132 }
133}
134
135struct ExecutionRuntime {
136 executor: Arc<dyn crate::AgentExecutor>,
137 task_store: Arc<dyn crate::TaskStore>,
138 push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
139 push_sender: Option<Arc<crate::HttpPushSender>>,
140 execution_manager: Arc<ExecutionManager>,
141}
142
143struct SubscriptionState {
144 receiver: broadcast::Receiver<ExecutionEvent>,
145 snapshot_task: Option<Task>,
146 min_sequence: u64,
147 done: bool,
148}
149
150fn is_terminal_event(event: &StreamResponse) -> bool {
151 match event {
152 StreamResponse::Task(task) => task.status.state.is_terminal(),
153 StreamResponse::StatusUpdate(update) => update.status.state.is_terminal(),
154 _ => false,
155 }
156}
157
158fn task_id_for_event(event: &StreamResponse) -> Option<&str> {
159 match event {
160 StreamResponse::Task(task) => Some(&task.id),
161 StreamResponse::StatusUpdate(update) => Some(&update.task_id),
162 StreamResponse::ArtifactUpdate(update) => Some(&update.task_id),
163 StreamResponse::Message(_) => None,
164 }
165}
166
167fn should_interrupt_non_streaming(
168 req: &SendMessageRequest,
169 event: &StreamResponse,
170) -> Option<TaskId> {
171 if req
172 .configuration
173 .as_ref()
174 .and_then(|config| config.return_immediately)
175 .unwrap_or(false)
176 {
177 return match event {
178 StreamResponse::Message(_) => None,
179 _ => task_id_for_event(event).map(ToOwned::to_owned),
180 };
181 }
182
183 match event {
184 StreamResponse::Task(task) if task.status.state == TaskState::AuthRequired => {
185 Some(task.id.clone())
186 }
187 StreamResponse::StatusUpdate(update) if update.status.state == TaskState::AuthRequired => {
188 Some(update.task_id.clone())
189 }
190 _ => None,
191 }
192}
193
194async fn send_push_notifications(
195 task_id: &str,
196 push_config_store: Option<&dyn crate::PushConfigStore>,
197 push_sender: Option<&crate::HttpPushSender>,
198 event: &StreamResponse,
199) -> Result<(), A2AError> {
200 let Some(push_config_store) = push_config_store else {
201 return Ok(());
202 };
203 let Some(push_sender) = push_sender else {
204 return Ok(());
205 };
206
207 let configs = push_config_store.list(task_id).await?;
208 for config in configs {
209 push_sender.send_push(&config, event.clone()).await?;
210 }
211
212 Ok(())
213}
214
215async fn save_task(task_store: &dyn crate::TaskStore, task: Task) -> Result<Task, A2AError> {
216 match task_store.update(task.clone()).await {
217 Ok(_) => Ok(task),
218 Err(error) if error.code == error_code::TASK_NOT_FOUND => {
219 task_store.create(task.clone()).await?;
220 Ok(task)
221 }
222 Err(error) => Err(error),
223 }
224}
225
226async fn apply_event_to_task(
227 task_store: &dyn crate::TaskStore,
228 current_task: Option<Task>,
229 event: &StreamResponse,
230) -> Result<Option<Task>, A2AError> {
231 match event {
232 StreamResponse::Task(task) => save_task(task_store, task.clone()).await.map(Some),
233 StreamResponse::StatusUpdate(update) => {
234 let mut task = current_task
235 .or(task_store.get(&update.task_id).await?)
236 .ok_or_else(|| A2AError::task_not_found(&update.task_id))?;
237 task.status = update.status.clone();
238 save_task(task_store, task).await.map(Some)
239 }
240 StreamResponse::ArtifactUpdate(update) => {
241 let mut task = current_task
242 .or(task_store.get(&update.task_id).await?)
243 .ok_or_else(|| A2AError::task_not_found(&update.task_id))?;
244 let artifacts = task.artifacts.get_or_insert_with(Vec::new);
245 artifacts.push(update.artifact.clone());
246 save_task(task_store, task).await.map(Some)
247 }
248 StreamResponse::Message(_) => Ok(current_task),
249 }
250}
251
252fn subscription_stream(
253 receiver: broadcast::Receiver<ExecutionEvent>,
254 snapshot_task: Option<Task>,
255 min_sequence: u64,
256) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
257 Box::pin(futures::stream::unfold(
258 SubscriptionState {
259 receiver,
260 snapshot_task,
261 min_sequence,
262 done: false,
263 },
264 |mut state| async move {
265 if state.done {
266 return None;
267 }
268
269 if let Some(task) = state.snapshot_task.take() {
270 if task.status.state.is_terminal() {
271 state.done = true;
272 }
273 return Some((Ok(StreamResponse::Task(task)), state));
274 }
275
276 loop {
277 match state.receiver.recv().await {
278 Ok(event) if event.sequence <= state.min_sequence => continue,
279 Ok(event) => {
280 state.min_sequence = event.sequence;
281 match &event.result {
282 Ok(item) if is_terminal_event(item) => state.done = true,
283 Err(_) => state.done = true,
284 _ => {}
285 }
286 return Some((event.result.clone(), state));
287 }
288 Err(broadcast::error::RecvError::Lagged(_)) => {
289 state.done = true;
290 return Some((
291 Err(A2AError::internal(
292 "subscription fell behind active execution",
293 )),
294 state,
295 ));
296 }
297 Err(broadcast::error::RecvError::Closed) => return None,
298 }
299 }
300 },
301 ))
302}
303
304async fn drive_execution(
305 runtime: ExecutionRuntime,
306 task_id: TaskId,
307 active: Arc<ActiveExecution>,
308 exec_ctx: crate::ExecutorContext,
309) {
310 let mut current_task = {
311 let state = active.state.read().await;
312 state.snapshot_task.clone()
313 };
314 let mut stream = runtime.executor.execute(exec_ctx);
315
316 while let Some(result) = stream.next().await {
317 if active.snapshot_is_terminal().await {
318 break;
319 }
320
321 match result {
322 Ok(event) => {
323 match apply_event_to_task(runtime.task_store.as_ref(), current_task.clone(), &event)
324 .await
325 {
326 Ok(updated_task) => {
327 current_task = updated_task.clone();
328 if let Err(error) = send_push_notifications(
329 &task_id,
330 runtime.push_config_store.as_deref(),
331 runtime.push_sender.as_deref(),
332 &event,
333 )
334 .await
335 {
336 active.publish(Err(&error), None).await;
337 break;
338 }
339 active.publish(Ok(&event), updated_task).await;
340 if is_terminal_event(&event) {
341 break;
342 }
343 }
344 Err(error) => {
345 active.publish(Err(&error), None).await;
346 break;
347 }
348 }
349 }
350 Err(error) => {
351 active.publish(Err(&error), None).await;
352 break;
353 }
354 }
355 }
356
357 runtime.execution_manager.finish(&task_id, &active).await;
358}
359
360#[async_trait]
367pub trait RequestHandler: Send + Sync + 'static {
368 async fn send_message(
369 &self,
370 params: &ServiceParams,
371 req: SendMessageRequest,
372 ) -> Result<SendMessageResponse, A2AError>;
373
374 async fn send_streaming_message(
375 &self,
376 params: &ServiceParams,
377 req: SendMessageRequest,
378 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
379
380 async fn get_task(&self, params: &ServiceParams, req: GetTaskRequest)
381 -> Result<Task, A2AError>;
382
383 async fn list_tasks(
384 &self,
385 params: &ServiceParams,
386 req: ListTasksRequest,
387 ) -> Result<ListTasksResponse, A2AError>;
388
389 async fn cancel_task(
390 &self,
391 params: &ServiceParams,
392 req: CancelTaskRequest,
393 ) -> Result<Task, A2AError>;
394
395 async fn subscribe_to_task(
396 &self,
397 params: &ServiceParams,
398 req: SubscribeToTaskRequest,
399 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
400
401 async fn create_push_config(
402 &self,
403 params: &ServiceParams,
404 req: TaskPushNotificationConfig,
405 ) -> Result<TaskPushNotificationConfig, A2AError>;
406
407 async fn get_push_config(
408 &self,
409 params: &ServiceParams,
410 req: GetTaskPushNotificationConfigRequest,
411 ) -> Result<TaskPushNotificationConfig, A2AError>;
412
413 async fn list_push_configs(
414 &self,
415 params: &ServiceParams,
416 req: ListTaskPushNotificationConfigsRequest,
417 ) -> Result<ListTaskPushNotificationConfigsResponse, A2AError>;
418
419 async fn delete_push_config(
420 &self,
421 params: &ServiceParams,
422 req: DeleteTaskPushNotificationConfigRequest,
423 ) -> Result<(), A2AError>;
424
425 async fn get_extended_agent_card(
426 &self,
427 params: &ServiceParams,
428 req: GetExtendedAgentCardRequest,
429 ) -> Result<AgentCard, A2AError>;
430}
431
432pub trait RequestAuthorizer: Send + Sync {
439 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError>;
440}
441
442impl<F> RequestAuthorizer for F
443where
444 F: Fn(&ServiceParams, Option<&str>) -> Result<(), A2AError> + Send + Sync,
445{
446 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
447 self(params, task_id)
448 }
449}
450
451pub struct DefaultRequestHandler {
454 executor: Arc<dyn crate::AgentExecutor>,
455 task_store: Arc<dyn crate::TaskStore>,
456 execution_manager: Arc<ExecutionManager>,
457 push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
458 push_sender: Option<Arc<crate::HttpPushSender>>,
459 capabilities: AgentCapabilities,
460 authorizer: Option<Arc<dyn RequestAuthorizer>>,
461 extended_agent_card: Option<Arc<dyn ExtendedAgentCardResolver>>,
462}
463
464pub trait ExtendedAgentCardResolver: Send + Sync {
473 fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError>;
474}
475
476impl<F> ExtendedAgentCardResolver for F
477where
478 F: Fn(&ServiceParams) -> Result<AgentCard, A2AError> + Send + Sync,
479{
480 fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError> {
481 self(params)
482 }
483}
484
485impl DefaultRequestHandler {
486 pub fn new(executor: impl crate::AgentExecutor, task_store: impl crate::TaskStore) -> Self {
487 DefaultRequestHandler {
488 executor: Arc::new(executor),
489 task_store: Arc::new(task_store),
490 execution_manager: Arc::new(ExecutionManager::default()),
491 push_config_store: None,
492 push_sender: None,
493 capabilities: AgentCapabilities::default(),
494 authorizer: None,
495 extended_agent_card: None,
496 }
497 }
498
499 pub fn with_extended_agent_card(mut self, card: AgentCard) -> Self {
505 self.capabilities.extended_agent_card.get_or_insert(true);
506 self.with_extended_agent_card_resolver(move |_: &ServiceParams| Ok(card.clone()))
507 }
508
509 pub fn with_extended_agent_card_resolver(
512 mut self,
513 resolver: impl ExtendedAgentCardResolver + 'static,
514 ) -> Self {
515 self.capabilities.extended_agent_card.get_or_insert(true);
516 self.extended_agent_card = Some(Arc::new(resolver));
517 self
518 }
519
520 pub fn with_authorizer(mut self, authorizer: impl RequestAuthorizer + 'static) -> Self {
522 self.authorizer = Some(Arc::new(authorizer));
523 self
524 }
525
526 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
527 match &self.authorizer {
528 Some(authorizer) => authorizer.authorize(params, task_id),
529 None => Ok(()),
530 }
531 }
532
533 pub fn with_push_config_store(
534 mut self,
535 push_config_store: impl crate::PushConfigStore,
536 ) -> Self {
537 self.push_sender = Some(Arc::new(crate::HttpPushSender::new(None)));
538 self.push_config_store = Some(Arc::new(push_config_store));
539 self.capabilities.push_notifications.get_or_insert(true);
543 self
544 }
545
546 pub fn with_push_notifications(
547 mut self,
548 push_config_store: impl crate::PushConfigStore,
549 push_sender: crate::HttpPushSender,
550 ) -> Self {
551 self.push_config_store = Some(Arc::new(push_config_store));
552 self.push_sender = Some(Arc::new(push_sender));
553 self.capabilities.push_notifications.get_or_insert(true);
554 self
555 }
556
557 pub fn with_capabilities(mut self, capabilities: AgentCapabilities) -> Self {
558 self.capabilities = capabilities;
559 self
560 }
561
562 fn require_capability(
569 declared: Option<bool>,
570 unsupported: impl FnOnce() -> A2AError,
571 ) -> Result<(), A2AError> {
572 match declared {
573 Some(false) => Err(unsupported()),
574 _ => Ok(()),
575 }
576 }
577
578 fn require_streaming(&self) -> Result<(), A2AError> {
579 Self::require_capability(self.capabilities.streaming, || {
580 A2AError::unsupported_operation("the agent card does not declare streaming")
581 })
582 }
583
584 fn require_push_notifications(&self) -> Result<(), A2AError> {
585 Self::require_capability(
586 self.capabilities.push_notifications,
587 A2AError::push_notification_not_supported,
588 )
589 }
590
591 fn push_config_store_checked(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
596 self.require_push_notifications()?;
597 self.push_config_store()
598 }
599
600 fn push_config_store(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
601 self.push_config_store
602 .as_deref()
603 .ok_or_else(A2AError::push_notification_not_supported)
604 }
605
606 async fn load_task(&self, task_id: &str) -> Result<Task, A2AError> {
607 self.task_store
608 .get(task_id)
609 .await?
610 .ok_or_else(|| A2AError::task_not_found(task_id))
611 }
612
613 async fn prepare_task_for_execution(
614 &self,
615 req: &SendMessageRequest,
616 ) -> Result<(Task, Option<Task>, String), A2AError> {
617 let task_id = req.message.task_id.clone().unwrap_or_else(new_task_id);
618 let stored = self.task_store.get(&task_id).await?;
619 let context_id = stored
620 .as_ref()
621 .map(|task| task.context_id.clone())
622 .or_else(|| req.message.context_id.clone())
623 .unwrap_or_else(new_context_id);
624
625 let task = if let Some(existing) = stored.clone() {
626 existing
627 } else {
628 let task = Task {
629 id: task_id,
630 context_id: context_id.clone(),
631 status: TaskStatus {
632 state: TaskState::Submitted,
633 message: None,
634 timestamp: Some(chrono::Utc::now()),
635 },
636 artifacts: None,
637 history: Some(vec![req.message.clone()]),
638 metadata: None,
639 };
640 self.task_store.create(task.clone()).await?;
641 task
642 };
643
644 Ok((task, stored, context_id))
645 }
646
647 async fn save_request_push_config(
648 &self,
649 task_id: &str,
650 req: &SendMessageRequest,
651 ) -> Result<(), A2AError> {
652 let Some(config) = req
653 .configuration
654 .as_ref()
655 .and_then(|configuration| configuration.task_push_notification_config.clone())
656 else {
657 return Ok(());
658 };
659
660 let mut config = config;
661 config.task_id = task_id.to_string();
662 if config.tenant.is_none() {
663 config.tenant = req.tenant.clone();
664 }
665 self.push_config_store_checked()?.save(config).await?;
666 Ok(())
667 }
668
669 async fn start_execution(
670 &self,
671 params: &ServiceParams,
672 req: SendMessageRequest,
673 include_task_snapshot_in_context: bool,
674 ) -> Result<(TaskId, BoxStream<'static, Result<StreamResponse, A2AError>>), A2AError> {
675 let (task, stored_task, context_id) = self.prepare_task_for_execution(&req).await?;
676 let task_id = task.id.clone();
677 self.save_request_push_config(&task_id, &req).await?;
678 let active = self.execution_manager.start(task.clone()).await?;
679 let receiver = active.subscribe();
680
681 let exec_ctx = crate::ExecutorContext {
682 message: Some(req.message),
683 task_id: task_id.clone(),
684 stored_task: if include_task_snapshot_in_context {
685 Some(task.clone())
686 } else {
687 stored_task
688 },
689 context_id,
690 metadata: req.metadata,
691 user: None,
692 service_params: params.clone(),
693 tenant: req.tenant,
694 };
695
696 let runtime = ExecutionRuntime {
697 executor: Arc::clone(&self.executor),
698 task_store: Arc::clone(&self.task_store),
699 push_config_store: self.push_config_store.clone(),
700 push_sender: self.push_sender.clone(),
701 execution_manager: Arc::clone(&self.execution_manager),
702 };
703 let active_execution = active.clone();
704 let execution_task_id = task_id.clone();
705
706 tokio::spawn(async move {
707 drive_execution(runtime, execution_task_id, active_execution, exec_ctx).await;
708 });
709
710 Ok((task_id, subscription_stream(receiver, None, 0)))
711 }
712}
713
714#[async_trait]
715impl RequestHandler for DefaultRequestHandler {
716 async fn send_message(
717 &self,
718 params: &ServiceParams,
719 req: SendMessageRequest,
720 ) -> Result<SendMessageResponse, A2AError> {
721 self.authorize(params, req.message.task_id.as_deref())?;
722 let (task_id, mut stream) = self.start_execution(params, req.clone(), true).await?;
723 let mut last_event = None;
724
725 while let Some(item) = stream.next().await {
726 let event = item?;
727
728 if let Some(interrupt_task_id) = should_interrupt_non_streaming(&req, &event) {
729 return Ok(SendMessageResponse::Task(
730 self.load_task(&interrupt_task_id).await?,
731 ));
732 }
733
734 match event {
735 StreamResponse::Message(message) => {
736 return Ok(SendMessageResponse::Message(message));
737 }
738 other => last_event = Some(other),
739 }
740 }
741
742 match last_event {
743 Some(StreamResponse::Task(task)) => Ok(SendMessageResponse::Task(task)),
744 Some(StreamResponse::StatusUpdate(update)) => Ok(SendMessageResponse::Task(
745 self.load_task(&update.task_id).await?,
746 )),
747 Some(StreamResponse::ArtifactUpdate(update)) => Ok(SendMessageResponse::Task(
748 self.load_task(&update.task_id).await?,
749 )),
750 Some(StreamResponse::Message(message)) => Ok(SendMessageResponse::Message(message)),
751 None => Ok(SendMessageResponse::Task(self.load_task(&task_id).await?)),
752 }
753 }
754
755 async fn send_streaming_message(
756 &self,
757 params: &ServiceParams,
758 req: SendMessageRequest,
759 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
760 self.authorize(params, req.message.task_id.as_deref())?;
761 self.require_streaming()?;
762 let (_, stream) = self.start_execution(params, req, false).await?;
763 Ok(stream)
764 }
765
766 async fn get_task(
767 &self,
768 params: &ServiceParams,
769 req: GetTaskRequest,
770 ) -> Result<Task, A2AError> {
771 self.authorize(params, Some(&req.id))?;
772 let mut task = self
773 .task_store
774 .get(&req.id)
775 .await?
776 .ok_or_else(|| A2AError::task_not_found(&req.id))?;
777 apply_history_length(&mut task, req.history_length);
778 Ok(task)
779 }
780
781 async fn list_tasks(
794 &self,
795 params: &ServiceParams,
796 req: ListTasksRequest,
797 ) -> Result<ListTasksResponse, A2AError> {
798 self.authorize(params, None)?;
799 let page_size = resolve_page_size(req.page_size);
800 let req = ListTasksRequest {
801 page_size: Some(page_size as i32),
802 ..req
803 };
804 let mut response = self.task_store.list(&req).await?;
805 response.tasks.truncate(page_size);
806 for task in &mut response.tasks {
807 apply_history_length(task, req.history_length);
808 }
809 Ok(response)
810 }
811
812 async fn cancel_task(
813 &self,
814 params: &ServiceParams,
815 req: CancelTaskRequest,
816 ) -> Result<Task, A2AError> {
817 self.authorize(params, Some(&req.id))?;
818 let task = self.load_task(&req.id).await?;
819
820 if task.status.state == TaskState::Canceled {
824 return Ok(task);
825 }
826 if task.status.state.is_terminal() {
827 return Err(A2AError::task_not_cancelable(&req.id));
828 }
829
830 let active_execution = self.execution_manager.get(&req.id).await;
831
832 let exec_ctx = crate::ExecutorContext {
833 message: None,
834 task_id: req.id.clone(),
835 stored_task: Some(task.clone()),
836 context_id: task.context_id.clone(),
837 metadata: req.metadata,
838 user: None,
839 service_params: params.clone(),
840 tenant: req.tenant,
841 };
842
843 let mut stream = self.executor.cancel(exec_ctx);
844 let mut current_task = Some(task);
845
846 while let Some(event) = stream.next().await {
847 match event {
848 Ok(event) => {
849 let updated_task =
850 apply_event_to_task(self.task_store.as_ref(), current_task.clone(), &event)
851 .await?;
852 current_task = updated_task.clone();
853
854 send_push_notifications(
855 &req.id,
856 self.push_config_store.as_deref(),
857 self.push_sender.as_deref(),
858 &event,
859 )
860 .await?;
861
862 if let Some(active) = &active_execution {
863 active.publish(Ok(&event), updated_task).await;
864 }
865
866 if is_terminal_event(&event) {
867 break;
868 }
869 }
870 Err(error) => {
871 if let Some(active) = &active_execution {
872 active.publish(Err(&error), None).await;
873 }
874 return Err(error);
875 }
876 }
877 }
878
879 if let Some(active) = active_execution {
880 self.execution_manager.finish(&req.id, &active).await;
881 }
882
883 self.load_task(&req.id).await
884 }
885
886 async fn subscribe_to_task(
887 &self,
888 params: &ServiceParams,
889 req: SubscribeToTaskRequest,
890 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
891 self.authorize(params, Some(&req.id))?;
892 self.require_streaming()?;
893 let (receiver, snapshot_task, sequence) = self
894 .execution_manager
895 .resubscribe(&req.id)
896 .await
897 .ok_or_else(|| A2AError::task_not_found(&req.id))?;
898 Ok(subscription_stream(receiver, snapshot_task, sequence))
899 }
900
901 async fn create_push_config(
902 &self,
903 params: &ServiceParams,
904 req: TaskPushNotificationConfig,
905 ) -> Result<TaskPushNotificationConfig, A2AError> {
906 self.authorize(params, Some(&req.task_id))?;
907 self.push_config_store_checked()?.save(req).await
908 }
909
910 async fn get_push_config(
911 &self,
912 params: &ServiceParams,
913 req: GetTaskPushNotificationConfigRequest,
914 ) -> Result<TaskPushNotificationConfig, A2AError> {
915 self.authorize(params, Some(&req.task_id))?;
916 self.push_config_store_checked()?
917 .get(&req.task_id, &req.id)
918 .await
919 }
920
921 async fn list_push_configs(
922 &self,
923 params: &ServiceParams,
924 req: ListTaskPushNotificationConfigsRequest,
925 ) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
926 self.authorize(params, Some(&req.task_id))?;
927 let mut configs = self.push_config_store_checked()?.list(&req.task_id).await?;
928 configs.sort_by(|left, right| left.id.cmp(&right.id));
929
930 let page_size = resolve_page_size(req.page_size);
935 let start = if let Some(ref token) = req.page_token {
936 token
937 .parse::<usize>()
938 .map_err(|_| A2AError::invalid_params("invalid page token"))?
939 } else {
940 0
941 };
942 let start = start.min(configs.len());
943 let end = (start + page_size).min(configs.len());
944 let next_page_token = (end < configs.len()).then(|| end.to_string());
945
946 Ok(ListTaskPushNotificationConfigsResponse {
947 configs: configs[start..end].to_vec(),
948 next_page_token,
949 })
950 }
951
952 async fn delete_push_config(
953 &self,
954 params: &ServiceParams,
955 req: DeleteTaskPushNotificationConfigRequest,
956 ) -> Result<(), A2AError> {
957 self.authorize(params, Some(&req.task_id))?;
958 self.push_config_store_checked()?
959 .delete(&req.task_id, &req.id)
960 .await
961 }
962
963 async fn get_extended_agent_card(
964 &self,
965 params: &ServiceParams,
966 _req: GetExtendedAgentCardRequest,
967 ) -> Result<AgentCard, A2AError> {
968 self.authorize(params, None)?;
969 Self::require_capability(self.capabilities.extended_agent_card, || {
970 A2AError::unsupported_operation(
971 "the agent card does not declare an extended agent card",
972 )
973 })?;
974 match &self.extended_agent_card {
975 Some(resolver) => resolver.resolve(params),
976 None => Err(A2AError::extended_card_not_configured()),
981 }
982 }
983}
984
985#[cfg(test)]
986mod tests {
987 use super::*;
988 use crate::executor::ExecutorContext;
989 use crate::pagination::{DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE};
990 use crate::push::{InMemoryPushConfigStore, PushConfigStore};
991 use crate::task_store::{InMemoryTaskStore, TaskStore, TaskVersion};
992
993 use crate::test_util::install_crypto_provider;
994 use axum::{
995 Router,
996 body::Bytes,
997 extract::State,
998 http::{HeaderMap, StatusCode, header},
999 routing::post,
1000 };
1001 use futures::stream;
1002 use std::sync::Arc;
1003 use std::sync::atomic::{AtomicI32, Ordering};
1004 use tokio::{
1005 net::TcpListener,
1006 sync::{Notify, mpsc, oneshot},
1007 time::{Duration, timeout},
1008 };
1009
1010 struct EchoExecutor;
1011
1012 struct PushEventExecutor;
1013
1014 struct ResumableExecutor {
1015 release: Arc<Notify>,
1016 }
1017
1018 #[derive(Debug)]
1019 struct CapturedPush {
1020 authorization: Option<String>,
1021 notification_token: Option<String>,
1022 event: StreamResponse,
1023 }
1024
1025 impl crate::AgentExecutor for EchoExecutor {
1026 fn execute(
1027 &self,
1028 ctx: ExecutorContext,
1029 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1030 let task = Task {
1031 id: ctx.task_id.clone(),
1032 context_id: ctx.context_id.clone(),
1033 status: TaskStatus {
1034 state: TaskState::Completed,
1035 message: ctx.message,
1036 timestamp: None,
1037 },
1038 artifacts: None,
1039 history: None,
1040 metadata: None,
1041 };
1042 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1043 }
1044
1045 fn cancel(
1046 &self,
1047 ctx: ExecutorContext,
1048 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1049 let task = Task {
1050 id: ctx.task_id.clone(),
1051 context_id: ctx.context_id.clone(),
1052 status: TaskStatus {
1053 state: TaskState::Canceled,
1054 message: None,
1055 timestamp: None,
1056 },
1057 artifacts: None,
1058 history: None,
1059 metadata: None,
1060 };
1061 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1062 }
1063 }
1064
1065 impl crate::AgentExecutor for PushEventExecutor {
1066 fn execute(
1067 &self,
1068 ctx: ExecutorContext,
1069 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1070 let working = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1071 task_id: ctx.task_id.clone(),
1072 context_id: ctx.context_id.clone(),
1073 status: TaskStatus {
1074 state: TaskState::Working,
1075 message: None,
1076 timestamp: None,
1077 },
1078 metadata: None,
1079 });
1080 let completed = StreamResponse::Task(Task {
1081 id: ctx.task_id.clone(),
1082 context_id: ctx.context_id.clone(),
1083 status: TaskStatus {
1084 state: TaskState::Completed,
1085 message: ctx.message,
1086 timestamp: None,
1087 },
1088 artifacts: None,
1089 history: ctx.stored_task.and_then(|task| task.history),
1090 metadata: None,
1091 });
1092
1093 Box::pin(stream::iter(vec![Ok(working), Ok(completed)]))
1094 }
1095
1096 fn cancel(
1097 &self,
1098 ctx: ExecutorContext,
1099 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1100 let task = Task {
1101 id: ctx.task_id.clone(),
1102 context_id: ctx.context_id.clone(),
1103 status: TaskStatus {
1104 state: TaskState::Canceled,
1105 message: None,
1106 timestamp: None,
1107 },
1108 artifacts: None,
1109 history: None,
1110 metadata: None,
1111 };
1112 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1113 }
1114 }
1115
1116 impl crate::AgentExecutor for ResumableExecutor {
1117 fn execute(
1118 &self,
1119 ctx: ExecutorContext,
1120 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1121 let release = self.release.clone();
1122 let working_task = Task {
1123 id: ctx.task_id.clone(),
1124 context_id: ctx.context_id.clone(),
1125 status: TaskStatus {
1126 state: TaskState::Working,
1127 message: None,
1128 timestamp: None,
1129 },
1130 artifacts: None,
1131 history: ctx.stored_task.and_then(|task| task.history),
1132 metadata: None,
1133 };
1134 let completed_task = Task {
1135 id: working_task.id.clone(),
1136 context_id: working_task.context_id.clone(),
1137 status: TaskStatus {
1138 state: TaskState::Completed,
1139 message: None,
1140 timestamp: None,
1141 },
1142 artifacts: None,
1143 history: working_task.history.clone(),
1144 metadata: None,
1145 };
1146
1147 Box::pin(stream::unfold(
1148 (Some(working_task), Some((release, completed_task))),
1149 |(working_task, completion)| async move {
1150 if let Some(task) = working_task {
1151 return Some((Ok(StreamResponse::Task(task)), (None, completion)));
1152 }
1153
1154 if let Some((release, task)) = completion {
1155 release.notified().await;
1156 return Some((Ok(StreamResponse::Task(task)), (None, None)));
1157 }
1158
1159 None
1160 },
1161 ))
1162 }
1163
1164 fn cancel(
1165 &self,
1166 ctx: ExecutorContext,
1167 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1168 let task = Task {
1169 id: ctx.task_id.clone(),
1170 context_id: ctx.context_id.clone(),
1171 status: TaskStatus {
1172 state: TaskState::Canceled,
1173 message: None,
1174 timestamp: None,
1175 },
1176 artifacts: None,
1177 history: None,
1178 metadata: None,
1179 };
1180 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1181 }
1182 }
1183
1184 struct PagingTaskStore {
1192 inner: InMemoryTaskStore,
1193 count: usize,
1194 honour_page_size: bool,
1198 seen_page_size: Arc<AtomicI32>,
1202 next_page_token: String,
1205 }
1206
1207 impl PagingTaskStore {
1208 fn new(count: usize, honour_page_size: bool) -> Self {
1209 Self {
1210 inner: InMemoryTaskStore::new(),
1211 count,
1212 honour_page_size,
1213 seen_page_size: Arc::new(AtomicI32::new(-1)),
1214 next_page_token: String::new(),
1215 }
1216 }
1217
1218 fn with_next_page_token(mut self, token: &str) -> Self {
1219 self.next_page_token = token.to_string();
1220 self
1221 }
1222 }
1223
1224 #[async_trait]
1225 impl TaskStore for PagingTaskStore {
1226 async fn create(&self, task: Task) -> Result<TaskVersion, A2AError> {
1227 self.inner.create(task).await
1228 }
1229
1230 async fn update(&self, task: Task) -> Result<TaskVersion, A2AError> {
1231 self.inner.update(task).await
1232 }
1233
1234 async fn get(&self, task_id: &str) -> Result<Option<Task>, A2AError> {
1235 self.inner.get(task_id).await
1236 }
1237
1238 async fn list(&self, req: &ListTasksRequest) -> Result<ListTasksResponse, A2AError> {
1239 self.seen_page_size
1240 .store(req.page_size.unwrap_or(-1), Ordering::SeqCst);
1241 let served = if self.honour_page_size {
1242 req.page_size.unwrap_or(i32::MAX).max(0) as usize
1243 } else {
1244 self.count
1245 }
1246 .min(self.count);
1247 let tasks = (0..served)
1248 .map(|i| Task {
1249 id: format!("t-{i:04}"),
1250 context_id: "c1".into(),
1251 status: TaskStatus {
1252 state: TaskState::Completed,
1253 message: None,
1254 timestamp: None,
1255 },
1256 artifacts: None,
1257 history: None,
1258 metadata: None,
1259 })
1260 .collect::<Vec<_>>();
1261 let total = tasks.len() as i32;
1262 Ok(ListTasksResponse {
1263 tasks,
1264 next_page_token: self.next_page_token.clone(),
1265 page_size: total,
1266 total_size: total,
1267 })
1268 }
1269 }
1270
1271 struct PagingPushConfigStore {
1277 inner: InMemoryPushConfigStore,
1278 count: usize,
1279 }
1280
1281 impl PagingPushConfigStore {
1282 fn new(count: usize) -> Self {
1283 Self {
1284 inner: InMemoryPushConfigStore::new(),
1285 count,
1286 }
1287 }
1288 }
1289
1290 #[async_trait]
1291 impl crate::PushConfigStore for PagingPushConfigStore {
1292 async fn save(
1293 &self,
1294 config: TaskPushNotificationConfig,
1295 ) -> Result<TaskPushNotificationConfig, A2AError> {
1296 self.inner.save(config).await
1297 }
1298
1299 async fn get(
1300 &self,
1301 task_id: &str,
1302 config_id: &str,
1303 ) -> Result<TaskPushNotificationConfig, A2AError> {
1304 self.inner.get(task_id, config_id).await
1305 }
1306
1307 async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError> {
1308 Ok((0..self.count)
1309 .map(|i| TaskPushNotificationConfig {
1310 task_id: task_id.to_string(),
1311 url: format!("https://example.com/hook/{i}"),
1312 id: Some(format!("cfg-{i:04}")),
1313 token: None,
1314 authentication: None,
1315 tenant: None,
1316 })
1317 .collect())
1318 }
1319
1320 async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError> {
1321 self.inner.delete(task_id, config_id).await
1322 }
1323
1324 async fn delete_all(&self, task_id: &str) -> Result<(), A2AError> {
1325 self.inner.delete_all(task_id).await
1326 }
1327 }
1328
1329 fn make_handler() -> DefaultRequestHandler {
1330 DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
1331 }
1332
1333 fn make_handler_with_push_configs() -> DefaultRequestHandler {
1334 DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
1335 .with_push_config_store(InMemoryPushConfigStore::new())
1336 }
1337
1338 fn make_push_delivery_handler() -> DefaultRequestHandler {
1339 DefaultRequestHandler::new(PushEventExecutor, InMemoryTaskStore::new())
1342 .with_push_notifications(
1343 InMemoryPushConfigStore::new(),
1344 crate::HttpPushSender::new(Some(crate::HttpPushSenderConfig {
1345 validate_urls: false,
1346 ..Default::default()
1347 })),
1348 )
1349 }
1350
1351 fn make_resumable_handler() -> (DefaultRequestHandler, Arc<Notify>) {
1352 let release = Arc::new(Notify::new());
1353 (
1354 DefaultRequestHandler::new(
1355 ResumableExecutor {
1356 release: release.clone(),
1357 },
1358 InMemoryTaskStore::new(),
1359 ),
1360 release,
1361 )
1362 }
1363
1364 fn make_message() -> Message {
1365 Message::new(Role::User, vec![Part::text("hello")])
1366 }
1367
1368 async fn capture_push(
1369 State(sender): State<mpsc::UnboundedSender<CapturedPush>>,
1370 headers: HeaderMap,
1371 body: Bytes,
1372 ) -> StatusCode {
1373 let captured = CapturedPush {
1374 authorization: headers
1375 .get(header::AUTHORIZATION)
1376 .and_then(|value| value.to_str().ok())
1377 .map(ToOwned::to_owned),
1378 notification_token: headers
1379 .get("A2A-Notification-Token")
1380 .and_then(|value| value.to_str().ok())
1381 .map(ToOwned::to_owned),
1382 event: serde_json::from_slice(&body).unwrap(),
1383 };
1384 sender.send(captured).unwrap();
1385 StatusCode::ACCEPTED
1386 }
1387
1388 async fn start_push_webhook() -> (
1389 String,
1390 mpsc::UnboundedReceiver<CapturedPush>,
1391 oneshot::Sender<()>,
1392 tokio::task::JoinHandle<()>,
1393 ) {
1394 let (sender, receiver) = mpsc::unbounded_channel();
1395 let app = Router::new()
1396 .route("/", post(capture_push))
1397 .with_state(sender);
1398 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1399 let address = listener.local_addr().unwrap();
1400 let (shutdown_tx, shutdown_rx) = oneshot::channel();
1401 let server = tokio::spawn(async move {
1402 axum::serve(listener, app)
1403 .with_graceful_shutdown(async {
1404 let _ = shutdown_rx.await;
1405 })
1406 .await
1407 .unwrap();
1408 });
1409
1410 (format!("http://{address}/"), receiver, shutdown_tx, server)
1411 }
1412
1413 async fn next_push(receiver: &mut mpsc::UnboundedReceiver<CapturedPush>) -> CapturedPush {
1414 timeout(Duration::from_secs(5), receiver.recv())
1415 .await
1416 .unwrap()
1417 .unwrap()
1418 }
1419
1420 #[tokio::test]
1421 async fn test_send_message_creates_task() {
1422 let handler = make_handler();
1423 let params = ServiceParams::new();
1424 let req = SendMessageRequest {
1425 message: make_message(),
1426 configuration: None,
1427 metadata: None,
1428 tenant: None,
1429 };
1430 let result = handler.send_message(¶ms, req).await.unwrap();
1431 match result {
1432 SendMessageResponse::Task(task) => {
1433 assert_eq!(task.status.state, TaskState::Completed);
1434 }
1435 _ => panic!("expected Task response"),
1436 }
1437 }
1438
1439 #[tokio::test]
1440 async fn test_send_streaming_message() {
1441 use futures::StreamExt;
1442 let handler = make_handler();
1443 let params = ServiceParams::new();
1444 let req = SendMessageRequest {
1445 message: make_message(),
1446 configuration: None,
1447 metadata: None,
1448 tenant: None,
1449 };
1450 let mut stream = handler.send_streaming_message(¶ms, req).await.unwrap();
1451 let event = stream.next().await.unwrap().unwrap();
1452 assert!(matches!(event, StreamResponse::Task(_)));
1453 }
1454
1455 #[tokio::test]
1456 async fn test_stream_disconnect_preserves_execution_for_resubscription() {
1457 use futures::StreamExt;
1458
1459 let (handler, release) = make_resumable_handler();
1460 let params = ServiceParams::new();
1461 let mut message = make_message();
1462 message.task_id = Some("t-disconnect".into());
1463 message.context_id = Some("c-disconnect".into());
1464
1465 let mut stream = handler
1466 .send_streaming_message(
1467 ¶ms,
1468 SendMessageRequest {
1469 message,
1470 configuration: None,
1471 metadata: None,
1472 tenant: None,
1473 },
1474 )
1475 .await
1476 .unwrap();
1477
1478 let initial = stream.next().await.unwrap().unwrap();
1479 match initial {
1480 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
1481 _ => panic!("expected initial working task"),
1482 }
1483 drop(stream);
1484
1485 let mut resumed = handler
1486 .subscribe_to_task(
1487 ¶ms,
1488 SubscribeToTaskRequest {
1489 id: "t-disconnect".into(),
1490 tenant: None,
1491 },
1492 )
1493 .await
1494 .unwrap();
1495
1496 let snapshot = resumed.next().await.unwrap().unwrap();
1497 match snapshot {
1498 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
1499 _ => panic!("expected working task snapshot"),
1500 }
1501
1502 release.notify_waiters();
1503
1504 let terminal = resumed.next().await.unwrap().unwrap();
1505 match terminal {
1506 StreamResponse::Task(task) => {
1507 assert_eq!(task.status.state, TaskState::Completed)
1508 }
1509 _ => panic!("expected terminal task"),
1510 }
1511 assert!(resumed.next().await.is_none());
1512 }
1513
1514 #[tokio::test]
1515 async fn test_active_execution_keeps_snapshot_without_subscribers() {
1516 let task = Task {
1517 id: "t-buffering".into(),
1518 context_id: "c-buffering".into(),
1519 status: TaskStatus {
1520 state: TaskState::Working,
1521 message: None,
1522 timestamp: None,
1523 },
1524 artifacts: None,
1525 history: None,
1526 metadata: None,
1527 };
1528 let active = ActiveExecution::new(task.clone());
1529
1530 let event = StreamResponse::Task(task.clone());
1531 active.publish(Ok(&event), Some(task)).await;
1532
1533 let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
1534 assert_eq!(snapshot_task.unwrap().id, "t-buffering");
1535 assert_eq!(sequence, 1);
1536 assert!(receiver.try_recv().is_err());
1537 }
1538
1539 #[tokio::test]
1540 async fn test_active_execution_publishes_errors_to_subscribers() {
1541 let task = Task {
1542 id: "t-buffering-error".into(),
1543 context_id: "c-buffering-error".into(),
1544 status: TaskStatus {
1545 state: TaskState::Working,
1546 message: None,
1547 timestamp: None,
1548 },
1549 artifacts: None,
1550 history: None,
1551 metadata: None,
1552 };
1553 let active = ActiveExecution::new(task);
1554 let mut receiver = active.subscribe();
1555 let error = A2AError::internal("execution failed");
1556
1557 active.publish(Err(&error), None).await;
1558
1559 let event = receiver.recv().await.unwrap();
1560 assert_eq!(event.sequence, 1);
1561 assert_eq!(event.result.unwrap_err().code, error.code);
1562 }
1563
1564 #[tokio::test]
1565 async fn test_active_execution_drops_errors_without_subscribers() {
1566 let task = Task {
1567 id: "t-error-no-sub".into(),
1568 context_id: "c-error-no-sub".into(),
1569 status: TaskStatus {
1570 state: TaskState::Working,
1571 message: None,
1572 timestamp: None,
1573 },
1574 artifacts: None,
1575 history: None,
1576 metadata: None,
1577 };
1578 let active = ActiveExecution::new(task.clone());
1579 let error = A2AError::internal("execution failed");
1580
1581 active.publish(Err(&error), Some(task.clone())).await;
1583
1584 let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
1586 assert_eq!(snapshot_task.unwrap().id, "t-error-no-sub");
1587 assert_eq!(sequence, 1);
1588 assert!(receiver.try_recv().is_err());
1589 }
1590
1591 #[tokio::test]
1592 async fn test_active_execution_publishes_events_with_snapshot_updates() {
1593 let initial_task = Task {
1594 id: "t-snapshot-update".into(),
1595 context_id: "c-snapshot-update".into(),
1596 status: TaskStatus {
1597 state: TaskState::Working,
1598 message: None,
1599 timestamp: None,
1600 },
1601 artifacts: None,
1602 history: None,
1603 metadata: None,
1604 };
1605 let active = ActiveExecution::new(initial_task.clone());
1606 let mut receiver = active.subscribe();
1607
1608 let updated_task = Task {
1609 id: "t-snapshot-update".into(),
1610 context_id: "c-snapshot-update".into(),
1611 status: TaskStatus {
1612 state: TaskState::Working,
1613 message: None,
1614 timestamp: None,
1615 },
1616 artifacts: None,
1617 history: None,
1618 metadata: None,
1619 };
1620
1621 let event = StreamResponse::Task(updated_task.clone());
1622 active.publish(Ok(&event), Some(updated_task.clone())).await;
1623
1624 let published = receiver.recv().await.unwrap();
1626 assert_eq!(published.sequence, 1);
1627 match published.result {
1628 Ok(StreamResponse::Task(task)) => {
1629 assert_eq!(task.id, "t-snapshot-update");
1630 }
1631 _ => panic!("expected Ok(Task) event"),
1632 }
1633
1634 let (_, snapshot_task, sequence) = active.resubscribe().await;
1636 assert_eq!(snapshot_task.unwrap().id, "t-snapshot-update");
1637 assert_eq!(sequence, 1);
1638 }
1639
1640 #[tokio::test]
1641 async fn test_active_execution_broadcasts_to_multiple_subscribers() {
1642 let task = Task {
1643 id: "t-multi-sub".into(),
1644 context_id: "c-multi-sub".into(),
1645 status: TaskStatus {
1646 state: TaskState::Working,
1647 message: None,
1648 timestamp: None,
1649 },
1650 artifacts: None,
1651 history: None,
1652 metadata: None,
1653 };
1654 let active = ActiveExecution::new(task.clone());
1655
1656 let mut sub1 = active.subscribe();
1658 let mut sub2 = active.subscribe();
1659
1660 let event = StreamResponse::Task(task);
1661 active.publish(Ok(&event), None).await;
1662
1663 let evt1 = sub1.recv().await.unwrap();
1665 let evt2 = sub2.recv().await.unwrap();
1666
1667 assert_eq!(evt1.sequence, 1);
1668 assert_eq!(evt2.sequence, 1);
1669 assert!(matches!(evt1.result, Ok(StreamResponse::Task(_))));
1670 assert!(matches!(evt2.result, Ok(StreamResponse::Task(_))));
1671 }
1672
1673 #[tokio::test]
1674 async fn test_get_task_not_found() {
1675 let handler = make_handler();
1676 let params = ServiceParams::new();
1677 let req = GetTaskRequest {
1678 id: "nonexistent".into(),
1679 history_length: None,
1680 tenant: None,
1681 };
1682 let result = handler.get_task(¶ms, req).await;
1683 assert!(result.is_err());
1684 }
1685
1686 #[tokio::test]
1687 async fn test_get_task_after_send() {
1688 let handler = make_handler();
1689 let params = ServiceParams::new();
1690 let mut msg = make_message();
1691 msg.task_id = Some("t1".into());
1692 msg.context_id = Some("c1".into());
1693 let req = SendMessageRequest {
1694 message: msg,
1695 configuration: None,
1696 metadata: None,
1697 tenant: None,
1698 };
1699 handler.send_message(¶ms, req).await.unwrap();
1700 let task = handler
1701 .get_task(
1702 ¶ms,
1703 GetTaskRequest {
1704 id: "t1".into(),
1705 history_length: None,
1706 tenant: None,
1707 },
1708 )
1709 .await
1710 .unwrap();
1711 assert_eq!(task.id, "t1");
1712 }
1713
1714 #[tokio::test]
1715 async fn test_list_tasks() {
1716 let handler = make_handler();
1717 let params = ServiceParams::new();
1718 let req = ListTasksRequest {
1719 context_id: None,
1720 status: None,
1721 page_size: None,
1722 page_token: None,
1723 history_length: None,
1724 status_timestamp_after: None,
1725 include_artifacts: None,
1726 tenant: None,
1727 };
1728 let resp = handler.list_tasks(¶ms, req).await.unwrap();
1729 assert!(resp.tasks.is_empty());
1730 }
1731
1732 #[tokio::test]
1733 async fn test_cancel_task_not_found() {
1734 let handler = make_handler();
1735 let params = ServiceParams::new();
1736 let req = CancelTaskRequest {
1737 id: "nonexistent".into(),
1738 metadata: None,
1739 tenant: None,
1740 };
1741 let result = handler.cancel_task(¶ms, req).await;
1742 assert!(result.is_err());
1743 }
1744
1745 #[tokio::test]
1746 async fn test_cancel_task_working() {
1747 let handler = make_handler();
1748 let params = ServiceParams::new();
1749 let task = Task {
1751 id: "t2".into(),
1752 context_id: "c2".into(),
1753 status: TaskStatus {
1754 state: TaskState::Working,
1755 message: None,
1756 timestamp: None,
1757 },
1758 artifacts: None,
1759 history: None,
1760 metadata: None,
1761 };
1762 handler.task_store.create(task).await.unwrap();
1763 let req = CancelTaskRequest {
1764 id: "t2".into(),
1765 metadata: None,
1766 tenant: None,
1767 };
1768 let result = handler.cancel_task(¶ms, req).await.unwrap();
1769 assert_eq!(result.status.state, TaskState::Canceled);
1770 }
1771
1772 #[tokio::test]
1773 async fn test_cancel_task_already_completed() {
1774 let handler = make_handler();
1775 let params = ServiceParams::new();
1776 let task = Task {
1777 id: "t3".into(),
1778 context_id: "c3".into(),
1779 status: TaskStatus {
1780 state: TaskState::Completed,
1781 message: None,
1782 timestamp: None,
1783 },
1784 artifacts: None,
1785 history: None,
1786 metadata: None,
1787 };
1788 handler.task_store.create(task).await.unwrap();
1789 let req = CancelTaskRequest {
1790 id: "t3".into(),
1791 metadata: None,
1792 tenant: None,
1793 };
1794 let result = handler.cancel_task(¶ms, req).await;
1795 match result {
1796 Ok(_) => panic!("expected cancel of COMPLETED to fail"),
1797 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
1798 }
1799 }
1800
1801 #[tokio::test]
1802 async fn test_cancel_task_already_canceled_is_idempotent() {
1803 let handler = make_handler();
1804 let params = ServiceParams::new();
1805 let task = Task {
1806 id: "t-canceled".into(),
1807 context_id: "c-canceled".into(),
1808 status: TaskStatus {
1809 state: TaskState::Canceled,
1810 message: None,
1811 timestamp: None,
1812 },
1813 artifacts: None,
1814 history: None,
1815 metadata: None,
1816 };
1817 handler.task_store.create(task).await.unwrap();
1818 let req = CancelTaskRequest {
1819 id: "t-canceled".into(),
1820 metadata: None,
1821 tenant: None,
1822 };
1823 let result = handler.cancel_task(¶ms, req).await.unwrap();
1824 assert_eq!(result.status.state, TaskState::Canceled);
1825 }
1826
1827 #[tokio::test]
1828 async fn test_cancel_task_failed_and_rejected_are_not_cancelable() {
1829 let handler = make_handler();
1830 let params = ServiceParams::new();
1831 for (id, state) in [
1832 ("t-failed", TaskState::Failed),
1833 ("t-rejected", TaskState::Rejected),
1834 ] {
1835 handler
1836 .task_store
1837 .create(Task {
1838 id: id.into(),
1839 context_id: format!("c-{id}"),
1840 status: TaskStatus {
1841 state,
1842 message: None,
1843 timestamp: None,
1844 },
1845 artifacts: None,
1846 history: None,
1847 metadata: None,
1848 })
1849 .await
1850 .unwrap();
1851 let result = handler
1852 .cancel_task(
1853 ¶ms,
1854 CancelTaskRequest {
1855 id: id.into(),
1856 metadata: None,
1857 tenant: None,
1858 },
1859 )
1860 .await;
1861 match result {
1862 Ok(_) => panic!("expected cancel of {id} to fail"),
1863 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
1864 }
1865 }
1866 }
1867
1868 fn task_with_history(id: &str, context_id: &str, n: usize) -> Task {
1869 Task {
1870 id: id.into(),
1871 context_id: context_id.into(),
1872 status: TaskStatus {
1873 state: TaskState::Completed,
1874 message: None,
1875 timestamp: None,
1876 },
1877 artifacts: None,
1878 history: Some(
1879 (0..n)
1880 .map(|i| Message::new(Role::User, vec![Part::text(format!("m{i}"))]))
1881 .collect(),
1882 ),
1883 metadata: None,
1884 }
1885 }
1886
1887 #[tokio::test]
1888 async fn test_get_task_applies_history_length() {
1889 let handler = make_handler();
1890 let params = ServiceParams::new();
1891 handler
1892 .task_store
1893 .create(task_with_history("hist", "c", 3))
1894 .await
1895 .unwrap();
1896
1897 let full = handler
1898 .get_task(
1899 ¶ms,
1900 GetTaskRequest {
1901 id: "hist".into(),
1902 history_length: None,
1903 tenant: None,
1904 },
1905 )
1906 .await
1907 .unwrap();
1908 assert_eq!(full.history.as_ref().unwrap().len(), 3);
1909
1910 let last = handler
1911 .get_task(
1912 ¶ms,
1913 GetTaskRequest {
1914 id: "hist".into(),
1915 history_length: Some(1),
1916 tenant: None,
1917 },
1918 )
1919 .await
1920 .unwrap();
1921 assert_eq!(last.history.as_ref().unwrap().len(), 1);
1922
1923 for hl in [0, -1] {
1924 let empty = handler
1925 .get_task(
1926 ¶ms,
1927 GetTaskRequest {
1928 id: "hist".into(),
1929 history_length: Some(hl),
1930 tenant: None,
1931 },
1932 )
1933 .await
1934 .unwrap();
1935 assert!(
1936 empty.history.as_ref().unwrap().is_empty(),
1937 "history_length={hl} must clear history"
1938 );
1939 }
1940 }
1941
1942 fn reject_unauthenticated(
1943 _params: &ServiceParams,
1944 _task_id: Option<&str>,
1945 ) -> Result<(), A2AError> {
1946 Err(A2AError::invalid_request("authentication required"))
1947 }
1948
1949 #[tokio::test]
1950 async fn test_authorizer_rejects_handler_methods() {
1951 let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
1952 .with_authorizer(reject_unauthenticated)
1953 .with_push_config_store(InMemoryPushConfigStore::new());
1954 let params = ServiceParams::new();
1955
1956 assert!(
1957 handler
1958 .send_message(
1959 ¶ms,
1960 SendMessageRequest {
1961 message: make_message(),
1962 configuration: None,
1963 metadata: None,
1964 tenant: None,
1965 },
1966 )
1967 .await
1968 .is_err()
1969 );
1970 assert!(
1971 handler
1972 .get_task(
1973 ¶ms,
1974 GetTaskRequest {
1975 id: "t".into(),
1976 history_length: None,
1977 tenant: None,
1978 },
1979 )
1980 .await
1981 .is_err()
1982 );
1983 assert!(
1984 handler
1985 .list_tasks(
1986 ¶ms,
1987 ListTasksRequest {
1988 context_id: None,
1989 status: None,
1990 page_size: None,
1991 page_token: None,
1992 history_length: None,
1993 status_timestamp_after: None,
1994 include_artifacts: None,
1995 tenant: None,
1996 },
1997 )
1998 .await
1999 .is_err()
2000 );
2001 assert!(
2002 handler
2003 .cancel_task(
2004 ¶ms,
2005 CancelTaskRequest {
2006 id: "t".into(),
2007 metadata: None,
2008 tenant: None,
2009 },
2010 )
2011 .await
2012 .is_err()
2013 );
2014 assert!(
2015 handler
2016 .subscribe_to_task(
2017 ¶ms,
2018 SubscribeToTaskRequest {
2019 id: "t".into(),
2020 tenant: None,
2021 },
2022 )
2023 .await
2024 .is_err()
2025 );
2026 assert!(
2027 handler
2028 .create_push_config(
2029 ¶ms,
2030 TaskPushNotificationConfig {
2031 url: "https://example.com/push".into(),
2032 id: None,
2033 task_id: "t".into(),
2034 token: None,
2035 authentication: None,
2036 tenant: None,
2037 },
2038 )
2039 .await
2040 .is_err()
2041 );
2042 assert!(
2043 handler
2044 .get_push_config(
2045 ¶ms,
2046 GetTaskPushNotificationConfigRequest {
2047 task_id: "t".into(),
2048 id: "cfg".into(),
2049 tenant: None,
2050 },
2051 )
2052 .await
2053 .is_err()
2054 );
2055 assert!(
2056 handler
2057 .list_push_configs(
2058 ¶ms,
2059 ListTaskPushNotificationConfigsRequest {
2060 task_id: "t".into(),
2061 page_size: None,
2062 page_token: None,
2063 tenant: None,
2064 },
2065 )
2066 .await
2067 .is_err()
2068 );
2069 assert!(
2070 handler
2071 .delete_push_config(
2072 ¶ms,
2073 DeleteTaskPushNotificationConfigRequest {
2074 task_id: "t".into(),
2075 id: "cfg".into(),
2076 tenant: None,
2077 },
2078 )
2079 .await
2080 .is_err()
2081 );
2082 assert!(
2083 handler
2084 .get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
2085 .await
2086 .is_err()
2087 );
2088 }
2089
2090 #[tokio::test]
2091 async fn test_subscribe_requires_active_execution() {
2092 let handler = make_handler();
2093 let params = ServiceParams::new();
2094 let req = SubscribeToTaskRequest {
2095 id: "t1".into(),
2096 tenant: None,
2097 };
2098 let result = handler.subscribe_to_task(¶ms, req).await;
2099 match result {
2100 Ok(_) => panic!("expected subscribe_to_task to fail"),
2101 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_FOUND),
2102 }
2103 }
2104
2105 #[tokio::test]
2106 async fn test_send_message_return_immediately_allows_resubscribe() {
2107 use futures::StreamExt;
2108
2109 let (handler, release) = make_resumable_handler();
2110 let params = ServiceParams::new();
2111 let mut message = make_message();
2112 message.task_id = Some("t-resume".into());
2113 message.context_id = Some("c-resume".into());
2114
2115 let response = handler
2116 .send_message(
2117 ¶ms,
2118 SendMessageRequest {
2119 message,
2120 configuration: Some(SendMessageConfiguration {
2121 accepted_output_modes: None,
2122 task_push_notification_config: None,
2123 history_length: None,
2124 return_immediately: Some(true),
2125 }),
2126 metadata: None,
2127 tenant: None,
2128 },
2129 )
2130 .await
2131 .unwrap();
2132
2133 match response {
2134 SendMessageResponse::Task(task) => {
2135 assert_eq!(task.status.state, TaskState::Working);
2136 }
2137 _ => panic!("expected Task response"),
2138 }
2139
2140 let mut subscription = handler
2141 .subscribe_to_task(
2142 ¶ms,
2143 SubscribeToTaskRequest {
2144 id: "t-resume".into(),
2145 tenant: None,
2146 },
2147 )
2148 .await
2149 .unwrap();
2150
2151 let snapshot = subscription.next().await.unwrap().unwrap();
2152 match snapshot {
2153 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
2154 _ => panic!("expected snapshot task"),
2155 }
2156
2157 release.notify_waiters();
2158
2159 let terminal = subscription.next().await.unwrap().unwrap();
2160 match terminal {
2161 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
2162 _ => panic!("expected terminal task"),
2163 }
2164
2165 assert!(subscription.next().await.is_none());
2166
2167 let result = handler
2168 .subscribe_to_task(
2169 ¶ms,
2170 SubscribeToTaskRequest {
2171 id: "t-resume".into(),
2172 tenant: None,
2173 },
2174 )
2175 .await;
2176 match result {
2177 Ok(_) => panic!("expected subscribe_to_task to fail after completion"),
2178 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_FOUND),
2179 }
2180 }
2181
2182 #[tokio::test]
2183 async fn test_push_config_not_supported_without_store() {
2184 let handler = make_handler();
2185 let params = ServiceParams::new();
2186 let result = handler
2187 .create_push_config(
2188 ¶ms,
2189 TaskPushNotificationConfig {
2190 task_id: "t1".into(),
2191 url: "http://example.com".into(),
2192 id: None,
2193 token: None,
2194 authentication: None,
2195 tenant: None,
2196 },
2197 )
2198 .await;
2199 assert!(result.is_err());
2200 }
2201
2202 #[tokio::test]
2203 async fn test_list_push_configs_invalid_page_token_errors() {
2204 install_crypto_provider();
2205 let handler = make_handler_with_push_configs();
2206 let params = ServiceParams::new();
2207 handler
2208 .create_push_config(
2209 ¶ms,
2210 TaskPushNotificationConfig {
2211 task_id: "t1".into(),
2212 url: "https://example.com/callback".into(),
2213 id: Some("cfg-1".into()),
2214 token: None,
2215 authentication: None,
2216 tenant: None,
2217 },
2218 )
2219 .await
2220 .unwrap();
2221
2222 let result = handler
2223 .list_push_configs(
2224 ¶ms,
2225 ListTaskPushNotificationConfigsRequest {
2226 task_id: "t1".into(),
2227 page_size: None,
2228 page_token: Some("not-a-number".into()),
2229 tenant: None,
2230 },
2231 )
2232 .await;
2233 assert!(result.is_err());
2234 assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
2235 }
2236
2237 #[tokio::test]
2238 async fn test_list_tasks_page_size_cap_bounds_a_store_that_ignores_it() {
2239 install_crypto_provider();
2240 let handler = DefaultRequestHandler::new(
2243 EchoExecutor,
2244 PagingTaskStore::new(MAX_PAGE_SIZE * 3, false),
2245 );
2246
2247 let resp = handler
2248 .list_tasks(
2249 &ServiceParams::new(),
2250 ListTasksRequest {
2251 context_id: None,
2252 status: None,
2253 page_size: Some(1_000_000),
2254 page_token: None,
2255 history_length: None,
2256 include_artifacts: None,
2257 status_timestamp_after: None,
2258 tenant: None,
2259 },
2260 )
2261 .await
2262 .unwrap();
2263
2264 assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
2265 }
2266
2267 #[tokio::test]
2268 async fn test_paging_stores_delegate_everything_except_list() {
2269 install_crypto_provider();
2270 let tasks = PagingTaskStore::new(0, true);
2274 let task = Task {
2275 id: "t-delegated".into(),
2276 context_id: "c1".into(),
2277 status: TaskStatus {
2278 state: TaskState::Submitted,
2279 message: None,
2280 timestamp: None,
2281 },
2282 artifacts: None,
2283 history: None,
2284 metadata: None,
2285 };
2286 tasks.create(task.clone()).await.unwrap();
2287 assert_eq!(
2288 tasks
2289 .get("t-delegated")
2290 .await
2291 .unwrap()
2292 .map(|t| t.status.state),
2293 Some(TaskState::Submitted)
2294 );
2295 let mut done = task.clone();
2296 done.status.state = TaskState::Completed;
2297 tasks.update(done).await.unwrap();
2298 assert_eq!(
2299 tasks
2300 .get("t-delegated")
2301 .await
2302 .unwrap()
2303 .map(|t| t.status.state),
2304 Some(TaskState::Completed)
2305 );
2306
2307 let configs = PagingPushConfigStore::new(0);
2308 configs
2309 .save(TaskPushNotificationConfig {
2310 task_id: "t1".into(),
2311 url: "https://example.com/callback".into(),
2312 id: Some("cfg-1".into()),
2313 token: None,
2314 authentication: None,
2315 tenant: None,
2316 })
2317 .await
2318 .unwrap();
2319 assert_eq!(
2320 configs.get("t1", "cfg-1").await.unwrap().id.as_deref(),
2321 Some("cfg-1")
2322 );
2323 configs.delete("t1", "cfg-1").await.unwrap();
2324 assert!(configs.get("t1", "cfg-1").await.is_err());
2325 configs.delete_all("t1").await.unwrap();
2326 }
2327
2328 #[tokio::test]
2329 async fn test_list_tasks_page_size_cap_bounds_a_store_that_honours_it() {
2330 install_crypto_provider();
2331 let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
2335 let seen = Arc::clone(&store.seen_page_size);
2336 let handler = DefaultRequestHandler::new(EchoExecutor, store);
2337
2338 let resp = handler
2339 .list_tasks(
2340 &ServiceParams::new(),
2341 ListTasksRequest {
2342 context_id: None,
2343 status: None,
2344 page_size: Some(1_000_000),
2345 page_token: None,
2346 history_length: None,
2347 include_artifacts: None,
2348 status_timestamp_after: None,
2349 tenant: None,
2350 },
2351 )
2352 .await
2353 .unwrap();
2354
2355 assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
2356 assert_eq!(seen.load(Ordering::SeqCst), MAX_PAGE_SIZE as i32);
2360 }
2361
2362 #[tokio::test]
2363 async fn test_list_tasks_truncation_keeps_the_store_continuation_token() {
2364 install_crypto_provider();
2365 let handler = DefaultRequestHandler::new(
2370 EchoExecutor,
2371 PagingTaskStore::new(MAX_PAGE_SIZE * 3, false)
2372 .with_next_page_token("cursor-from-store"),
2373 );
2374
2375 let resp = handler
2376 .list_tasks(
2377 &ServiceParams::new(),
2378 ListTasksRequest {
2379 context_id: None,
2380 status: None,
2381 page_size: None,
2382 page_token: None,
2383 history_length: None,
2384 include_artifacts: None,
2385 status_timestamp_after: None,
2386 tenant: None,
2387 },
2388 )
2389 .await
2390 .unwrap();
2391
2392 assert_eq!(resp.tasks.len(), DEFAULT_PAGE_SIZE);
2393 assert_eq!(resp.next_page_token, "cursor-from-store");
2394 }
2395
2396 #[tokio::test]
2397 async fn test_list_tasks_honours_a_page_size_under_the_cap() {
2398 install_crypto_provider();
2399 let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
2402 let seen = Arc::clone(&store.seen_page_size);
2403 let handler = DefaultRequestHandler::new(EchoExecutor, store);
2404
2405 let resp = handler
2406 .list_tasks(
2407 &ServiceParams::new(),
2408 ListTasksRequest {
2409 context_id: None,
2410 status: None,
2411 page_size: Some(7),
2412 page_token: None,
2413 history_length: None,
2414 include_artifacts: None,
2415 status_timestamp_after: None,
2416 tenant: None,
2417 },
2418 )
2419 .await
2420 .unwrap();
2421
2422 assert_eq!(seen.load(Ordering::SeqCst), 7);
2423 assert_eq!(resp.tasks.len(), 7);
2424 }
2425
2426 #[tokio::test]
2427 async fn test_list_push_configs_page_size_cap_bounds_a_custom_store() {
2428 install_crypto_provider();
2429 let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
2433 .with_push_config_store(PagingPushConfigStore::new(MAX_PAGE_SIZE * 3));
2434
2435 let resp = handler
2436 .list_push_configs(
2437 &ServiceParams::new(),
2438 ListTaskPushNotificationConfigsRequest {
2439 task_id: "t1".into(),
2440 page_size: Some(1_000_000),
2441 page_token: None,
2442 tenant: None,
2443 },
2444 )
2445 .await
2446 .unwrap();
2447
2448 assert_eq!(resp.configs.len(), MAX_PAGE_SIZE);
2449 assert_eq!(resp.next_page_token.as_deref(), Some("100"));
2451 }
2452
2453 #[tokio::test]
2454 async fn test_list_push_configs_page_size_is_capped() {
2455 install_crypto_provider();
2456 let handler = make_handler_with_push_configs();
2457 let params = ServiceParams::new();
2458 for i in 0..40 {
2461 handler
2462 .create_push_config(
2463 ¶ms,
2464 TaskPushNotificationConfig {
2465 task_id: "t1".into(),
2466 url: format!("https://example.com/callback/{i}"),
2467 id: Some(format!("cfg-{i:03}")),
2468 token: None,
2469 authentication: None,
2470 tenant: None,
2471 },
2472 )
2473 .await
2474 .unwrap();
2475 }
2476
2477 let resp = handler
2483 .list_push_configs(
2484 ¶ms,
2485 ListTaskPushNotificationConfigsRequest {
2486 task_id: "t1".into(),
2487 page_size: Some(1000),
2488 page_token: None,
2489 tenant: None,
2490 },
2491 )
2492 .await
2493 .unwrap();
2494 assert_eq!(resp.configs.len(), 40);
2495 }
2496
2497 #[tokio::test]
2498 async fn test_push_config_crud_with_store() {
2499 install_crypto_provider();
2500 let handler = make_handler_with_push_configs();
2501 let params = ServiceParams::new();
2502
2503 let created = handler
2504 .create_push_config(
2505 ¶ms,
2506 TaskPushNotificationConfig {
2507 task_id: "t1".into(),
2508 url: "https://example.com/first".into(),
2509 id: Some("cfg-1".into()),
2510 token: None,
2511 authentication: None,
2512 tenant: Some("tenant-a".into()),
2513 },
2514 )
2515 .await
2516 .unwrap();
2517 assert_eq!(created.id.as_deref(), Some("cfg-1"));
2518
2519 handler
2520 .create_push_config(
2521 ¶ms,
2522 TaskPushNotificationConfig {
2523 task_id: "t1".into(),
2524 url: "https://example.com/second".into(),
2525 id: Some("cfg-2".into()),
2526 token: None,
2527 authentication: None,
2528 tenant: Some("tenant-a".into()),
2529 },
2530 )
2531 .await
2532 .unwrap();
2533
2534 let fetched = handler
2535 .get_push_config(
2536 ¶ms,
2537 GetTaskPushNotificationConfigRequest {
2538 task_id: "t1".into(),
2539 id: "cfg-1".into(),
2540 tenant: Some("tenant-a".into()),
2541 },
2542 )
2543 .await
2544 .unwrap();
2545 assert_eq!(fetched.url, "https://example.com/first");
2546
2547 let first_page = handler
2548 .list_push_configs(
2549 ¶ms,
2550 ListTaskPushNotificationConfigsRequest {
2551 task_id: "t1".into(),
2552 page_size: Some(1),
2553 page_token: Some("0".into()),
2554 tenant: Some("tenant-a".into()),
2555 },
2556 )
2557 .await
2558 .unwrap();
2559 assert_eq!(first_page.configs.len(), 1);
2560 assert_eq!(first_page.configs[0].id.as_deref(), Some("cfg-1"));
2561 assert_eq!(first_page.next_page_token.as_deref(), Some("1"));
2562
2563 let default_page = handler
2564 .list_push_configs(
2565 ¶ms,
2566 ListTaskPushNotificationConfigsRequest {
2567 task_id: "t1".into(),
2568 page_size: Some(0),
2569 page_token: None,
2570 tenant: Some("tenant-a".into()),
2571 },
2572 )
2573 .await
2574 .unwrap();
2575 assert_eq!(default_page.configs.len(), 2);
2576 assert!(default_page.next_page_token.is_none());
2577
2578 handler
2579 .delete_push_config(
2580 ¶ms,
2581 DeleteTaskPushNotificationConfigRequest {
2582 task_id: "t1".into(),
2583 id: "cfg-1".into(),
2584 tenant: Some("tenant-a".into()),
2585 },
2586 )
2587 .await
2588 .unwrap();
2589
2590 let remaining = handler
2591 .list_push_configs(
2592 ¶ms,
2593 ListTaskPushNotificationConfigsRequest {
2594 task_id: "t1".into(),
2595 page_size: None,
2596 page_token: None,
2597 tenant: Some("tenant-a".into()),
2598 },
2599 )
2600 .await
2601 .unwrap();
2602 assert_eq!(remaining.configs.len(), 1);
2603 assert_eq!(remaining.configs[0].id.as_deref(), Some("cfg-2"));
2604 }
2605
2606 #[tokio::test]
2607 async fn test_send_message_request_push_config_delivers_webhooks() {
2608 install_crypto_provider();
2609 let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
2610 let handler = make_push_delivery_handler();
2611 let params = ServiceParams::new();
2612 let mut message = make_message();
2613 message.task_id = Some("t-push-request".into());
2614 message.context_id = Some("c-push-request".into());
2615
2616 let response = handler
2617 .send_message(
2618 ¶ms,
2619 SendMessageRequest {
2620 message,
2621 configuration: Some(SendMessageConfiguration {
2622 accepted_output_modes: None,
2623 task_push_notification_config: Some(TaskPushNotificationConfig {
2624 task_id: String::new(),
2625 url: url.clone(),
2626 id: Some("cfg-request".into()),
2627 token: Some("notify-token".into()),
2628 authentication: Some(AuthenticationInfo {
2629 scheme: "bearer".into(),
2630 credentials: Some("secret-token".into()),
2631 }),
2632 tenant: None,
2633 }),
2634 history_length: None,
2635 return_immediately: None,
2636 }),
2637 metadata: None,
2638 tenant: None,
2639 },
2640 )
2641 .await
2642 .unwrap();
2643
2644 match response {
2645 SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
2646 _ => panic!("expected Task response"),
2647 }
2648
2649 let saved = handler
2650 .get_push_config(
2651 ¶ms,
2652 GetTaskPushNotificationConfigRequest {
2653 task_id: "t-push-request".into(),
2654 id: "cfg-request".into(),
2655 tenant: None,
2656 },
2657 )
2658 .await
2659 .unwrap();
2660 assert_eq!(saved.url, url);
2661
2662 let first = next_push(&mut receiver).await;
2663 assert_eq!(first.authorization.as_deref(), Some("Bearer secret-token"));
2664 assert_eq!(first.notification_token.as_deref(), Some("notify-token"));
2665 match first.event {
2666 StreamResponse::StatusUpdate(update) => {
2667 assert_eq!(update.task_id, "t-push-request");
2668 assert_eq!(update.status.state, TaskState::Working);
2669 }
2670 _ => panic!("expected status update push"),
2671 }
2672
2673 let second = next_push(&mut receiver).await;
2674 assert_eq!(second.authorization.as_deref(), Some("Bearer secret-token"));
2675 assert_eq!(second.notification_token.as_deref(), Some("notify-token"));
2676 match second.event {
2677 StreamResponse::Task(task) => {
2678 assert_eq!(task.id, "t-push-request");
2679 assert_eq!(task.status.state, TaskState::Completed);
2680 }
2681 _ => panic!("expected final task push"),
2682 }
2683
2684 shutdown_tx.send(()).unwrap();
2685 server.await.unwrap();
2686 }
2687
2688 #[tokio::test]
2689 async fn test_send_message_stored_push_config_delivers_webhooks() {
2690 install_crypto_provider();
2691 let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
2692 let handler = make_push_delivery_handler();
2693 let params = ServiceParams::new();
2694
2695 handler
2696 .create_push_config(
2697 ¶ms,
2698 TaskPushNotificationConfig {
2699 task_id: "t-push-stored".into(),
2700 url: url.clone(),
2701 id: Some("cfg-stored".into()),
2702 token: Some("stored-token".into()),
2703 authentication: None,
2704 tenant: None,
2705 },
2706 )
2707 .await
2708 .unwrap();
2709
2710 let mut message = make_message();
2711 message.task_id = Some("t-push-stored".into());
2712 message.context_id = Some("c-push-stored".into());
2713 let response = handler
2714 .send_message(
2715 ¶ms,
2716 SendMessageRequest {
2717 message,
2718 configuration: None,
2719 metadata: None,
2720 tenant: None,
2721 },
2722 )
2723 .await
2724 .unwrap();
2725
2726 match response {
2727 SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
2728 _ => panic!("expected Task response"),
2729 }
2730
2731 let first = next_push(&mut receiver).await;
2732 assert_eq!(first.notification_token.as_deref(), Some("stored-token"));
2733 match first.event {
2734 StreamResponse::StatusUpdate(update) => {
2735 assert_eq!(update.task_id, "t-push-stored");
2736 assert_eq!(update.status.state, TaskState::Working);
2737 }
2738 _ => panic!("expected status update push"),
2739 }
2740
2741 let second = next_push(&mut receiver).await;
2742 assert_eq!(second.notification_token.as_deref(), Some("stored-token"));
2743 match second.event {
2744 StreamResponse::Task(task) => {
2745 assert_eq!(task.id, "t-push-stored");
2746 assert_eq!(task.status.state, TaskState::Completed);
2747 }
2748 _ => panic!("expected final task push"),
2749 }
2750
2751 shutdown_tx.send(()).unwrap();
2752 server.await.unwrap();
2753 }
2754
2755 #[tokio::test]
2756 async fn test_extended_agent_card_not_configured() {
2757 let handler = make_handler();
2758 let params = ServiceParams::new();
2759 let result = handler
2760 .get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
2761 .await;
2762 assert!(result.is_err());
2763 }
2764
2765 #[test]
2766 fn test_with_capabilities() {
2767 let handler = make_handler().with_capabilities(AgentCapabilities {
2768 streaming: Some(true),
2769 push_notifications: None,
2770 extensions: None,
2771 extended_agent_card: None,
2772 });
2773 assert_eq!(handler.capabilities.streaming, Some(true));
2774 }
2775
2776 #[test]
2777 fn test_with_push_config_store_enables_push_capability() {
2778 install_crypto_provider();
2779 let handler = make_handler_with_push_configs();
2780 assert_eq!(handler.capabilities.push_notifications, Some(true));
2781 }
2782
2783 fn caps(streaming: Option<bool>, push: Option<bool>) -> AgentCapabilities {
2786 AgentCapabilities {
2787 streaming,
2788 push_notifications: push,
2789 extensions: None,
2790 extended_agent_card: None,
2791 }
2792 }
2793
2794 fn send_request() -> SendMessageRequest {
2795 SendMessageRequest {
2796 message: Message::new(Role::User, vec![Part::text("hi")]),
2797 configuration: None,
2798 metadata: None,
2799 tenant: None,
2800 }
2801 }
2802
2803 fn push_config(task_id: &str) -> TaskPushNotificationConfig {
2804 TaskPushNotificationConfig {
2805 task_id: task_id.into(),
2806 url: "https://example.com/callback".into(),
2807 id: Some("cfg-1".into()),
2808 token: None,
2809 authentication: None,
2810 tenant: None,
2811 }
2812 }
2813
2814 #[tokio::test]
2815 async fn test_streaming_refused_when_the_card_declares_it_absent() {
2816 let handler = make_handler().with_capabilities(caps(Some(false), None));
2817 let params = ServiceParams::new();
2818
2819 match handler
2822 .send_streaming_message(¶ms, send_request())
2823 .await
2824 {
2825 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
2826 Ok(_) => panic!("streaming must be refused when the card declares it absent"),
2827 }
2828
2829 match handler
2830 .subscribe_to_task(
2831 ¶ms,
2832 SubscribeToTaskRequest {
2833 id: "t1".into(),
2834 tenant: None,
2835 },
2836 )
2837 .await
2838 {
2839 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
2840 Ok(_) => panic!("subscribe must be refused when the card declares it absent"),
2841 }
2842 }
2843
2844 #[tokio::test]
2845 async fn test_push_refused_when_the_card_declares_it_absent() {
2846 let handler = make_handler_with_push_configs().with_capabilities(caps(None, Some(false)));
2849 let params = ServiceParams::new();
2850
2851 let error = handler
2852 .create_push_config(¶ms, push_config("t1"))
2853 .await
2854 .unwrap_err();
2855 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
2856
2857 let error = handler
2858 .list_push_configs(
2859 ¶ms,
2860 ListTaskPushNotificationConfigsRequest {
2861 task_id: "t1".into(),
2862 page_size: None,
2863 page_token: None,
2864 tenant: None,
2865 },
2866 )
2867 .await
2868 .unwrap_err();
2869 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
2870 }
2871
2872 #[tokio::test]
2876 async fn test_undeclared_capabilities_stay_permissive() {
2877 let handler = make_handler_with_push_configs();
2878 let params = ServiceParams::new();
2879
2880 assert!(
2881 handler
2882 .send_streaming_message(¶ms, send_request())
2883 .await
2884 .is_ok(),
2885 "an undeclared capability must not be treated as refused"
2886 );
2887 assert!(
2888 handler
2889 .create_push_config(¶ms, push_config("t1"))
2890 .await
2891 .is_ok()
2892 );
2893 }
2894
2895 #[tokio::test]
2899 async fn test_builder_order_does_not_change_the_declared_capability() {
2900 let params = ServiceParams::new();
2901
2902 for handler in [
2903 make_handler()
2904 .with_capabilities(caps(None, Some(false)))
2905 .with_push_config_store(InMemoryPushConfigStore::new()),
2906 make_handler()
2907 .with_push_config_store(InMemoryPushConfigStore::new())
2908 .with_capabilities(caps(None, Some(false))),
2909 ] {
2910 let error = handler
2911 .create_push_config(¶ms, push_config("t1"))
2912 .await
2913 .unwrap_err();
2914 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
2915 }
2916 }
2917
2918 #[tokio::test]
2923 async fn test_with_push_notifications_defaults_the_capability() {
2924 let params = ServiceParams::new();
2925
2926 let handler = make_handler().with_push_notifications(
2927 InMemoryPushConfigStore::new(),
2928 crate::HttpPushSender::new(None),
2929 );
2930 assert_eq!(handler.capabilities.push_notifications, Some(true));
2931 assert!(
2932 handler
2933 .create_push_config(¶ms, push_config("t1"))
2934 .await
2935 .is_ok()
2936 );
2937
2938 let handler = make_handler()
2939 .with_capabilities(caps(None, Some(false)))
2940 .with_push_notifications(
2941 InMemoryPushConfigStore::new(),
2942 crate::HttpPushSender::new(None),
2943 );
2944 assert_eq!(handler.capabilities.push_notifications, Some(false));
2945 let error = handler
2946 .create_push_config(¶ms, push_config("t1"))
2947 .await
2948 .unwrap_err();
2949 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
2950 }
2951
2952 #[tokio::test]
2955 async fn test_missing_store_still_reports_push_unsupported() {
2956 let handler = make_handler();
2957 let error = handler
2958 .create_push_config(&ServiceParams::new(), push_config("t1"))
2959 .await
2960 .unwrap_err();
2961 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
2962 }
2963
2964 #[tokio::test]
2970 async fn test_artifacts_do_not_aggregate_across_tasks_in_one_context() {
2971 let handler = make_handler();
2972 let params = ServiceParams::new();
2973 let context_id = "ctx-shared";
2974
2975 for (task_id, artifact_id) in [("t-first", "a-1"), ("t-second", "a-2")] {
2976 handler
2977 .task_store
2978 .create(Task {
2979 id: task_id.into(),
2980 context_id: context_id.into(),
2981 status: TaskStatus {
2982 state: TaskState::Working,
2983 message: None,
2984 timestamp: None,
2985 },
2986 artifacts: None,
2987 history: None,
2988 metadata: None,
2989 })
2990 .await
2991 .unwrap();
2992
2993 apply_event_to_task(
2994 handler.task_store.as_ref(),
2995 None,
2996 &StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
2997 task_id: task_id.into(),
2998 context_id: context_id.into(),
2999 artifact: Artifact {
3000 artifact_id: artifact_id.into(),
3001 name: None,
3002 description: None,
3003 parts: vec![Part::text("payload")],
3004 metadata: None,
3005 extensions: None,
3006 },
3007 append: None,
3008 last_chunk: None,
3009 metadata: None,
3010 }),
3011 )
3012 .await
3013 .unwrap();
3014 }
3015
3016 for (task_id, expected) in [("t-first", "a-1"), ("t-second", "a-2")] {
3017 let task = handler
3018 .get_task(
3019 ¶ms,
3020 GetTaskRequest {
3021 id: task_id.into(),
3022 history_length: None,
3023 tenant: None,
3024 },
3025 )
3026 .await
3027 .unwrap();
3028 let ids: Vec<&str> = task
3029 .artifacts
3030 .as_deref()
3031 .unwrap_or_default()
3032 .iter()
3033 .map(|a| a.artifact_id.as_str())
3034 .collect();
3035 assert_eq!(
3036 ids,
3037 vec![expected],
3038 "{task_id} must carry only its own artifact"
3039 );
3040 }
3041 }
3042 fn extended_card(name: &str) -> AgentCard {
3045 AgentCard {
3046 name: name.into(),
3047 description: "extended".into(),
3048 version: "1.0".into(),
3049 supported_interfaces: vec![],
3050 capabilities: AgentCapabilities::default(),
3051 default_input_modes: vec![],
3052 default_output_modes: vec![],
3053 skills: vec![],
3054 provider: None,
3055 documentation_url: None,
3056 icon_url: None,
3057 security_schemes: None,
3058 security_requirements: None,
3059 signatures: None,
3060 }
3061 }
3062
3063 fn extended_request() -> GetExtendedAgentCardRequest {
3064 GetExtendedAgentCardRequest { tenant: None }
3065 }
3066
3067 #[tokio::test]
3068 async fn test_extended_agent_card_is_served_when_configured() {
3069 let handler = make_handler().with_extended_agent_card(extended_card("secret skills"));
3070 let card = handler
3071 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3072 .await
3073 .unwrap();
3074 assert_eq!(card.name, "secret skills");
3075 }
3076
3077 #[tokio::test]
3080 async fn test_configuring_an_extended_card_declares_the_capability() {
3081 let handler = make_handler().with_extended_agent_card(extended_card("x"));
3082 assert_eq!(handler.capabilities.extended_agent_card, Some(true));
3083 }
3084
3085 #[tokio::test]
3089 async fn test_unconfigured_extended_card_reports_its_own_code() {
3090 let handler = make_handler().with_capabilities(AgentCapabilities {
3091 streaming: None,
3092 push_notifications: None,
3093 extensions: None,
3094 extended_agent_card: Some(true),
3095 });
3096 let error = handler
3097 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3098 .await
3099 .unwrap_err();
3100 assert_eq!(error.code, error_code::EXTENDED_CARD_NOT_CONFIGURED);
3101 assert_ne!(error.code, error_code::UNSUPPORTED_OPERATION);
3102 }
3103
3104 #[tokio::test]
3108 async fn test_capability_gate_precedes_the_not_configured_check() {
3109 let handler = make_handler()
3110 .with_extended_agent_card(extended_card("x"))
3111 .with_capabilities(AgentCapabilities {
3112 streaming: None,
3113 push_notifications: None,
3114 extensions: None,
3115 extended_agent_card: Some(false),
3116 });
3117 let error = handler
3118 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3119 .await
3120 .unwrap_err();
3121 assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION);
3122 }
3123
3124 #[tokio::test]
3127 async fn test_resolver_sees_the_caller_params() {
3128 let handler = make_handler().with_extended_agent_card_resolver(|params: &ServiceParams| {
3129 let privileged = params.contains_key("x-privileged");
3130 Ok(extended_card(if privileged { "full" } else { "limited" }))
3131 });
3132
3133 let card = handler
3134 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3135 .await
3136 .unwrap();
3137 assert_eq!(card.name, "limited");
3138
3139 let mut params = ServiceParams::new();
3140 params.insert("x-privileged".into(), vec!["yes".into()]);
3141 let card = handler
3142 .get_extended_agent_card(¶ms, extended_request())
3143 .await
3144 .unwrap();
3145 assert_eq!(card.name, "full");
3146 }
3147}