Skip to main content

mlua_swarm/store/task/
inmemory.rs

1//! `InMemoryTaskStore` — a process-volatile `TaskStore` used by the current
2//! default.
3
4use super::{Inner, SharedInner, TaskId, TaskRecord, TaskRecordStatus, TaskStore, TaskStoreError};
5use async_trait::async_trait;
6use std::sync::Mutex;
7
8/// Process-volatile [`TaskStore`] used as the current default. Entries are
9/// lost on restart; persistent backends (SQLite / Git / mini-app / …) are
10/// future carries.
11#[derive(Default)]
12pub struct InMemoryTaskStore {
13    inner: SharedInner,
14}
15
16impl InMemoryTaskStore {
17    /// Create an empty store.
18    pub fn new() -> Self {
19        Self {
20            inner: Mutex::new(Inner::default()),
21        }
22    }
23}
24
25#[async_trait]
26impl TaskStore for InMemoryTaskStore {
27    fn name(&self) -> &str {
28        "in-memory"
29    }
30
31    async fn create(&self, record: TaskRecord) -> Result<(), TaskStoreError> {
32        let mut inner = self.inner.lock().unwrap();
33        if inner.records.contains_key(&record.id) {
34            return Err(TaskStoreError::Duplicate(record.id));
35        }
36        inner.order.push(record.id.clone());
37        inner.records.insert(record.id.clone(), record);
38        Ok(())
39    }
40
41    async fn get(&self, id: &TaskId) -> Result<TaskRecord, TaskStoreError> {
42        let inner = self.inner.lock().unwrap();
43        inner
44            .records
45            .get(id)
46            .cloned()
47            .ok_or_else(|| TaskStoreError::NotFound(id.clone()))
48    }
49
50    async fn list(&self) -> Result<Vec<TaskRecord>, TaskStoreError> {
51        let inner = self.inner.lock().unwrap();
52        let mut records: Vec<TaskRecord> = inner
53            .order
54            .iter()
55            .filter_map(|id| inner.records.get(id).cloned())
56            .collect();
57        records.sort_by_key(|r| std::cmp::Reverse(r.created_at));
58        Ok(records)
59    }
60
61    async fn update_status(
62        &self,
63        id: &TaskId,
64        status: TaskRecordStatus,
65    ) -> Result<(), TaskStoreError> {
66        let mut inner = self.inner.lock().unwrap();
67        let record = inner
68            .records
69            .get_mut(id)
70            .ok_or_else(|| TaskStoreError::NotFound(id.clone()))?;
71        record.status = status;
72        record.updated_at = crate::types::now_unix();
73        Ok(())
74    }
75}
76
77// ──────────────────────────────────────────────────────────────────────────
78// tests
79// ──────────────────────────────────────────────────────────────────────────
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use serde_json::json;
85
86    fn mk(id: &str, created_at: u64) -> TaskRecord {
87        TaskRecord {
88            id: TaskId::parse(id).unwrap(),
89            goal: format!("goal for {id}"),
90            blueprint_ref: json!({"id": "bp-1"}),
91            input_ctx: json!({"k": "v"}),
92            task_input_spec: None,
93            status: TaskRecordStatus::Pending,
94            created_at,
95            updated_at: created_at,
96        }
97    }
98
99    #[tokio::test]
100    async fn create_then_get() {
101        let s = InMemoryTaskStore::new();
102        s.create(mk("T-1", 100)).await.unwrap();
103        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
104        assert_eq!(got.id, TaskId::parse("T-1").unwrap());
105        assert_eq!(got.goal, "goal for T-1");
106        assert_eq!(got.status, TaskRecordStatus::Pending);
107    }
108
109    #[tokio::test]
110    async fn duplicate_create_rejected() {
111        let s = InMemoryTaskStore::new();
112        s.create(mk("T-1", 100)).await.unwrap();
113        let err = s.create(mk("T-1", 200)).await.unwrap_err();
114        assert!(matches!(err, TaskStoreError::Duplicate(_)));
115    }
116
117    #[tokio::test]
118    async fn get_missing_returns_not_found() {
119        let s = InMemoryTaskStore::new();
120        let err = s.get(&TaskId::parse("T-nope").unwrap()).await.unwrap_err();
121        assert!(matches!(err, TaskStoreError::NotFound(_)));
122    }
123
124    #[tokio::test]
125    async fn list_returns_newest_first() {
126        let s = InMemoryTaskStore::new();
127        s.create(mk("T-1", 100)).await.unwrap();
128        s.create(mk("T-2", 300)).await.unwrap();
129        s.create(mk("T-3", 200)).await.unwrap();
130        let list = s.list().await.unwrap();
131        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
132        assert_eq!(ids, vec!["T-2", "T-3", "T-1"]);
133    }
134
135    #[tokio::test]
136    async fn update_status_bumps_updated_at_and_persists_status() {
137        let s = InMemoryTaskStore::new();
138        s.create(mk("T-1", 100)).await.unwrap();
139        s.update_status(&TaskId::parse("T-1").unwrap(), TaskRecordStatus::Running)
140            .await
141            .unwrap();
142        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
143        assert_eq!(got.status, TaskRecordStatus::Running);
144        assert!(got.updated_at >= got.created_at);
145    }
146
147    #[tokio::test]
148    async fn update_status_unknown_fails() {
149        let s = InMemoryTaskStore::new();
150        let err = s
151            .update_status(&TaskId::parse("T-nope").unwrap(), TaskRecordStatus::Done)
152            .await
153            .unwrap_err();
154        assert!(matches!(err, TaskStoreError::NotFound(_)));
155    }
156
157    #[tokio::test]
158    async fn name_is_in_memory() {
159        assert_eq!(InMemoryTaskStore::new().name(), "in-memory");
160    }
161
162    #[tokio::test]
163    async fn task_input_spec_round_trips_through_create_and_get() {
164        // Issue #19 ST4: `TaskRecord.task_input_spec` is a plain field on
165        // this backend (no encode/decode step), but a regression guard
166        // here still catches an accidental drop/clear in `create`/`get`.
167        let s = InMemoryTaskStore::new();
168        let mut record = mk("T-1", 100);
169        record.task_input_spec = Some(json!({"project_root": "/repo"}));
170        s.create(record).await.unwrap();
171        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
172        assert_eq!(got.task_input_spec, Some(json!({"project_root": "/repo"})));
173    }
174}