1use crate::contract::event::AgentEvent;
13use crate::contract::finding::Finding;
14use crate::contract::ids::{AgentId, PhaseId, RunId};
15use serde::{Deserialize, Serialize};
16use std::collections::HashMap;
17use std::fs::{self, File, OpenOptions};
18use std::io::{BufRead, BufReader, Write};
19use std::path::{Path, PathBuf};
20use std::sync::{Arc, RwLock};
21use std::time::{SystemTime, UNIX_EPOCH};
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct RunCheckpoint {
26 pub run_id: RunId,
27 pub task: String,
28 pub status: CheckpointStatus,
29 pub current_phase: u32,
30 pub completed_phases: Vec<PhaseSummary>,
31 pub agent_results: HashMap<AgentId, AgentResultCache>,
32 #[serde(default)]
35 pub agent_sessions: HashMap<AgentId, AgentSessionCheckpoint>,
36 pub findings: Vec<Finding>,
37 pub total_tokens: u64,
38 pub created_at: u64,
39 pub updated_at: u64,
40 #[serde(default)]
41 pub workflow_meta: Option<serde_json::Value>,
42 #[serde(default)]
45 pub started_agent_ids: Vec<AgentId>,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
49#[serde(rename_all = "lowercase")]
50pub enum CheckpointStatus {
51 Running,
52 Completed,
53 Failed,
54 Cancelled,
55}
56
57impl std::fmt::Display for CheckpointStatus {
58 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59 let s = match self {
60 CheckpointStatus::Running => "Running",
61 CheckpointStatus::Completed => "Completed",
62 CheckpointStatus::Failed => "Failed",
63 CheckpointStatus::Cancelled => "Cancelled",
64 };
65 f.write_str(s)
66 }
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct PhaseSummary {
71 pub phase_id: PhaseId,
72 pub label: String,
73 pub planned: usize,
74 pub ok: usize,
75 pub failed: usize,
76 #[serde(default)]
77 pub description: Option<String>,
78 #[serde(default)]
79 pub role: Option<String>,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct AgentResultCache {
84 pub agent_id: AgentId,
85 pub phase_id: PhaseId,
86 pub status: String,
87 pub output: serde_json::Value,
88 pub findings: Vec<Finding>,
89 pub tokens: u64,
90 pub completed_at: u64,
91 #[serde(default)]
94 pub cache_key_hash: Option<String>,
95 #[serde(default)]
96 pub description: Option<String>,
97 #[serde(default)]
98 pub role: Option<String>,
99}
100
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct AgentSessionCheckpoint {
103 pub agent_id: AgentId,
104 #[serde(default)]
106 pub backend_id: Option<String>,
107 #[serde(default)]
110 pub protocol_session_id: Option<String>,
111 pub session_id: String,
112 pub status: String,
113 pub updated_at: u64,
114 #[serde(default)]
115 pub resumable: bool,
116}
117
118#[derive(Debug)]
120pub struct RunStore {
121 run_dir: PathBuf,
122 checkpoint: RwLock<Option<RunCheckpoint>>,
123 events_file: RwLock<Option<File>>,
124}
125
126impl RunStore {
127 pub fn new(run_dir: &Path) -> Result<Arc<Self>, std::io::Error> {
129 tracing::debug!(path = %run_dir.display(), "creating RunStore");
130 fs::create_dir_all(run_dir)?;
131
132 let store = Arc::new(Self {
133 run_dir: run_dir.to_path_buf(),
134 checkpoint: RwLock::new(None),
135 events_file: RwLock::new(None),
136 });
137
138 Ok(store)
139 }
140
141 pub fn upsert_agent_result(&self, cache: &AgentResultCache) -> Result<(), std::io::Error> {
144 let mut guard = self.checkpoint.write().unwrap();
145 if let Some(ref mut checkpoint) = *guard {
146 checkpoint
147 .agent_results
148 .insert(cache.agent_id, cache.clone());
149 checkpoint.updated_at = current_timestamp();
150 let cp = checkpoint.clone();
151 drop(guard);
152 let cp_path = self.run_dir.join("checkpoint.json");
153 let content = serde_json::to_string_pretty(&cp)
154 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
155 fs::write(&cp_path, content)?;
156 }
157 Ok(())
158 }
159
160 pub fn upsert_agent_session(
162 &self,
163 session: &AgentSessionCheckpoint,
164 ) -> Result<(), std::io::Error> {
165 let mut guard = self.checkpoint.write().unwrap();
166 if let Some(ref mut checkpoint) = *guard {
167 checkpoint
168 .agent_sessions
169 .insert(session.agent_id, session.clone());
170 checkpoint.updated_at = current_timestamp();
171 let cp = checkpoint.clone();
172 drop(guard);
173 let cp_path = self.run_dir.join("checkpoint.json");
174 let content = serde_json::to_string_pretty(&cp)
175 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
176 fs::write(&cp_path, content)?;
177 }
178 Ok(())
179 }
180
181 pub fn init_run(&self, run_id: RunId, task: &str) -> Result<(), std::io::Error> {
183 tracing::info!(%run_id, %task, "initializing run store");
184 let checkpoint = RunCheckpoint {
185 run_id,
186 task: task.to_string(),
187 status: CheckpointStatus::Running,
188 current_phase: 0,
189 completed_phases: vec![],
190 agent_results: HashMap::new(),
191 agent_sessions: HashMap::new(),
192 findings: vec![],
193 total_tokens: 0,
194 created_at: current_timestamp(),
195 updated_at: current_timestamp(),
196 workflow_meta: None,
197 started_agent_ids: vec![],
198 };
199
200 self.save_checkpoint(&checkpoint)?;
202
203 let events_path = self.run_dir.join("events.jsonl");
205 let events_file = OpenOptions::new()
206 .create(true)
207 .append(true)
208 .open(events_path)?;
209
210 let mut checkpoint_guard = self.checkpoint.write().unwrap();
211 *checkpoint_guard = Some(checkpoint);
212
213 let mut events_guard = self.events_file.write().unwrap();
214 *events_guard = Some(events_file);
215
216 Ok(())
217 }
218
219 pub fn init_run_with_meta(
221 &self,
222 run_id: RunId,
223 task: &str,
224 workflow_meta: serde_json::Value,
225 ) -> Result<(), std::io::Error> {
226 tracing::info!(%run_id, %task, "initializing run store with meta");
227 let checkpoint = RunCheckpoint {
228 run_id,
229 task: task.to_string(),
230 status: CheckpointStatus::Running,
231 current_phase: 0,
232 completed_phases: vec![],
233 agent_results: HashMap::new(),
234 agent_sessions: HashMap::new(),
235 findings: vec![],
236 total_tokens: 0,
237 created_at: current_timestamp(),
238 updated_at: current_timestamp(),
239 workflow_meta: Some(workflow_meta),
240 started_agent_ids: vec![],
241 };
242
243 self.save_checkpoint(&checkpoint)?;
244
245 let events_path = self.run_dir.join("events.jsonl");
246 let events_file = OpenOptions::new()
247 .create(true)
248 .append(true)
249 .open(events_path)?;
250
251 let mut checkpoint_guard = self.checkpoint.write().unwrap();
252 *checkpoint_guard = Some(checkpoint);
253
254 let mut events_guard = self.events_file.write().unwrap();
255 *events_guard = Some(events_file);
256
257 Ok(())
258 }
259
260 pub fn open_run(&self, _run_id: RunId) -> Result<Option<RunCheckpoint>, std::io::Error> {
262 tracing::debug!(%_run_id, "opening existing run");
263 let checkpoint_path = self.run_dir.join("checkpoint.json");
264
265 if !checkpoint_path.exists() {
266 return Ok(None);
267 }
268
269 let content = fs::read_to_string(&checkpoint_path)?;
270 let checkpoint: RunCheckpoint = serde_json::from_str(&content)
271 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
272
273 let events_path = self.run_dir.join("events.jsonl");
278 let events_file = OpenOptions::new()
279 .read(true)
280 .append(true)
281 .open(events_path)?;
282
283 let mut checkpoint_guard = self.checkpoint.write().unwrap();
284 *checkpoint_guard = Some(checkpoint.clone());
285
286 let mut events_guard = self.events_file.write().unwrap();
287 *events_guard = Some(events_file);
288
289 Ok(Some(checkpoint))
290 }
291
292 pub fn append_event(&self, event: &AgentEvent) -> Result<(), std::io::Error> {
294 let json = serde_json::to_string(event)
295 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
296
297 let mut events_guard = self.events_file.write().unwrap();
298 if let Some(ref mut file) = *events_guard {
299 writeln!(file, "{}", json)?;
300 file.flush()?;
301 }
302
303 self.update_from_event(event);
305
306 Ok(())
307 }
308
309 fn update_from_event(&self, event: &AgentEvent) {
311 let mut checkpoint_guard = self.checkpoint.write().unwrap();
312 if let Some(ref mut checkpoint) = *checkpoint_guard {
313 match event {
314 AgentEvent::AgentDone {
315 agent_id,
316 status,
317 tokens,
318 ..
319 } => {
320 let existing = checkpoint.agent_results.get(agent_id);
321 let cache = AgentResultCache {
322 agent_id: *agent_id,
323 phase_id: existing.map(|c| c.phase_id).unwrap_or(0),
324 status: status.as_str().to_string(),
325 output: existing
326 .map(|c| c.output.clone())
327 .unwrap_or(serde_json::Value::Null),
328 findings: existing.map(|c| c.findings.clone()).unwrap_or_default(),
329 tokens: tokens.total(),
330 completed_at: existing
331 .map(|c| c.completed_at)
332 .unwrap_or(current_timestamp()),
333 cache_key_hash: existing.and_then(|c| c.cache_key_hash.clone()),
334 description: existing.and_then(|c| c.description.clone()),
335 role: existing.and_then(|c| c.role.clone()),
336 };
337 checkpoint.agent_results.insert(*agent_id, cache);
338 checkpoint.total_tokens += tokens.total();
339 }
340 AgentEvent::AgentStarted { agent_id, .. } => {
341 if !checkpoint.started_agent_ids.contains(agent_id) {
342 checkpoint.started_agent_ids.push(*agent_id);
343 }
344 }
345 AgentEvent::PhaseDone { phase_id, .. } => {
346 if *phase_id > 0 {
347 checkpoint.current_phase = *phase_id;
348 }
349 }
350 AgentEvent::RunDone {
351 status,
352 total_tokens,
353 ..
354 } => {
355 let cancelled_on_disk =
359 fs::read_to_string(self.run_dir.join("checkpoint.json"))
360 .ok()
361 .and_then(|content| {
362 serde_json::from_str::<RunCheckpoint>(&content).ok()
363 })
364 .is_some_and(|disk_checkpoint| {
365 disk_checkpoint.status == CheckpointStatus::Cancelled
366 });
367 if cancelled_on_disk {
368 checkpoint.status = CheckpointStatus::Cancelled;
369 } else if checkpoint.status != CheckpointStatus::Cancelled {
370 checkpoint.status = match status {
371 crate::contract::event::RunStatus::Completed => {
372 CheckpointStatus::Completed
373 }
374 crate::contract::event::RunStatus::Failed => CheckpointStatus::Failed,
375 crate::contract::event::RunStatus::Cancelled => {
376 CheckpointStatus::Cancelled
377 }
378 crate::contract::event::RunStatus::Partial => CheckpointStatus::Running,
379 };
380 }
381 let t = total_tokens.total();
384 if t > 0 {
385 checkpoint.total_tokens = t;
386 }
387 }
388 _ => {}
389 }
390 checkpoint.updated_at = current_timestamp();
391
392 if let Err(e) = self.write_checkpoint_to_disk(checkpoint) {
394 tracing::warn!(error = %e, "failed to save checkpoint");
395 }
396 }
397 }
398
399 fn write_checkpoint_to_disk(&self, checkpoint: &RunCheckpoint) -> Result<(), std::io::Error> {
401 let checkpoint_path = self.run_dir.join("checkpoint.json");
402 let temp_path = self.run_dir.join("checkpoint.json.tmp");
403 let mut checkpoint_to_write = checkpoint.clone();
408 if let Ok(existing) = std::fs::read_to_string(&checkpoint_path) {
409 if serde_json::from_str::<RunCheckpoint>(&existing)
410 .ok()
411 .is_some_and(|cp| cp.status == CheckpointStatus::Cancelled)
412 {
413 checkpoint_to_write.status = CheckpointStatus::Cancelled;
414 }
415 }
416 let content = serde_json::to_string_pretty(&checkpoint_to_write)
417 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
418 std::fs::write(&temp_path, &content)?;
419 std::fs::rename(&temp_path, &checkpoint_path)?;
420 Ok(())
421 }
422
423 pub fn save_checkpoint(&self, checkpoint: &RunCheckpoint) -> Result<(), std::io::Error> {
425 let checkpoint_path = self.run_dir.join("checkpoint.json");
426 let temp_path = self.run_dir.join("checkpoint.json.tmp");
427 let content = serde_json::to_string_pretty(checkpoint)
428 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
429 std::fs::write(&temp_path, &content)?;
430 std::fs::rename(&temp_path, &checkpoint_path)?;
431
432 let mut checkpoint_guard = self.checkpoint.write().unwrap();
433 *checkpoint_guard = Some(checkpoint.clone());
434
435 Ok(())
436 }
437
438 pub fn get_checkpoint(&self) -> Option<RunCheckpoint> {
440 let guard = self.checkpoint.read().unwrap();
441 guard.clone()
442 }
443
444 pub fn get_findings(&self) -> Vec<Finding> {
446 let guard = self.checkpoint.read().unwrap();
447 guard
448 .as_ref()
449 .map(|c| c.findings.clone())
450 .unwrap_or_default()
451 }
452
453 pub fn get_event_log(&self) -> Result<Vec<AgentEvent>, std::io::Error> {
455 let events_path = self.run_dir.join("events.jsonl");
456 let file = File::open(events_path)?;
457 let reader = BufReader::new(file);
458 let mut events = Vec::new();
459
460 for line in reader.lines() {
461 let line = line?;
462 if line.trim().is_empty() {
463 continue;
464 }
465 let event: AgentEvent = serde_json::from_str(&line)
466 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
467 events.push(event);
468 }
469
470 Ok(events)
471 }
472
473 pub fn can_resume(&self) -> bool {
475 let guard = self.checkpoint.read().unwrap();
476 matches!(
477 guard.as_ref().map(|c| c.status.clone()),
478 Some(CheckpointStatus::Running)
479 )
480 }
481
482 pub fn reset_status_to_running(&self) -> Result<(), std::io::Error> {
487 let mut guard = self.checkpoint.write().unwrap();
488
489 if guard.is_none() {
490 let checkpoint_path = self.run_dir.join("checkpoint.json");
491 if !checkpoint_path.exists() {
492 return Ok(());
493 }
494 let content = fs::read_to_string(&checkpoint_path)?;
495 let cp: RunCheckpoint = serde_json::from_str(&content)
496 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
497 *guard = Some(cp);
498 }
499
500 if let Some(ref mut checkpoint) = *guard {
501 checkpoint.status = CheckpointStatus::Running;
502 checkpoint.updated_at = current_timestamp();
503 let cp_clone = checkpoint.clone();
504 drop(guard);
505 self.write_checkpoint_to_disk(&cp_clone)?;
506 }
507 Ok(())
508 }
509
510 pub fn cancel(&self) -> Result<(), std::io::Error> {
512 tracing::info!("cancelling run");
513 let mut guard = self.checkpoint.write().unwrap();
514
515 if guard.is_none() {
523 let checkpoint_path = self.run_dir.join("checkpoint.json");
524 if !checkpoint_path.exists() {
525 return Ok(());
526 }
527 let content = fs::read_to_string(&checkpoint_path)?;
528 let cp: RunCheckpoint = serde_json::from_str(&content)
529 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
530 *guard = Some(cp);
531 }
532
533 if let Some(ref mut checkpoint) = *guard {
534 checkpoint.status = CheckpointStatus::Cancelled;
535 checkpoint.updated_at = current_timestamp();
536 let checkpoint = checkpoint.clone();
537 drop(guard);
538 self.write_checkpoint_to_disk(&checkpoint)?;
542 }
543 Ok(())
544 }
545}
546
547fn current_timestamp() -> u64 {
549 SystemTime::now()
550 .duration_since(UNIX_EPOCH)
551 .map(|d| d.as_secs())
552 .unwrap_or(0)
553}
554
555use std::sync::OnceLock;
560
561static RUN_STORES: OnceLock<dashmap::DashMap<String, Arc<RunStore>>> = OnceLock::new();
562
563fn get_run_stores() -> &'static dashmap::DashMap<String, Arc<RunStore>> {
565 RUN_STORES.get_or_init(dashmap::DashMap::new)
566}
567
568pub fn get_run_store(run_dir_name: &str, base_dir: &Path) -> Result<Arc<RunStore>, std::io::Error> {
570 let stores = get_run_stores();
571
572 if let Some(store) = stores.get(run_dir_name) {
573 return Ok(store.clone());
574 }
575
576 let run_dir = base_dir.join(run_dir_name);
577 let store = RunStore::new(&run_dir)?;
578 stores.insert(run_dir_name.to_string(), store.clone());
579
580 Ok(store)
581}
582
583pub fn list_runs(base_dir: &Path) -> Result<Vec<String>, std::io::Error> {
585 if !base_dir.exists() {
586 return Ok(vec![]);
587 }
588
589 let mut run_dirs = Vec::new();
590 for entry in fs::read_dir(base_dir)? {
591 let entry = entry?;
592 let path = entry.path();
593 if path.is_dir() {
594 if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
595 run_dirs.push(name.to_string());
596 }
597 }
598 }
599
600 run_dirs.sort();
601 Ok(run_dirs)
602}
603
604#[cfg(test)]
605mod tests {
606 use super::*;
607 use tempfile::tempdir;
608
609 #[test]
610 fn test_run_store_init() {
611 let dir = tempdir().unwrap();
612 let run_id = uuid::Uuid::now_v7();
613 let store = RunStore::new(dir.path()).unwrap();
614 store.init_run(run_id, "Test task").unwrap();
615
616 let checkpoint = store.get_checkpoint().unwrap();
617 assert_eq!(checkpoint.run_id, run_id);
618 assert_eq!(checkpoint.task, "Test task");
619 assert_eq!(checkpoint.status, CheckpointStatus::Running);
620 }
621
622 #[test]
623 fn test_run_store_resume() {
624 let dir = tempdir().unwrap();
625 let run_id = uuid::Uuid::now_v7();
626 let store = RunStore::new(dir.path()).unwrap();
627 store.init_run(run_id, "Test task").unwrap();
628
629 let store2 = RunStore::new(dir.path()).unwrap();
631 let checkpoint = store2.open_run(run_id).unwrap().unwrap();
632 assert_eq!(checkpoint.run_id, run_id);
633 assert_eq!(checkpoint.task, "Test task");
634 }
635
636 #[test]
637 fn test_can_resume() {
638 let dir = tempdir().unwrap();
639 let run_id = uuid::Uuid::now_v7();
640 let store = RunStore::new(dir.path()).unwrap();
641 store.init_run(run_id, "Test task").unwrap();
642
643 assert!(store.can_resume());
644 }
645
646 #[test]
647 fn test_resume_appends_events() {
648 let dir = tempdir().unwrap();
652 let run_id = uuid::Uuid::now_v7();
653 let store = RunStore::new(dir.path()).unwrap();
654 store.init_run(run_id, "Test task").unwrap();
655
656 let store2 = RunStore::new(dir.path()).unwrap();
657 store2.open_run(run_id).unwrap().unwrap();
658
659 let evt = AgentEvent::Log {
661 run_id,
662 agent_id: None,
663 level: crate::contract::event::LogLevel::Info,
664 msg: "resume smoke test".to_string(),
665 };
666 store2
667 .append_event(&evt)
668 .expect("append_event after resume must succeed");
669
670 let log = store2.get_event_log().expect("read events.jsonl");
671 assert!(
672 log.iter().any(|e| matches!(
673 e,
674 AgentEvent::Log { msg, .. } if msg == "resume smoke test"
675 )),
676 "event written after open_run must appear in events.jsonl"
677 );
678 }
679
680 use crate::contract::backend::AgentStatus;
691 use crate::contract::ids::TokenUsage;
692 use std::collections::HashSet;
693
694 fn sample_token_usage() -> TokenUsage {
695 TokenUsage {
696 input: 10,
697 output: 5,
698 cache_read: 0,
699 cache_write: 0,
700 }
701 }
702
703 fn build_agent_done(
704 run_id: RunId,
705 agent_id: AgentId,
706 status: AgentStatus,
707 tokens: TokenUsage,
708 ) -> AgentEvent {
709 AgentEvent::AgentDone {
710 run_id,
711 agent_id,
712 status,
713 tokens,
714 elapsed_ms: 0,
715 name: None,
716 agent_seq: 0,
717 output: serde_json::Value::Null,
718 findings: vec![],
719 prompt: String::new(),
720 retry_count: 0,
721 ts: Default::default(),
722 }
723 }
724
725 fn read_raw_checkpoint(run_dir: &Path) -> serde_json::Value {
726 let path = run_dir.join("checkpoint.json");
727 let content =
728 std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read checkpoint.json: {e}"));
729 serde_json::from_str(&content).unwrap_or_else(|e| panic!("parse checkpoint.json: {e}"))
730 }
731
732 #[test]
735 fn upsert_agent_result_persists_to_disk() {
736 let dir = tempdir().unwrap();
740 let run_id = uuid::Uuid::now_v7();
741 let store = RunStore::new(dir.path()).unwrap();
742 store.init_run(run_id, "upsert test").unwrap();
743
744 let agent_id = uuid::Uuid::now_v7();
745 let cache = AgentResultCache {
746 agent_id,
747 phase_id: 1,
748 status: "ok".into(),
749 output: serde_json::json!({"v": 42}),
750 findings: vec![],
751 tokens: 100,
752 completed_at: 1_700_000_000,
753 cache_key_hash: Some("deadbeef".into()),
754 description: None,
755 role: None,
756 };
757 store.upsert_agent_result(&cache).unwrap();
758
759 let cp = store.get_checkpoint().expect("checkpoint present");
761 let cached = cp
762 .agent_results
763 .get(&agent_id)
764 .expect("agent_id indexed after upsert");
765 assert_eq!(cached.tokens, 100);
766 assert_eq!(cached.status, "ok");
767
768 let raw = read_raw_checkpoint(dir.path());
770 let ar = raw
771 .get("agent_results")
772 .and_then(|v| v.as_object())
773 .expect("agent_results object");
774 assert_eq!(ar.len(), 1, "exactly one agent cached on disk");
775 let entry = ar.values().next().expect("non-empty agent_results on disk");
776 assert_eq!(entry.get("tokens").and_then(|v| v.as_u64()), Some(100));
777 assert_eq!(entry.get("status").and_then(|v| v.as_str()), Some("ok"));
778 assert_eq!(
779 entry.get("cache_key_hash").and_then(|v| v.as_str()),
780 Some("deadbeef")
781 );
782
783 drop(store);
785 let reopened = RunStore::new(dir.path()).unwrap();
786 let restored = reopened.open_run(run_id).unwrap().unwrap();
787 assert!(
788 restored.agent_results.contains_key(&agent_id),
789 "upserted entry must survive close+reopen"
790 );
791 assert_eq!(restored.agent_results[&agent_id].tokens, 100);
792 }
793
794 #[test]
795 fn upsert_agent_session_persists_and_reloads() {
796 let dir = tempdir().unwrap();
797 let run_id = uuid::Uuid::now_v7();
798 let store = RunStore::new(dir.path()).unwrap();
799 store.init_run(run_id, "session checkpoint test").unwrap();
800
801 let agent_id = uuid::Uuid::now_v7();
802 store
803 .upsert_agent_session(&AgentSessionCheckpoint {
804 agent_id,
805 backend_id: Some("mock".into()),
806 protocol_session_id: Some("luft-session-1".into()),
807 session_id: "luft-session-1".into(),
808 status: "ok".into(),
809 updated_at: 1_700_000_000,
810 resumable: true,
811 })
812 .unwrap();
813
814 let raw = read_raw_checkpoint(dir.path());
815 assert_eq!(
816 raw["agent_sessions"][agent_id.to_string()]["session_id"],
817 "luft-session-1"
818 );
819
820 drop(store);
821 let reopened = RunStore::new(dir.path()).unwrap();
822 let restored = reopened.open_run(run_id).unwrap().unwrap();
823 let session = restored.agent_sessions.get(&agent_id).unwrap();
824 assert_eq!(session.session_id, "luft-session-1");
825 assert!(session.resumable);
826 }
827
828 #[test]
829 fn upsert_agent_result_updates_existing_entry() {
830 let dir = tempdir().unwrap();
833 let run_id = uuid::Uuid::now_v7();
834 let store = RunStore::new(dir.path()).unwrap();
835 store.init_run(run_id, "overwrite test").unwrap();
836
837 let agent_id = uuid::Uuid::now_v7();
838 let first = AgentResultCache {
839 agent_id,
840 phase_id: 1,
841 status: "ok".into(),
842 output: serde_json::json!("first"),
843 findings: vec![],
844 tokens: 10,
845 completed_at: 1,
846 cache_key_hash: None,
847 description: None,
848 role: None,
849 };
850 let second = AgentResultCache {
851 agent_id,
852 phase_id: 1,
853 status: "error".into(),
854 output: serde_json::json!("second"),
855 findings: vec![],
856 tokens: 99,
857 completed_at: 2,
858 cache_key_hash: None,
859 description: None,
860 role: None,
861 };
862 store.upsert_agent_result(&first).unwrap();
863 store.upsert_agent_result(&second).unwrap();
864
865 let cp = store.get_checkpoint().unwrap();
866 assert_eq!(cp.agent_results.len(), 1, "no duplicate entries");
867 let cached = &cp.agent_results[&agent_id];
868 assert_eq!(cached.status, "error");
869 assert_eq!(cached.tokens, 99);
870 assert_eq!(cached.completed_at, 2);
871
872 let raw = read_raw_checkpoint(dir.path());
874 let ar = raw
875 .get("agent_results")
876 .and_then(|v| v.as_object())
877 .unwrap();
878 assert_eq!(ar.len(), 1);
879 let entry = ar.values().next().unwrap();
880 assert_eq!(entry.get("tokens").and_then(|v| v.as_u64()), Some(99));
881 assert_eq!(entry.get("status").and_then(|v| v.as_str()), Some("error"));
882 }
883
884 #[test]
885 fn upsert_agent_result_noop_when_uninitialized() {
886 let dir = tempdir().unwrap();
890 let store = RunStore::new(dir.path()).unwrap();
891 assert!(store.get_checkpoint().is_none());
892 let cp_path = dir.path().join("checkpoint.json");
893 assert!(!cp_path.exists(), "no checkpoint.json before init");
894
895 let cache = AgentResultCache {
896 agent_id: uuid::Uuid::now_v7(),
897 phase_id: 1,
898 status: "ok".into(),
899 output: serde_json::json!(null),
900 findings: vec![],
901 tokens: 0,
902 completed_at: 0,
903 cache_key_hash: None,
904 description: None,
905 role: None,
906 };
907 store.upsert_agent_result(&cache).unwrap();
908 assert!(
909 !cp_path.exists(),
910 "upsert_agent_result must not create checkpoint.json before init_run"
911 );
912 assert!(store.get_checkpoint().is_none());
913 }
914
915 #[test]
916 fn upsert_agent_result_advances_updated_at() {
917 let dir = tempdir().unwrap();
920 let run_id = uuid::Uuid::now_v7();
921 let store = RunStore::new(dir.path()).unwrap();
922 store.init_run(run_id, "ts test").unwrap();
923 let before = store.get_checkpoint().unwrap().updated_at;
924
925 std::thread::sleep(std::time::Duration::from_millis(1100));
926
927 let cache = AgentResultCache {
928 agent_id: uuid::Uuid::now_v7(),
929 phase_id: 1,
930 status: "ok".into(),
931 output: serde_json::json!(null),
932 findings: vec![],
933 tokens: 0,
934 completed_at: 0,
935 cache_key_hash: None,
936 description: None,
937 role: None,
938 };
939 store.upsert_agent_result(&cache).unwrap();
940 let after = store.get_checkpoint().unwrap().updated_at;
941 assert!(
942 after > before,
943 "updated_at must advance after upsert (before={before}, after={after})"
944 );
945 }
946
947 #[test]
950 fn cancel_persists_cancelled_status_to_disk() {
951 let dir = tempdir().unwrap();
956 let run_id = uuid::Uuid::now_v7();
957 let store = RunStore::new(dir.path()).unwrap();
958 store.init_run(run_id, "cancel me").unwrap();
959 assert!(store.can_resume());
960
961 store.cancel().unwrap();
962
963 let cp = store.get_checkpoint().unwrap();
965 assert_eq!(cp.status, CheckpointStatus::Cancelled);
966 assert!(!store.can_resume());
967
968 let raw = read_raw_checkpoint(dir.path());
970 assert_eq!(
971 raw.get("status").and_then(|v| v.as_str()),
972 Some("cancelled")
973 );
974
975 drop(store);
977 let reopened = RunStore::new(dir.path()).unwrap();
978 let restored = reopened.open_run(run_id).unwrap().unwrap();
979 assert_eq!(restored.status, CheckpointStatus::Cancelled);
980 assert!(!reopened.can_resume());
981 }
982
983 #[test]
984 fn cancel_is_idempotent() {
985 let dir = tempdir().unwrap();
990 let run_id = uuid::Uuid::now_v7();
991 let store = RunStore::new(dir.path()).unwrap();
992 store.init_run(run_id, "double cancel").unwrap();
993
994 store.cancel().unwrap();
995 let after_first = store.get_checkpoint().unwrap().updated_at;
996 std::thread::sleep(std::time::Duration::from_millis(1100));
997
998 store.cancel().expect("second cancel must succeed");
999 let after_second = store.get_checkpoint().unwrap().updated_at;
1000
1001 assert_eq!(
1002 store.get_checkpoint().unwrap().status,
1003 CheckpointStatus::Cancelled
1004 );
1005 assert!(
1006 after_second >= after_first,
1007 "updated_at must not regress (was {after_first}, now {after_second})"
1008 );
1009
1010 let raw = read_raw_checkpoint(dir.path());
1011 assert_eq!(
1012 raw.get("status").and_then(|v| v.as_str()),
1013 Some("cancelled")
1014 );
1015 }
1016
1017 #[test]
1018 fn cancel_before_init_is_safe_noop() {
1019 let dir = tempdir().unwrap();
1023 let store = RunStore::new(dir.path()).unwrap();
1024 assert!(store.get_checkpoint().is_none());
1025 store.cancel().expect("cancel before init must succeed");
1026 assert!(store.get_checkpoint().is_none());
1027 assert!(
1028 !dir.path().join("checkpoint.json").exists(),
1029 "cancel before init must not create checkpoint.json"
1030 );
1031 }
1032
1033 #[test]
1034 fn cancel_loads_checkpoint_from_disk_when_cache_is_empty() {
1035 let dir = tempdir().unwrap();
1041 let run_id = uuid::Uuid::now_v7();
1042
1043 let cp = RunCheckpoint {
1046 run_id,
1047 task: "started elsewhere".into(),
1048 status: CheckpointStatus::Running,
1049 current_phase: 0,
1050 completed_phases: vec![],
1051 agent_results: HashMap::new(),
1052 agent_sessions: HashMap::new(),
1053 findings: vec![],
1054 total_tokens: 0,
1055 created_at: current_timestamp(),
1056 updated_at: current_timestamp(),
1057 workflow_meta: None,
1058 started_agent_ids: vec![],
1059 };
1060 let cp_path = dir.path().join("checkpoint.json");
1061 std::fs::write(&cp_path, serde_json::to_string_pretty(&cp).unwrap()).unwrap();
1062
1063 let store = RunStore::new(dir.path()).unwrap();
1065 assert!(store.get_checkpoint().is_none(), "cache must start empty");
1066
1067 store.cancel().expect("cancel must succeed cross-process");
1068
1069 let cached = store
1071 .get_checkpoint()
1072 .expect("cache populated after cancel");
1073 assert_eq!(cached.status, CheckpointStatus::Cancelled);
1074
1075 let raw = read_raw_checkpoint(dir.path());
1077 assert_eq!(
1078 raw.get("status").and_then(|v| v.as_str()),
1079 Some("cancelled")
1080 );
1081 assert_eq!(
1082 raw.get("run_id").and_then(|v| v.as_str()),
1083 Some(run_id.to_string()).as_deref()
1084 );
1085 }
1086
1087 #[test]
1088 fn cancel_preserves_agent_results_and_findings() {
1089 let dir = tempdir().unwrap();
1093 let run_id = uuid::Uuid::now_v7();
1094 let store = RunStore::new(dir.path()).unwrap();
1095 store.init_run(run_id, "preserve").unwrap();
1096
1097 let agent_id = uuid::Uuid::now_v7();
1098 let cache = AgentResultCache {
1099 agent_id,
1100 phase_id: 1,
1101 status: "ok".into(),
1102 output: serde_json::json!({"x": 1}),
1103 findings: vec![],
1104 tokens: 250,
1105 completed_at: 7,
1106 cache_key_hash: Some("hash-1".into()),
1107 description: None,
1108 role: None,
1109 };
1110 store.upsert_agent_result(&cache).unwrap();
1111 let before = store.get_checkpoint().unwrap();
1112
1113 store.cancel().unwrap();
1114 let after = store.get_checkpoint().unwrap();
1115
1116 assert_eq!(after.status, CheckpointStatus::Cancelled);
1117 assert_eq!(after.agent_results.len(), 1);
1118 assert_eq!(after.agent_results[&agent_id].tokens, 250);
1119 assert_eq!(
1120 after.agent_results[&agent_id].cache_key_hash.as_deref(),
1121 Some("hash-1")
1122 );
1123 assert_eq!(after.total_tokens, before.total_tokens);
1124 }
1125
1126 #[test]
1127 fn late_event_from_another_process_cannot_resurrect_cancelled_run() {
1128 let dir = tempdir().unwrap();
1129 let run_id = uuid::Uuid::now_v7();
1130 let writer = RunStore::new(dir.path()).unwrap();
1131 writer.init_run(run_id, "cross process race").unwrap();
1132
1133 let canceller = RunStore::new(dir.path()).unwrap();
1136 canceller.cancel().unwrap();
1137
1138 writer
1139 .append_event(&AgentEvent::RunDone {
1140 run_id,
1141 status: crate::contract::event::RunStatus::Completed,
1142 total_tokens: TokenUsage::default(),
1143 report: serde_json::Value::Null,
1144 ts: chrono::Utc::now(),
1145 })
1146 .unwrap();
1147
1148 let restored = RunStore::new(dir.path())
1149 .unwrap()
1150 .open_run(run_id)
1151 .unwrap()
1152 .unwrap();
1153 assert_eq!(restored.status, CheckpointStatus::Cancelled);
1154 }
1155
1156 #[test]
1159 fn agent_done_persists_snake_case_status_for_each_variant() {
1160 let dir = tempdir().unwrap();
1165 let run_id = uuid::Uuid::now_v7();
1166 let store = RunStore::new(dir.path()).unwrap();
1167 store.init_run(run_id, "F5 variants").unwrap();
1168
1169 let cases: Vec<(AgentStatus, &str)> = vec![
1170 (AgentStatus::Ok, "ok"),
1171 (AgentStatus::Error, "error"),
1172 (AgentStatus::Cancelled, "cancelled"),
1173 (AgentStatus::TimedOut, "timed_out"),
1174 ];
1175 for (status, expected) in &cases {
1176 let agent_id = uuid::Uuid::now_v7();
1177 let evt = build_agent_done(run_id, agent_id, status.clone(), sample_token_usage());
1178 store.append_event(&evt).unwrap();
1179
1180 let raw = read_raw_checkpoint(dir.path());
1181 let ar = raw
1182 .get("agent_results")
1183 .and_then(|v| v.as_object())
1184 .expect("agent_results object");
1185 let entry = ar
1186 .values()
1187 .find(|v| {
1188 v.get("agent_id").and_then(|id| id.as_str()) == Some(&agent_id.to_string())
1189 })
1190 .unwrap_or_else(|| panic!("entry for {agent_id} missing"));
1191 let persisted = entry
1192 .get("status")
1193 .and_then(|v| v.as_str())
1194 .unwrap_or_else(|| panic!("status missing for {status:?}"));
1195 assert_eq!(
1196 persisted, *expected,
1197 "AgentDone({status:?}) must persist status={expected:?} (snake_case); \
1198 got {persisted:?}. If this fails with \"timedout\" for TimedOut, \
1199 F5 has regressed to Debug formatting."
1200 );
1201 }
1202 }
1203
1204 #[test]
1205 fn agent_done_timed_out_persists_with_underscore_not_collapsed() {
1206 let dir = tempdir().unwrap();
1211 let run_id = uuid::Uuid::now_v7();
1212 let store = RunStore::new(dir.path()).unwrap();
1213 store.init_run(run_id, "timed-out guard").unwrap();
1214
1215 let agent_id = uuid::Uuid::now_v7();
1216 let evt = build_agent_done(
1217 run_id,
1218 agent_id,
1219 AgentStatus::TimedOut,
1220 sample_token_usage(),
1221 );
1222 store.append_event(&evt).unwrap();
1223
1224 let raw = read_raw_checkpoint(dir.path());
1225 let ar = raw
1226 .get("agent_results")
1227 .and_then(|v| v.as_object())
1228 .unwrap();
1229 let entry = ar.values().next().expect("entry exists");
1230 let persisted = entry.get("status").and_then(|v| v.as_str()).unwrap();
1231
1232 assert_eq!(
1233 persisted, "timed_out",
1234 "AgentDone(TimedOut) must persist \"timed_out\" with an underscore; got {persisted:?}"
1235 );
1236 assert_ne!(
1237 persisted, "timedout",
1238 "AgentDone(TimedOut) must NOT collapse to Debug-lowercased \"timedout\""
1239 );
1240 }
1241
1242 #[test]
1243 fn agent_done_then_reopen_restores_snake_case_status() {
1244 let dir = tempdir().unwrap();
1249 let run_id = uuid::Uuid::now_v7();
1250 let store = RunStore::new(dir.path()).unwrap();
1251 store.init_run(run_id, "round-trip").unwrap();
1252
1253 let agent_id = uuid::Uuid::now_v7();
1254 let evt = build_agent_done(
1255 run_id,
1256 agent_id,
1257 AgentStatus::Cancelled,
1258 TokenUsage {
1259 input: 1,
1260 output: 2,
1261 cache_read: 0,
1262 cache_write: 0,
1263 },
1264 );
1265 store.append_event(&evt).unwrap();
1266 drop(store);
1267
1268 let reopened = RunStore::new(dir.path()).unwrap();
1269 let cp = reopened.open_run(run_id).unwrap().unwrap();
1270 let cached = cp
1271 .agent_results
1272 .get(&agent_id)
1273 .expect("agent cached on disk");
1274 assert_eq!(cached.status, "cancelled");
1275 assert_eq!(cached.tokens, 3);
1276 }
1277
1278 #[test]
1281 fn open_run_with_corrupt_checkpoint_returns_invalid_data() {
1282 let dir = tempdir().unwrap();
1286 std::fs::create_dir_all(dir.path()).unwrap();
1287 std::fs::write(
1288 dir.path().join("checkpoint.json"),
1289 b"{ this is not valid json",
1290 )
1291 .unwrap();
1292
1293 let store = RunStore::new(dir.path()).unwrap();
1294 let err = store
1295 .open_run(uuid::Uuid::now_v7())
1296 .expect_err("corrupt JSON must surface as an io::Error");
1297 assert_eq!(
1298 err.kind(),
1299 std::io::ErrorKind::InvalidData,
1300 "corrupt checkpoint must map to InvalidData via serde_to_io; got {:?}",
1301 err.kind()
1302 );
1303 }
1304
1305 #[test]
1306 fn open_run_with_wrong_typed_checkpoint_returns_invalid_data() {
1307 let dir = tempdir().unwrap();
1310 std::fs::create_dir_all(dir.path()).unwrap();
1311 std::fs::write(
1314 dir.path().join("checkpoint.json"),
1315 br#"{"run_id":"00000000-0000-0000-0000-000000000000","status":"running"}"#,
1316 )
1317 .unwrap();
1318
1319 let store = RunStore::new(dir.path()).unwrap();
1320 let err = store
1321 .open_run(uuid::Uuid::now_v7())
1322 .expect_err("missing-field JSON must surface as an io::Error");
1323 assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
1324 }
1325
1326 #[test]
1327 fn get_event_log_with_corrupt_line_returns_invalid_data() {
1328 let dir = tempdir().unwrap();
1333 std::fs::create_dir_all(dir.path()).unwrap();
1334 std::fs::write(dir.path().join("events.jsonl"), b"not-json\n").unwrap();
1335
1336 let store = RunStore::new(dir.path()).unwrap();
1337 let err = store
1338 .get_event_log()
1339 .expect_err("corrupt event line must surface as an io::Error");
1340 assert_eq!(
1341 err.kind(),
1342 std::io::ErrorKind::InvalidData,
1343 "corrupt event line must map to InvalidData via serde_to_io; got {:?}",
1344 err.kind()
1345 );
1346 }
1347
1348 #[test]
1351 fn as_str_variants_round_trip_through_checkpoint_pipeline() {
1352 let dir = tempdir().unwrap();
1357 let run_id = uuid::Uuid::now_v7();
1358 let store = RunStore::new(dir.path()).unwrap();
1359 store.init_run(run_id, "round-trip property").unwrap();
1360
1361 let variants = [
1362 AgentStatus::Ok,
1363 AgentStatus::Error,
1364 AgentStatus::Cancelled,
1365 AgentStatus::TimedOut,
1366 ];
1367 let mut seen: HashSet<String> = HashSet::new();
1368
1369 for variant in &variants {
1370 let agent_id = uuid::Uuid::now_v7();
1371 let evt = build_agent_done(run_id, agent_id, variant.clone(), sample_token_usage());
1372 store.append_event(&evt).unwrap();
1373
1374 let cp = store.get_checkpoint().unwrap();
1375 let cached = cp
1376 .agent_results
1377 .get(&agent_id)
1378 .expect("entry for {agent_id}");
1379 assert_eq!(
1380 cached.status,
1381 variant.as_str(),
1382 "{variant:?}.as_str() must round-trip via append_event→update_from_event"
1383 );
1384 assert!(
1386 seen.insert(cached.status.clone()),
1387 "duplicate status {cached_status:?} persisted for {variant:?}",
1388 cached_status = cached.status
1389 );
1390 }
1391 }
1392}