luft_core/
in_memory_backend.rs

1//! In-memory `CheckpointBackend` for tests and lightweight usage.
2//!
3//! No file I/O, no SQLite — just `RwLock` state in memory.
4
5use crate::contract::event::AgentEvent;
6use crate::contract::finding::Finding;
7use crate::contract::ids::RunId;
8use crate::state::{
9    AgentResultCache, AgentSessionCheckpoint, CheckpointBackend, CheckpointStatus, RunCheckpoint,
10};
11use std::collections::HashMap;
12use std::sync::RwLock;
13
14#[derive(Debug, Default)]
15pub struct InMemoryBackend {
16    checkpoint: RwLock<Option<RunCheckpoint>>,
17    events: RwLock<Vec<AgentEvent>>,
18}
19
20impl InMemoryBackend {
21    pub fn new() -> Self {
22        Self::default()
23    }
24}
25
26impl CheckpointBackend for InMemoryBackend {
27    fn init_run(&self, run_id: RunId, task: &str, _run_dir: &str) -> anyhow::Result<()> {
28        let now = crate::state::current_timestamp();
29        let cp = RunCheckpoint {
30            run_id,
31            task: task.to_string(),
32            status: CheckpointStatus::Running,
33            current_phase: 0,
34            completed_phases: vec![],
35            agent_results: HashMap::new(),
36            agent_sessions: HashMap::new(),
37            findings: vec![],
38            total_tokens: 0,
39            created_at: now,
40            updated_at: now,
41            workflow_meta: None,
42            started_agent_ids: vec![],
43        };
44        *self.checkpoint.write().unwrap() = Some(cp);
45        Ok(())
46    }
47
48    fn init_run_with_meta(
49        &self,
50        run_id: RunId,
51        task: &str,
52        _run_dir: &str,
53        workflow_meta: serde_json::Value,
54    ) -> anyhow::Result<()> {
55        self.init_run(run_id, task, "")?;
56        if let Some(ref mut cp) = *self.checkpoint.write().unwrap() {
57            cp.workflow_meta = Some(workflow_meta);
58        }
59        Ok(())
60    }
61
62    fn open_run(&self, run_id: RunId) -> anyhow::Result<Option<RunCheckpoint>> {
63        let cp = self.checkpoint.read().unwrap().clone();
64        Ok(cp.filter(|c| c.run_id == run_id))
65    }
66
67    fn append_event(&self, event: &AgentEvent) -> anyhow::Result<()> {
68        self.events.write().unwrap().push(event.clone());
69        Ok(())
70    }
71
72    fn upsert_agent_result(&self, cache: &AgentResultCache) -> anyhow::Result<()> {
73        if let Some(ref mut cp) = *self.checkpoint.write().unwrap() {
74            cp.agent_results.insert(cache.agent_id, cache.clone());
75        }
76        Ok(())
77    }
78
79    fn upsert_agent_session(&self, session: &AgentSessionCheckpoint) -> anyhow::Result<()> {
80        if let Some(ref mut cp) = *self.checkpoint.write().unwrap() {
81            cp.agent_sessions.insert(session.agent_id, session.clone());
82        }
83        Ok(())
84    }
85
86    fn get_checkpoint(&self) -> Option<RunCheckpoint> {
87        self.checkpoint.read().unwrap().clone()
88    }
89
90    fn get_findings(&self) -> Vec<Finding> {
91        self.checkpoint
92            .read()
93            .unwrap()
94            .as_ref()
95            .map(|cp| cp.findings.clone())
96            .unwrap_or_default()
97    }
98
99    fn get_event_log(&self) -> anyhow::Result<Vec<AgentEvent>> {
100        Ok(self.events.read().unwrap().clone())
101    }
102
103    fn can_resume(&self) -> bool {
104        self.checkpoint
105            .read()
106            .unwrap()
107            .as_ref()
108            .map(|cp| {
109                matches!(
110                    cp.status,
111                    CheckpointStatus::Running
112                        | CheckpointStatus::Failed
113                        | CheckpointStatus::Cancelled
114                )
115            })
116            .unwrap_or(false)
117    }
118
119    fn reset_status_to_running(&self) -> anyhow::Result<()> {
120        if let Some(ref mut cp) = *self.checkpoint.write().unwrap() {
121            cp.status = CheckpointStatus::Running;
122        }
123        Ok(())
124    }
125
126    fn cancel(&self) -> anyhow::Result<()> {
127        if let Some(ref mut cp) = *self.checkpoint.write().unwrap() {
128            cp.status = CheckpointStatus::Cancelled;
129        }
130        Ok(())
131    }
132
133    fn save_checkpoint(&self, checkpoint: &RunCheckpoint) -> anyhow::Result<()> {
134        *self.checkpoint.write().unwrap() = Some(checkpoint.clone());
135        Ok(())
136    }
137}
luft_core/in_memory_backend.rs

luft_core/
in_memory_backend.rs
