Skip to main content

mlua_swarm/store/run/
inmemory.rs

1//! `InMemoryRunStore` — a process-volatile `RunStore` used by the current
2//! default.
3
4use super::{
5    DegradationEntry, Inner, RunId, RunRecord, RunStatus, RunStore, RunStoreError, SharedInner,
6    StepEntry, TaskId,
7};
8use async_trait::async_trait;
9use std::sync::Mutex;
10
11/// Process-volatile [`RunStore`] used as the current default. Entries are
12/// lost on restart; persistent backends (SQLite / Git / mini-app / …) are
13/// future carries.
14#[derive(Default)]
15pub struct InMemoryRunStore {
16    inner: SharedInner,
17}
18
19impl InMemoryRunStore {
20    /// Create an empty store.
21    pub fn new() -> Self {
22        Self {
23            inner: Mutex::new(Inner::default()),
24        }
25    }
26}
27
28#[async_trait]
29impl RunStore for InMemoryRunStore {
30    fn name(&self) -> &str {
31        "in-memory"
32    }
33
34    async fn create(&self, record: RunRecord) -> Result<(), RunStoreError> {
35        let mut inner = self.inner.lock().unwrap();
36        if inner.records.contains_key(&record.id) {
37            return Err(RunStoreError::Duplicate(record.id));
38        }
39        inner.order.push(record.id.clone());
40        inner.records.insert(record.id.clone(), record);
41        Ok(())
42    }
43
44    async fn get(&self, id: &RunId) -> Result<RunRecord, RunStoreError> {
45        let inner = self.inner.lock().unwrap();
46        inner
47            .records
48            .get(id)
49            .cloned()
50            .ok_or_else(|| RunStoreError::NotFound(id.clone()))
51    }
52
53    async fn list_by_task(&self, task_id: &TaskId) -> Result<Vec<RunRecord>, RunStoreError> {
54        let inner = self.inner.lock().unwrap();
55        let mut records: Vec<RunRecord> = inner
56            .order
57            .iter()
58            .filter_map(|id| inner.records.get(id).cloned())
59            .filter(|r| &r.task_id == task_id)
60            .collect();
61        records.sort_by_key(|r| r.created_at);
62        Ok(records)
63    }
64
65    async fn append_step_entry(&self, id: &RunId, entry: StepEntry) -> Result<(), RunStoreError> {
66        let mut inner = self.inner.lock().unwrap();
67        let record = inner
68            .records
69            .get_mut(id)
70            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
71        record.step_entries.push(entry);
72        record.updated_at = crate::types::now_unix();
73        Ok(())
74    }
75
76    async fn append_degradation(
77        &self,
78        id: &RunId,
79        entry: DegradationEntry,
80    ) -> Result<(), RunStoreError> {
81        let mut inner = self.inner.lock().unwrap();
82        let record = inner
83            .records
84            .get_mut(id)
85            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
86        record.degradations.push(entry);
87        record.updated_at = crate::types::now_unix();
88        Ok(())
89    }
90
91    async fn update_status(&self, id: &RunId, status: RunStatus) -> Result<(), RunStoreError> {
92        let mut inner = self.inner.lock().unwrap();
93        let record = inner
94            .records
95            .get_mut(id)
96            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
97        record.status = status;
98        record.updated_at = crate::types::now_unix();
99        Ok(())
100    }
101
102    async fn set_result(
103        &self,
104        id: &RunId,
105        result_ref: serde_json::Value,
106    ) -> Result<(), RunStoreError> {
107        let mut inner = self.inner.lock().unwrap();
108        let record = inner
109            .records
110            .get_mut(id)
111            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
112        record.result_ref = Some(result_ref);
113        record.updated_at = crate::types::now_unix();
114        Ok(())
115    }
116
117    async fn list_running(&self) -> Result<Vec<RunRecord>, RunStoreError> {
118        let inner = self.inner.lock().unwrap();
119        let records: Vec<RunRecord> = inner
120            .order
121            .iter()
122            .filter_map(|id| inner.records.get(id).cloned())
123            .filter(|r| r.status == RunStatus::Running)
124            .collect();
125        Ok(records)
126    }
127}
128
129// ──────────────────────────────────────────────────────────────────────────
130// tests
131// ──────────────────────────────────────────────────────────────────────────
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use serde_json::json;
137
138    fn mk(id: &str, task_id: &str, created_at: u64) -> RunRecord {
139        RunRecord {
140            id: RunId::parse(id).unwrap(),
141            task_id: TaskId::parse(task_id).unwrap(),
142            status: RunStatus::Pending,
143            step_entries: vec![],
144            degradations: vec![],
145            operator_sid: None,
146            result_ref: None,
147            created_at,
148            updated_at: created_at,
149        }
150    }
151
152    fn mk_degradation(tool: &str, at: u64) -> DegradationEntry {
153        DegradationEntry {
154            tool: tool.to_string(),
155            error: "boom".to_string(),
156            fallback: "cached-default".to_string(),
157            note: None,
158            step_ref: Some("worker".to_string()),
159            attempt: Some(1),
160            at,
161        }
162    }
163
164    #[tokio::test]
165    async fn create_then_get() {
166        let s = InMemoryRunStore::new();
167        s.create(mk("R-1", "T-1", 100)).await.unwrap();
168        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
169        assert_eq!(got.task_id, TaskId::parse("T-1").unwrap());
170        assert_eq!(got.status, RunStatus::Pending);
171        assert!(got.step_entries.is_empty());
172    }
173
174    #[tokio::test]
175    async fn duplicate_create_rejected() {
176        let s = InMemoryRunStore::new();
177        s.create(mk("R-1", "T-1", 100)).await.unwrap();
178        let err = s.create(mk("R-1", "T-1", 200)).await.unwrap_err();
179        assert!(matches!(err, RunStoreError::Duplicate(_)));
180    }
181
182    #[tokio::test]
183    async fn get_missing_returns_not_found() {
184        let s = InMemoryRunStore::new();
185        let err = s.get(&RunId::parse("R-nope").unwrap()).await.unwrap_err();
186        assert!(matches!(err, RunStoreError::NotFound(_)));
187    }
188
189    #[tokio::test]
190    async fn list_by_task_filters_and_orders_ascending() {
191        let s = InMemoryRunStore::new();
192        s.create(mk("R-1", "T-1", 300)).await.unwrap();
193        s.create(mk("R-2", "T-2", 50)).await.unwrap();
194        s.create(mk("R-3", "T-1", 100)).await.unwrap();
195        let list = s
196            .list_by_task(&TaskId::parse("T-1").unwrap())
197            .await
198            .unwrap();
199        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
200        assert_eq!(ids, vec!["R-3", "R-1"]);
201    }
202
203    #[tokio::test]
204    async fn append_step_entry_accumulates_in_order() {
205        let s = InMemoryRunStore::new();
206        s.create(mk("R-1", "T-1", 100)).await.unwrap();
207        s.append_step_entry(
208            &RunId::parse("R-1").unwrap(),
209            StepEntry {
210                step_id: crate::types::StepId::parse("ST-1").unwrap(),
211                step_ref: Some("step-a".into()),
212                status: Some("dispatched".into()),
213                at: 101,
214            },
215        )
216        .await
217        .unwrap();
218        s.append_step_entry(
219            &RunId::parse("R-1").unwrap(),
220            StepEntry {
221                step_id: crate::types::StepId::parse("ST-2").unwrap(),
222                step_ref: Some("step-b".into()),
223                status: Some("passed".into()),
224                at: 102,
225            },
226        )
227        .await
228        .unwrap();
229        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
230        assert_eq!(got.step_entries.len(), 2);
231        assert_eq!(got.step_entries[0].step_ref, Some("step-a".into()));
232        assert_eq!(got.step_entries[1].step_ref, Some("step-b".into()));
233        assert!(got.updated_at >= got.created_at);
234    }
235
236    #[tokio::test]
237    async fn append_degradation_accumulates_in_order() {
238        let s = InMemoryRunStore::new();
239        s.create(mk("R-1", "T-1", 100)).await.unwrap();
240        s.append_degradation(
241            &RunId::parse("R-1").unwrap(),
242            mk_degradation("web_search", 101),
243        )
244        .await
245        .unwrap();
246        s.append_degradation(
247            &RunId::parse("R-1").unwrap(),
248            mk_degradation("code_exec", 102),
249        )
250        .await
251        .unwrap();
252        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
253        assert_eq!(got.degradations.len(), 2);
254        assert_eq!(got.degradations[0].tool, "web_search");
255        assert_eq!(got.degradations[1].tool, "code_exec");
256        assert!(got.updated_at >= got.created_at);
257    }
258
259    #[tokio::test]
260    async fn append_degradation_unknown_run_fails() {
261        let s = InMemoryRunStore::new();
262        let err = s
263            .append_degradation(
264                &RunId::parse("R-nope").unwrap(),
265                mk_degradation("web_search", 1),
266            )
267            .await
268            .unwrap_err();
269        assert!(matches!(err, RunStoreError::NotFound(_)));
270    }
271
272    #[tokio::test]
273    async fn append_degradation_bumps_updated_at() {
274        let s = InMemoryRunStore::new();
275        s.create(mk("R-1", "T-1", 100)).await.unwrap();
276        s.append_degradation(
277            &RunId::parse("R-1").unwrap(),
278            mk_degradation("web_search", 200),
279        )
280        .await
281        .unwrap();
282        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
283        assert!(got.updated_at > 100);
284    }
285
286    #[tokio::test]
287    async fn append_step_entry_unknown_run_fails() {
288        let s = InMemoryRunStore::new();
289        let err = s
290            .append_step_entry(
291                &RunId::parse("R-nope").unwrap(),
292                StepEntry {
293                    step_id: crate::types::StepId::parse("ST-1").unwrap(),
294                    step_ref: None,
295                    status: None,
296                    at: 1,
297                },
298            )
299            .await
300            .unwrap_err();
301        assert!(matches!(err, RunStoreError::NotFound(_)));
302    }
303
304    #[tokio::test]
305    async fn update_status_persists() {
306        let s = InMemoryRunStore::new();
307        s.create(mk("R-1", "T-1", 100)).await.unwrap();
308        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Running)
309            .await
310            .unwrap();
311        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
312        assert_eq!(got.status, RunStatus::Running);
313    }
314
315    #[tokio::test]
316    async fn set_result_persists() {
317        let s = InMemoryRunStore::new();
318        s.create(mk("R-1", "T-1", 100)).await.unwrap();
319        s.set_result(&RunId::parse("R-1").unwrap(), json!({"ok": true}))
320            .await
321            .unwrap();
322        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
323        assert_eq!(got.result_ref, Some(json!({"ok": true})));
324    }
325
326    #[tokio::test]
327    async fn name_is_in_memory() {
328        assert_eq!(InMemoryRunStore::new().name(), "in-memory");
329    }
330
331    #[tokio::test]
332    async fn list_running_filters_by_status() {
333        let s = InMemoryRunStore::new();
334        s.create(mk("R-1", "T-1", 100)).await.unwrap();
335        s.create(mk("R-2", "T-2", 200)).await.unwrap();
336        s.create(mk("R-3", "T-3", 300)).await.unwrap();
337        s.update_status(&RunId::parse("R-2").unwrap(), RunStatus::Running)
338            .await
339            .unwrap();
340        s.update_status(&RunId::parse("R-3").unwrap(), RunStatus::Done)
341            .await
342            .unwrap();
343        let running = s.list_running().await.unwrap();
344        assert_eq!(running.len(), 1);
345        assert_eq!(running[0].id, RunId::parse("R-2").unwrap());
346        assert_eq!(running[0].status, RunStatus::Running);
347    }
348}