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