1use std::collections::{BTreeSet, HashMap};
103use std::fmt::Write as _;
104use std::sync::atomic::{AtomicBool, Ordering};
105use std::sync::{Arc, RwLock};
106use std::time::{Duration, Instant};
107
108use async_trait::async_trait;
109
110use crate::error::JsonRpcError;
111use crate::protocol::{CallToolResult, InputRequests, InputResponses, TaskObject, TaskStatus};
112
113const DEFAULT_TTL_MS: u64 = 300_000;
118
119const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
121
122#[derive(Debug)]
124pub struct Task {
125 pub id: String,
127 pub tool_name: String,
129 pub arguments: serde_json::Value,
131 pub status: TaskStatus,
133 pub created_at: Instant,
135 pub created_at_str: String,
137 pub last_updated_at_str: String,
139 pub ttl: u64,
141 pub poll_interval: u64,
143 pub status_message: Option<String>,
145 pub result: Option<CallToolResult>,
147 pub error: Option<JsonRpcError>,
153 pub owner: TaskOwner,
158 pub input_requests: InputRequests,
160 pub answered_input_keys: BTreeSet<String>,
162 pub superseded_input_keys: BTreeSet<String>,
165 pub cancellation_token: CancellationToken,
167 pub completed_at: Option<Instant>,
169 pub completion_notify: Arc<tokio::sync::Notify>,
171}
172
173impl Task {
174 fn new(
176 id: String,
177 tool_name: String,
178 arguments: serde_json::Value,
179 ttl: Option<u64>,
180 owner: TaskOwner,
181 ) -> Self {
182 let cancelled = Arc::new(AtomicBool::new(false));
183 let now_str = chrono_now_iso8601();
184 Self {
185 id,
186 tool_name,
187 arguments,
188 status: TaskStatus::Working,
189 created_at: Instant::now(),
190 created_at_str: now_str.clone(),
191 last_updated_at_str: now_str,
192 ttl: ttl.unwrap_or(DEFAULT_TTL_MS),
193 poll_interval: DEFAULT_POLL_INTERVAL_MS,
194 status_message: Some("Task started".to_string()),
195 result: None,
196 error: None,
197 owner,
198 input_requests: InputRequests::new(),
199 answered_input_keys: BTreeSet::new(),
200 superseded_input_keys: BTreeSet::new(),
201 cancellation_token: CancellationToken { cancelled },
202 completed_at: None,
203 completion_notify: Arc::new(tokio::sync::Notify::new()),
204 }
205 }
206
207 pub fn to_task_object(&self) -> TaskObject {
209 TaskObject {
210 task_id: self.id.clone(),
211 status: self.status,
212 status_message: self.status_message.clone(),
213 created_at: self.created_at_str.clone(),
214 last_updated_at: self.last_updated_at_str.clone(),
215 ttl: Some(self.ttl),
216 poll_interval: Some(self.poll_interval),
217 result: None,
218 error: None,
219 meta: None,
220 }
221 }
222
223 pub fn is_expired(&self) -> bool {
230 self.created_at.elapsed() > Duration::from_millis(self.ttl)
231 }
232
233 pub fn outstanding_input_requests(&self) -> &InputRequests {
235 &self.input_requests
236 }
237
238 pub fn is_cancelled(&self) -> bool {
240 self.cancellation_token.is_cancelled()
241 }
242}
243
244pub fn generate_task_id() -> String {
257 let mut bytes = [0u8; 16];
258 getrandom::fill(&mut bytes).expect("system entropy source unavailable for task ID generation");
259 let mut id = String::with_capacity(2 * bytes.len());
260 for byte in bytes {
261 let _ = write!(id, "{byte:02x}");
262 }
263 id
264}
265
266pub type TaskOwner = Option<String>;
275
276pub fn owner_matches(owner: &TaskOwner, principal: Option<&str>) -> bool {
283 owner.as_deref() == principal
284}
285
286#[derive(Debug, Clone, Default, PartialEq, Eq)]
292pub struct AppliedInputResponses {
293 pub accepted: BTreeSet<String>,
295 pub ignored: BTreeSet<String>,
298 pub still_outstanding: BTreeSet<String>,
300}
301
302impl AppliedInputResponses {
303 pub fn is_complete(&self) -> bool {
305 self.still_outstanding.is_empty()
306 }
307}
308
309#[derive(Debug, Clone)]
311pub struct CancellationToken {
312 cancelled: Arc<AtomicBool>,
313}
314
315impl CancellationToken {
316 pub fn is_cancelled(&self) -> bool {
318 self.cancelled.load(Ordering::Relaxed)
319 }
320
321 pub fn cancel(&self) {
323 self.cancelled.store(true, Ordering::Relaxed);
324 }
325}
326
327#[derive(Debug, thiserror::Error)]
335#[non_exhaustive]
336pub enum TaskStoreError {
337 #[error("encode error: {0}")]
339 Encode(String),
340 #[error("decode error: {0}")]
342 Decode(String),
343 #[error("backend error: {0}")]
345 Backend(String),
346}
347
348pub type Result<T> = std::result::Result<T, TaskStoreError>;
350
351pub type TaskSnapshot = (TaskObject, Option<CallToolResult>, Option<JsonRpcError>);
357
358#[async_trait]
380pub trait TaskStore: Send + Sync + 'static {
381 async fn create_task(
387 &self,
388 tool_name: &str,
389 arguments: serde_json::Value,
390 ttl: Option<u64>,
391 owner: TaskOwner,
392 ) -> Result<(String, CancellationToken)>;
393
394 async fn task_owner(&self, task_id: &str) -> Result<Option<TaskOwner>>;
400
401 async fn get_task(&self, task_id: &str) -> Result<Option<TaskObject>>;
403
404 async fn get_task_result(&self, task_id: &str) -> Result<Option<TaskSnapshot>>;
406
407 async fn wait_for_completion(&self, task_id: &str) -> Result<Option<TaskSnapshot>>;
413
414 async fn list_tasks(&self, status_filter: Option<TaskStatus>) -> Result<Vec<TaskObject>>;
416
417 async fn require_input(
426 &self,
427 task_id: &str,
428 requests: InputRequests,
429 message: Option<&str>,
430 ) -> Result<bool>;
431
432 async fn outstanding_input_requests(&self, task_id: &str) -> Result<Option<InputRequests>>;
437
438 async fn apply_input_responses(
446 &self,
447 task_id: &str,
448 responses: InputResponses,
449 ) -> Result<Option<AppliedInputResponses>>;
450
451 async fn set_ttl(&self, task_id: &str, ttl_ms: u64) -> Result<bool>;
456
457 async fn complete_task(&self, task_id: &str, result: CallToolResult) -> Result<bool>;
466
467 async fn fail_task(&self, task_id: &str, error: JsonRpcError) -> Result<bool>;
472
473 async fn cancel_task(&self, task_id: &str, reason: Option<&str>) -> Result<Option<TaskObject>>;
479}
480
481#[derive(Debug, Clone)]
489pub struct MemoryTaskStore {
490 tasks: Arc<RwLock<HashMap<String, Task>>>,
491}
492
493impl Default for MemoryTaskStore {
494 fn default() -> Self {
495 Self::new()
496 }
497}
498
499impl MemoryTaskStore {
500 pub fn new() -> Self {
502 Self {
503 tasks: Arc::new(RwLock::new(HashMap::new())),
504 }
505 }
506
507 pub fn cleanup_expired(&self) -> usize {
515 if let Ok(mut tasks) = self.tasks.write() {
516 let before = tasks.len();
517 tasks.retain(|_, t| !t.is_expired());
518 before - tasks.len()
519 } else {
520 0
521 }
522 }
523
524 #[cfg(test)]
526 pub fn len(&self) -> usize {
527 if let Ok(tasks) = self.tasks.read() {
528 tasks.len()
529 } else {
530 0
531 }
532 }
533
534 #[cfg(test)]
536 pub fn is_empty(&self) -> bool {
537 self.len() == 0
538 }
539}
540
541#[async_trait]
542impl TaskStore for MemoryTaskStore {
543 async fn create_task(
544 &self,
545 tool_name: &str,
546 arguments: serde_json::Value,
547 ttl: Option<u64>,
548 owner: TaskOwner,
549 ) -> Result<(String, CancellationToken)> {
550 let id = generate_task_id();
551 let task = Task::new(id.clone(), tool_name.to_string(), arguments, ttl, owner);
552 let token = task.cancellation_token.clone();
553
554 if let Ok(mut tasks) = self.tasks.write() {
555 tasks.insert(id.clone(), task);
556 }
557
558 Ok((id, token))
559 }
560
561 async fn get_task(&self, task_id: &str) -> Result<Option<TaskObject>> {
562 Ok(if let Ok(tasks) = self.tasks.read() {
563 tasks
564 .get(task_id)
565 .filter(|t| !t.is_expired())
566 .map(|t| t.to_task_object())
567 } else {
568 None
569 })
570 }
571
572 async fn task_owner(&self, task_id: &str) -> Result<Option<TaskOwner>> {
573 Ok(if let Ok(tasks) = self.tasks.read() {
574 tasks
575 .get(task_id)
576 .filter(|t| !t.is_expired())
577 .map(|t| t.owner.clone())
578 } else {
579 None
580 })
581 }
582
583 async fn get_task_result(&self, task_id: &str) -> Result<Option<TaskSnapshot>> {
584 Ok(if let Ok(tasks) = self.tasks.read() {
585 tasks
586 .get(task_id)
587 .filter(|t| !t.is_expired())
588 .map(|t| (t.to_task_object(), t.result.clone(), t.error.clone()))
589 } else {
590 None
591 })
592 }
593
594 async fn wait_for_completion(&self, task_id: &str) -> Result<Option<TaskSnapshot>> {
595 let notify = {
597 let Ok(tasks) = self.tasks.read() else {
598 return Ok(None);
599 };
600 let Some(task) = tasks.get(task_id).filter(|t| !t.is_expired()) else {
601 return Ok(None);
602 };
603 if task.status.is_terminal() {
604 return Ok(Some((
605 task.to_task_object(),
606 task.result.clone(),
607 task.error.clone(),
608 )));
609 }
610 task.completion_notify.clone()
611 };
612
613 notify.notified().await;
615
616 self.get_task_result(task_id).await
618 }
619
620 async fn list_tasks(&self, status_filter: Option<TaskStatus>) -> Result<Vec<TaskObject>> {
621 Ok(if let Ok(tasks) = self.tasks.read() {
622 tasks
623 .values()
624 .filter(|t| !t.is_expired())
625 .filter(|t| status_filter.is_none() || status_filter == Some(t.status))
626 .map(|t| t.to_task_object())
627 .collect()
628 } else {
629 vec![]
630 })
631 }
632
633 async fn require_input(
634 &self,
635 task_id: &str,
636 requests: InputRequests,
637 message: Option<&str>,
638 ) -> Result<bool> {
639 let Ok(mut tasks) = self.tasks.write() else {
640 return Ok(false);
641 };
642 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
643 return Ok(false);
644 };
645 if task.status.is_terminal() {
646 return Ok(false);
647 }
648
649 for key in std::mem::take(&mut task.input_requests).into_keys() {
651 if !requests.contains_key(&key) {
652 task.superseded_input_keys.insert(key);
653 }
654 }
655 for key in requests.keys() {
657 task.answered_input_keys.remove(key);
658 task.superseded_input_keys.remove(key);
659 }
660
661 task.input_requests = requests;
662 task.status = TaskStatus::InputRequired;
663 task.status_message = Some(
664 message
665 .map(str::to_string)
666 .unwrap_or_else(|| "Awaiting client input".to_string()),
667 );
668 task.last_updated_at_str = chrono_now_iso8601();
669 Ok(true)
670 }
671
672 async fn outstanding_input_requests(&self, task_id: &str) -> Result<Option<InputRequests>> {
673 Ok(if let Ok(tasks) = self.tasks.read() {
674 tasks
675 .get(task_id)
676 .filter(|t| !t.is_expired())
677 .map(|t| t.input_requests.clone())
678 } else {
679 None
680 })
681 }
682
683 async fn apply_input_responses(
684 &self,
685 task_id: &str,
686 responses: InputResponses,
687 ) -> Result<Option<AppliedInputResponses>> {
688 let Ok(mut tasks) = self.tasks.write() else {
689 return Ok(None);
690 };
691 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
692 return Ok(None);
693 };
694 if task.status.is_terminal() {
695 return Ok(None);
696 }
697
698 let mut applied = AppliedInputResponses::default();
699 for key in responses.into_keys() {
700 if task.input_requests.remove(&key).is_some() {
701 task.answered_input_keys.insert(key.clone());
702 applied.accepted.insert(key);
703 } else {
704 applied.ignored.insert(key);
708 }
709 }
710 applied.still_outstanding = task.input_requests.keys().cloned().collect();
711
712 if !applied.accepted.is_empty() {
713 task.last_updated_at_str = chrono_now_iso8601();
714 }
715 if applied.is_complete() && task.status == TaskStatus::InputRequired {
716 task.status = TaskStatus::Working;
717 task.status_message = Some("Task resumed".to_string());
718 }
719 Ok(Some(applied))
720 }
721
722 async fn set_ttl(&self, task_id: &str, ttl_ms: u64) -> Result<bool> {
723 let Ok(mut tasks) = self.tasks.write() else {
724 return Ok(false);
725 };
726 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
727 return Ok(false);
728 };
729 task.ttl = ttl_ms;
730 task.last_updated_at_str = chrono_now_iso8601();
731 Ok(true)
732 }
733
734 async fn complete_task(&self, task_id: &str, result: CallToolResult) -> Result<bool> {
735 let Ok(mut tasks) = self.tasks.write() else {
736 return Ok(false);
737 };
738 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
739 return Ok(false);
740 };
741 if task.status.is_terminal() {
742 return Ok(false);
743 }
744 task.status = TaskStatus::Completed;
745 task.status_message = Some("Task completed".to_string());
746 task.result = Some(result);
747 task.input_requests.clear();
748 task.completed_at = Some(Instant::now());
749 task.last_updated_at_str = chrono_now_iso8601();
750 task.completion_notify.notify_waiters();
751 Ok(true)
752 }
753
754 async fn fail_task(&self, task_id: &str, error: JsonRpcError) -> Result<bool> {
755 let Ok(mut tasks) = self.tasks.write() else {
756 return Ok(false);
757 };
758 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
759 return Ok(false);
760 };
761 if task.status.is_terminal() {
762 return Ok(false);
763 }
764 task.status = TaskStatus::Failed;
765 task.status_message = Some(format!("Task failed: {}", error.message));
766 task.error = Some(error);
767 task.input_requests.clear();
768 task.completed_at = Some(Instant::now());
769 task.last_updated_at_str = chrono_now_iso8601();
770 task.completion_notify.notify_waiters();
771 Ok(true)
772 }
773
774 async fn cancel_task(&self, task_id: &str, reason: Option<&str>) -> Result<Option<TaskObject>> {
775 let Ok(mut tasks) = self.tasks.write() else {
776 return Ok(None);
777 };
778 let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
779 return Ok(None);
780 };
781
782 task.cancellation_token.cancel();
784
785 if !task.status.is_terminal() {
787 task.input_requests.clear();
788 task.status = TaskStatus::Cancelled;
789 task.status_message = Some(
790 reason
791 .map(|r| format!("Cancelled: {}", r))
792 .unwrap_or_else(|| "Task cancelled".to_string()),
793 );
794 task.completed_at = Some(Instant::now());
795 task.last_updated_at_str = chrono_now_iso8601();
796 task.completion_notify.notify_waiters();
797 }
798 Ok(Some(task.to_task_object()))
799 }
800}
801
802pub fn tasks_extension() -> crate::ExtensionDeclaration {
807 crate::ExtensionDeclaration::empty(crate::protocol::TASKS_EXTENSION_ID)
808 .expect("the built-in Tasks extension declaration is valid")
809}
810
811impl crate::McpRouter {
812 pub fn with_tasks(self) -> Self {
820 self.with_protocol_extension(tasks_extension())
821 }
822}
823
824impl crate::McpClientBuilder {
825 pub fn with_tasks(self) -> Self {
827 self.with_protocol_extension(tasks_extension())
828 }
829}
830
831impl crate::RequestContext {
832 pub fn supports_tasks(&self) -> bool {
838 self.negotiated_extensions()
839 .is_some_and(|extensions| extensions.contains(crate::protocol::TASKS_EXTENSION_ID))
840 }
841}
842
843fn chrono_now_iso8601() -> String {
845 use std::time::SystemTime;
846
847 let now = SystemTime::now();
848 let duration = now
849 .duration_since(SystemTime::UNIX_EPOCH)
850 .unwrap_or_default();
851
852 let secs = duration.as_secs();
853 let millis = duration.subsec_millis();
854
855 let days = secs / 86400;
858 let remaining = secs % 86400;
859 let hours = remaining / 3600;
860 let remaining = remaining % 3600;
861 let minutes = remaining / 60;
862 let seconds = remaining % 60;
863
864 let mut year = 1970i32;
867 let mut remaining_days = days as i32;
868
869 loop {
870 let days_in_year = if is_leap_year(year) { 366 } else { 365 };
871 if remaining_days < days_in_year {
872 break;
873 }
874 remaining_days -= days_in_year;
875 year += 1;
876 }
877
878 let days_in_months: [i32; 12] = if is_leap_year(year) {
879 [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
880 } else {
881 [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
882 };
883
884 let mut month = 1;
885 for days_in_month in days_in_months.iter() {
886 if remaining_days < *days_in_month {
887 break;
888 }
889 remaining_days -= days_in_month;
890 month += 1;
891 }
892
893 let day = remaining_days + 1;
894
895 format!(
896 "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:03}Z",
897 year, month, day, hours, minutes, seconds, millis
898 )
899}
900
901fn is_leap_year(year: i32) -> bool {
902 (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0)
903}
904
905#[cfg(test)]
906mod tests {
907 use super::*;
908 use crate::protocol::{
909 ElicitAction, ElicitResult, InputRequest, InputResponse, ListRootsParams,
910 };
911
912 #[tokio::test]
913 async fn test_create_task() {
914 let store = MemoryTaskStore::new();
915 let (id, token) = store
916 .create_task("test-tool", serde_json::json!({"a": 1}), None, None)
917 .await
918 .unwrap();
919
920 assert!(!id.is_empty());
921 assert!(!token.is_cancelled());
922
923 let info = store
924 .get_task(&id)
925 .await
926 .unwrap()
927 .expect("task should exist");
928 assert_eq!(info.task_id, id);
929 assert_eq!(info.status, TaskStatus::Working);
930 }
931
932 #[tokio::test]
933 async fn test_task_lifecycle() {
934 let store = MemoryTaskStore::new();
935 let (id, _) = store
936 .create_task("test-tool", serde_json::json!({}), None, None)
937 .await
938 .unwrap();
939
940 assert!(
942 store
943 .complete_task(&id, CallToolResult::text("Done"))
944 .await
945 .unwrap()
946 );
947
948 let info = store.get_task(&id).await.unwrap().unwrap();
949 assert_eq!(info.status, TaskStatus::Completed);
950 }
951
952 #[tokio::test]
953 async fn test_task_cancellation() {
954 let store = MemoryTaskStore::new();
955 let (id, token) = store
956 .create_task("test-tool", serde_json::json!({}), None, None)
957 .await
958 .unwrap();
959
960 assert!(!token.is_cancelled());
961
962 let task_obj = store
963 .cancel_task(&id, Some("User requested"))
964 .await
965 .unwrap();
966 assert!(task_obj.is_some());
967 assert_eq!(task_obj.unwrap().status, TaskStatus::Cancelled);
968 assert!(token.is_cancelled());
969
970 let info = store.get_task(&id).await.unwrap().unwrap();
971 assert_eq!(info.status, TaskStatus::Cancelled);
972 }
973
974 #[tokio::test]
975 async fn test_task_failure() {
976 let store = MemoryTaskStore::new();
977 let (id, _) = store
978 .create_task("test-tool", serde_json::json!({}), None, None)
979 .await
980 .unwrap();
981
982 assert!(
983 store
984 .fail_task(&id, JsonRpcError::internal_error("Something went wrong"))
985 .await
986 .unwrap()
987 );
988
989 let info = store.get_task(&id).await.unwrap().unwrap();
990 assert_eq!(info.status, TaskStatus::Failed);
991 assert!(info.status_message.as_ref().unwrap().contains("failed"));
992 }
993
994 #[tokio::test]
995 async fn test_list_tasks() {
996 let store = MemoryTaskStore::new();
997 store
998 .create_task("tool1", serde_json::json!({}), None, None)
999 .await
1000 .unwrap();
1001 store
1002 .create_task("tool2", serde_json::json!({}), None, None)
1003 .await
1004 .unwrap();
1005 let (id3, _) = store
1006 .create_task("tool3", serde_json::json!({}), None, None)
1007 .await
1008 .unwrap();
1009
1010 store
1012 .complete_task(&id3, CallToolResult::text("Done"))
1013 .await
1014 .unwrap();
1015
1016 let all = store.list_tasks(None).await.unwrap();
1018 assert_eq!(all.len(), 3);
1019
1020 let working = store.list_tasks(Some(TaskStatus::Working)).await.unwrap();
1022 assert_eq!(working.len(), 2);
1023
1024 let completed = store.list_tasks(Some(TaskStatus::Completed)).await.unwrap();
1026 assert_eq!(completed.len(), 1);
1027 }
1028
1029 #[tokio::test]
1030 async fn test_terminal_state_immutable() {
1031 let store = MemoryTaskStore::new();
1032 let (id, _) = store
1033 .create_task("test-tool", serde_json::json!({}), None, None)
1034 .await
1035 .unwrap();
1036
1037 store
1039 .complete_task(&id, CallToolResult::text("Done"))
1040 .await
1041 .unwrap();
1042
1043 assert!(
1045 !store
1046 .fail_task(&id, JsonRpcError::internal_error("Error"))
1047 .await
1048 .unwrap()
1049 );
1050
1051 let info = store.get_task(&id).await.unwrap().unwrap();
1053 assert_eq!(info.status, TaskStatus::Completed);
1054 }
1055
1056 #[tokio::test]
1057 async fn test_task_ids_unique() {
1058 let store = MemoryTaskStore::new();
1059 let (id1, _) = store
1060 .create_task("tool", serde_json::json!({}), None, None)
1061 .await
1062 .unwrap();
1063 let (id2, _) = store
1064 .create_task("tool", serde_json::json!({}), None, None)
1065 .await
1066 .unwrap();
1067 let (id3, _) = store
1068 .create_task("tool", serde_json::json!({}), None, None)
1069 .await
1070 .unwrap();
1071
1072 assert_ne!(id1, id2);
1073 assert_ne!(id2, id3);
1074 assert_ne!(id1, id3);
1075 }
1076
1077 #[tokio::test]
1078 async fn test_get_task_result() {
1079 let store = MemoryTaskStore::new();
1080 let (id, _) = store
1081 .create_task("test-tool", serde_json::json!({}), None, None)
1082 .await
1083 .unwrap();
1084
1085 let result = CallToolResult::text("The result");
1087 store.complete_task(&id, result).await.unwrap();
1088
1089 let (task_obj, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1090 assert_eq!(task_obj.status, TaskStatus::Completed);
1091 assert!(result.is_some());
1092 assert!(error.is_none());
1093 }
1094
1095 #[tokio::test]
1096 async fn test_wait_for_completion_returns_terminal_snapshot() {
1097 let store = MemoryTaskStore::new();
1098 let (id, _) = store
1099 .create_task("test-tool", serde_json::json!({}), None, None)
1100 .await
1101 .unwrap();
1102
1103 let waiter_store = store.clone();
1105 let waiter_id = id.clone();
1106 let waiter =
1107 tokio::spawn(async move { waiter_store.wait_for_completion(&waiter_id).await });
1108
1109 tokio::time::sleep(Duration::from_millis(10)).await;
1110 store
1111 .complete_task(&id, CallToolResult::text("Done"))
1112 .await
1113 .unwrap();
1114
1115 let (task_obj, result, error) = waiter.await.unwrap().unwrap().unwrap();
1116 assert_eq!(task_obj.status, TaskStatus::Completed);
1117 assert!(result.is_some());
1118 assert!(error.is_none());
1119 }
1120
1121 #[tokio::test]
1122 async fn dyn_task_store_object_safe() {
1123 let store: Arc<dyn TaskStore> = Arc::new(MemoryTaskStore::new());
1125 let (id, _) = store
1126 .create_task("tool", serde_json::json!({}), None, None)
1127 .await
1128 .unwrap();
1129 assert!(store.get_task(&id).await.unwrap().is_some());
1130 }
1131
1132 #[test]
1133 fn test_iso8601_timestamp() {
1134 let ts = chrono_now_iso8601();
1135 assert!(ts.ends_with('Z'));
1137 assert!(ts.contains('T'));
1138 assert_eq!(ts.len(), 24); }
1140
1141 #[test]
1142 fn test_task_status_display() {
1143 assert_eq!(TaskStatus::Working.to_string(), "working");
1144 assert_eq!(TaskStatus::InputRequired.to_string(), "input_required");
1145 assert_eq!(TaskStatus::Completed.to_string(), "completed");
1146 assert_eq!(TaskStatus::Failed.to_string(), "failed");
1147 assert_eq!(TaskStatus::Cancelled.to_string(), "cancelled");
1148 }
1149
1150 #[test]
1151 fn test_task_status_is_terminal() {
1152 assert!(!TaskStatus::Working.is_terminal());
1153 assert!(!TaskStatus::InputRequired.is_terminal());
1154 assert!(TaskStatus::Completed.is_terminal());
1155 assert!(TaskStatus::Failed.is_terminal());
1156 assert!(TaskStatus::Cancelled.is_terminal());
1157 }
1158
1159 fn requests(keys: &[&str]) -> InputRequests {
1160 keys.iter()
1161 .map(|k| {
1162 (
1163 k.to_string(),
1164 InputRequest::ListRoots(ListRootsParams { meta: None }),
1165 )
1166 })
1167 .collect()
1168 }
1169
1170 fn accept(key: &str) -> (String, InputResponse) {
1171 (
1172 key.to_string(),
1173 InputResponse::Elicit(ElicitResult {
1174 action: ElicitAction::Accept,
1175 content: None,
1176 meta: None,
1177 }),
1178 )
1179 }
1180
1181 async fn working_task(store: &MemoryTaskStore, ttl: Option<u64>) -> String {
1182 store
1183 .create_task("tool", serde_json::json!({}), ttl, None)
1184 .await
1185 .unwrap()
1186 .0
1187 }
1188
1189 #[tokio::test]
1190 async fn task_ids_are_unguessable_not_sequential() {
1191 let store = MemoryTaskStore::new();
1192 let mut ids = BTreeSet::new();
1193 for _ in 0..64 {
1194 ids.insert(working_task(&store, None).await);
1195 }
1196 assert_eq!(ids.len(), 64, "task IDs collided");
1197
1198 for id in &ids {
1199 assert_eq!(id.len(), 32, "expected 128 bits of hex: {id}");
1200 assert!(id.chars().all(|c| c.is_ascii_hexdigit()), "{id}");
1201 assert!(!id.starts_with("task-"), "sequential-looking ID: {id}");
1202 }
1203
1204 let leading: BTreeSet<&str> = ids.iter().map(|id| &id[..2]).collect();
1207 assert!(
1208 leading.len() > 32,
1209 "only {} distinct leading bytes across 64 IDs",
1210 leading.len()
1211 );
1212 }
1213
1214 #[tokio::test]
1215 async fn ttl_runs_from_creation_and_expired_tasks_read_as_absent() {
1216 let store = MemoryTaskStore::new();
1217 let id = working_task(&store, Some(0)).await;
1218
1219 tokio::time::sleep(Duration::from_millis(5)).await;
1222
1223 assert!(store.get_task(&id).await.unwrap().is_none());
1224 assert!(store.get_task_result(&id).await.unwrap().is_none());
1225 assert!(store.list_tasks(None).await.unwrap().is_empty());
1226 assert!(
1227 store
1228 .outstanding_input_requests(&id)
1229 .await
1230 .unwrap()
1231 .is_none()
1232 );
1233 assert!(store.cancel_task(&id, None).await.unwrap().is_none());
1234 assert!(!store.set_ttl(&id, 60_000).await.unwrap());
1235 assert!(
1236 !store
1237 .complete_task(&id, CallToolResult::text("late"))
1238 .await
1239 .unwrap()
1240 );
1241 }
1242
1243 #[tokio::test]
1244 async fn ttl_is_mutable_over_the_task_lifetime() {
1245 let store = MemoryTaskStore::new();
1246 let id = working_task(&store, Some(60_000)).await;
1247
1248 assert!(store.set_ttl(&id, 120_000).await.unwrap());
1249 let task = store.get_task(&id).await.unwrap().unwrap();
1250 assert_eq!(task.ttl, Some(120_000));
1251
1252 assert!(store.set_ttl(&id, 0).await.unwrap());
1254 tokio::time::sleep(Duration::from_millis(5)).await;
1255 assert!(store.get_task(&id).await.unwrap().is_none());
1256 }
1257
1258 #[tokio::test]
1259 async fn require_input_records_requests_and_exposes_them() {
1260 let store = MemoryTaskStore::new();
1261 let id = working_task(&store, None).await;
1262
1263 assert!(
1264 store
1265 .require_input(&id, requests(&["approval", "region"]), Some("need input"))
1266 .await
1267 .unwrap()
1268 );
1269
1270 let task = store.get_task(&id).await.unwrap().unwrap();
1271 assert_eq!(task.status, TaskStatus::InputRequired);
1272 assert_eq!(task.status_message.as_deref(), Some("need input"));
1273
1274 let outstanding = store
1275 .outstanding_input_requests(&id)
1276 .await
1277 .unwrap()
1278 .unwrap();
1279 assert_eq!(
1280 outstanding.keys().collect::<Vec<_>>(),
1281 vec!["approval", "region"],
1282 "every outstanding request must be exposed, not just the newest"
1283 );
1284 }
1285
1286 #[tokio::test]
1287 async fn partial_input_responses_leave_the_rest_outstanding() {
1288 let store = MemoryTaskStore::new();
1289 let id = working_task(&store, None).await;
1290 store
1291 .require_input(&id, requests(&["approval", "region"]), None)
1292 .await
1293 .unwrap();
1294
1295 let applied = store
1296 .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1297 .await
1298 .unwrap()
1299 .unwrap();
1300
1301 assert_eq!(applied.accepted, ["approval".to_string()].into());
1302 assert!(applied.ignored.is_empty());
1303 assert_eq!(applied.still_outstanding, ["region".to_string()].into());
1304 assert!(!applied.is_complete());
1305
1306 let task = store.get_task(&id).await.unwrap().unwrap();
1308 assert_eq!(task.status, TaskStatus::InputRequired);
1309 assert_eq!(
1310 store
1311 .outstanding_input_requests(&id)
1312 .await
1313 .unwrap()
1314 .unwrap()
1315 .keys()
1316 .collect::<Vec<_>>(),
1317 vec!["region"]
1318 );
1319
1320 let applied = store
1322 .apply_input_responses(&id, [accept("region")].into_iter().collect())
1323 .await
1324 .unwrap()
1325 .unwrap();
1326 assert!(applied.is_complete());
1327 assert_eq!(
1328 store.get_task(&id).await.unwrap().unwrap().status,
1329 TaskStatus::Working
1330 );
1331 }
1332
1333 #[tokio::test]
1334 async fn unknown_answered_and_superseded_response_keys_are_ignored() {
1335 let store = MemoryTaskStore::new();
1336 let id = working_task(&store, None).await;
1337 store
1338 .require_input(&id, requests(&["approval", "stale"]), None)
1339 .await
1340 .unwrap();
1341
1342 store
1344 .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1345 .await
1346 .unwrap()
1347 .unwrap();
1348 store
1349 .require_input(&id, requests(&["region"]), None)
1350 .await
1351 .unwrap();
1352
1353 let applied = store
1354 .apply_input_responses(
1355 &id,
1356 [accept("never-issued"), accept("approval"), accept("stale")]
1357 .into_iter()
1358 .collect(),
1359 )
1360 .await
1361 .unwrap()
1362 .unwrap();
1363
1364 assert!(
1365 applied.accepted.is_empty(),
1366 "none of these keys are outstanding"
1367 );
1368 assert_eq!(
1369 applied.ignored,
1370 [
1371 "never-issued".to_string(),
1372 "approval".to_string(),
1373 "stale".to_string()
1374 ]
1375 .into(),
1376 "unknown, already-answered, and superseded keys are all ignored"
1377 );
1378 assert_eq!(applied.still_outstanding, ["region".to_string()].into());
1379 assert_eq!(
1380 store.get_task(&id).await.unwrap().unwrap().status,
1381 TaskStatus::InputRequired,
1382 "ignoring a stale update must not resume or fail the task"
1383 );
1384 }
1385
1386 #[tokio::test]
1387 async fn reissued_key_becomes_a_fresh_question() {
1388 let store = MemoryTaskStore::new();
1389 let id = working_task(&store, None).await;
1390 store
1391 .require_input(&id, requests(&["approval"]), None)
1392 .await
1393 .unwrap();
1394 store
1395 .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1396 .await
1397 .unwrap()
1398 .unwrap();
1399
1400 store
1403 .require_input(&id, requests(&["approval"]), None)
1404 .await
1405 .unwrap();
1406 let applied = store
1407 .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1408 .await
1409 .unwrap()
1410 .unwrap();
1411 assert_eq!(applied.accepted, ["approval".to_string()].into());
1412 assert!(applied.is_complete());
1413 }
1414
1415 #[tokio::test]
1416 async fn failed_tasks_preserve_the_structured_error() {
1417 let store = MemoryTaskStore::new();
1418 let id = working_task(&store, None).await;
1419
1420 let mut error = JsonRpcError::invalid_params("bad region");
1421 error.data = Some(serde_json::json!({"field": "region"}));
1422 assert!(store.fail_task(&id, error).await.unwrap());
1423
1424 let (_, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1425 assert!(result.is_none());
1426 let error = error.expect("structured error must survive the store");
1427 assert_eq!(
1428 error.code, -32602,
1429 "the original code must not be flattened"
1430 );
1431 assert_eq!(error.message, "bad region");
1432 assert_eq!(error.data.unwrap()["field"], "region");
1433 }
1434
1435 #[tokio::test]
1436 async fn tool_error_results_complete_the_task() {
1437 let store = MemoryTaskStore::new();
1438 let id = working_task(&store, None).await;
1439
1440 let mut result = CallToolResult::text("domain failure");
1441 result.is_error = true;
1442 assert!(store.complete_task(&id, result).await.unwrap());
1443
1444 let (task, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1445 assert_eq!(
1446 task.status,
1447 TaskStatus::Completed,
1448 "isError is a domain error, not an execution failure"
1449 );
1450 assert!(result.unwrap().is_error);
1451 assert!(error.is_none(), "no JSON-RPC error accompanies isError");
1452 }
1453
1454 #[tokio::test]
1455 async fn tasks_record_their_creating_principal() {
1456 let store = MemoryTaskStore::new();
1457 let (owned, _) = store
1458 .create_task("tool", serde_json::json!({}), None, Some("alice".into()))
1459 .await
1460 .unwrap();
1461 let (unowned, _) = store
1462 .create_task("tool", serde_json::json!({}), None, None)
1463 .await
1464 .unwrap();
1465
1466 assert_eq!(
1467 store.task_owner(&owned).await.unwrap(),
1468 Some(Some("alice".to_string()))
1469 );
1470 assert_eq!(store.task_owner(&unowned).await.unwrap(), Some(None));
1471 assert_eq!(
1472 store.task_owner("does-not-exist").await.unwrap(),
1473 None,
1474 "an unknown task has no owner record at all"
1475 );
1476
1477 let wire = serde_json::to_value(store.get_task(&owned).await.unwrap().unwrap()).unwrap();
1479 assert!(
1480 wire.get("owner").is_none(),
1481 "owner leaked to the wire: {wire}"
1482 );
1483 assert!(!wire.to_string().contains("alice"));
1484 }
1485
1486 #[test]
1487 fn owner_matching_is_equality_not_leniency() {
1488 assert!(owner_matches(&None, None), "no auth configured");
1489 assert!(owner_matches(&Some("alice".into()), Some("alice")));
1490
1491 assert!(
1492 !owner_matches(&Some("alice".into()), Some("bob")),
1493 "a different principal must not inherit the task"
1494 );
1495 assert!(
1496 !owner_matches(&Some("alice".into()), None),
1497 "dropping the token must not grant access"
1498 );
1499 assert!(
1500 !owner_matches(&None, Some("alice")),
1501 "an unowned task belongs to a different security context"
1502 );
1503 }
1504
1505 #[tokio::test]
1506 async fn terminal_states_clear_outstanding_requests() {
1507 for (label, terminate) in [("completed", true), ("cancelled", false)] {
1508 let store = MemoryTaskStore::new();
1509 let id = working_task(&store, None).await;
1510 store
1511 .require_input(&id, requests(&["approval"]), None)
1512 .await
1513 .unwrap();
1514
1515 if terminate {
1516 store
1517 .complete_task(&id, CallToolResult::text("done"))
1518 .await
1519 .unwrap();
1520 } else {
1521 store.cancel_task(&id, None).await.unwrap();
1522 }
1523
1524 assert!(
1525 store
1526 .outstanding_input_requests(&id)
1527 .await
1528 .unwrap()
1529 .unwrap()
1530 .is_empty(),
1531 "{label} task still advertises outstanding input requests"
1532 );
1533 assert!(
1534 store
1535 .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1536 .await
1537 .unwrap()
1538 .is_none(),
1539 "{label} task accepted a late input response"
1540 );
1541 }
1542 }
1543}