a2a_server/task_store/
inmemory.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3use a2a::*;
4use async_trait::async_trait;
5use std::collections::HashMap;
6use tokio::sync::RwLock;
7
8use super::store::{TaskStore, TaskVersion};
9
10struct StoredEntry {
11    task: Task,
12    version: TaskVersion,
13}
14
15/// In-memory task store. Contents do not survive restarts.
16pub struct InMemoryTaskStore {
17    tasks: RwLock<HashMap<TaskId, StoredEntry>>,
18}
19
20impl InMemoryTaskStore {
21    pub fn new() -> Self {
22        InMemoryTaskStore {
23            tasks: RwLock::new(HashMap::new()),
24        }
25    }
26}
27
28impl Default for InMemoryTaskStore {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34#[async_trait]
35impl TaskStore for InMemoryTaskStore {
36    async fn create(&self, task: Task) -> Result<TaskVersion, A2AError> {
37        let mut store = self.tasks.write().await;
38        if store.contains_key(&task.id) {
39            return Err(A2AError::internal("task already exists"));
40        }
41        let id = task.id.clone();
42        store.insert(id, StoredEntry { task, version: 1 });
43        Ok(1)
44    }
45
46    async fn update(&self, task: Task) -> Result<TaskVersion, A2AError> {
47        let mut store = self.tasks.write().await;
48        let entry = store
49            .get_mut(&task.id)
50            .ok_or_else(|| A2AError::task_not_found(&task.id))?;
51        entry.version += 1;
52        entry.task = task;
53        Ok(entry.version)
54    }
55
56    async fn get(&self, task_id: &str) -> Result<Option<Task>, A2AError> {
57        let store = self.tasks.read().await;
58        Ok(store.get(task_id).map(|e| e.task.clone()))
59    }
60
61    async fn list(&self, req: &ListTasksRequest) -> Result<ListTasksResponse, A2AError> {
62        let store = self.tasks.read().await;
63        let mut tasks: Vec<Task> = store
64            .values()
65            .filter(|entry| {
66                // Filter by context_id if provided
67                if let Some(ref ctx_id) = req.context_id {
68                    if entry.task.context_id != *ctx_id {
69                        return false;
70                    }
71                }
72                // Filter by status if provided
73                if let Some(ref status) = req.status {
74                    if entry.task.status.state != *status {
75                        return false;
76                    }
77                }
78                true
79            })
80            .map(|e| e.task.clone())
81            .collect();
82
83        // Sort by ID for deterministic output
84        tasks.sort_by(|a, b| a.id.cmp(&b.id));
85
86        // Apply pagination
87        let page_size = match req.page_size {
88            Some(size) if size > 0 => size as usize,
89            _ => 50,
90        };
91        let start = if let Some(ref token) = req.page_token {
92            // Simple offset-based pagination
93            token.parse::<usize>().unwrap_or(0)
94        } else {
95            0
96        };
97
98        let total_size = tasks.len();
99        let end = (start + page_size).min(total_size);
100        let page = tasks[start..end].to_vec();
101
102        let next_page_token = if end < total_size {
103            Some(end.to_string())
104        } else {
105            None
106        };
107
108        // Apply history length truncation
109        let page = page
110            .into_iter()
111            .map(|mut task| {
112                if let Some(ref hl) = req.history_length {
113                    let hl = *hl as usize;
114                    if let Some(ref mut history) = task.history {
115                        if hl == 0 {
116                            *history = Vec::new();
117                        } else if history.len() > hl {
118                            let start = history.len() - hl;
119                            *history = history[start..].to_vec();
120                        }
121                    }
122                }
123                task
124            })
125            .collect();
126
127        Ok(ListTasksResponse {
128            tasks: page,
129            next_page_token: next_page_token.unwrap_or_default(),
130            page_size: page_size as i32,
131            total_size: total_size as i32,
132        })
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    fn make_task(id: &str, ctx: &str, state: TaskState) -> Task {
141        Task {
142            id: id.to_string(),
143            context_id: ctx.to_string(),
144            status: TaskStatus {
145                state,
146                message: None,
147                timestamp: None,
148            },
149            artifacts: None,
150            history: None,
151            metadata: None,
152        }
153    }
154
155    #[tokio::test]
156    async fn test_create_and_get() {
157        let store = InMemoryTaskStore::new();
158        let task = make_task("t1", "c1", TaskState::Submitted);
159        let ver = store.create(task.clone()).await.unwrap();
160        assert_eq!(ver, 1);
161
162        let got = store.get("t1").await.unwrap().unwrap();
163        assert_eq!(got.id, "t1");
164        assert_eq!(got.context_id, "c1");
165    }
166
167    #[tokio::test]
168    async fn test_create_duplicate() {
169        let store = InMemoryTaskStore::new();
170        let task = make_task("t1", "c1", TaskState::Submitted);
171        store.create(task.clone()).await.unwrap();
172
173        let result = store.create(task).await;
174        assert!(result.is_err());
175    }
176
177    #[tokio::test]
178    async fn test_update() {
179        let store = InMemoryTaskStore::new();
180        let task = make_task("t1", "c1", TaskState::Submitted);
181        store.create(task).await.unwrap();
182
183        let updated = make_task("t1", "c1", TaskState::Working);
184        let ver = store.update(updated).await.unwrap();
185        assert_eq!(ver, 2);
186
187        let got = store.get("t1").await.unwrap().unwrap();
188        assert_eq!(got.status.state, TaskState::Working);
189    }
190
191    #[tokio::test]
192    async fn test_update_not_found() {
193        let store = InMemoryTaskStore::new();
194        let task = make_task("t1", "c1", TaskState::Working);
195        let result = store.update(task).await;
196        assert!(result.is_err());
197    }
198
199    #[tokio::test]
200    async fn test_get_not_found() {
201        let store = InMemoryTaskStore::new();
202        let got = store.get("nonexistent").await.unwrap();
203        assert!(got.is_none());
204    }
205
206    #[tokio::test]
207    async fn test_list_filter_by_context() {
208        let store = InMemoryTaskStore::new();
209        store
210            .create(make_task("t1", "c1", TaskState::Submitted))
211            .await
212            .unwrap();
213        store
214            .create(make_task("t2", "c2", TaskState::Working))
215            .await
216            .unwrap();
217        store
218            .create(make_task("t3", "c1", TaskState::Completed))
219            .await
220            .unwrap();
221
222        let req = ListTasksRequest {
223            context_id: Some("c1".to_string()),
224            status: None,
225            page_size: None,
226            page_token: None,
227            history_length: None,
228            status_timestamp_after: None,
229            include_artifacts: None,
230            tenant: None,
231        };
232        let resp = store.list(&req).await.unwrap();
233        assert_eq!(resp.tasks.len(), 2);
234        assert_eq!(resp.total_size, 2);
235    }
236
237    #[tokio::test]
238    async fn test_list_filter_by_status() {
239        let store = InMemoryTaskStore::new();
240        store
241            .create(make_task("t1", "c1", TaskState::Submitted))
242            .await
243            .unwrap();
244        store
245            .create(make_task("t2", "c1", TaskState::Working))
246            .await
247            .unwrap();
248
249        let req = ListTasksRequest {
250            context_id: None,
251            status: Some(TaskState::Working),
252            page_size: None,
253            page_token: None,
254            history_length: None,
255            status_timestamp_after: None,
256            include_artifacts: None,
257            tenant: None,
258        };
259        let resp = store.list(&req).await.unwrap();
260        assert_eq!(resp.tasks.len(), 1);
261        assert_eq!(resp.tasks[0].id, "t2");
262    }
263
264    #[tokio::test]
265    async fn test_list_pagination() {
266        let store = InMemoryTaskStore::new();
267        for i in 0..5 {
268            store
269                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
270                .await
271                .unwrap();
272        }
273
274        let req = ListTasksRequest {
275            context_id: None,
276            status: None,
277            page_size: Some(2),
278            page_token: None,
279            history_length: None,
280            status_timestamp_after: None,
281            include_artifacts: None,
282            tenant: None,
283        };
284        let resp = store.list(&req).await.unwrap();
285        assert_eq!(resp.tasks.len(), 2);
286        assert_eq!(resp.total_size, 5);
287        assert!(!resp.next_page_token.is_empty());
288
289        // Get second page
290        let req2 = ListTasksRequest {
291            page_token: Some(resp.next_page_token),
292            page_size: Some(2),
293            ..req
294        };
295        let resp2 = store.list(&req2).await.unwrap();
296        assert_eq!(resp2.tasks.len(), 2);
297    }
298
299    #[tokio::test]
300    async fn test_list_zero_page_size_uses_default_window() {
301        let store = InMemoryTaskStore::new();
302        for i in 0..3 {
303            store
304                .create(make_task(&format!("t{i}"), "c1", TaskState::Submitted))
305                .await
306                .unwrap();
307        }
308
309        let req = ListTasksRequest {
310            context_id: None,
311            status: None,
312            page_size: Some(0),
313            page_token: None,
314            history_length: None,
315            status_timestamp_after: None,
316            include_artifacts: None,
317            tenant: None,
318        };
319        let resp = store.list(&req).await.unwrap();
320        assert_eq!(resp.tasks.len(), 3);
321        assert_eq!(resp.page_size, 50);
322    }
323
324    #[tokio::test]
325    async fn test_list_history_truncation() {
326        let store = InMemoryTaskStore::new();
327        let mut task = make_task("t1", "c1", TaskState::Working);
328        task.history = Some(vec![
329            Message::new(Role::User, vec![Part::text("1")]),
330            Message::new(Role::Agent, vec![Part::text("2")]),
331            Message::new(Role::User, vec![Part::text("3")]),
332        ]);
333        store.create(task).await.unwrap();
334
335        let req = ListTasksRequest {
336            context_id: None,
337            status: None,
338            page_size: None,
339            page_token: None,
340            history_length: Some(1),
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[0].history.as_ref().unwrap().len(), 1);
347    }
348}
a2a_server/task_store/inmemory.rs

a2a_server/task_store/
inmemory.rs
