a2a_server/task_store/
inmemory.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// Copyright A2A Contributors (https://github.com/a2aproject)
3// SPDX-License-Identifier: Apache-2.0
4use a2a::*;
5use async_trait::async_trait;
6use std::collections::HashMap;
7use tokio::sync::RwLock;
8
9use super::apply_history_length;
10use super::store::{TaskStore, TaskVersion};
11use crate::pagination::resolve_page_size;
12
13struct StoredEntry {
14    task: Task,
15    version: TaskVersion,
16}
17
18/// In-memory task store. Contents do not survive restarts.
19pub struct InMemoryTaskStore {
20    tasks: RwLock<HashMap<TaskId, StoredEntry>>,
21}
22
23impl InMemoryTaskStore {
24    pub fn new() -> Self {
25        InMemoryTaskStore {
26            tasks: RwLock::new(HashMap::new()),
27        }
28    }
29}
30
31impl Default for InMemoryTaskStore {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37#[async_trait]
38impl TaskStore for InMemoryTaskStore {
39    async fn create(&self, task: Task) -> Result<TaskVersion, A2AError> {
40        let mut store = self.tasks.write().await;
41        if store.contains_key(&task.id) {
42            return Err(crate::sanitized_internal_error("task already exists"));
43        }
44        let id = task.id.clone();
45        store.insert(id, StoredEntry { task, version: 1 });
46        Ok(1)
47    }
48
49    async fn update(&self, task: Task) -> Result<TaskVersion, A2AError> {
50        let mut store = self.tasks.write().await;
51        let entry = store
52            .get_mut(&task.id)
53            .ok_or_else(|| A2AError::task_not_found(&task.id))?;
54        // A terminal task's state is immutable. Re-applying the same
55        // terminal state is allowed so idempotent re-emission (e.g. cancel
56        // flows) keeps working, but a terminal state can never be overwritten
57        // by a different state.
58        if entry.task.status.state.is_terminal() && entry.task.status.state != task.status.state {
59            return Err(A2AError::invalid_request(format!(
60                "task {} is in terminal state {:?} and cannot transition to {:?}",
61                task.id, entry.task.status.state, task.status.state
62            )));
63        }
64        entry.version += 1;
65        entry.task = task;
66        Ok(entry.version)
67    }
68
69    async fn begin_cancel(&self, task_id: &str) -> Result<Task, A2AError> {
70        // Hold the write lock across the check and the state transition so two
71        // concurrent cancel requests cannot both pass the check-then-act window
72        // window.
73        let mut store = self.tasks.write().await;
74        let entry = store
75            .get_mut(task_id)
76            .ok_or_else(|| A2AError::task_not_found(task_id))?;
77        if entry.task.status.state.is_terminal() {
78            return Err(A2AError::task_not_cancelable(task_id));
79        }
80        entry.version += 1;
81        entry.task.status.state = TaskState::Canceled;
82        entry.task.status.timestamp = Some(chrono::Utc::now());
83        Ok(entry.task.clone())
84    }
85
86    async fn get(&self, task_id: &str) -> Result<Option<Task>, A2AError> {
87        let store = self.tasks.read().await;
88        Ok(store.get(task_id).map(|e| e.task.clone()))
89    }
90
91    async fn list(&self, req: &ListTasksRequest) -> Result<ListTasksResponse, A2AError> {
92        let store = self.tasks.read().await;
93        let mut tasks: Vec<Task> = store
94            .values()
95            .filter(|entry| {
96                // Filter by context_id if provided
97                if let Some(ref ctx_id) = req.context_id {
98                    if entry.task.context_id != *ctx_id {
99                        return false;
100                    }
101                }
102                // Filter by status if provided
103                if let Some(ref status) = req.status {
104                    if entry.task.status.state != *status {
105                        return false;
106                    }
107                }
108                // A task without a status timestamp cannot satisfy an
109                // "after" filter. Keep this predicate in the store so the
110                // result count and pagination are based on the filtered set.
111                if let Some(after) = req.status_timestamp_after {
112                    if entry
113                        .task
114                        .status
115                        .timestamp
116                        .is_none_or(|timestamp| timestamp <= after)
117                    {
118                        return false;
119                    }
120                }
121                true
122            })
123            .map(|e| e.task.clone())
124            .collect();
125
126        // Sort by ID for deterministic output
127        tasks.sort_by(|a, b| a.id.cmp(&b.id));
128
129        // Apply pagination
130        // Defence in depth: DefaultRequestHandler clamps before calling a
131        // store, but a store reached by another route must still bound its
132        // own page.
133        let page_size = resolve_page_size(req.page_size);
134        let start = if let Some(ref token) = req.page_token {
135            // Simple offset-based pagination
136            token
137                .parse::<usize>()
138                .map_err(|_| A2AError::invalid_params("invalid page token"))?
139        } else {
140            0
141        };
142
143        let total_size = tasks.len();
144        // Make sure start is within bounds
145        let start = start.min(total_size);
146        let end = start.saturating_add(page_size).min(total_size);
147        let page = tasks[start..end].to_vec();
148
149        let next_page_token = if end < total_size {
150            Some(end.to_string())
151        } else {
152            None
153        };
154
155        let page = page
156            .into_iter()
157            .map(|mut task| {
158                apply_history_length(&mut task, req.history_length);
159                task
160            })
161            .collect();
162
163        Ok(ListTasksResponse {
164            tasks: page,
165            next_page_token: next_page_token.unwrap_or_default(),
166            page_size: page_size as i32,
167            total_size: total_size as i32,
168        })
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    fn make_task(id: &str, ctx: &str, state: TaskState) -> Task {
177        Task {
178            id: id.to_string(),
179            context_id: ctx.to_string(),
180            status: TaskStatus {
181                state,
182                message: None,
183                timestamp: None,
184            },
185            artifacts: None,
186            history: None,
187            metadata: None,
188        }
189    }
190
191    #[tokio::test]
192    async fn test_create_and_get() {
193        let store = InMemoryTaskStore::new();
194        let task = make_task("t1", "c1", TaskState::Submitted);
195        let ver = store.create(task.clone()).await.unwrap();
196        assert_eq!(ver, 1);
197
198        let got = store.get("t1").await.unwrap().unwrap();
199        assert_eq!(got.id, "t1");
200        assert_eq!(got.context_id, "c1");
201    }
202
203    #[tokio::test]
204    async fn test_create_duplicate() {
205        let store = InMemoryTaskStore::new();
206        let task = make_task("t1", "c1", TaskState::Submitted);
207        store.create(task.clone()).await.unwrap();
208
209        let result = store.create(task).await;
210        assert!(result.is_err());
211    }
212
213    #[tokio::test]
214    async fn test_update() {
215        let store = InMemoryTaskStore::new();
216        let task = make_task("t1", "c1", TaskState::Submitted);
217        store.create(task).await.unwrap();
218
219        let updated = make_task("t1", "c1", TaskState::Working);
220        let ver = store.update(updated).await.unwrap();
221        assert_eq!(ver, 2);
222
223        let got = store.get("t1").await.unwrap().unwrap();
224        assert_eq!(got.status.state, TaskState::Working);
225    }
226
227    #[tokio::test]
228    async fn test_update_not_found() {
229        let store = InMemoryTaskStore::new();
230        let task = make_task("t1", "c1", TaskState::Working);
231        let result = store.update(task).await;
232        assert!(result.is_err());
233    }
234
235    #[tokio::test]
236    async fn test_update_terminal_state_rejected() {
237        let store = InMemoryTaskStore::new();
238        store
239            .create(make_task("t1", "c1", TaskState::Completed))
240            .await
241            .unwrap();
242
243        // A terminal task cannot transition to a different state.
244        for state in [
245            TaskState::Submitted,
246            TaskState::Working,
247            TaskState::Canceled,
248            TaskState::Failed,
249        ] {
250            let result = store.update(make_task("t1", "c1", state.clone())).await;
251            assert!(result.is_err(), "expected {state:?} update to be rejected");
252            assert_eq!(
253                result.unwrap_err().code,
254                error_code::INVALID_REQUEST,
255                "unexpected code for {state:?}"
256            );
257        }
258    }
259
260    #[tokio::test]
261    async fn test_update_terminal_same_state_allowed() {
262        let store = InMemoryTaskStore::new();
263        let mut task = make_task("t1", "c1", TaskState::Canceled);
264        store.create(task.clone()).await.unwrap();
265
266        // Re-applying the same terminal state (idempotent re-emission) is fine.
267        task.status.message = Some(Message::new(Role::Agent, vec![Part::text("after cancel")]));
268        let ver = store.update(task).await.unwrap();
269        assert_eq!(ver, 2);
270        let got = store.get("t1").await.unwrap().unwrap();
271        assert_eq!(got.status.state, TaskState::Canceled);
272        assert!(got.status.message.is_some());
273    }
274
275    #[tokio::test]
276    async fn test_begin_cancel_transitions_non_terminal() {
277        let store = InMemoryTaskStore::new();
278        store
279            .create(make_task("t1", "c1", TaskState::Working))
280            .await
281            .unwrap();
282
283        let task = store.begin_cancel("t1").await.unwrap();
284        assert_eq!(task.status.state, TaskState::Canceled);
285
286        // A second cancel on the now-terminal task fails.
287        let result = store.begin_cancel("t1").await;
288        assert!(result.is_err());
289        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_CANCELABLE);
290    }
291
292    #[tokio::test]
293    async fn test_begin_cancel_terminal_rejected() {
294        let store = InMemoryTaskStore::new();
295        store
296            .create(make_task("t1", "c1", TaskState::Completed))
297            .await
298            .unwrap();
299        let result = store.begin_cancel("t1").await;
300        assert!(result.is_err());
301        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_CANCELABLE);
302    }
303
304    #[tokio::test]
305    async fn test_begin_cancel_not_found() {
306        let store = InMemoryTaskStore::new();
307        let result = store.begin_cancel("nonexistent").await;
308        assert!(result.is_err());
309        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_FOUND);
310    }
311
312    #[tokio::test]
313    async fn test_get_not_found() {
314        let store = InMemoryTaskStore::new();
315        let got = store.get("nonexistent").await.unwrap();
316        assert!(got.is_none());
317    }
318
319    #[tokio::test]
320    async fn test_list_filter_by_context() {
321        let store = InMemoryTaskStore::new();
322        store
323            .create(make_task("t1", "c1", TaskState::Submitted))
324            .await
325            .unwrap();
326        store
327            .create(make_task("t2", "c2", TaskState::Working))
328            .await
329            .unwrap();
330        store
331            .create(make_task("t3", "c1", TaskState::Completed))
332            .await
333            .unwrap();
334
335        let req = ListTasksRequest {
336            context_id: Some("c1".to_string()),
337            status: None,
338            page_size: None,
339            page_token: None,
340            history_length: None,
341            status_timestamp_after: None,
342            include_artifacts: None,
343            tenant: None,
344        };
345        let resp = store.list(&req).await.unwrap();
346        assert_eq!(resp.tasks.len(), 2);
347        assert_eq!(resp.total_size, 2);
348    }
349
350    #[tokio::test]
351    async fn test_list_filter_by_status() {
352        let store = InMemoryTaskStore::new();
353        store
354            .create(make_task("t1", "c1", TaskState::Submitted))
355            .await
356            .unwrap();
357        store
358            .create(make_task("t2", "c1", TaskState::Working))
359            .await
360            .unwrap();
361
362        let req = ListTasksRequest {
363            context_id: None,
364            status: Some(TaskState::Working),
365            page_size: None,
366            page_token: None,
367            history_length: None,
368            status_timestamp_after: None,
369            include_artifacts: None,
370            tenant: None,
371        };
372        let resp = store.list(&req).await.unwrap();
373        assert_eq!(resp.tasks.len(), 1);
374        assert_eq!(resp.tasks[0].id, "t2");
375    }
376
377    #[tokio::test]
378    async fn test_list_filter_by_status_timestamp_after() {
379        let store = InMemoryTaskStore::new();
380        let cutoff = chrono::DateTime::parse_from_rfc3339("2025-01-01T00:00:00Z")
381            .unwrap()
382            .with_timezone(&chrono::Utc);
383
384        let mut old = make_task("old", "c1", TaskState::Completed);
385        old.status.timestamp = Some(cutoff - chrono::Duration::seconds(1));
386        store.create(old).await.unwrap();
387
388        let mut recent = make_task("recent", "c1", TaskState::Completed);
389        recent.status.timestamp = Some(cutoff + chrono::Duration::seconds(1));
390        store.create(recent).await.unwrap();
391
392        // Missing timestamps do not satisfy an "after" filter.
393        store
394            .create(make_task("missing", "c1", TaskState::Completed))
395            .await
396            .unwrap();
397
398        let resp = store
399            .list(&ListTasksRequest {
400                context_id: None,
401                status: None,
402                page_size: None,
403                page_token: None,
404                history_length: None,
405                status_timestamp_after: Some(cutoff),
406                include_artifacts: None,
407                tenant: None,
408            })
409            .await
410            .unwrap();
411
412        assert_eq!(resp.total_size, 1);
413        assert_eq!(
414            resp.tasks
415                .iter()
416                .map(|task| task.id.as_str())
417                .collect::<Vec<_>>(),
418            ["recent"]
419        );
420    }
421
422    #[tokio::test]
423    async fn test_list_pagination() {
424        let store = InMemoryTaskStore::new();
425        for i in 0..5 {
426            store
427                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
428                .await
429                .unwrap();
430        }
431
432        let req = ListTasksRequest {
433            context_id: None,
434            status: None,
435            page_size: Some(2),
436            page_token: None,
437            history_length: None,
438            status_timestamp_after: None,
439            include_artifacts: None,
440            tenant: None,
441        };
442        let resp = store.list(&req).await.unwrap();
443        assert_eq!(resp.tasks.len(), 2);
444        assert_eq!(resp.total_size, 5);
445        assert!(!resp.next_page_token.is_empty());
446
447        // Get second page
448        let req2 = ListTasksRequest {
449            page_token: Some(resp.next_page_token),
450            page_size: Some(2),
451            ..req
452        };
453        let resp2 = store.list(&req2).await.unwrap();
454        assert_eq!(resp2.tasks.len(), 2);
455    }
456
457    #[tokio::test]
458    async fn test_list_zero_page_size_uses_default_window() {
459        let store = InMemoryTaskStore::new();
460        for i in 0..3 {
461            store
462                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
463                .await
464                .unwrap();
465        }
466
467        let req = ListTasksRequest {
468            context_id: None,
469            status: None,
470            page_size: Some(0),
471            page_token: None,
472            history_length: None,
473            status_timestamp_after: None,
474            include_artifacts: None,
475            tenant: None,
476        };
477        let resp = store.list(&req).await.unwrap();
478        assert_eq!(resp.tasks.len(), 3);
479        assert_eq!(resp.page_size, 50);
480    }
481
482    #[tokio::test]
483    async fn test_list_history_truncation() {
484        let store = InMemoryTaskStore::new();
485        let mut task = make_task("t1", "c1", TaskState::Working);
486        task.history = Some(vec![
487            Message::new(Role::User, vec![Part::text("1")]),
488            Message::new(Role::Agent, vec![Part::text("2")]),
489            Message::new(Role::User, vec![Part::text("3")]),
490        ]);
491        store.create(task).await.unwrap();
492
493        let req = ListTasksRequest {
494            context_id: None,
495            status: None,
496            page_size: None,
497            page_token: None,
498            history_length: Some(1),
499            status_timestamp_after: None,
500            include_artifacts: None,
501            tenant: None,
502        };
503        let resp = store.list(&req).await.unwrap();
504        assert_eq!(resp.tasks[0].history.as_ref().unwrap().len(), 1);
505
506        let empty = store
507            .list(&ListTasksRequest {
508                history_length: Some(-1),
509                ..req
510            })
511            .await
512            .unwrap();
513        assert!(empty.tasks[0].history.as_ref().unwrap().is_empty());
514    }
515
516    fn make_task_with_history(id: &str, messages: Vec<&str>) -> Task {
517        let mut task = make_task(id, "c1", TaskState::Working);
518        task.history = Some(
519            messages
520                .into_iter()
521                .map(|m| Message::new(Role::User, vec![Part::text(m)]))
522                .collect(),
523        );
524        task
525    }
526
527    #[tokio::test]
528    async fn test_list_negative_history_length_returns_empty_history() {
529        let store = InMemoryTaskStore::new();
530        store
531            .create(make_task_with_history("t1", vec!["1", "2", "3"]))
532            .await
533            .unwrap();
534
535        let req = ListTasksRequest {
536            context_id: None,
537            status: None,
538            page_size: None,
539            page_token: None,
540            history_length: Some(-1),
541            status_timestamp_after: None,
542            include_artifacts: None,
543            tenant: None,
544        };
545        let resp = store.list(&req).await.unwrap();
546        assert_eq!(resp.tasks.len(), 1);
547        assert_eq!(resp.tasks[0].history.as_ref().unwrap().len(), 0);
548    }
549
550    #[tokio::test]
551    async fn test_list_page_size_is_capped() {
552        let store = InMemoryTaskStore::new();
553        for i in 0..150 {
554            store
555                .create(make_task(&format!("t{i:03}"), "c1", TaskState::Submitted))
556                .await
557                .unwrap();
558        }
559
560        let req = ListTasksRequest {
561            context_id: None,
562            status: None,
563            page_size: Some(1000),
564            page_token: None,
565            history_length: None,
566            status_timestamp_after: None,
567            include_artifacts: None,
568            tenant: None,
569        };
570        let resp = store.list(&req).await.unwrap();
571        assert_eq!(resp.tasks.len(), 100);
572        assert_eq!(resp.page_size, 100);
573        assert_eq!(resp.total_size, 150);
574    }
575
576    #[tokio::test]
577    async fn test_list_invalid_page_token_errors() {
578        let store = InMemoryTaskStore::new();
579        store
580            .create(make_task("t1", "c1", TaskState::Submitted))
581            .await
582            .unwrap();
583
584        let req = ListTasksRequest {
585            context_id: None,
586            status: None,
587            page_size: None,
588            page_token: Some("not-a-number".into()),
589            history_length: None,
590            status_timestamp_after: None,
591            include_artifacts: None,
592            tenant: None,
593        };
594        let result = store.list(&req).await;
595        assert!(result.is_err());
596        assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
597    }
598
599    #[tokio::test]
600    async fn test_list_out_of_range_page_token_returns_empty_page() {
601        let store = InMemoryTaskStore::new();
602        for i in 0..3 {
603            store
604                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
605                .await
606                .unwrap();
607        }
608
609        let req = ListTasksRequest {
610            context_id: None,
611            status: None,
612            page_size: None,
613            page_token: Some("999".to_string()),
614            history_length: None,
615            status_timestamp_after: None,
616            include_artifacts: None,
617            tenant: None,
618        };
619        let resp = store.list(&req).await.unwrap();
620
621        assert!(resp.tasks.is_empty());
622        assert!(
623            resp.next_page_token.is_empty(),
624            "must not hand back another token"
625        );
626        assert_eq!(resp.total_size, 3);
627    }
628}
a2a_server/task_store/inmemory.rs

a2a_server/task_store/
inmemory.rs
