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                true
109            })
110            .map(|e| e.task.clone())
111            .collect();
112
113        // Sort by ID for deterministic output
114        tasks.sort_by(|a, b| a.id.cmp(&b.id));
115
116        // Apply pagination
117        // Defence in depth: DefaultRequestHandler clamps before calling a
118        // store, but a store reached by another route must still bound its
119        // own page.
120        let page_size = resolve_page_size(req.page_size);
121        let start = if let Some(ref token) = req.page_token {
122            // Simple offset-based pagination
123            token
124                .parse::<usize>()
125                .map_err(|_| A2AError::invalid_params("invalid page token"))?
126        } else {
127            0
128        };
129
130        let total_size = tasks.len();
131        // Make sure start is within bounds
132        let start = start.min(total_size);
133        let end = start.saturating_add(page_size).min(total_size);
134        let page = tasks[start..end].to_vec();
135
136        let next_page_token = if end < total_size {
137            Some(end.to_string())
138        } else {
139            None
140        };
141
142        let page = page
143            .into_iter()
144            .map(|mut task| {
145                apply_history_length(&mut task, req.history_length);
146                task
147            })
148            .collect();
149
150        Ok(ListTasksResponse {
151            tasks: page,
152            next_page_token: next_page_token.unwrap_or_default(),
153            page_size: page_size as i32,
154            total_size: total_size as i32,
155        })
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    fn make_task(id: &str, ctx: &str, state: TaskState) -> Task {
164        Task {
165            id: id.to_string(),
166            context_id: ctx.to_string(),
167            status: TaskStatus {
168                state,
169                message: None,
170                timestamp: None,
171            },
172            artifacts: None,
173            history: None,
174            metadata: None,
175        }
176    }
177
178    #[tokio::test]
179    async fn test_create_and_get() {
180        let store = InMemoryTaskStore::new();
181        let task = make_task("t1", "c1", TaskState::Submitted);
182        let ver = store.create(task.clone()).await.unwrap();
183        assert_eq!(ver, 1);
184
185        let got = store.get("t1").await.unwrap().unwrap();
186        assert_eq!(got.id, "t1");
187        assert_eq!(got.context_id, "c1");
188    }
189
190    #[tokio::test]
191    async fn test_create_duplicate() {
192        let store = InMemoryTaskStore::new();
193        let task = make_task("t1", "c1", TaskState::Submitted);
194        store.create(task.clone()).await.unwrap();
195
196        let result = store.create(task).await;
197        assert!(result.is_err());
198    }
199
200    #[tokio::test]
201    async fn test_update() {
202        let store = InMemoryTaskStore::new();
203        let task = make_task("t1", "c1", TaskState::Submitted);
204        store.create(task).await.unwrap();
205
206        let updated = make_task("t1", "c1", TaskState::Working);
207        let ver = store.update(updated).await.unwrap();
208        assert_eq!(ver, 2);
209
210        let got = store.get("t1").await.unwrap().unwrap();
211        assert_eq!(got.status.state, TaskState::Working);
212    }
213
214    #[tokio::test]
215    async fn test_update_not_found() {
216        let store = InMemoryTaskStore::new();
217        let task = make_task("t1", "c1", TaskState::Working);
218        let result = store.update(task).await;
219        assert!(result.is_err());
220    }
221
222    #[tokio::test]
223    async fn test_update_terminal_state_rejected() {
224        let store = InMemoryTaskStore::new();
225        store
226            .create(make_task("t1", "c1", TaskState::Completed))
227            .await
228            .unwrap();
229
230        // A terminal task cannot transition to a different state.
231        for state in [
232            TaskState::Submitted,
233            TaskState::Working,
234            TaskState::Canceled,
235            TaskState::Failed,
236        ] {
237            let result = store.update(make_task("t1", "c1", state.clone())).await;
238            assert!(result.is_err(), "expected {state:?} update to be rejected");
239            assert_eq!(
240                result.unwrap_err().code,
241                error_code::INVALID_REQUEST,
242                "unexpected code for {state:?}"
243            );
244        }
245    }
246
247    #[tokio::test]
248    async fn test_update_terminal_same_state_allowed() {
249        let store = InMemoryTaskStore::new();
250        let mut task = make_task("t1", "c1", TaskState::Canceled);
251        store.create(task.clone()).await.unwrap();
252
253        // Re-applying the same terminal state (idempotent re-emission) is fine.
254        task.status.message = Some(Message::new(Role::Agent, vec![Part::text("after cancel")]));
255        let ver = store.update(task).await.unwrap();
256        assert_eq!(ver, 2);
257        let got = store.get("t1").await.unwrap().unwrap();
258        assert_eq!(got.status.state, TaskState::Canceled);
259        assert!(got.status.message.is_some());
260    }
261
262    #[tokio::test]
263    async fn test_begin_cancel_transitions_non_terminal() {
264        let store = InMemoryTaskStore::new();
265        store
266            .create(make_task("t1", "c1", TaskState::Working))
267            .await
268            .unwrap();
269
270        let task = store.begin_cancel("t1").await.unwrap();
271        assert_eq!(task.status.state, TaskState::Canceled);
272
273        // A second cancel on the now-terminal task fails.
274        let result = store.begin_cancel("t1").await;
275        assert!(result.is_err());
276        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_CANCELABLE);
277    }
278
279    #[tokio::test]
280    async fn test_begin_cancel_terminal_rejected() {
281        let store = InMemoryTaskStore::new();
282        store
283            .create(make_task("t1", "c1", TaskState::Completed))
284            .await
285            .unwrap();
286        let result = store.begin_cancel("t1").await;
287        assert!(result.is_err());
288        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_CANCELABLE);
289    }
290
291    #[tokio::test]
292    async fn test_begin_cancel_not_found() {
293        let store = InMemoryTaskStore::new();
294        let result = store.begin_cancel("nonexistent").await;
295        assert!(result.is_err());
296        assert_eq!(result.unwrap_err().code, error_code::TASK_NOT_FOUND);
297    }
298
299    #[tokio::test]
300    async fn test_get_not_found() {
301        let store = InMemoryTaskStore::new();
302        let got = store.get("nonexistent").await.unwrap();
303        assert!(got.is_none());
304    }
305
306    #[tokio::test]
307    async fn test_list_filter_by_context() {
308        let store = InMemoryTaskStore::new();
309        store
310            .create(make_task("t1", "c1", TaskState::Submitted))
311            .await
312            .unwrap();
313        store
314            .create(make_task("t2", "c2", TaskState::Working))
315            .await
316            .unwrap();
317        store
318            .create(make_task("t3", "c1", TaskState::Completed))
319            .await
320            .unwrap();
321
322        let req = ListTasksRequest {
323            context_id: Some("c1".to_string()),
324            status: None,
325            page_size: None,
326            page_token: None,
327            history_length: None,
328            status_timestamp_after: None,
329            include_artifacts: None,
330            tenant: None,
331        };
332        let resp = store.list(&req).await.unwrap();
333        assert_eq!(resp.tasks.len(), 2);
334        assert_eq!(resp.total_size, 2);
335    }
336
337    #[tokio::test]
338    async fn test_list_filter_by_status() {
339        let store = InMemoryTaskStore::new();
340        store
341            .create(make_task("t1", "c1", TaskState::Submitted))
342            .await
343            .unwrap();
344        store
345            .create(make_task("t2", "c1", TaskState::Working))
346            .await
347            .unwrap();
348
349        let req = ListTasksRequest {
350            context_id: None,
351            status: Some(TaskState::Working),
352            page_size: None,
353            page_token: None,
354            history_length: None,
355            status_timestamp_after: None,
356            include_artifacts: None,
357            tenant: None,
358        };
359        let resp = store.list(&req).await.unwrap();
360        assert_eq!(resp.tasks.len(), 1);
361        assert_eq!(resp.tasks[0].id, "t2");
362    }
363
364    #[tokio::test]
365    async fn test_list_pagination() {
366        let store = InMemoryTaskStore::new();
367        for i in 0..5 {
368            store
369                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
370                .await
371                .unwrap();
372        }
373
374        let req = ListTasksRequest {
375            context_id: None,
376            status: None,
377            page_size: Some(2),
378            page_token: None,
379            history_length: None,
380            status_timestamp_after: None,
381            include_artifacts: None,
382            tenant: None,
383        };
384        let resp = store.list(&req).await.unwrap();
385        assert_eq!(resp.tasks.len(), 2);
386        assert_eq!(resp.total_size, 5);
387        assert!(!resp.next_page_token.is_empty());
388
389        // Get second page
390        let req2 = ListTasksRequest {
391            page_token: Some(resp.next_page_token),
392            page_size: Some(2),
393            ..req
394        };
395        let resp2 = store.list(&req2).await.unwrap();
396        assert_eq!(resp2.tasks.len(), 2);
397    }
398
399    #[tokio::test]
400    async fn test_list_zero_page_size_uses_default_window() {
401        let store = InMemoryTaskStore::new();
402        for i in 0..3 {
403            store
404                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
405                .await
406                .unwrap();
407        }
408
409        let req = ListTasksRequest {
410            context_id: None,
411            status: None,
412            page_size: Some(0),
413            page_token: None,
414            history_length: None,
415            status_timestamp_after: None,
416            include_artifacts: None,
417            tenant: None,
418        };
419        let resp = store.list(&req).await.unwrap();
420        assert_eq!(resp.tasks.len(), 3);
421        assert_eq!(resp.page_size, 50);
422    }
423
424    #[tokio::test]
425    async fn test_list_history_truncation() {
426        let store = InMemoryTaskStore::new();
427        let mut task = make_task("t1", "c1", TaskState::Working);
428        task.history = Some(vec![
429            Message::new(Role::User, vec![Part::text("1")]),
430            Message::new(Role::Agent, vec![Part::text("2")]),
431            Message::new(Role::User, vec![Part::text("3")]),
432        ]);
433        store.create(task).await.unwrap();
434
435        let req = ListTasksRequest {
436            context_id: None,
437            status: None,
438            page_size: None,
439            page_token: None,
440            history_length: Some(1),
441            status_timestamp_after: None,
442            include_artifacts: None,
443            tenant: None,
444        };
445        let resp = store.list(&req).await.unwrap();
446        assert_eq!(resp.tasks[0].history.as_ref().unwrap().len(), 1);
447
448        let empty = store
449            .list(&ListTasksRequest {
450                history_length: Some(-1),
451                ..req
452            })
453            .await
454            .unwrap();
455        assert!(empty.tasks[0].history.as_ref().unwrap().is_empty());
456    }
457
458    fn make_task_with_history(id: &str, messages: Vec<&str>) -> Task {
459        let mut task = make_task(id, "c1", TaskState::Working);
460        task.history = Some(
461            messages
462                .into_iter()
463                .map(|m| Message::new(Role::User, vec![Part::text(m)]))
464                .collect(),
465        );
466        task
467    }
468
469    #[tokio::test]
470    async fn test_list_negative_history_length_returns_empty_history() {
471        let store = InMemoryTaskStore::new();
472        store
473            .create(make_task_with_history("t1", vec!["1", "2", "3"]))
474            .await
475            .unwrap();
476
477        let req = ListTasksRequest {
478            context_id: None,
479            status: None,
480            page_size: None,
481            page_token: None,
482            history_length: Some(-1),
483            status_timestamp_after: None,
484            include_artifacts: None,
485            tenant: None,
486        };
487        let resp = store.list(&req).await.unwrap();
488        assert_eq!(resp.tasks.len(), 1);
489        assert_eq!(resp.tasks[0].history.as_ref().unwrap().len(), 0);
490    }
491
492    #[tokio::test]
493    async fn test_list_page_size_is_capped() {
494        let store = InMemoryTaskStore::new();
495        for i in 0..150 {
496            store
497                .create(make_task(&format!("t{i:03}"), "c1", TaskState::Submitted))
498                .await
499                .unwrap();
500        }
501
502        let req = ListTasksRequest {
503            context_id: None,
504            status: None,
505            page_size: Some(1000),
506            page_token: None,
507            history_length: None,
508            status_timestamp_after: None,
509            include_artifacts: None,
510            tenant: None,
511        };
512        let resp = store.list(&req).await.unwrap();
513        assert_eq!(resp.tasks.len(), 100);
514        assert_eq!(resp.page_size, 100);
515        assert_eq!(resp.total_size, 150);
516    }
517
518    #[tokio::test]
519    async fn test_list_invalid_page_token_errors() {
520        let store = InMemoryTaskStore::new();
521        store
522            .create(make_task("t1", "c1", TaskState::Submitted))
523            .await
524            .unwrap();
525
526        let req = ListTasksRequest {
527            context_id: None,
528            status: None,
529            page_size: None,
530            page_token: Some("not-a-number".into()),
531            history_length: None,
532            status_timestamp_after: None,
533            include_artifacts: None,
534            tenant: None,
535        };
536        let result = store.list(&req).await;
537        assert!(result.is_err());
538        assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
539    }
540
541    #[tokio::test]
542    async fn test_list_out_of_range_page_token_returns_empty_page() {
543        let store = InMemoryTaskStore::new();
544        for i in 0..3 {
545            store
546                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
547                .await
548                .unwrap();
549        }
550
551        let req = ListTasksRequest {
552            context_id: None,
553            status: None,
554            page_size: None,
555            page_token: Some("999".to_string()),
556            history_length: None,
557            status_timestamp_after: None,
558            include_artifacts: None,
559            tenant: None,
560        };
561        let resp = store.list(&req).await.unwrap();
562
563        assert!(resp.tasks.is_empty());
564        assert!(
565            resp.next_page_token.is_empty(),
566            "must not hand back another token"
567        );
568        assert_eq!(resp.total_size, 3);
569    }
570}
a2a_server/task_store/inmemory.rs

a2a_server/task_store/
inmemory.rs
