1use a2a::*;
5use async_trait::async_trait;
6use futures::{StreamExt, stream::BoxStream};
7use std::{collections::HashMap, sync::Arc};
8use tokio::sync::{RwLock, broadcast};
9
10use crate::middleware::ServiceParams;
11use crate::pagination::resolve_page_size;
12use crate::task_store::apply_history_length;
13
14const EXECUTION_BUFFER_CAPACITY: usize = 32;
15
16#[derive(Clone)]
17struct ExecutionEvent {
18 sequence: u64,
19 result: Result<StreamResponse, A2AError>,
20}
21
22#[derive(Default)]
23struct ActiveExecutionState {
24 sequence: u64,
25 snapshot_task: Option<Task>,
26}
27
28struct ActiveExecution {
29 sender: broadcast::Sender<ExecutionEvent>,
30 state: RwLock<ActiveExecutionState>,
31}
32
33impl ActiveExecution {
34 fn new(task: Task) -> Arc<Self> {
35 let (sender, _) = broadcast::channel(EXECUTION_BUFFER_CAPACITY);
36 Arc::new(Self {
37 sender,
38 state: RwLock::new(ActiveExecutionState {
39 sequence: 0,
40 snapshot_task: Some(task),
41 }),
42 })
43 }
44
45 fn subscribe(&self) -> broadcast::Receiver<ExecutionEvent> {
46 self.sender.subscribe()
47 }
48
49 async fn resubscribe(&self) -> (broadcast::Receiver<ExecutionEvent>, Option<Task>, u64) {
50 let receiver = self.sender.subscribe();
51 let state = self.state.read().await;
52 (receiver, state.snapshot_task.clone(), state.sequence)
53 }
54
55 async fn publish(
56 &self,
57 result: Result<&StreamResponse, &A2AError>,
58 snapshot_task: Option<Task>,
59 ) {
60 let sequence = {
61 let mut state = self.state.write().await;
62 if let Some(task) = snapshot_task {
63 state.snapshot_task = Some(task);
64 }
65 state.sequence += 1;
66 state.sequence
67 };
68
69 if self.sender.receiver_count() == 0 {
72 return;
73 }
74
75 let event = ExecutionEvent {
76 sequence,
77 result: result.cloned().map_err(Clone::clone),
78 };
79 let _ = self.sender.send(event);
80 }
81
82 async fn snapshot_is_terminal(&self) -> bool {
83 let state = self.state.read().await;
84 state
85 .snapshot_task
86 .as_ref()
87 .is_some_and(|task| task.status.state.is_terminal())
88 }
89}
90
91#[derive(Default)]
92struct ExecutionManager {
93 executions: RwLock<HashMap<TaskId, Arc<ActiveExecution>>>,
94}
95
96impl ExecutionManager {
97 async fn start(&self, task: Task) -> Result<Arc<ActiveExecution>, A2AError> {
98 let mut executions = self.executions.write().await;
99 if executions.contains_key(&task.id) {
100 return Err(A2AError::invalid_request(format!(
101 "task execution is already in progress: {}",
102 task.id
103 )));
104 }
105
106 let active = ActiveExecution::new(task.clone());
107 executions.insert(task.id.clone(), active.clone());
108 Ok(active)
109 }
110
111 async fn get(&self, task_id: &str) -> Option<Arc<ActiveExecution>> {
112 let executions = self.executions.read().await;
113 executions.get(task_id).cloned()
114 }
115
116 async fn resubscribe(
117 &self,
118 task_id: &str,
119 ) -> Option<(broadcast::Receiver<ExecutionEvent>, Option<Task>, u64)> {
120 let active = self.get(task_id).await?;
121 let (receiver, snapshot_task, sequence) = active.resubscribe().await;
122 Some((receiver, snapshot_task, sequence))
123 }
124
125 async fn finish(&self, task_id: &str, active: &Arc<ActiveExecution>) {
126 let mut executions = self.executions.write().await;
127 let should_remove = executions
128 .get(task_id)
129 .is_some_and(|current| Arc::ptr_eq(current, active));
130 if should_remove {
131 executions.remove(task_id);
132 }
133 }
134}
135
136struct ExecutionRuntime {
137 executor: Arc<dyn crate::AgentExecutor>,
138 task_store: Arc<dyn crate::TaskStore>,
139 push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
140 push_sender: Option<Arc<crate::HttpPushSender>>,
141 execution_manager: Arc<ExecutionManager>,
142}
143
144struct SubscriptionState {
145 receiver: broadcast::Receiver<ExecutionEvent>,
146 snapshot_task: Option<Task>,
147 min_sequence: u64,
148 done: bool,
149}
150
151fn is_terminal_event(event: &StreamResponse) -> bool {
152 match event {
153 StreamResponse::Task(task) => task.status.state.is_terminal(),
154 StreamResponse::StatusUpdate(update) => update.status.state.is_terminal(),
155 _ => false,
156 }
157}
158
159fn task_id_for_event(event: &StreamResponse) -> Option<&str> {
160 match event {
161 StreamResponse::Task(task) => Some(&task.id),
162 StreamResponse::StatusUpdate(update) => Some(&update.task_id),
163 StreamResponse::ArtifactUpdate(update) => Some(&update.task_id),
164 StreamResponse::Message(_) => None,
165 }
166}
167
168fn should_interrupt_non_streaming(
169 req: &SendMessageRequest,
170 event: &StreamResponse,
171) -> Option<TaskId> {
172 if req
173 .configuration
174 .as_ref()
175 .and_then(|config| config.return_immediately)
176 .unwrap_or(false)
177 {
178 return match event {
179 StreamResponse::Message(_) => None,
180 _ => task_id_for_event(event).map(ToOwned::to_owned),
181 };
182 }
183
184 match event {
185 StreamResponse::Task(task) if task.status.state == TaskState::AuthRequired => {
186 Some(task.id.clone())
187 }
188 StreamResponse::StatusUpdate(update) if update.status.state == TaskState::AuthRequired => {
189 Some(update.task_id.clone())
190 }
191 _ => None,
192 }
193}
194
195async fn send_push_notifications(
196 task_id: &str,
197 push_config_store: Option<&dyn crate::PushConfigStore>,
198 push_sender: Option<&crate::HttpPushSender>,
199 event: &StreamResponse,
200) -> Result<(), A2AError> {
201 let Some(push_config_store) = push_config_store else {
202 return Ok(());
203 };
204 let Some(push_sender) = push_sender else {
205 return Ok(());
206 };
207
208 let configs = push_config_store.list(task_id).await?;
209 for config in configs {
210 push_sender.send_push(&config, event.clone()).await?;
211 }
212
213 Ok(())
214}
215
216async fn save_task(task_store: &dyn crate::TaskStore, mut task: Task) -> Result<Task, A2AError> {
217 if task.status.timestamp.is_none() {
222 task.status.timestamp = Some(chrono::Utc::now());
223 }
224 match task_store.update(task.clone()).await {
225 Ok(_) => Ok(task),
226 Err(error) if error.code == error_code::TASK_NOT_FOUND => {
227 task_store.create(task.clone()).await?;
228 Ok(task)
229 }
230 Err(error) => Err(error),
231 }
232}
233
234fn backfill_status_timestamp(event: &mut StreamResponse) {
242 let status = match event {
243 StreamResponse::Task(task) => &mut task.status,
244 StreamResponse::StatusUpdate(update) => &mut update.status,
245 StreamResponse::ArtifactUpdate(_) | StreamResponse::Message(_) => return,
246 };
247 if status.timestamp.is_none() {
248 status.timestamp = Some(chrono::Utc::now());
249 }
250}
251
252async fn apply_event_to_task(
253 task_store: &dyn crate::TaskStore,
254 current_task: Option<Task>,
255 event: &StreamResponse,
256) -> Result<Option<Task>, A2AError> {
257 match event {
258 StreamResponse::Task(task) => save_task(task_store, task.clone()).await.map(Some),
259 StreamResponse::StatusUpdate(update) => {
260 let mut task = match current_task {
261 Some(task) => task,
262 None => task_store
263 .get(&update.task_id)
264 .await?
265 .ok_or_else(|| A2AError::task_not_found(&update.task_id))?,
266 };
267 task.status = update.status.clone();
268 save_task(task_store, task).await.map(Some)
269 }
270 StreamResponse::ArtifactUpdate(update) => {
271 let mut task = match current_task {
272 Some(task) => task,
273 None => task_store
274 .get(&update.task_id)
275 .await?
276 .ok_or_else(|| A2AError::task_not_found(&update.task_id))?,
277 };
278 let artifacts = task.artifacts.get_or_insert_with(Vec::new);
279 let existing = artifacts
280 .iter_mut()
281 .find(|artifact| artifact.artifact_id == update.artifact.artifact_id);
282 if update.append.unwrap_or(false) {
283 let artifact = existing.ok_or_else(A2AError::invalid_agent_response)?;
284 artifact.parts.extend(update.artifact.parts.iter().cloned());
285 if let Some(metadata) = &update.artifact.metadata {
286 artifact
287 .metadata
288 .get_or_insert_with(Default::default)
289 .extend(metadata.clone());
290 }
291 } else if let Some(artifact) = existing {
292 *artifact = update.artifact.clone();
293 } else {
294 artifacts.push(update.artifact.clone());
295 }
296 save_task(task_store, task).await.map(Some)
297 }
298 StreamResponse::Message(_) => Ok(current_task),
299 }
300}
301
302fn subscription_stream(
303 receiver: broadcast::Receiver<ExecutionEvent>,
304 snapshot_task: Option<Task>,
305 min_sequence: u64,
306) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
307 Box::pin(futures::stream::unfold(
308 SubscriptionState {
309 receiver,
310 snapshot_task,
311 min_sequence,
312 done: false,
313 },
314 |mut state| async move {
315 if state.done {
316 return None;
317 }
318
319 if let Some(task) = state.snapshot_task.take() {
320 if task.status.state.is_terminal() {
321 state.done = true;
322 }
323 return Some((Ok(StreamResponse::Task(task)), state));
324 }
325
326 loop {
327 match state.receiver.recv().await {
328 Ok(event) if event.sequence <= state.min_sequence => continue,
329 Ok(event) => {
330 state.min_sequence = event.sequence;
331 match &event.result {
332 Ok(item) if is_terminal_event(item) => state.done = true,
333 Err(_) => state.done = true,
334 _ => {}
335 }
336 return Some((event.result.clone(), state));
337 }
338 Err(broadcast::error::RecvError::Lagged(_)) => {
339 state.done = true;
340 return Some((
341 Err(crate::sanitized_internal_error(
342 "subscription fell behind active execution",
343 )),
344 state,
345 ));
346 }
347 Err(broadcast::error::RecvError::Closed) => return None,
348 }
349 }
350 },
351 ))
352}
353
354async fn drive_execution(
355 runtime: ExecutionRuntime,
356 task_id: TaskId,
357 active: Arc<ActiveExecution>,
358 exec_ctx: crate::ExecutorContext,
359) {
360 let mut current_task = {
361 let state = active.state.read().await;
362 state.snapshot_task.clone()
363 };
364 let mut stream = runtime.executor.execute(exec_ctx);
365
366 while let Some(result) = stream.next().await {
367 if active.snapshot_is_terminal().await {
368 break;
369 }
370
371 match result {
372 Ok(mut event) => {
373 backfill_status_timestamp(&mut event);
374 match apply_event_to_task(runtime.task_store.as_ref(), current_task.clone(), &event)
375 .await
376 {
377 Ok(updated_task) => {
378 current_task = updated_task.clone();
379 if let Err(error) = send_push_notifications(
380 &task_id,
381 runtime.push_config_store.as_deref(),
382 runtime.push_sender.as_deref(),
383 &event,
384 )
385 .await
386 {
387 active.publish(Err(&error), None).await;
388 break;
389 }
390 active.publish(Ok(&event), updated_task).await;
391 if is_terminal_event(&event) {
392 break;
393 }
394 }
395 Err(error) => {
396 active.publish(Err(&error), None).await;
397 break;
398 }
399 }
400 }
401 Err(error) => {
402 active.publish(Err(&error), None).await;
403 break;
404 }
405 }
406 }
407
408 runtime.execution_manager.finish(&task_id, &active).await;
409}
410
411#[async_trait]
418pub trait RequestHandler: Send + Sync + 'static {
419 async fn send_message(
420 &self,
421 params: &ServiceParams,
422 req: SendMessageRequest,
423 ) -> Result<SendMessageResponse, A2AError>;
424
425 async fn send_streaming_message(
426 &self,
427 params: &ServiceParams,
428 req: SendMessageRequest,
429 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
430
431 async fn get_task(&self, params: &ServiceParams, req: GetTaskRequest)
432 -> Result<Task, A2AError>;
433
434 async fn list_tasks(
435 &self,
436 params: &ServiceParams,
437 req: ListTasksRequest,
438 ) -> Result<ListTasksResponse, A2AError>;
439
440 async fn cancel_task(
441 &self,
442 params: &ServiceParams,
443 req: CancelTaskRequest,
444 ) -> Result<Task, A2AError>;
445
446 async fn subscribe_to_task(
447 &self,
448 params: &ServiceParams,
449 req: SubscribeToTaskRequest,
450 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
451
452 async fn create_push_config(
453 &self,
454 params: &ServiceParams,
455 req: TaskPushNotificationConfig,
456 ) -> Result<TaskPushNotificationConfig, A2AError>;
457
458 async fn get_push_config(
459 &self,
460 params: &ServiceParams,
461 req: GetTaskPushNotificationConfigRequest,
462 ) -> Result<TaskPushNotificationConfig, A2AError>;
463
464 async fn list_push_configs(
465 &self,
466 params: &ServiceParams,
467 req: ListTaskPushNotificationConfigsRequest,
468 ) -> Result<ListTaskPushNotificationConfigsResponse, A2AError>;
469
470 async fn delete_push_config(
471 &self,
472 params: &ServiceParams,
473 req: DeleteTaskPushNotificationConfigRequest,
474 ) -> Result<(), A2AError>;
475
476 async fn get_extended_agent_card(
477 &self,
478 params: &ServiceParams,
479 req: GetExtendedAgentCardRequest,
480 ) -> Result<AgentCard, A2AError>;
481}
482
483pub trait RequestAuthorizer: Send + Sync {
490 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError>;
491}
492
493impl<F> RequestAuthorizer for F
494where
495 F: Fn(&ServiceParams, Option<&str>) -> Result<(), A2AError> + Send + Sync,
496{
497 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
498 self(params, task_id)
499 }
500}
501
502pub struct DefaultRequestHandler {
505 executor: Arc<dyn crate::AgentExecutor>,
506 task_store: Arc<dyn crate::TaskStore>,
507 execution_manager: Arc<ExecutionManager>,
508 push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
509 push_sender: Option<Arc<crate::HttpPushSender>>,
510 capabilities: AgentCapabilities,
511 default_input_modes: Vec<String>,
512 authorizer: Option<Arc<dyn RequestAuthorizer>>,
513 extended_agent_card: Option<Arc<dyn ExtendedAgentCardResolver>>,
514}
515
516pub trait ExtendedAgentCardResolver: Send + Sync {
525 fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError>;
526}
527
528impl<F> ExtendedAgentCardResolver for F
529where
530 F: Fn(&ServiceParams) -> Result<AgentCard, A2AError> + Send + Sync,
531{
532 fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError> {
533 self(params)
534 }
535}
536
537impl DefaultRequestHandler {
538 pub fn new(executor: impl crate::AgentExecutor, task_store: impl crate::TaskStore) -> Self {
539 DefaultRequestHandler {
540 executor: Arc::new(executor),
541 task_store: Arc::new(task_store),
542 execution_manager: Arc::new(ExecutionManager::default()),
543 push_config_store: None,
544 push_sender: None,
545 capabilities: AgentCapabilities::default(),
546 default_input_modes: Vec::new(),
547 authorizer: None,
548 extended_agent_card: None,
549 }
550 }
551
552 pub fn with_extended_agent_card(mut self, card: AgentCard) -> Self {
558 self.capabilities.extended_agent_card.get_or_insert(true);
559 self.with_extended_agent_card_resolver(move |_: &ServiceParams| Ok(card.clone()))
560 }
561
562 pub fn with_extended_agent_card_resolver(
565 mut self,
566 resolver: impl ExtendedAgentCardResolver + 'static,
567 ) -> Self {
568 self.capabilities.extended_agent_card.get_or_insert(true);
569 self.extended_agent_card = Some(Arc::new(resolver));
570 self
571 }
572
573 pub fn with_authorizer(mut self, authorizer: impl RequestAuthorizer + 'static) -> Self {
575 self.authorizer = Some(Arc::new(authorizer));
576 self
577 }
578
579 fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
580 match &self.authorizer {
581 Some(authorizer) => authorizer.authorize(params, task_id),
582 None => Ok(()),
583 }
584 }
585
586 pub fn with_push_config_store(
587 mut self,
588 push_config_store: impl crate::PushConfigStore,
589 ) -> Self {
590 self.push_sender = Some(Arc::new(crate::HttpPushSender::new(None)));
591 self.push_config_store = Some(Arc::new(push_config_store));
592 self.capabilities.push_notifications.get_or_insert(true);
596 self
597 }
598
599 pub fn with_push_notifications(
600 mut self,
601 push_config_store: impl crate::PushConfigStore,
602 push_sender: crate::HttpPushSender,
603 ) -> Self {
604 self.push_config_store = Some(Arc::new(push_config_store));
605 self.push_sender = Some(Arc::new(push_sender));
606 self.capabilities.push_notifications.get_or_insert(true);
607 self
608 }
609
610 pub fn with_capabilities(mut self, capabilities: AgentCapabilities) -> Self {
611 self.capabilities = capabilities;
612 self
613 }
614
615 pub fn with_default_input_modes(mut self, modes: impl Into<Vec<String>>) -> Self {
619 self.default_input_modes = modes.into();
620 self
621 }
622
623 fn validate_content_types(&self, message: &Message) -> Result<(), A2AError> {
627 if self.default_input_modes.is_empty() {
628 return Ok(());
629 }
630 let supported = message.parts.iter().all(|part| {
631 part.media_type.as_deref().is_none_or(|media_type| {
632 media_type.is_empty() || self.default_input_modes.iter().any(|m| m == media_type)
633 })
634 });
635 if supported {
636 Ok(())
637 } else {
638 Err(A2AError::content_type_not_supported())
639 }
640 }
641
642 fn require_capability(
649 declared: Option<bool>,
650 unsupported: impl FnOnce() -> A2AError,
651 ) -> Result<(), A2AError> {
652 match declared {
653 Some(false) => Err(unsupported()),
654 _ => Ok(()),
655 }
656 }
657
658 fn require_streaming(&self) -> Result<(), A2AError> {
659 Self::require_capability(self.capabilities.streaming, || {
660 A2AError::unsupported_operation("the agent card does not declare streaming")
661 })
662 }
663
664 fn require_push_notifications(&self) -> Result<(), A2AError> {
665 Self::require_capability(
666 self.capabilities.push_notifications,
667 A2AError::push_notification_not_supported,
668 )
669 }
670
671 fn push_config_store_checked(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
676 self.require_push_notifications()?;
677 self.push_config_store()
678 }
679
680 fn push_config_store(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
681 self.push_config_store
682 .as_deref()
683 .ok_or_else(A2AError::push_notification_not_supported)
684 }
685
686 async fn load_task(&self, task_id: &str) -> Result<Task, A2AError> {
687 self.task_store
688 .get(task_id)
689 .await?
690 .ok_or_else(|| A2AError::task_not_found(task_id))
691 }
692
693 async fn prepare_task_for_execution(
694 &self,
695 req: &SendMessageRequest,
696 ) -> Result<(Task, Option<Task>, String), A2AError> {
697 let task_id = match req.message.task_id.clone() {
701 Some(id) => id,
702 None => new_task_id(),
703 };
704 let stored = self.task_store.get(&task_id).await?;
705 if stored.is_none() && req.message.task_id.is_some() {
706 return Err(A2AError::task_not_found(&task_id));
707 }
708 if stored
709 .as_ref()
710 .is_some_and(|task| task.status.state.is_terminal())
711 {
712 return Err(A2AError::unsupported_operation(format!(
713 "task {task_id} has already reached a terminal state"
714 )));
715 }
716
717 let context_id = match stored.as_ref() {
720 Some(task) => {
721 if let Some(requested) = req.message.context_id.as_deref() {
722 if requested != task.context_id {
723 return Err(A2AError::invalid_params(format!(
724 "contextId {requested} does not match task {task_id}, \
725 which belongs to context {}",
726 task.context_id
727 )));
728 }
729 }
730 task.context_id.clone()
731 }
732 None => req
733 .message
734 .context_id
735 .clone()
736 .unwrap_or_else(new_context_id),
737 };
738
739 let task = if let Some(mut existing) = stored.clone() {
740 existing
747 .history
748 .get_or_insert_with(Vec::new)
749 .push(req.message.clone());
750 existing
751 } else {
752 let task = Task {
753 id: task_id,
754 context_id: context_id.clone(),
755 status: TaskStatus {
756 state: TaskState::Submitted,
757 message: None,
758 timestamp: Some(chrono::Utc::now()),
759 },
760 artifacts: None,
761 history: Some(vec![req.message.clone()]),
762 metadata: None,
763 };
764 self.task_store.create(task.clone()).await?;
765 task
766 };
767
768 Ok((task, stored, context_id))
769 }
770
771 async fn save_request_push_config(
772 &self,
773 task_id: &str,
774 req: &SendMessageRequest,
775 ) -> Result<(), A2AError> {
776 let Some(config) = req
777 .configuration
778 .as_ref()
779 .and_then(|configuration| configuration.task_push_notification_config.clone())
780 else {
781 return Ok(());
782 };
783
784 let mut config = config;
785 config.task_id = task_id.to_string();
786 if config.tenant.is_none() {
787 config.tenant = req.tenant.clone();
788 }
789 self.push_config_store_checked()?.save(config).await?;
790 Ok(())
791 }
792
793 async fn start_execution(
794 &self,
795 params: &ServiceParams,
796 req: SendMessageRequest,
797 include_task_snapshot_in_context: bool,
798 ) -> Result<(TaskId, BoxStream<'static, Result<StreamResponse, A2AError>>), A2AError> {
799 self.validate_content_types(&req.message)?;
800 let (task, stored_task, context_id) = self.prepare_task_for_execution(&req).await?;
801 let task_id = task.id.clone();
802 self.save_request_push_config(&task_id, &req).await?;
803 let active = self.execution_manager.start(task.clone()).await?;
804 let receiver = active.subscribe();
805
806 let exec_ctx = crate::ExecutorContext {
807 message: Some(req.message),
808 task_id: task_id.clone(),
809 stored_task: if include_task_snapshot_in_context {
810 Some(task.clone())
811 } else {
812 stored_task
813 },
814 context_id,
815 metadata: req.metadata,
816 user: None,
817 service_params: params.clone(),
818 tenant: req.tenant,
819 };
820
821 let runtime = ExecutionRuntime {
822 executor: Arc::clone(&self.executor),
823 task_store: Arc::clone(&self.task_store),
824 push_config_store: self.push_config_store.clone(),
825 push_sender: self.push_sender.clone(),
826 execution_manager: Arc::clone(&self.execution_manager),
827 };
828 let active_execution = active.clone();
829 let execution_task_id = task_id.clone();
830
831 tokio::spawn(async move {
832 drive_execution(runtime, execution_task_id, active_execution, exec_ctx).await;
833 });
834
835 Ok((task_id, subscription_stream(receiver, None, 0)))
836 }
837}
838
839#[async_trait]
840impl RequestHandler for DefaultRequestHandler {
841 async fn send_message(
842 &self,
843 params: &ServiceParams,
844 req: SendMessageRequest,
845 ) -> Result<SendMessageResponse, A2AError> {
846 self.authorize(params, req.message.task_id.as_deref())?;
847 let history_length = req
851 .configuration
852 .as_ref()
853 .and_then(|configuration| configuration.history_length);
854 let (task_id, mut stream) = self.start_execution(params, req.clone(), true).await?;
855 let mut last_event = None;
856
857 while let Some(item) = stream.next().await {
858 let event = item?;
859
860 if let Some(interrupt_task_id) = should_interrupt_non_streaming(&req, &event) {
861 let mut task = self.load_task(&interrupt_task_id).await?;
862 apply_history_length(&mut task, history_length);
863 return Ok(SendMessageResponse::Task(task));
864 }
865
866 match event {
867 StreamResponse::Message(message) => {
868 return Ok(SendMessageResponse::Message(message));
869 }
870 other => last_event = Some(other),
871 }
872 }
873
874 match last_event {
875 Some(StreamResponse::Task(mut task)) => {
876 apply_history_length(&mut task, history_length);
877 Ok(SendMessageResponse::Task(task))
878 }
879 Some(StreamResponse::StatusUpdate(update)) => {
880 let mut task = self.load_task(&update.task_id).await?;
881 apply_history_length(&mut task, history_length);
882 Ok(SendMessageResponse::Task(task))
883 }
884 Some(StreamResponse::ArtifactUpdate(update)) => {
885 let mut task = self.load_task(&update.task_id).await?;
886 apply_history_length(&mut task, history_length);
887 Ok(SendMessageResponse::Task(task))
888 }
889 Some(StreamResponse::Message(message)) => Ok(SendMessageResponse::Message(message)),
890 None => {
891 let mut task = self.load_task(&task_id).await?;
892 apply_history_length(&mut task, history_length);
893 Ok(SendMessageResponse::Task(task))
894 }
895 }
896 }
897
898 async fn send_streaming_message(
899 &self,
900 params: &ServiceParams,
901 req: SendMessageRequest,
902 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
903 self.authorize(params, req.message.task_id.as_deref())?;
904 self.require_streaming()?;
905 let (_, stream) = self.start_execution(params, req, false).await?;
906 Ok(stream)
907 }
908
909 async fn get_task(
910 &self,
911 params: &ServiceParams,
912 req: GetTaskRequest,
913 ) -> Result<Task, A2AError> {
914 self.authorize(params, Some(&req.id))?;
915 let mut task = self
916 .task_store
917 .get(&req.id)
918 .await?
919 .ok_or_else(|| A2AError::task_not_found(&req.id))?;
920 apply_history_length(&mut task, req.history_length);
921 Ok(task)
922 }
923
924 async fn list_tasks(
937 &self,
938 params: &ServiceParams,
939 req: ListTasksRequest,
940 ) -> Result<ListTasksResponse, A2AError> {
941 self.authorize(params, None)?;
942 let page_size = resolve_page_size(req.page_size);
943 let req = ListTasksRequest {
944 page_size: Some(page_size as i32),
945 ..req
946 };
947 let mut response = self.task_store.list(&req).await?;
948 response.tasks.truncate(page_size);
949 for task in &mut response.tasks {
950 apply_history_length(task, req.history_length);
951 }
952 Ok(response)
953 }
954
955 async fn cancel_task(
956 &self,
957 params: &ServiceParams,
958 req: CancelTaskRequest,
959 ) -> Result<Task, A2AError> {
960 self.authorize(params, Some(&req.id))?;
961 let task = self.load_task(&req.id).await?;
962
963 if task.status.state == TaskState::Canceled {
967 return Ok(task);
968 }
969 if task.status.state.is_terminal() {
970 return Err(A2AError::task_not_cancelable(&req.id));
971 }
972
973 let task = match self.task_store.begin_cancel(&req.id).await {
976 Ok(task) => task,
977 Err(error) if error.code == error_code::TASK_NOT_CANCELABLE => {
978 let raced = self.load_task(&req.id).await?;
983 if raced.status.state == TaskState::Canceled {
984 return Ok(raced);
985 }
986 return Err(error);
987 }
988 Err(error) => return Err(error),
989 };
990
991 let active_execution = self.execution_manager.get(&req.id).await;
992
993 let exec_ctx = crate::ExecutorContext {
994 message: None,
995 task_id: req.id.clone(),
996 stored_task: Some(task.clone()),
997 context_id: task.context_id.clone(),
998 metadata: req.metadata,
999 user: None,
1000 service_params: params.clone(),
1001 tenant: req.tenant,
1002 };
1003
1004 let mut stream = self.executor.cancel(exec_ctx);
1005 let mut current_task = Some(task);
1006
1007 while let Some(event) = stream.next().await {
1008 match event {
1009 Ok(mut event) => {
1010 backfill_status_timestamp(&mut event);
1011 let updated_task =
1012 apply_event_to_task(self.task_store.as_ref(), current_task.clone(), &event)
1013 .await?;
1014 current_task = updated_task.clone();
1015
1016 send_push_notifications(
1017 &req.id,
1018 self.push_config_store.as_deref(),
1019 self.push_sender.as_deref(),
1020 &event,
1021 )
1022 .await?;
1023
1024 if let Some(active) = &active_execution {
1025 active.publish(Ok(&event), updated_task).await;
1026 }
1027
1028 if is_terminal_event(&event) {
1029 break;
1030 }
1031 }
1032 Err(error) => {
1033 if let Some(active) = &active_execution {
1034 active.publish(Err(&error), None).await;
1035 }
1036 return Err(error);
1037 }
1038 }
1039 }
1040
1041 if let Some(active) = active_execution {
1042 self.execution_manager.finish(&req.id, &active).await;
1043 }
1044
1045 self.load_task(&req.id).await
1046 }
1047
1048 async fn subscribe_to_task(
1049 &self,
1050 params: &ServiceParams,
1051 req: SubscribeToTaskRequest,
1052 ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
1053 self.authorize(params, Some(&req.id))?;
1054 self.require_streaming()?;
1055 if let Some((receiver, snapshot_task, sequence)) =
1056 self.execution_manager.resubscribe(&req.id).await
1057 {
1058 return Ok(subscription_stream(receiver, snapshot_task, sequence));
1059 }
1060
1061 if let Some(task) = self.task_store.get(&req.id).await? {
1066 if task.status.state.is_terminal() {
1067 return Err(A2AError::unsupported_operation(format!(
1068 "task {} has already reached a terminal state",
1069 req.id
1070 )));
1071 }
1072 }
1073 Err(A2AError::task_not_found(&req.id))
1074 }
1075
1076 async fn create_push_config(
1077 &self,
1078 params: &ServiceParams,
1079 req: TaskPushNotificationConfig,
1080 ) -> Result<TaskPushNotificationConfig, A2AError> {
1081 self.authorize(params, Some(&req.task_id))?;
1082 self.push_config_store_checked()?.save(req).await
1083 }
1084
1085 async fn get_push_config(
1086 &self,
1087 params: &ServiceParams,
1088 req: GetTaskPushNotificationConfigRequest,
1089 ) -> Result<TaskPushNotificationConfig, A2AError> {
1090 self.authorize(params, Some(&req.task_id))?;
1091 self.push_config_store_checked()?
1092 .get(&req.task_id, &req.id)
1093 .await
1094 }
1095
1096 async fn list_push_configs(
1097 &self,
1098 params: &ServiceParams,
1099 req: ListTaskPushNotificationConfigsRequest,
1100 ) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
1101 self.authorize(params, Some(&req.task_id))?;
1102 let mut configs = self.push_config_store_checked()?.list(&req.task_id).await?;
1103 configs.sort_by(|left, right| left.id.cmp(&right.id));
1104
1105 let page_size = resolve_page_size(req.page_size);
1110 let start = if let Some(ref token) = req.page_token {
1111 token
1112 .parse::<usize>()
1113 .map_err(|_| A2AError::invalid_params("invalid page token"))?
1114 } else {
1115 0
1116 };
1117 let start = start.min(configs.len());
1118 let end = (start + page_size).min(configs.len());
1119 let next_page_token = (end < configs.len()).then(|| end.to_string());
1120
1121 Ok(ListTaskPushNotificationConfigsResponse {
1122 configs: configs[start..end].to_vec(),
1123 next_page_token,
1124 })
1125 }
1126
1127 async fn delete_push_config(
1128 &self,
1129 params: &ServiceParams,
1130 req: DeleteTaskPushNotificationConfigRequest,
1131 ) -> Result<(), A2AError> {
1132 self.authorize(params, Some(&req.task_id))?;
1133 self.push_config_store_checked()?
1134 .delete(&req.task_id, &req.id)
1135 .await
1136 }
1137
1138 async fn get_extended_agent_card(
1139 &self,
1140 params: &ServiceParams,
1141 _req: GetExtendedAgentCardRequest,
1142 ) -> Result<AgentCard, A2AError> {
1143 self.authorize(params, None)?;
1144 Self::require_capability(self.capabilities.extended_agent_card, || {
1145 A2AError::unsupported_operation(
1146 "the agent card does not declare an extended agent card",
1147 )
1148 })?;
1149 match &self.extended_agent_card {
1150 Some(resolver) => resolver.resolve(params),
1151 None => Err(A2AError::extended_card_not_configured()),
1156 }
1157 }
1158}
1159
1160#[cfg(test)]
1161mod tests {
1162 use super::*;
1163 use crate::executor::ExecutorContext;
1164 use crate::pagination::{DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE};
1165 use crate::push::{InMemoryPushConfigStore, PushConfigStore};
1166 use crate::task_store::{InMemoryTaskStore, TaskStore, TaskVersion};
1167
1168 use crate::test_util::install_crypto_provider;
1169 use axum::{
1170 Router,
1171 body::Bytes,
1172 extract::State,
1173 http::{HeaderMap, StatusCode, header},
1174 routing::post,
1175 };
1176 use futures::stream;
1177 use std::sync::Arc;
1178 use std::sync::atomic::{AtomicI32, Ordering};
1179 use tokio::{
1180 net::TcpListener,
1181 sync::{Notify, mpsc, oneshot},
1182 time::{Duration, timeout},
1183 };
1184
1185 struct EchoExecutor;
1186
1187 struct PushEventExecutor;
1188
1189 struct ResumableExecutor {
1190 release: Arc<Notify>,
1191 }
1192
1193 #[derive(Debug)]
1194 struct CapturedPush {
1195 authorization: Option<String>,
1196 notification_token: Option<String>,
1197 event: StreamResponse,
1198 }
1199
1200 impl crate::AgentExecutor for EchoExecutor {
1201 fn execute(
1202 &self,
1203 ctx: ExecutorContext,
1204 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1205 let task = Task {
1206 id: ctx.task_id.clone(),
1207 context_id: ctx.context_id.clone(),
1208 status: TaskStatus {
1209 state: TaskState::Completed,
1210 message: ctx.message,
1211 timestamp: None,
1212 },
1213 artifacts: None,
1214 history: None,
1215 metadata: None,
1216 };
1217 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1218 }
1219
1220 fn cancel(
1221 &self,
1222 ctx: ExecutorContext,
1223 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1224 let task = Task {
1225 id: ctx.task_id.clone(),
1226 context_id: ctx.context_id.clone(),
1227 status: TaskStatus {
1228 state: TaskState::Canceled,
1229 message: None,
1230 timestamp: None,
1231 },
1232 artifacts: None,
1233 history: None,
1234 metadata: None,
1235 };
1236 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1237 }
1238 }
1239
1240 impl crate::AgentExecutor for PushEventExecutor {
1241 fn execute(
1242 &self,
1243 ctx: ExecutorContext,
1244 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1245 let working = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1246 task_id: ctx.task_id.clone(),
1247 context_id: ctx.context_id.clone(),
1248 status: TaskStatus {
1249 state: TaskState::Working,
1250 message: None,
1251 timestamp: None,
1252 },
1253 metadata: None,
1254 });
1255 let completed = StreamResponse::Task(Task {
1256 id: ctx.task_id.clone(),
1257 context_id: ctx.context_id.clone(),
1258 status: TaskStatus {
1259 state: TaskState::Completed,
1260 message: ctx.message,
1261 timestamp: None,
1262 },
1263 artifacts: None,
1264 history: ctx.stored_task.and_then(|task| task.history),
1265 metadata: None,
1266 });
1267
1268 Box::pin(stream::iter(vec![Ok(working), Ok(completed)]))
1269 }
1270
1271 fn cancel(
1272 &self,
1273 ctx: ExecutorContext,
1274 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1275 let task = Task {
1276 id: ctx.task_id.clone(),
1277 context_id: ctx.context_id.clone(),
1278 status: TaskStatus {
1279 state: TaskState::Canceled,
1280 message: None,
1281 timestamp: None,
1282 },
1283 artifacts: None,
1284 history: None,
1285 metadata: None,
1286 };
1287 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1288 }
1289 }
1290
1291 impl crate::AgentExecutor for ResumableExecutor {
1292 fn execute(
1293 &self,
1294 ctx: ExecutorContext,
1295 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1296 let release = self.release.clone();
1297 let working_task = Task {
1298 id: ctx.task_id.clone(),
1299 context_id: ctx.context_id.clone(),
1300 status: TaskStatus {
1301 state: TaskState::Working,
1302 message: None,
1303 timestamp: None,
1304 },
1305 artifacts: None,
1306 history: ctx.stored_task.and_then(|task| task.history),
1307 metadata: None,
1308 };
1309 let completed_task = Task {
1310 id: working_task.id.clone(),
1311 context_id: working_task.context_id.clone(),
1312 status: TaskStatus {
1313 state: TaskState::Completed,
1314 message: None,
1315 timestamp: None,
1316 },
1317 artifacts: None,
1318 history: working_task.history.clone(),
1319 metadata: None,
1320 };
1321
1322 Box::pin(stream::unfold(
1323 (Some(working_task), Some((release, completed_task))),
1324 |(working_task, completion)| async move {
1325 if let Some(task) = working_task {
1326 return Some((Ok(StreamResponse::Task(task)), (None, completion)));
1327 }
1328
1329 if let Some((release, task)) = completion {
1330 release.notified().await;
1331 return Some((Ok(StreamResponse::Task(task)), (None, None)));
1332 }
1333
1334 None
1335 },
1336 ))
1337 }
1338
1339 fn cancel(
1340 &self,
1341 ctx: ExecutorContext,
1342 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1343 let task = Task {
1344 id: ctx.task_id.clone(),
1345 context_id: ctx.context_id.clone(),
1346 status: TaskStatus {
1347 state: TaskState::Canceled,
1348 message: None,
1349 timestamp: None,
1350 },
1351 artifacts: None,
1352 history: None,
1353 metadata: None,
1354 };
1355 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1356 }
1357 }
1358
1359 struct MultiTurnExecutor;
1364
1365 impl crate::AgentExecutor for MultiTurnExecutor {
1366 fn execute(
1367 &self,
1368 ctx: ExecutorContext,
1369 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1370 let done = ctx
1371 .message
1372 .as_ref()
1373 .and_then(Message::text)
1374 .is_some_and(|text| text.trim().eq_ignore_ascii_case("done"));
1375 let state = if done {
1376 TaskState::Completed
1377 } else {
1378 TaskState::InputRequired
1379 };
1380 let update = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1381 task_id: ctx.task_id.clone(),
1382 context_id: ctx.context_id.clone(),
1383 status: TaskStatus {
1384 state,
1385 message: None,
1386 timestamp: None,
1387 },
1388 metadata: None,
1389 });
1390 Box::pin(stream::once(async move { Ok(update) }))
1391 }
1392
1393 fn cancel(
1394 &self,
1395 ctx: ExecutorContext,
1396 ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
1397 let task = Task {
1398 id: ctx.task_id.clone(),
1399 context_id: ctx.context_id.clone(),
1400 status: TaskStatus {
1401 state: TaskState::Canceled,
1402 message: None,
1403 timestamp: None,
1404 },
1405 artifacts: None,
1406 history: None,
1407 metadata: None,
1408 };
1409 Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
1410 }
1411 }
1412
1413 struct PagingTaskStore {
1421 inner: InMemoryTaskStore,
1422 count: usize,
1423 honour_page_size: bool,
1427 seen_page_size: Arc<AtomicI32>,
1431 next_page_token: String,
1434 }
1435
1436 impl PagingTaskStore {
1437 fn new(count: usize, honour_page_size: bool) -> Self {
1438 Self {
1439 inner: InMemoryTaskStore::new(),
1440 count,
1441 honour_page_size,
1442 seen_page_size: Arc::new(AtomicI32::new(-1)),
1443 next_page_token: String::new(),
1444 }
1445 }
1446
1447 fn with_next_page_token(mut self, token: &str) -> Self {
1448 self.next_page_token = token.to_string();
1449 self
1450 }
1451 }
1452
1453 #[async_trait]
1454 impl TaskStore for PagingTaskStore {
1455 async fn create(&self, task: Task) -> Result<TaskVersion, A2AError> {
1456 self.inner.create(task).await
1457 }
1458
1459 async fn update(&self, task: Task) -> Result<TaskVersion, A2AError> {
1460 self.inner.update(task).await
1461 }
1462
1463 async fn get(&self, task_id: &str) -> Result<Option<Task>, A2AError> {
1464 self.inner.get(task_id).await
1465 }
1466
1467 async fn begin_cancel(&self, task_id: &str) -> Result<Task, A2AError> {
1468 self.inner.begin_cancel(task_id).await
1469 }
1470
1471 async fn list(&self, req: &ListTasksRequest) -> Result<ListTasksResponse, A2AError> {
1472 self.seen_page_size
1473 .store(req.page_size.unwrap_or(-1), Ordering::SeqCst);
1474 let served = if self.honour_page_size {
1475 req.page_size.unwrap_or(i32::MAX).max(0) as usize
1476 } else {
1477 self.count
1478 }
1479 .min(self.count);
1480 let tasks = (0..served)
1481 .map(|i| Task {
1482 id: format!("t-{i:04}"),
1483 context_id: "c1".into(),
1484 status: TaskStatus {
1485 state: TaskState::Completed,
1486 message: None,
1487 timestamp: None,
1488 },
1489 artifacts: None,
1490 history: None,
1491 metadata: None,
1492 })
1493 .collect::<Vec<_>>();
1494 let total = tasks.len() as i32;
1495 Ok(ListTasksResponse {
1496 tasks,
1497 next_page_token: self.next_page_token.clone(),
1498 page_size: total,
1499 total_size: total,
1500 })
1501 }
1502 }
1503
1504 struct PagingPushConfigStore {
1510 inner: InMemoryPushConfigStore,
1511 count: usize,
1512 }
1513
1514 impl PagingPushConfigStore {
1515 fn new(count: usize) -> Self {
1516 Self {
1517 inner: InMemoryPushConfigStore::new(),
1518 count,
1519 }
1520 }
1521 }
1522
1523 #[async_trait]
1524 impl crate::PushConfigStore for PagingPushConfigStore {
1525 async fn save(
1526 &self,
1527 config: TaskPushNotificationConfig,
1528 ) -> Result<TaskPushNotificationConfig, A2AError> {
1529 self.inner.save(config).await
1530 }
1531
1532 async fn get(
1533 &self,
1534 task_id: &str,
1535 config_id: &str,
1536 ) -> Result<TaskPushNotificationConfig, A2AError> {
1537 self.inner.get(task_id, config_id).await
1538 }
1539
1540 async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError> {
1541 Ok((0..self.count)
1542 .map(|i| TaskPushNotificationConfig {
1543 task_id: task_id.to_string(),
1544 url: format!("https://example.com/hook/{i}"),
1545 id: Some(format!("cfg-{i:04}")),
1546 token: None,
1547 authentication: None,
1548 tenant: None,
1549 })
1550 .collect())
1551 }
1552
1553 async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError> {
1554 self.inner.delete(task_id, config_id).await
1555 }
1556
1557 async fn delete_all(&self, task_id: &str) -> Result<(), A2AError> {
1558 self.inner.delete_all(task_id).await
1559 }
1560 }
1561
1562 async fn seed_task(handler: &DefaultRequestHandler, task_id: &str, context_id: &str) {
1565 handler
1566 .task_store
1567 .create(Task {
1568 id: task_id.into(),
1569 context_id: context_id.into(),
1570 status: TaskStatus {
1571 state: TaskState::Submitted,
1572 message: None,
1573 timestamp: None,
1574 },
1575 artifacts: None,
1576 history: None,
1577 metadata: None,
1578 })
1579 .await
1580 .unwrap();
1581 }
1582
1583 fn make_handler() -> DefaultRequestHandler {
1584 DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
1585 }
1586
1587 fn make_handler_with_push_configs() -> DefaultRequestHandler {
1588 DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
1589 .with_push_config_store(InMemoryPushConfigStore::new())
1590 }
1591
1592 fn make_push_delivery_handler() -> DefaultRequestHandler {
1593 DefaultRequestHandler::new(PushEventExecutor, InMemoryTaskStore::new())
1596 .with_push_notifications(
1597 InMemoryPushConfigStore::new(),
1598 crate::HttpPushSender::new(Some(crate::HttpPushSenderConfig {
1599 validate_urls: false,
1600 ..Default::default()
1601 })),
1602 )
1603 }
1604
1605 fn make_resumable_handler() -> (DefaultRequestHandler, Arc<Notify>) {
1606 let release = Arc::new(Notify::new());
1607 (
1608 DefaultRequestHandler::new(
1609 ResumableExecutor {
1610 release: release.clone(),
1611 },
1612 InMemoryTaskStore::new(),
1613 ),
1614 release,
1615 )
1616 }
1617
1618 fn make_message() -> Message {
1619 Message::new(Role::User, vec![Part::text("hello")])
1620 }
1621
1622 async fn capture_push(
1623 State(sender): State<mpsc::UnboundedSender<CapturedPush>>,
1624 headers: HeaderMap,
1625 body: Bytes,
1626 ) -> StatusCode {
1627 let captured = CapturedPush {
1628 authorization: headers
1629 .get(header::AUTHORIZATION)
1630 .and_then(|value| value.to_str().ok())
1631 .map(ToOwned::to_owned),
1632 notification_token: headers
1633 .get("A2A-Notification-Token")
1634 .and_then(|value| value.to_str().ok())
1635 .map(ToOwned::to_owned),
1636 event: serde_json::from_slice(&body).unwrap(),
1637 };
1638 sender.send(captured).unwrap();
1639 StatusCode::ACCEPTED
1640 }
1641
1642 async fn start_push_webhook() -> (
1643 String,
1644 mpsc::UnboundedReceiver<CapturedPush>,
1645 oneshot::Sender<()>,
1646 tokio::task::JoinHandle<()>,
1647 ) {
1648 let (sender, receiver) = mpsc::unbounded_channel();
1649 let app = Router::new()
1650 .route("/", post(capture_push))
1651 .with_state(sender);
1652 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1653 let address = listener.local_addr().unwrap();
1654 let (shutdown_tx, shutdown_rx) = oneshot::channel();
1655 let server = tokio::spawn(async move {
1656 axum::serve(listener, app)
1657 .with_graceful_shutdown(async {
1658 let _ = shutdown_rx.await;
1659 })
1660 .await
1661 .unwrap();
1662 });
1663
1664 (format!("http://{address}/"), receiver, shutdown_tx, server)
1665 }
1666
1667 async fn next_push(receiver: &mut mpsc::UnboundedReceiver<CapturedPush>) -> CapturedPush {
1668 timeout(Duration::from_secs(5), receiver.recv())
1669 .await
1670 .unwrap()
1671 .unwrap()
1672 }
1673
1674 #[tokio::test]
1675 async fn test_send_message_creates_task() {
1676 let handler = make_handler();
1677 let params = ServiceParams::new();
1678 let req = SendMessageRequest {
1679 message: make_message(),
1680 configuration: None,
1681 metadata: None,
1682 tenant: None,
1683 };
1684 let result = handler.send_message(¶ms, req).await.unwrap();
1685 match result {
1686 SendMessageResponse::Task(task) => {
1687 assert_eq!(task.status.state, TaskState::Completed);
1688 }
1689 _ => panic!("expected Task response"),
1690 }
1691 }
1692
1693 #[tokio::test]
1694 async fn test_send_streaming_message() {
1695 use futures::StreamExt;
1696 let handler = make_handler();
1697 let params = ServiceParams::new();
1698 let req = SendMessageRequest {
1699 message: make_message(),
1700 configuration: None,
1701 metadata: None,
1702 tenant: None,
1703 };
1704 let mut stream = handler.send_streaming_message(¶ms, req).await.unwrap();
1705 let event = stream.next().await.unwrap().unwrap();
1706 assert!(matches!(event, StreamResponse::Task(_)));
1707 }
1708
1709 #[tokio::test]
1714 async fn test_streamed_task_event_has_a_backfilled_timestamp() {
1715 use futures::StreamExt;
1716 let handler = make_handler();
1717 let params = ServiceParams::new();
1718 let req = SendMessageRequest {
1719 message: make_message(),
1720 configuration: None,
1721 metadata: None,
1722 tenant: None,
1723 };
1724 let mut stream = handler.send_streaming_message(¶ms, req).await.unwrap();
1725 let event = stream.next().await.unwrap().unwrap();
1726 match event {
1727 StreamResponse::Task(task) => assert!(
1728 task.status.timestamp.is_some(),
1729 "a live-streamed task event must carry a backfilled timestamp"
1730 ),
1731 other => panic!("unexpected event: {other:?}"),
1732 }
1733 }
1734
1735 #[test]
1738 fn test_backfill_status_timestamp_ignores_events_without_a_status() {
1739 let mut artifact_update = StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
1740 task_id: "t1".into(),
1741 context_id: "c1".into(),
1742 artifact: Artifact {
1743 artifact_id: "a1".into(),
1744 name: None,
1745 description: None,
1746 parts: vec![Part::text("payload")],
1747 metadata: None,
1748 extensions: None,
1749 },
1750 append: None,
1751 last_chunk: None,
1752 metadata: None,
1753 });
1754 let before = artifact_update.clone();
1755 backfill_status_timestamp(&mut artifact_update);
1756 assert_eq!(artifact_update, before);
1757
1758 let mut message = StreamResponse::Message(make_message());
1759 let before = message.clone();
1760 backfill_status_timestamp(&mut message);
1761 assert_eq!(message, before);
1762 }
1763
1764 #[tokio::test]
1768 async fn test_status_update_backfills_missing_timestamp() {
1769 let store = InMemoryTaskStore::new();
1770 store
1771 .create(Task {
1772 id: "t-backfill".into(),
1773 context_id: "ctx".into(),
1774 status: TaskStatus {
1775 state: TaskState::Submitted,
1776 message: None,
1777 timestamp: Some(chrono::Utc::now()),
1778 },
1779 artifacts: None,
1780 history: None,
1781 metadata: None,
1782 })
1783 .await
1784 .unwrap();
1785
1786 let updated = apply_event_to_task(
1787 &store,
1788 None,
1789 &StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1790 task_id: "t-backfill".into(),
1791 context_id: "ctx".into(),
1792 status: TaskStatus {
1793 state: TaskState::Working,
1794 message: None,
1795 timestamp: None,
1796 },
1797 metadata: None,
1798 }),
1799 )
1800 .await
1801 .unwrap()
1802 .expect("status update yields a task");
1803
1804 assert_eq!(updated.status.state, TaskState::Working);
1805 assert!(
1806 updated.status.timestamp.is_some(),
1807 "missing status timestamp must be backfilled"
1808 );
1809 }
1810
1811 #[tokio::test]
1814 async fn test_status_update_preserves_supplied_timestamp() {
1815 let store = InMemoryTaskStore::new();
1816 store
1817 .create(Task {
1818 id: "t-preserve".into(),
1819 context_id: "ctx".into(),
1820 status: TaskStatus {
1821 state: TaskState::Submitted,
1822 message: None,
1823 timestamp: None,
1824 },
1825 artifacts: None,
1826 history: None,
1827 metadata: None,
1828 })
1829 .await
1830 .unwrap();
1831
1832 let supplied = chrono::DateTime::parse_from_rfc3339("2020-01-02T03:04:05Z")
1833 .unwrap()
1834 .with_timezone(&chrono::Utc);
1835
1836 let updated = apply_event_to_task(
1837 &store,
1838 None,
1839 &StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1840 task_id: "t-preserve".into(),
1841 context_id: "ctx".into(),
1842 status: TaskStatus {
1843 state: TaskState::Working,
1844 message: None,
1845 timestamp: Some(supplied),
1846 },
1847 metadata: None,
1848 }),
1849 )
1850 .await
1851 .unwrap()
1852 .expect("status update yields a task");
1853
1854 assert_eq!(updated.status.timestamp, Some(supplied));
1855 }
1856
1857 #[tokio::test]
1861 async fn test_task_snapshot_backfills_missing_timestamp() {
1862 let store = InMemoryTaskStore::new();
1863
1864 let returned = apply_event_to_task(
1865 &store,
1866 None,
1867 &StreamResponse::Task(Task {
1868 id: "t-snapshot".into(),
1869 context_id: "ctx".into(),
1870 status: TaskStatus {
1871 state: TaskState::Completed,
1872 message: None,
1873 timestamp: None,
1874 },
1875 artifacts: None,
1876 history: None,
1877 metadata: None,
1878 }),
1879 )
1880 .await
1881 .unwrap()
1882 .expect("task event yields a task");
1883
1884 assert!(
1885 returned.status.timestamp.is_some(),
1886 "missing status timestamp must be backfilled on the task-snapshot path"
1887 );
1888 let persisted = store.get("t-snapshot").await.unwrap().unwrap();
1889 assert!(
1890 persisted.status.timestamp.is_some(),
1891 "the persisted task must carry a timestamp"
1892 );
1893 }
1894
1895 #[tokio::test]
1898 async fn test_task_snapshot_preserves_supplied_timestamp() {
1899 let store = InMemoryTaskStore::new();
1900
1901 let supplied = chrono::DateTime::parse_from_rfc3339("2021-06-07T08:09:10Z")
1902 .unwrap()
1903 .with_timezone(&chrono::Utc);
1904
1905 let returned = apply_event_to_task(
1906 &store,
1907 None,
1908 &StreamResponse::Task(Task {
1909 id: "t-snapshot-kept".into(),
1910 context_id: "ctx".into(),
1911 status: TaskStatus {
1912 state: TaskState::Completed,
1913 message: None,
1914 timestamp: Some(supplied),
1915 },
1916 artifacts: None,
1917 history: None,
1918 metadata: None,
1919 }),
1920 )
1921 .await
1922 .unwrap()
1923 .expect("task event yields a task");
1924
1925 assert_eq!(returned.status.timestamp, Some(supplied));
1926 let persisted = store.get("t-snapshot-kept").await.unwrap().unwrap();
1927 assert_eq!(persisted.status.timestamp, Some(supplied));
1928 }
1929
1930 #[tokio::test]
1931 async fn test_stream_disconnect_preserves_execution_for_resubscription() {
1932 use futures::StreamExt;
1933
1934 let (handler, release) = make_resumable_handler();
1935 let params = ServiceParams::new();
1936 let mut message = make_message();
1937 message.task_id = Some("t-disconnect".into());
1938 message.context_id = Some("c-disconnect".into());
1939 seed_task(&handler, "t-disconnect", "c-disconnect").await;
1940
1941 let mut stream = handler
1942 .send_streaming_message(
1943 ¶ms,
1944 SendMessageRequest {
1945 message,
1946 configuration: None,
1947 metadata: None,
1948 tenant: None,
1949 },
1950 )
1951 .await
1952 .unwrap();
1953
1954 let initial = stream.next().await.unwrap().unwrap();
1955 match initial {
1956 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
1957 _ => panic!("expected initial working task"),
1958 }
1959 drop(stream);
1960
1961 let mut resumed = handler
1962 .subscribe_to_task(
1963 ¶ms,
1964 SubscribeToTaskRequest {
1965 id: "t-disconnect".into(),
1966 tenant: None,
1967 },
1968 )
1969 .await
1970 .unwrap();
1971
1972 let snapshot = resumed.next().await.unwrap().unwrap();
1973 match snapshot {
1974 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
1975 _ => panic!("expected working task snapshot"),
1976 }
1977
1978 release.notify_waiters();
1979
1980 let terminal = resumed.next().await.unwrap().unwrap();
1981 match terminal {
1982 StreamResponse::Task(task) => {
1983 assert_eq!(task.status.state, TaskState::Completed)
1984 }
1985 _ => panic!("expected terminal task"),
1986 }
1987 assert!(resumed.next().await.is_none());
1988 }
1989
1990 #[tokio::test]
1991 async fn test_active_execution_keeps_snapshot_without_subscribers() {
1992 let task = Task {
1993 id: "t-buffering".into(),
1994 context_id: "c-buffering".into(),
1995 status: TaskStatus {
1996 state: TaskState::Working,
1997 message: None,
1998 timestamp: None,
1999 },
2000 artifacts: None,
2001 history: None,
2002 metadata: None,
2003 };
2004 let active = ActiveExecution::new(task.clone());
2005
2006 let event = StreamResponse::Task(task.clone());
2007 active.publish(Ok(&event), Some(task)).await;
2008
2009 let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
2010 assert_eq!(snapshot_task.unwrap().id, "t-buffering");
2011 assert_eq!(sequence, 1);
2012 assert!(receiver.try_recv().is_err());
2013 }
2014
2015 #[tokio::test]
2016 async fn test_active_execution_publishes_errors_to_subscribers() {
2017 let task = Task {
2018 id: "t-buffering-error".into(),
2019 context_id: "c-buffering-error".into(),
2020 status: TaskStatus {
2021 state: TaskState::Working,
2022 message: None,
2023 timestamp: None,
2024 },
2025 artifacts: None,
2026 history: None,
2027 metadata: None,
2028 };
2029 let active = ActiveExecution::new(task);
2030 let mut receiver = active.subscribe();
2031 let error = A2AError::internal("execution failed");
2032
2033 active.publish(Err(&error), None).await;
2034
2035 let event = receiver.recv().await.unwrap();
2036 assert_eq!(event.sequence, 1);
2037 assert_eq!(event.result.unwrap_err().code, error.code);
2038 }
2039
2040 #[tokio::test]
2041 async fn test_active_execution_drops_errors_without_subscribers() {
2042 let task = Task {
2043 id: "t-error-no-sub".into(),
2044 context_id: "c-error-no-sub".into(),
2045 status: TaskStatus {
2046 state: TaskState::Working,
2047 message: None,
2048 timestamp: None,
2049 },
2050 artifacts: None,
2051 history: None,
2052 metadata: None,
2053 };
2054 let active = ActiveExecution::new(task.clone());
2055 let error = A2AError::internal("execution failed");
2056
2057 active.publish(Err(&error), Some(task.clone())).await;
2059
2060 let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
2062 assert_eq!(snapshot_task.unwrap().id, "t-error-no-sub");
2063 assert_eq!(sequence, 1);
2064 assert!(receiver.try_recv().is_err());
2065 }
2066
2067 #[tokio::test]
2068 async fn test_active_execution_publishes_events_with_snapshot_updates() {
2069 let initial_task = Task {
2070 id: "t-snapshot-update".into(),
2071 context_id: "c-snapshot-update".into(),
2072 status: TaskStatus {
2073 state: TaskState::Working,
2074 message: None,
2075 timestamp: None,
2076 },
2077 artifacts: None,
2078 history: None,
2079 metadata: None,
2080 };
2081 let active = ActiveExecution::new(initial_task.clone());
2082 let mut receiver = active.subscribe();
2083
2084 let updated_task = Task {
2085 id: "t-snapshot-update".into(),
2086 context_id: "c-snapshot-update".into(),
2087 status: TaskStatus {
2088 state: TaskState::Working,
2089 message: None,
2090 timestamp: None,
2091 },
2092 artifacts: None,
2093 history: None,
2094 metadata: None,
2095 };
2096
2097 let event = StreamResponse::Task(updated_task.clone());
2098 active.publish(Ok(&event), Some(updated_task.clone())).await;
2099
2100 let published = receiver.recv().await.unwrap();
2102 assert_eq!(published.sequence, 1);
2103 match published.result {
2104 Ok(StreamResponse::Task(task)) => {
2105 assert_eq!(task.id, "t-snapshot-update");
2106 }
2107 _ => panic!("expected Ok(Task) event"),
2108 }
2109
2110 let (_, snapshot_task, sequence) = active.resubscribe().await;
2112 assert_eq!(snapshot_task.unwrap().id, "t-snapshot-update");
2113 assert_eq!(sequence, 1);
2114 }
2115
2116 #[tokio::test]
2117 async fn test_active_execution_broadcasts_to_multiple_subscribers() {
2118 let task = Task {
2119 id: "t-multi-sub".into(),
2120 context_id: "c-multi-sub".into(),
2121 status: TaskStatus {
2122 state: TaskState::Working,
2123 message: None,
2124 timestamp: None,
2125 },
2126 artifacts: None,
2127 history: None,
2128 metadata: None,
2129 };
2130 let active = ActiveExecution::new(task.clone());
2131
2132 let mut sub1 = active.subscribe();
2134 let mut sub2 = active.subscribe();
2135
2136 let event = StreamResponse::Task(task);
2137 active.publish(Ok(&event), None).await;
2138
2139 let evt1 = sub1.recv().await.unwrap();
2141 let evt2 = sub2.recv().await.unwrap();
2142
2143 assert_eq!(evt1.sequence, 1);
2144 assert_eq!(evt2.sequence, 1);
2145 assert!(matches!(evt1.result, Ok(StreamResponse::Task(_))));
2146 assert!(matches!(evt2.result, Ok(StreamResponse::Task(_))));
2147 }
2148
2149 #[tokio::test]
2150 async fn test_get_task_not_found() {
2151 let handler = make_handler();
2152 let params = ServiceParams::new();
2153 let req = GetTaskRequest {
2154 id: "nonexistent".into(),
2155 history_length: None,
2156 tenant: None,
2157 };
2158 let result = handler.get_task(¶ms, req).await;
2159 assert!(result.is_err());
2160 }
2161
2162 fn send_req(task_id: Option<&str>, context_id: Option<&str>) -> SendMessageRequest {
2165 let mut msg = make_message();
2166 msg.task_id = task_id.map(Into::into);
2167 msg.context_id = context_id.map(Into::into);
2168 SendMessageRequest {
2169 message: msg,
2170 configuration: None,
2171 metadata: None,
2172 tenant: None,
2173 }
2174 }
2175
2176 #[tokio::test]
2177 async fn test_send_message_rejects_unsupported_content_type() {
2178 install_crypto_provider();
2179 let handler = make_handler().with_default_input_modes(vec!["text/plain".into()]);
2180 let mut req = send_req(None, None);
2181 req.message.parts = vec![
2182 Part::data(serde_json::json!({"value": "unsupported"}))
2183 .with_media_type("application/x-unsupported-type-12345"),
2184 ];
2185 let err = handler
2186 .send_message(&ServiceParams::new(), req)
2187 .await
2188 .expect_err("an undeclared content type must be rejected");
2189 assert_eq!(err.code, error_code::CONTENT_TYPE_NOT_SUPPORTED);
2190 }
2191
2192 #[tokio::test]
2193 async fn test_send_message_accepts_a_declared_content_type() {
2194 install_crypto_provider();
2195 let handler = make_handler().with_default_input_modes(vec!["text/plain".into()]);
2196 let mut req = send_req(None, None);
2197 req.message.parts = vec![Part::text("hello").with_media_type("text/plain")];
2198 handler
2199 .send_message(&ServiceParams::new(), req)
2200 .await
2201 .expect("a declared content type must be accepted");
2202 }
2203
2204 #[tokio::test]
2205 async fn test_send_message_allows_any_content_type_when_none_declared() {
2206 install_crypto_provider();
2207 let handler = make_handler();
2208 let mut req = send_req(None, None);
2209 req.message.parts = vec![
2210 Part::data(serde_json::json!({"value": "x"})).with_media_type("application/x-anything"),
2211 ];
2212 handler
2213 .send_message(&ServiceParams::new(), req)
2214 .await
2215 .expect("no declared modes means no restriction");
2216 }
2217
2218 #[tokio::test]
2219 async fn test_send_message_unknown_task_id_is_not_found() {
2220 install_crypto_provider();
2221 let handler = make_handler();
2222 let err = handler
2226 .send_message(&ServiceParams::new(), send_req(Some("t-absent"), None))
2227 .await
2228 .expect_err("an unknown taskId must not create a task");
2229 assert_eq!(err.code, error_code::TASK_NOT_FOUND);
2230 assert!(
2231 handler.task_store.get("t-absent").await.unwrap().is_none(),
2232 "the rejected id must not have been created"
2233 );
2234 }
2235
2236 #[tokio::test]
2237 async fn test_send_message_to_terminal_task_is_rejected() {
2238 install_crypto_provider();
2239 let handler = make_handler();
2240 handler
2241 .task_store
2242 .create(Task {
2243 id: "t-done".into(),
2244 context_id: "c-done".into(),
2245 status: TaskStatus {
2246 state: TaskState::Completed,
2247 message: None,
2248 timestamp: None,
2249 },
2250 artifacts: None,
2251 history: None,
2252 metadata: None,
2253 })
2254 .await
2255 .unwrap();
2256 let err = handler
2257 .send_message(&ServiceParams::new(), send_req(Some("t-done"), None))
2258 .await
2259 .expect_err("a follow-up to a terminal task must not be accepted");
2260 assert_eq!(err.code, error_code::UNSUPPORTED_OPERATION);
2261 }
2262
2263 #[tokio::test]
2264 async fn test_send_message_mismatched_context_id_is_rejected() {
2265 install_crypto_provider();
2266 let handler = make_handler();
2267 seed_task(&handler, "t-ctx", "c-real").await;
2268 let err = handler
2271 .send_message(
2272 &ServiceParams::new(),
2273 send_req(Some("t-ctx"), Some("c-wrong")),
2274 )
2275 .await
2276 .expect_err("a mismatching contextId must be rejected");
2277 assert_eq!(err.code, error_code::INVALID_PARAMS);
2278 assert!(
2279 err.message.contains("c-wrong"),
2280 "error names the bad id: {}",
2281 err.message
2282 );
2283 }
2284
2285 #[tokio::test]
2286 async fn test_send_message_matching_context_id_is_accepted() {
2287 install_crypto_provider();
2288 let handler = make_handler();
2289 seed_task(&handler, "t-ok", "c-ok").await;
2290 handler
2291 .send_message(&ServiceParams::new(), send_req(Some("t-ok"), Some("c-ok")))
2292 .await
2293 .expect("a matching pair must be accepted");
2294 }
2295
2296 #[tokio::test]
2297 async fn test_send_message_accumulates_history_across_turns() {
2298 install_crypto_provider();
2304 let handler = DefaultRequestHandler::new(MultiTurnExecutor, InMemoryTaskStore::new());
2305
2306 let mut turn1 = make_message();
2307 turn1.parts = vec![Part::text("turn one")];
2308 let resp1 = handler
2309 .send_message(
2310 &ServiceParams::new(),
2311 SendMessageRequest {
2312 message: turn1,
2313 configuration: None,
2314 metadata: None,
2315 tenant: None,
2316 },
2317 )
2318 .await
2319 .unwrap();
2320 let task_id = match resp1 {
2321 SendMessageResponse::Task(t) => t.id,
2322 SendMessageResponse::Message(_) => panic!("expected a task"),
2323 };
2324
2325 for text in ["turn two", "done"] {
2326 let mut turn = make_message();
2327 turn.task_id = Some(task_id.clone());
2328 turn.parts = vec![Part::text(text)];
2329 handler
2330 .send_message(
2331 &ServiceParams::new(),
2332 SendMessageRequest {
2333 message: turn,
2334 configuration: None,
2335 metadata: None,
2336 tenant: None,
2337 },
2338 )
2339 .await
2340 .unwrap();
2341 }
2342
2343 let task = handler.task_store.get(&task_id).await.unwrap().unwrap();
2344 let history = task.history.expect("history must accumulate across turns");
2345 let texts: Vec<_> = history.iter().map(|m| m.text()).collect();
2346 assert_eq!(
2347 texts,
2348 vec![Some("turn one"), Some("turn two"), Some("done")]
2349 );
2350 }
2351
2352 #[tokio::test]
2353 async fn test_send_message_applies_configuration_history_length_to_its_own_response() {
2354 install_crypto_provider();
2359 let handler = DefaultRequestHandler::new(MultiTurnExecutor, InMemoryTaskStore::new());
2360 let mut message = make_message();
2361 message.parts = vec![Part::text("done")];
2362 let resp = handler
2363 .send_message(
2364 &ServiceParams::new(),
2365 SendMessageRequest {
2366 message,
2367 configuration: Some(SendMessageConfiguration {
2368 accepted_output_modes: None,
2369 task_push_notification_config: None,
2370 history_length: Some(0),
2371 return_immediately: None,
2372 }),
2373 metadata: None,
2374 tenant: None,
2375 },
2376 )
2377 .await
2378 .unwrap();
2379 let task = match resp {
2380 SendMessageResponse::Task(t) => t,
2381 SendMessageResponse::Message(_) => panic!("expected a task"),
2382 };
2383 assert!(
2384 task.history.is_none_or(|h| h.is_empty()),
2385 "historyLength=0 must clear history on SendMessage's own response"
2386 );
2387 }
2388
2389 #[tokio::test]
2390 async fn test_send_message_infers_context_id_from_the_task() {
2391 install_crypto_provider();
2392 let handler = make_handler();
2393 seed_task(&handler, "t-infer", "c-infer").await;
2394 handler
2396 .send_message(&ServiceParams::new(), send_req(Some("t-infer"), None))
2397 .await
2398 .unwrap();
2399 let task = handler.task_store.get("t-infer").await.unwrap().unwrap();
2400 assert_eq!(task.context_id, "c-infer");
2401 }
2402
2403 #[tokio::test]
2404 async fn test_send_message_without_task_id_still_creates_one() {
2405 install_crypto_provider();
2406 let handler = make_handler();
2407 let resp = handler
2410 .send_message(&ServiceParams::new(), send_req(None, Some("c-new")))
2411 .await
2412 .expect("a message with no taskId must still create a task");
2413 let _ = resp;
2414 let listed = handler
2415 .task_store
2416 .list(&ListTasksRequest {
2417 context_id: Some("c-new".into()),
2418 status: None,
2419 page_size: None,
2420 page_token: None,
2421 history_length: None,
2422 include_artifacts: None,
2423 status_timestamp_after: None,
2424 tenant: None,
2425 })
2426 .await
2427 .unwrap();
2428 assert_eq!(listed.tasks.len(), 1);
2429 assert_eq!(listed.tasks[0].context_id, "c-new");
2430 }
2431
2432 #[tokio::test]
2433 async fn test_get_task_after_send() {
2434 let handler = make_handler();
2435 let params = ServiceParams::new();
2436 let mut msg = make_message();
2437 msg.task_id = Some("t1".into());
2438 msg.context_id = Some("c1".into());
2439 seed_task(&handler, "t1", "c1").await;
2440 let req = SendMessageRequest {
2441 message: msg,
2442 configuration: None,
2443 metadata: None,
2444 tenant: None,
2445 };
2446 handler.send_message(¶ms, req).await.unwrap();
2447 let task = handler
2448 .get_task(
2449 ¶ms,
2450 GetTaskRequest {
2451 id: "t1".into(),
2452 history_length: None,
2453 tenant: None,
2454 },
2455 )
2456 .await
2457 .unwrap();
2458 assert_eq!(task.id, "t1");
2459 }
2460
2461 #[tokio::test]
2462 async fn test_list_tasks() {
2463 let handler = make_handler();
2464 let params = ServiceParams::new();
2465 let req = ListTasksRequest {
2466 context_id: None,
2467 status: None,
2468 page_size: None,
2469 page_token: None,
2470 history_length: None,
2471 status_timestamp_after: None,
2472 include_artifacts: None,
2473 tenant: None,
2474 };
2475 let resp = handler.list_tasks(¶ms, req).await.unwrap();
2476 assert!(resp.tasks.is_empty());
2477 }
2478
2479 #[tokio::test]
2480 async fn test_cancel_task_not_found() {
2481 let handler = make_handler();
2482 let params = ServiceParams::new();
2483 let req = CancelTaskRequest {
2484 id: "nonexistent".into(),
2485 metadata: None,
2486 tenant: None,
2487 };
2488 let result = handler.cancel_task(¶ms, req).await;
2489 assert!(result.is_err());
2490 }
2491
2492 #[tokio::test]
2493 async fn test_cancel_task_working() {
2494 let handler = make_handler();
2495 let params = ServiceParams::new();
2496 let task = Task {
2498 id: "t2".into(),
2499 context_id: "c2".into(),
2500 status: TaskStatus {
2501 state: TaskState::Working,
2502 message: None,
2503 timestamp: None,
2504 },
2505 artifacts: None,
2506 history: None,
2507 metadata: None,
2508 };
2509 handler.task_store.create(task).await.unwrap();
2510 let req = CancelTaskRequest {
2511 id: "t2".into(),
2512 metadata: None,
2513 tenant: None,
2514 };
2515 let result = handler.cancel_task(¶ms, req).await.unwrap();
2516 assert_eq!(result.status.state, TaskState::Canceled);
2517 }
2518
2519 #[tokio::test]
2520 async fn test_cancel_task_concurrent_both_succeed_idempotently() {
2521 let handler = Arc::new(make_handler());
2522 let params = ServiceParams::new();
2523 let task = Task {
2524 id: "t-race".into(),
2525 context_id: "c-race".into(),
2526 status: TaskStatus {
2527 state: TaskState::Working,
2528 message: None,
2529 timestamp: None,
2530 },
2531 artifacts: None,
2532 history: None,
2533 metadata: None,
2534 };
2535 handler.task_store.create(task).await.unwrap();
2536
2537 let handlers = vec![Arc::clone(&handler), Arc::clone(&handler)];
2538 let mut results = Vec::new();
2539 for handler in handlers {
2540 let params = params.clone();
2541 let req = CancelTaskRequest {
2542 id: "t-race".into(),
2543 metadata: None,
2544 tenant: None,
2545 };
2546 results.push(tokio::spawn(async move {
2547 handler.cancel_task(¶ms, req).await
2548 }));
2549 }
2550 for result in results {
2551 let task = result.await.unwrap().unwrap();
2552 assert_eq!(task.status.state, TaskState::Canceled);
2553 }
2554 }
2555
2556 #[tokio::test]
2557 async fn test_cancel_task_already_completed() {
2558 let handler = make_handler();
2559 let params = ServiceParams::new();
2560 let task = Task {
2561 id: "t3".into(),
2562 context_id: "c3".into(),
2563 status: TaskStatus {
2564 state: TaskState::Completed,
2565 message: None,
2566 timestamp: None,
2567 },
2568 artifacts: None,
2569 history: None,
2570 metadata: None,
2571 };
2572 handler.task_store.create(task).await.unwrap();
2573 let req = CancelTaskRequest {
2574 id: "t3".into(),
2575 metadata: None,
2576 tenant: None,
2577 };
2578 let result = handler.cancel_task(¶ms, req).await;
2579 match result {
2580 Ok(_) => panic!("expected cancel of COMPLETED to fail"),
2581 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
2582 }
2583 }
2584
2585 #[tokio::test]
2586 async fn test_cancel_task_already_canceled_is_idempotent() {
2587 let handler = make_handler();
2588 let params = ServiceParams::new();
2589 let task = Task {
2590 id: "t-canceled".into(),
2591 context_id: "c-canceled".into(),
2592 status: TaskStatus {
2593 state: TaskState::Canceled,
2594 message: None,
2595 timestamp: None,
2596 },
2597 artifacts: None,
2598 history: None,
2599 metadata: None,
2600 };
2601 handler.task_store.create(task).await.unwrap();
2602 let req = CancelTaskRequest {
2603 id: "t-canceled".into(),
2604 metadata: None,
2605 tenant: None,
2606 };
2607 let result = handler.cancel_task(¶ms, req).await.unwrap();
2608 assert_eq!(result.status.state, TaskState::Canceled);
2609 }
2610
2611 #[tokio::test]
2612 async fn test_cancel_task_failed_and_rejected_are_not_cancelable() {
2613 let handler = make_handler();
2614 let params = ServiceParams::new();
2615 for (id, state) in [
2616 ("t-failed", TaskState::Failed),
2617 ("t-rejected", TaskState::Rejected),
2618 ] {
2619 handler
2620 .task_store
2621 .create(Task {
2622 id: id.into(),
2623 context_id: format!("c-{id}"),
2624 status: TaskStatus {
2625 state,
2626 message: None,
2627 timestamp: None,
2628 },
2629 artifacts: None,
2630 history: None,
2631 metadata: None,
2632 })
2633 .await
2634 .unwrap();
2635 let result = handler
2636 .cancel_task(
2637 ¶ms,
2638 CancelTaskRequest {
2639 id: id.into(),
2640 metadata: None,
2641 tenant: None,
2642 },
2643 )
2644 .await;
2645 match result {
2646 Ok(_) => panic!("expected cancel of {id} to fail"),
2647 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
2648 }
2649 }
2650 }
2651
2652 fn task_with_history(id: &str, context_id: &str, n: usize) -> Task {
2653 Task {
2654 id: id.into(),
2655 context_id: context_id.into(),
2656 status: TaskStatus {
2657 state: TaskState::Completed,
2658 message: None,
2659 timestamp: None,
2660 },
2661 artifacts: None,
2662 history: Some(
2663 (0..n)
2664 .map(|i| Message::new(Role::User, vec![Part::text(format!("m{i}"))]))
2665 .collect(),
2666 ),
2667 metadata: None,
2668 }
2669 }
2670
2671 #[tokio::test]
2672 async fn test_get_task_applies_history_length() {
2673 let handler = make_handler();
2674 let params = ServiceParams::new();
2675 handler
2676 .task_store
2677 .create(task_with_history("hist", "c", 3))
2678 .await
2679 .unwrap();
2680
2681 let full = handler
2682 .get_task(
2683 ¶ms,
2684 GetTaskRequest {
2685 id: "hist".into(),
2686 history_length: None,
2687 tenant: None,
2688 },
2689 )
2690 .await
2691 .unwrap();
2692 assert_eq!(full.history.as_ref().unwrap().len(), 3);
2693
2694 let last = handler
2695 .get_task(
2696 ¶ms,
2697 GetTaskRequest {
2698 id: "hist".into(),
2699 history_length: Some(1),
2700 tenant: None,
2701 },
2702 )
2703 .await
2704 .unwrap();
2705 assert_eq!(last.history.as_ref().unwrap().len(), 1);
2706
2707 for hl in [0, -1] {
2708 let empty = handler
2709 .get_task(
2710 ¶ms,
2711 GetTaskRequest {
2712 id: "hist".into(),
2713 history_length: Some(hl),
2714 tenant: None,
2715 },
2716 )
2717 .await
2718 .unwrap();
2719 assert!(
2720 empty.history.as_ref().unwrap().is_empty(),
2721 "history_length={hl} must clear history"
2722 );
2723 }
2724 }
2725
2726 fn reject_unauthenticated(
2727 _params: &ServiceParams,
2728 _task_id: Option<&str>,
2729 ) -> Result<(), A2AError> {
2730 Err(A2AError::invalid_request("authentication required"))
2731 }
2732
2733 #[tokio::test]
2734 async fn test_authorizer_rejects_handler_methods() {
2735 let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
2736 .with_authorizer(reject_unauthenticated)
2737 .with_push_config_store(InMemoryPushConfigStore::new());
2738 let params = ServiceParams::new();
2739
2740 assert!(
2741 handler
2742 .send_message(
2743 ¶ms,
2744 SendMessageRequest {
2745 message: make_message(),
2746 configuration: None,
2747 metadata: None,
2748 tenant: None,
2749 },
2750 )
2751 .await
2752 .is_err()
2753 );
2754 assert!(
2755 handler
2756 .get_task(
2757 ¶ms,
2758 GetTaskRequest {
2759 id: "t".into(),
2760 history_length: None,
2761 tenant: None,
2762 },
2763 )
2764 .await
2765 .is_err()
2766 );
2767 assert!(
2768 handler
2769 .list_tasks(
2770 ¶ms,
2771 ListTasksRequest {
2772 context_id: None,
2773 status: None,
2774 page_size: None,
2775 page_token: None,
2776 history_length: None,
2777 status_timestamp_after: None,
2778 include_artifacts: None,
2779 tenant: None,
2780 },
2781 )
2782 .await
2783 .is_err()
2784 );
2785 assert!(
2786 handler
2787 .cancel_task(
2788 ¶ms,
2789 CancelTaskRequest {
2790 id: "t".into(),
2791 metadata: None,
2792 tenant: None,
2793 },
2794 )
2795 .await
2796 .is_err()
2797 );
2798 assert!(
2799 handler
2800 .subscribe_to_task(
2801 ¶ms,
2802 SubscribeToTaskRequest {
2803 id: "t".into(),
2804 tenant: None,
2805 },
2806 )
2807 .await
2808 .is_err()
2809 );
2810 assert!(
2811 handler
2812 .create_push_config(
2813 ¶ms,
2814 TaskPushNotificationConfig {
2815 url: "https://example.com/push".into(),
2816 id: None,
2817 task_id: "t".into(),
2818 token: None,
2819 authentication: None,
2820 tenant: None,
2821 },
2822 )
2823 .await
2824 .is_err()
2825 );
2826 assert!(
2827 handler
2828 .get_push_config(
2829 ¶ms,
2830 GetTaskPushNotificationConfigRequest {
2831 task_id: "t".into(),
2832 id: "cfg".into(),
2833 tenant: None,
2834 },
2835 )
2836 .await
2837 .is_err()
2838 );
2839 assert!(
2840 handler
2841 .list_push_configs(
2842 ¶ms,
2843 ListTaskPushNotificationConfigsRequest {
2844 task_id: "t".into(),
2845 page_size: None,
2846 page_token: None,
2847 tenant: None,
2848 },
2849 )
2850 .await
2851 .is_err()
2852 );
2853 assert!(
2854 handler
2855 .delete_push_config(
2856 ¶ms,
2857 DeleteTaskPushNotificationConfigRequest {
2858 task_id: "t".into(),
2859 id: "cfg".into(),
2860 tenant: None,
2861 },
2862 )
2863 .await
2864 .is_err()
2865 );
2866 assert!(
2867 handler
2868 .get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
2869 .await
2870 .is_err()
2871 );
2872 }
2873
2874 #[tokio::test]
2875 async fn test_subscribe_requires_active_execution() {
2876 let handler = make_handler();
2877 let params = ServiceParams::new();
2878 let req = SubscribeToTaskRequest {
2879 id: "t1".into(),
2880 tenant: None,
2881 };
2882 let result = handler.subscribe_to_task(¶ms, req).await;
2883 match result {
2884 Ok(_) => panic!("expected subscribe_to_task to fail"),
2885 Err(error) => assert_eq!(error.code, error_code::TASK_NOT_FOUND),
2886 }
2887 }
2888
2889 #[tokio::test]
2890 async fn test_subscribe_to_terminal_task_is_unsupported_not_not_found() {
2891 let handler = make_handler();
2892 handler
2893 .task_store
2894 .create(Task {
2895 id: "t-done".into(),
2896 context_id: "c-done".into(),
2897 status: TaskStatus {
2898 state: TaskState::Completed,
2899 message: None,
2900 timestamp: None,
2901 },
2902 artifacts: None,
2903 history: None,
2904 metadata: None,
2905 })
2906 .await
2907 .unwrap();
2908 let result = handler
2909 .subscribe_to_task(
2910 &ServiceParams::new(),
2911 SubscribeToTaskRequest {
2912 id: "t-done".into(),
2913 tenant: None,
2914 },
2915 )
2916 .await;
2917 match result {
2918 Ok(_) => panic!("expected subscribe_to_task to fail"),
2919 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
2920 }
2921 }
2922
2923 #[tokio::test]
2924 async fn test_send_message_return_immediately_allows_resubscribe() {
2925 use futures::StreamExt;
2926
2927 let (handler, release) = make_resumable_handler();
2928 let params = ServiceParams::new();
2929 let mut message = make_message();
2930 message.task_id = Some("t-resume".into());
2931 message.context_id = Some("c-resume".into());
2932 seed_task(&handler, "t-resume", "c-resume").await;
2933
2934 let response = handler
2935 .send_message(
2936 ¶ms,
2937 SendMessageRequest {
2938 message,
2939 configuration: Some(SendMessageConfiguration {
2940 accepted_output_modes: None,
2941 task_push_notification_config: None,
2942 history_length: None,
2943 return_immediately: Some(true),
2944 }),
2945 metadata: None,
2946 tenant: None,
2947 },
2948 )
2949 .await
2950 .unwrap();
2951
2952 match response {
2953 SendMessageResponse::Task(task) => {
2954 assert_eq!(task.status.state, TaskState::Working);
2955 }
2956 _ => panic!("expected Task response"),
2957 }
2958
2959 let mut subscription = handler
2960 .subscribe_to_task(
2961 ¶ms,
2962 SubscribeToTaskRequest {
2963 id: "t-resume".into(),
2964 tenant: None,
2965 },
2966 )
2967 .await
2968 .unwrap();
2969
2970 let snapshot = subscription.next().await.unwrap().unwrap();
2971 match snapshot {
2972 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
2973 _ => panic!("expected snapshot task"),
2974 }
2975
2976 release.notify_waiters();
2977
2978 let terminal = subscription.next().await.unwrap().unwrap();
2979 match terminal {
2980 StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
2981 _ => panic!("expected terminal task"),
2982 }
2983
2984 assert!(subscription.next().await.is_none());
2985
2986 let result = handler
2987 .subscribe_to_task(
2988 ¶ms,
2989 SubscribeToTaskRequest {
2990 id: "t-resume".into(),
2991 tenant: None,
2992 },
2993 )
2994 .await;
2995 match result {
2996 Ok(_) => panic!("expected subscribe_to_task to fail after completion"),
2997 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
2998 }
2999 }
3000
3001 #[tokio::test]
3002 async fn test_push_config_not_supported_without_store() {
3003 let handler = make_handler();
3004 let params = ServiceParams::new();
3005 let result = handler
3006 .create_push_config(
3007 ¶ms,
3008 TaskPushNotificationConfig {
3009 task_id: "t1".into(),
3010 url: "http://example.com".into(),
3011 id: None,
3012 token: None,
3013 authentication: None,
3014 tenant: None,
3015 },
3016 )
3017 .await;
3018 assert!(result.is_err());
3019 }
3020
3021 #[tokio::test]
3022 async fn test_list_push_configs_invalid_page_token_errors() {
3023 install_crypto_provider();
3024 let handler = make_handler_with_push_configs();
3025 let params = ServiceParams::new();
3026 handler
3027 .create_push_config(
3028 ¶ms,
3029 TaskPushNotificationConfig {
3030 task_id: "t1".into(),
3031 url: "https://example.com/callback".into(),
3032 id: Some("cfg-1".into()),
3033 token: None,
3034 authentication: None,
3035 tenant: None,
3036 },
3037 )
3038 .await
3039 .unwrap();
3040
3041 let result = handler
3042 .list_push_configs(
3043 ¶ms,
3044 ListTaskPushNotificationConfigsRequest {
3045 task_id: "t1".into(),
3046 page_size: None,
3047 page_token: Some("not-a-number".into()),
3048 tenant: None,
3049 },
3050 )
3051 .await;
3052 assert!(result.is_err());
3053 assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
3054 }
3055
3056 #[tokio::test]
3057 async fn test_list_tasks_page_size_cap_bounds_a_store_that_ignores_it() {
3058 install_crypto_provider();
3059 let handler = DefaultRequestHandler::new(
3062 EchoExecutor,
3063 PagingTaskStore::new(MAX_PAGE_SIZE * 3, false),
3064 );
3065
3066 let resp = handler
3067 .list_tasks(
3068 &ServiceParams::new(),
3069 ListTasksRequest {
3070 context_id: None,
3071 status: None,
3072 page_size: Some(1_000_000),
3073 page_token: None,
3074 history_length: None,
3075 include_artifacts: None,
3076 status_timestamp_after: None,
3077 tenant: None,
3078 },
3079 )
3080 .await
3081 .unwrap();
3082
3083 assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
3084 }
3085
3086 #[tokio::test]
3087 async fn test_paging_stores_delegate_everything_except_list() {
3088 install_crypto_provider();
3089 let tasks = PagingTaskStore::new(0, true);
3093 let task = Task {
3094 id: "t-delegated".into(),
3095 context_id: "c1".into(),
3096 status: TaskStatus {
3097 state: TaskState::Submitted,
3098 message: None,
3099 timestamp: None,
3100 },
3101 artifacts: None,
3102 history: None,
3103 metadata: None,
3104 };
3105 tasks.create(task.clone()).await.unwrap();
3106 assert_eq!(
3107 tasks
3108 .get("t-delegated")
3109 .await
3110 .unwrap()
3111 .map(|t| t.status.state),
3112 Some(TaskState::Submitted)
3113 );
3114 let mut done = task.clone();
3115 done.status.state = TaskState::Completed;
3116 tasks.update(done).await.unwrap();
3117 assert_eq!(
3118 tasks
3119 .get("t-delegated")
3120 .await
3121 .unwrap()
3122 .map(|t| t.status.state),
3123 Some(TaskState::Completed)
3124 );
3125
3126 let mut cancelable = task.clone();
3130 cancelable.id = "t-cancelable".into();
3131 cancelable.status.state = TaskState::Working;
3132 tasks.create(cancelable).await.unwrap();
3133 assert_eq!(
3134 tasks
3135 .begin_cancel("t-cancelable")
3136 .await
3137 .unwrap()
3138 .status
3139 .state,
3140 TaskState::Canceled
3141 );
3142 assert_eq!(
3143 tasks.begin_cancel("t-delegated").await.unwrap_err().code,
3144 error_code::TASK_NOT_CANCELABLE
3145 );
3146
3147 let configs = PagingPushConfigStore::new(0);
3148 configs
3149 .save(TaskPushNotificationConfig {
3150 task_id: "t1".into(),
3151 url: "https://example.com/callback".into(),
3152 id: Some("cfg-1".into()),
3153 token: None,
3154 authentication: None,
3155 tenant: None,
3156 })
3157 .await
3158 .unwrap();
3159 assert_eq!(
3160 configs.get("t1", "cfg-1").await.unwrap().id.as_deref(),
3161 Some("cfg-1")
3162 );
3163 configs.delete("t1", "cfg-1").await.unwrap();
3164 assert!(configs.get("t1", "cfg-1").await.is_err());
3165 configs.delete_all("t1").await.unwrap();
3166 }
3167
3168 #[tokio::test]
3169 async fn test_list_tasks_page_size_cap_bounds_a_store_that_honours_it() {
3170 install_crypto_provider();
3171 let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
3175 let seen = Arc::clone(&store.seen_page_size);
3176 let handler = DefaultRequestHandler::new(EchoExecutor, store);
3177
3178 let resp = handler
3179 .list_tasks(
3180 &ServiceParams::new(),
3181 ListTasksRequest {
3182 context_id: None,
3183 status: None,
3184 page_size: Some(1_000_000),
3185 page_token: None,
3186 history_length: None,
3187 include_artifacts: None,
3188 status_timestamp_after: None,
3189 tenant: None,
3190 },
3191 )
3192 .await
3193 .unwrap();
3194
3195 assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
3196 assert_eq!(seen.load(Ordering::SeqCst), MAX_PAGE_SIZE as i32);
3200 }
3201
3202 #[tokio::test]
3203 async fn test_list_tasks_truncation_keeps_the_store_continuation_token() {
3204 install_crypto_provider();
3205 let handler = DefaultRequestHandler::new(
3210 EchoExecutor,
3211 PagingTaskStore::new(MAX_PAGE_SIZE * 3, false)
3212 .with_next_page_token("cursor-from-store"),
3213 );
3214
3215 let resp = handler
3216 .list_tasks(
3217 &ServiceParams::new(),
3218 ListTasksRequest {
3219 context_id: None,
3220 status: None,
3221 page_size: None,
3222 page_token: None,
3223 history_length: None,
3224 include_artifacts: None,
3225 status_timestamp_after: None,
3226 tenant: None,
3227 },
3228 )
3229 .await
3230 .unwrap();
3231
3232 assert_eq!(resp.tasks.len(), DEFAULT_PAGE_SIZE);
3233 assert_eq!(resp.next_page_token, "cursor-from-store");
3234 }
3235
3236 #[tokio::test]
3237 async fn test_list_tasks_honours_a_page_size_under_the_cap() {
3238 install_crypto_provider();
3239 let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
3242 let seen = Arc::clone(&store.seen_page_size);
3243 let handler = DefaultRequestHandler::new(EchoExecutor, store);
3244
3245 let resp = handler
3246 .list_tasks(
3247 &ServiceParams::new(),
3248 ListTasksRequest {
3249 context_id: None,
3250 status: None,
3251 page_size: Some(7),
3252 page_token: None,
3253 history_length: None,
3254 include_artifacts: None,
3255 status_timestamp_after: None,
3256 tenant: None,
3257 },
3258 )
3259 .await
3260 .unwrap();
3261
3262 assert_eq!(seen.load(Ordering::SeqCst), 7);
3263 assert_eq!(resp.tasks.len(), 7);
3264 }
3265
3266 #[tokio::test]
3267 async fn test_list_push_configs_page_size_cap_bounds_a_custom_store() {
3268 install_crypto_provider();
3269 let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
3273 .with_push_config_store(PagingPushConfigStore::new(MAX_PAGE_SIZE * 3));
3274
3275 let resp = handler
3276 .list_push_configs(
3277 &ServiceParams::new(),
3278 ListTaskPushNotificationConfigsRequest {
3279 task_id: "t1".into(),
3280 page_size: Some(1_000_000),
3281 page_token: None,
3282 tenant: None,
3283 },
3284 )
3285 .await
3286 .unwrap();
3287
3288 assert_eq!(resp.configs.len(), MAX_PAGE_SIZE);
3289 assert_eq!(resp.next_page_token.as_deref(), Some("100"));
3291 }
3292
3293 #[tokio::test]
3294 async fn test_list_push_configs_page_size_is_capped() {
3295 install_crypto_provider();
3296 let handler = make_handler_with_push_configs();
3297 let params = ServiceParams::new();
3298 for i in 0..40 {
3301 handler
3302 .create_push_config(
3303 ¶ms,
3304 TaskPushNotificationConfig {
3305 task_id: "t1".into(),
3306 url: format!("https://example.com/callback/{i}"),
3307 id: Some(format!("cfg-{i:03}")),
3308 token: None,
3309 authentication: None,
3310 tenant: None,
3311 },
3312 )
3313 .await
3314 .unwrap();
3315 }
3316
3317 let resp = handler
3323 .list_push_configs(
3324 ¶ms,
3325 ListTaskPushNotificationConfigsRequest {
3326 task_id: "t1".into(),
3327 page_size: Some(1000),
3328 page_token: None,
3329 tenant: None,
3330 },
3331 )
3332 .await
3333 .unwrap();
3334 assert_eq!(resp.configs.len(), 40);
3335 }
3336
3337 #[tokio::test]
3338 async fn test_push_config_crud_with_store() {
3339 install_crypto_provider();
3340 let handler = make_handler_with_push_configs();
3341 let params = ServiceParams::new();
3342
3343 let created = handler
3344 .create_push_config(
3345 ¶ms,
3346 TaskPushNotificationConfig {
3347 task_id: "t1".into(),
3348 url: "https://example.com/first".into(),
3349 id: Some("cfg-1".into()),
3350 token: None,
3351 authentication: None,
3352 tenant: Some("tenant-a".into()),
3353 },
3354 )
3355 .await
3356 .unwrap();
3357 assert_eq!(created.id.as_deref(), Some("cfg-1"));
3358
3359 handler
3360 .create_push_config(
3361 ¶ms,
3362 TaskPushNotificationConfig {
3363 task_id: "t1".into(),
3364 url: "https://example.com/second".into(),
3365 id: Some("cfg-2".into()),
3366 token: None,
3367 authentication: None,
3368 tenant: Some("tenant-a".into()),
3369 },
3370 )
3371 .await
3372 .unwrap();
3373
3374 let fetched = handler
3375 .get_push_config(
3376 ¶ms,
3377 GetTaskPushNotificationConfigRequest {
3378 task_id: "t1".into(),
3379 id: "cfg-1".into(),
3380 tenant: Some("tenant-a".into()),
3381 },
3382 )
3383 .await
3384 .unwrap();
3385 assert_eq!(fetched.url, "https://example.com/first");
3386
3387 let first_page = handler
3388 .list_push_configs(
3389 ¶ms,
3390 ListTaskPushNotificationConfigsRequest {
3391 task_id: "t1".into(),
3392 page_size: Some(1),
3393 page_token: Some("0".into()),
3394 tenant: Some("tenant-a".into()),
3395 },
3396 )
3397 .await
3398 .unwrap();
3399 assert_eq!(first_page.configs.len(), 1);
3400 assert_eq!(first_page.configs[0].id.as_deref(), Some("cfg-1"));
3401 assert_eq!(first_page.next_page_token.as_deref(), Some("1"));
3402
3403 let default_page = handler
3404 .list_push_configs(
3405 ¶ms,
3406 ListTaskPushNotificationConfigsRequest {
3407 task_id: "t1".into(),
3408 page_size: Some(0),
3409 page_token: None,
3410 tenant: Some("tenant-a".into()),
3411 },
3412 )
3413 .await
3414 .unwrap();
3415 assert_eq!(default_page.configs.len(), 2);
3416 assert!(default_page.next_page_token.is_none());
3417
3418 handler
3419 .delete_push_config(
3420 ¶ms,
3421 DeleteTaskPushNotificationConfigRequest {
3422 task_id: "t1".into(),
3423 id: "cfg-1".into(),
3424 tenant: Some("tenant-a".into()),
3425 },
3426 )
3427 .await
3428 .unwrap();
3429
3430 let remaining = handler
3431 .list_push_configs(
3432 ¶ms,
3433 ListTaskPushNotificationConfigsRequest {
3434 task_id: "t1".into(),
3435 page_size: None,
3436 page_token: None,
3437 tenant: Some("tenant-a".into()),
3438 },
3439 )
3440 .await
3441 .unwrap();
3442 assert_eq!(remaining.configs.len(), 1);
3443 assert_eq!(remaining.configs[0].id.as_deref(), Some("cfg-2"));
3444 }
3445
3446 #[tokio::test]
3447 async fn test_send_message_request_push_config_delivers_webhooks() {
3448 install_crypto_provider();
3449 let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
3450 let handler = make_push_delivery_handler();
3451 let params = ServiceParams::new();
3452 let mut message = make_message();
3453 message.task_id = Some("t-push-request".into());
3454 message.context_id = Some("c-push-request".into());
3455 seed_task(&handler, "t-push-request", "c-push-request").await;
3456
3457 let response = handler
3458 .send_message(
3459 ¶ms,
3460 SendMessageRequest {
3461 message,
3462 configuration: Some(SendMessageConfiguration {
3463 accepted_output_modes: None,
3464 task_push_notification_config: Some(TaskPushNotificationConfig {
3465 task_id: String::new(),
3466 url: url.clone(),
3467 id: Some("cfg-request".into()),
3468 token: Some("notify-token".into()),
3469 authentication: Some(AuthenticationInfo {
3470 scheme: "bearer".into(),
3471 credentials: Some("secret-token".into()),
3472 }),
3473 tenant: None,
3474 }),
3475 history_length: None,
3476 return_immediately: None,
3477 }),
3478 metadata: None,
3479 tenant: None,
3480 },
3481 )
3482 .await
3483 .unwrap();
3484
3485 match response {
3486 SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
3487 _ => panic!("expected Task response"),
3488 }
3489
3490 let saved = handler
3491 .get_push_config(
3492 ¶ms,
3493 GetTaskPushNotificationConfigRequest {
3494 task_id: "t-push-request".into(),
3495 id: "cfg-request".into(),
3496 tenant: None,
3497 },
3498 )
3499 .await
3500 .unwrap();
3501 assert_eq!(saved.url, url);
3502
3503 let first = next_push(&mut receiver).await;
3504 assert_eq!(first.authorization.as_deref(), Some("Bearer secret-token"));
3505 assert_eq!(first.notification_token.as_deref(), Some("notify-token"));
3506 match first.event {
3507 StreamResponse::StatusUpdate(update) => {
3508 assert_eq!(update.task_id, "t-push-request");
3509 assert_eq!(update.status.state, TaskState::Working);
3510 }
3511 _ => panic!("expected status update push"),
3512 }
3513
3514 let second = next_push(&mut receiver).await;
3515 assert_eq!(second.authorization.as_deref(), Some("Bearer secret-token"));
3516 assert_eq!(second.notification_token.as_deref(), Some("notify-token"));
3517 match second.event {
3518 StreamResponse::Task(task) => {
3519 assert_eq!(task.id, "t-push-request");
3520 assert_eq!(task.status.state, TaskState::Completed);
3521 }
3522 _ => panic!("expected final task push"),
3523 }
3524
3525 shutdown_tx.send(()).unwrap();
3526 server.await.unwrap();
3527 }
3528
3529 #[tokio::test]
3530 async fn test_send_message_stored_push_config_delivers_webhooks() {
3531 install_crypto_provider();
3532 let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
3533 let handler = make_push_delivery_handler();
3534 let params = ServiceParams::new();
3535
3536 handler
3537 .create_push_config(
3538 ¶ms,
3539 TaskPushNotificationConfig {
3540 task_id: "t-push-stored".into(),
3541 url: url.clone(),
3542 id: Some("cfg-stored".into()),
3543 token: Some("stored-token".into()),
3544 authentication: None,
3545 tenant: None,
3546 },
3547 )
3548 .await
3549 .unwrap();
3550
3551 let mut message = make_message();
3552 message.task_id = Some("t-push-stored".into());
3553 message.context_id = Some("c-push-stored".into());
3554 seed_task(&handler, "t-push-stored", "c-push-stored").await;
3555 let response = handler
3556 .send_message(
3557 ¶ms,
3558 SendMessageRequest {
3559 message,
3560 configuration: None,
3561 metadata: None,
3562 tenant: None,
3563 },
3564 )
3565 .await
3566 .unwrap();
3567
3568 match response {
3569 SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
3570 _ => panic!("expected Task response"),
3571 }
3572
3573 let first = next_push(&mut receiver).await;
3574 assert_eq!(first.notification_token.as_deref(), Some("stored-token"));
3575 match first.event {
3576 StreamResponse::StatusUpdate(update) => {
3577 assert_eq!(update.task_id, "t-push-stored");
3578 assert_eq!(update.status.state, TaskState::Working);
3579 }
3580 _ => panic!("expected status update push"),
3581 }
3582
3583 let second = next_push(&mut receiver).await;
3584 assert_eq!(second.notification_token.as_deref(), Some("stored-token"));
3585 match second.event {
3586 StreamResponse::Task(task) => {
3587 assert_eq!(task.id, "t-push-stored");
3588 assert_eq!(task.status.state, TaskState::Completed);
3589 }
3590 _ => panic!("expected final task push"),
3591 }
3592
3593 shutdown_tx.send(()).unwrap();
3594 server.await.unwrap();
3595 }
3596
3597 #[tokio::test]
3598 async fn test_extended_agent_card_not_configured() {
3599 let handler = make_handler();
3600 let params = ServiceParams::new();
3601 let result = handler
3602 .get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
3603 .await;
3604 assert!(result.is_err());
3605 }
3606
3607 #[test]
3608 fn test_with_capabilities() {
3609 let handler = make_handler().with_capabilities(AgentCapabilities {
3610 streaming: Some(true),
3611 push_notifications: None,
3612 extensions: None,
3613 extended_agent_card: None,
3614 });
3615 assert_eq!(handler.capabilities.streaming, Some(true));
3616 }
3617
3618 #[test]
3619 fn test_with_push_config_store_enables_push_capability() {
3620 install_crypto_provider();
3621 let handler = make_handler_with_push_configs();
3622 assert_eq!(handler.capabilities.push_notifications, Some(true));
3623 }
3624
3625 fn caps(streaming: Option<bool>, push: Option<bool>) -> AgentCapabilities {
3628 AgentCapabilities {
3629 streaming,
3630 push_notifications: push,
3631 extensions: None,
3632 extended_agent_card: None,
3633 }
3634 }
3635
3636 fn send_request() -> SendMessageRequest {
3637 SendMessageRequest {
3638 message: Message::new(Role::User, vec![Part::text("hi")]),
3639 configuration: None,
3640 metadata: None,
3641 tenant: None,
3642 }
3643 }
3644
3645 fn push_config(task_id: &str) -> TaskPushNotificationConfig {
3646 TaskPushNotificationConfig {
3647 task_id: task_id.into(),
3648 url: "https://example.com/callback".into(),
3649 id: Some("cfg-1".into()),
3650 token: None,
3651 authentication: None,
3652 tenant: None,
3653 }
3654 }
3655
3656 #[tokio::test]
3657 async fn test_streaming_refused_when_the_card_declares_it_absent() {
3658 let handler = make_handler().with_capabilities(caps(Some(false), None));
3659 let params = ServiceParams::new();
3660
3661 match handler
3664 .send_streaming_message(¶ms, send_request())
3665 .await
3666 {
3667 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
3668 Ok(_) => panic!("streaming must be refused when the card declares it absent"),
3669 }
3670
3671 match handler
3672 .subscribe_to_task(
3673 ¶ms,
3674 SubscribeToTaskRequest {
3675 id: "t1".into(),
3676 tenant: None,
3677 },
3678 )
3679 .await
3680 {
3681 Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
3682 Ok(_) => panic!("subscribe must be refused when the card declares it absent"),
3683 }
3684 }
3685
3686 #[tokio::test]
3687 async fn test_push_refused_when_the_card_declares_it_absent() {
3688 let handler = make_handler_with_push_configs().with_capabilities(caps(None, Some(false)));
3691 let params = ServiceParams::new();
3692
3693 let error = handler
3694 .create_push_config(¶ms, push_config("t1"))
3695 .await
3696 .unwrap_err();
3697 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
3698
3699 let error = handler
3700 .list_push_configs(
3701 ¶ms,
3702 ListTaskPushNotificationConfigsRequest {
3703 task_id: "t1".into(),
3704 page_size: None,
3705 page_token: None,
3706 tenant: None,
3707 },
3708 )
3709 .await
3710 .unwrap_err();
3711 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
3712 }
3713
3714 #[tokio::test]
3718 async fn test_undeclared_capabilities_stay_permissive() {
3719 let handler = make_handler_with_push_configs();
3720 let params = ServiceParams::new();
3721
3722 assert!(
3723 handler
3724 .send_streaming_message(¶ms, send_request())
3725 .await
3726 .is_ok(),
3727 "an undeclared capability must not be treated as refused"
3728 );
3729 assert!(
3730 handler
3731 .create_push_config(¶ms, push_config("t1"))
3732 .await
3733 .is_ok()
3734 );
3735 }
3736
3737 #[tokio::test]
3741 async fn test_builder_order_does_not_change_the_declared_capability() {
3742 let params = ServiceParams::new();
3743
3744 for handler in [
3745 make_handler()
3746 .with_capabilities(caps(None, Some(false)))
3747 .with_push_config_store(InMemoryPushConfigStore::new()),
3748 make_handler()
3749 .with_push_config_store(InMemoryPushConfigStore::new())
3750 .with_capabilities(caps(None, Some(false))),
3751 ] {
3752 let error = handler
3753 .create_push_config(¶ms, push_config("t1"))
3754 .await
3755 .unwrap_err();
3756 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
3757 }
3758 }
3759
3760 #[tokio::test]
3765 async fn test_with_push_notifications_defaults_the_capability() {
3766 let params = ServiceParams::new();
3767
3768 let handler = make_handler().with_push_notifications(
3769 InMemoryPushConfigStore::new(),
3770 crate::HttpPushSender::new(None),
3771 );
3772 assert_eq!(handler.capabilities.push_notifications, Some(true));
3773 assert!(
3774 handler
3775 .create_push_config(¶ms, push_config("t1"))
3776 .await
3777 .is_ok()
3778 );
3779
3780 let handler = make_handler()
3781 .with_capabilities(caps(None, Some(false)))
3782 .with_push_notifications(
3783 InMemoryPushConfigStore::new(),
3784 crate::HttpPushSender::new(None),
3785 );
3786 assert_eq!(handler.capabilities.push_notifications, Some(false));
3787 let error = handler
3788 .create_push_config(¶ms, push_config("t1"))
3789 .await
3790 .unwrap_err();
3791 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
3792 }
3793
3794 #[tokio::test]
3797 async fn test_missing_store_still_reports_push_unsupported() {
3798 let handler = make_handler();
3799 let error = handler
3800 .create_push_config(&ServiceParams::new(), push_config("t1"))
3801 .await
3802 .unwrap_err();
3803 assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
3804 }
3805
3806 #[tokio::test]
3812 async fn test_artifacts_do_not_aggregate_across_tasks_in_one_context() {
3813 let handler = make_handler();
3814 let params = ServiceParams::new();
3815 let context_id = "ctx-shared";
3816
3817 for (task_id, artifact_id) in [("t-first", "a-1"), ("t-second", "a-2")] {
3818 handler
3819 .task_store
3820 .create(Task {
3821 id: task_id.into(),
3822 context_id: context_id.into(),
3823 status: TaskStatus {
3824 state: TaskState::Working,
3825 message: None,
3826 timestamp: None,
3827 },
3828 artifacts: None,
3829 history: None,
3830 metadata: None,
3831 })
3832 .await
3833 .unwrap();
3834
3835 apply_event_to_task(
3836 handler.task_store.as_ref(),
3837 None,
3838 &StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
3839 task_id: task_id.into(),
3840 context_id: context_id.into(),
3841 artifact: Artifact {
3842 artifact_id: artifact_id.into(),
3843 name: None,
3844 description: None,
3845 parts: vec![Part::text("payload")],
3846 metadata: None,
3847 extensions: None,
3848 },
3849 append: None,
3850 last_chunk: None,
3851 metadata: None,
3852 }),
3853 )
3854 .await
3855 .unwrap();
3856 }
3857
3858 for (task_id, expected) in [("t-first", "a-1"), ("t-second", "a-2")] {
3859 let task = handler
3860 .get_task(
3861 ¶ms,
3862 GetTaskRequest {
3863 id: task_id.into(),
3864 history_length: None,
3865 tenant: None,
3866 },
3867 )
3868 .await
3869 .unwrap();
3870 let ids: Vec<&str> = task
3871 .artifacts
3872 .as_deref()
3873 .unwrap_or_default()
3874 .iter()
3875 .map(|a| a.artifact_id.as_str())
3876 .collect();
3877 assert_eq!(
3878 ids,
3879 vec![expected],
3880 "{task_id} must carry only its own artifact"
3881 );
3882 }
3883 }
3884 fn extended_card(name: &str) -> AgentCard {
3887 AgentCard {
3888 name: name.into(),
3889 description: "extended".into(),
3890 version: "1.0".into(),
3891 supported_interfaces: vec![],
3892 capabilities: AgentCapabilities::default(),
3893 default_input_modes: vec![],
3894 default_output_modes: vec![],
3895 skills: vec![],
3896 provider: None,
3897 documentation_url: None,
3898 icon_url: None,
3899 security_schemes: None,
3900 security_requirements: None,
3901 signatures: None,
3902 }
3903 }
3904
3905 fn extended_request() -> GetExtendedAgentCardRequest {
3906 GetExtendedAgentCardRequest { tenant: None }
3907 }
3908
3909 #[tokio::test]
3910 async fn test_extended_agent_card_is_served_when_configured() {
3911 let handler = make_handler().with_extended_agent_card(extended_card("secret skills"));
3912 let card = handler
3913 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3914 .await
3915 .unwrap();
3916 assert_eq!(card.name, "secret skills");
3917 }
3918
3919 #[tokio::test]
3922 async fn test_configuring_an_extended_card_declares_the_capability() {
3923 let handler = make_handler().with_extended_agent_card(extended_card("x"));
3924 assert_eq!(handler.capabilities.extended_agent_card, Some(true));
3925 }
3926
3927 #[tokio::test]
3931 async fn test_unconfigured_extended_card_reports_its_own_code() {
3932 let handler = make_handler().with_capabilities(AgentCapabilities {
3933 streaming: None,
3934 push_notifications: None,
3935 extensions: None,
3936 extended_agent_card: Some(true),
3937 });
3938 let error = handler
3939 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3940 .await
3941 .unwrap_err();
3942 assert_eq!(error.code, error_code::EXTENDED_CARD_NOT_CONFIGURED);
3943 assert_ne!(error.code, error_code::UNSUPPORTED_OPERATION);
3944 }
3945
3946 #[tokio::test]
3950 async fn test_capability_gate_precedes_the_not_configured_check() {
3951 let handler = make_handler()
3952 .with_extended_agent_card(extended_card("x"))
3953 .with_capabilities(AgentCapabilities {
3954 streaming: None,
3955 push_notifications: None,
3956 extensions: None,
3957 extended_agent_card: Some(false),
3958 });
3959 let error = handler
3960 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3961 .await
3962 .unwrap_err();
3963 assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION);
3964 }
3965
3966 #[tokio::test]
3969 async fn test_resolver_sees_the_caller_params() {
3970 let handler = make_handler().with_extended_agent_card_resolver(|params: &ServiceParams| {
3971 let privileged = params.contains_key("x-privileged");
3972 Ok(extended_card(if privileged { "full" } else { "limited" }))
3973 });
3974
3975 let card = handler
3976 .get_extended_agent_card(&ServiceParams::new(), extended_request())
3977 .await
3978 .unwrap();
3979 assert_eq!(card.name, "limited");
3980
3981 let mut params = ServiceParams::new();
3982 params.insert("x-privileged".into(), vec!["yes".into()]);
3983 let card = handler
3984 .get_extended_agent_card(¶ms, extended_request())
3985 .await
3986 .unwrap();
3987 assert_eq!(card.name, "full");
3988 }
3989}
3990
3991#[cfg(test)]
3992mod lag_coverage_tests {
3993 use super::*;
3994 use futures::StreamExt;
3995
3996 #[tokio::test]
3997 async fn test_subscription_lag_is_reported_as_a_sanitized_internal_error() {
3998 let (sender, receiver) = broadcast::channel(1);
4002
4003 for sequence in 0..4 {
4004 let _ = sender.send(ExecutionEvent {
4005 sequence,
4006 result: Ok(StreamResponse::Message(a2a::Message {
4007 message_id: "lag-probe".to_string(),
4008 role: a2a::Role::Agent,
4009 parts: Vec::new(),
4010 context_id: None,
4011 task_id: None,
4012 metadata: None,
4013 extensions: None,
4014 reference_task_ids: None,
4015 })),
4016 });
4017 }
4018
4019 let mut stream = subscription_stream(receiver, None, 0);
4020 let item = stream.next().await.expect("lagged item");
4021
4022 let error = item.expect_err("lag must surface as an error");
4023 assert_eq!(error.code, a2a::error_code::INTERNAL_ERROR);
4024 assert_eq!(error.message, "Internal error");
4025
4026 assert!(stream.next().await.is_none());
4028 }
4029}