1use 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#[derive(Default)]
15pub struct InMemoryRunStore {
16 inner: SharedInner,
17}
18
19impl InMemoryRunStore {
20 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#[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}