lc_langgraph/
checkpointer.rs

1// crates/lc-langgraph/src/checkpointer.rs
2//! Checkpointing for state persistence
3
4use crate::errors::{GraphError, GraphResult};
5use crate::state::StateSchema;
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use tokio::sync::Mutex;
10use uuid::Uuid;
11
12/// Checkpointer trait for state persistence
13#[async_trait]
14pub trait Checkpointer<S: StateSchema>: Send + Sync {
15    async fn save(&self, state: &S) -> GraphResult<String>;
16    async fn load(&self, checkpoint_id: &str) -> GraphResult<S>;
17    async fn list(&self) -> GraphResult<Vec<String>>;
18    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()>;
19}
20
21/// Checkpoint data structure
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(bound = "S: StateSchema")]
24pub struct CheckpointData<S: StateSchema> {
25    pub id: String,
26    pub state: S,
27    pub timestamp: i64,
28    pub metadata: HashMap<String, serde_json::Value>,
29}
30
31impl<S: StateSchema> CheckpointData<S> {
32    pub fn new(state: S) -> Self {
33        Self {
34            id: Uuid::new_v4().to_string(),
35            state,
36            timestamp: chrono::Utc::now().timestamp(),
37            metadata: HashMap::new(),
38        }
39    }
40}
41
42/// In-memory checkpointer for development
43pub struct MemoryCheckpointer<S: StateSchema> {
44    checkpoints: Mutex<HashMap<String, CheckpointData<S>>>,
45}
46
47impl<S: StateSchema> MemoryCheckpointer<S> {
48    pub fn new() -> Self {
49        Self {
50            checkpoints: Mutex::new(HashMap::new()),
51        }
52    }
53}
54
55impl<S: StateSchema> Default for MemoryCheckpointer<S> {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61#[async_trait]
62impl<S: StateSchema> Checkpointer<S> for MemoryCheckpointer<S> {
63    async fn save(&self, state: &S) -> GraphResult<String> {
64        let data = CheckpointData::new(state.clone());
65        let id = data.id.clone();
66        self.checkpoints.lock().await.insert(id.clone(), data);
67        Ok(id)
68    }
69
70    async fn load(&self, checkpoint_id: &str) -> GraphResult<S> {
71        self.checkpoints
72            .lock()
73            .await
74            .get(checkpoint_id)
75            .map(|d| d.state.clone())
76            .ok_or_else(|| {
77                GraphError::CheckpointError(format!("Checkpoint '{}' not found", checkpoint_id))
78            })
79    }
80
81    async fn list(&self) -> GraphResult<Vec<String>> {
82        Ok(self.checkpoints.lock().await.keys().cloned().collect())
83    }
84
85    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()> {
86        self.checkpoints.lock().await.remove(checkpoint_id);
87        Ok(())
88    }
89}
90
91/// Thread-safe memory checkpointer
92pub struct ThreadSafeMemoryCheckpointer<S: StateSchema> {
93    checkpoints: Mutex<HashMap<String, CheckpointData<S>>>,
94}
95
96impl<S: StateSchema> ThreadSafeMemoryCheckpointer<S> {
97    pub fn new() -> Self {
98        Self {
99            checkpoints: Mutex::new(HashMap::new()),
100        }
101    }
102}
103
104impl<S: StateSchema> Default for ThreadSafeMemoryCheckpointer<S> {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110#[async_trait]
111impl<S: StateSchema> Checkpointer<S> for ThreadSafeMemoryCheckpointer<S> {
112    async fn save(&self, state: &S) -> GraphResult<String> {
113        let data = CheckpointData::new(state.clone());
114        let id = data.id.clone();
115        self.checkpoints.lock().await.insert(id.clone(), data);
116        Ok(id)
117    }
118
119    async fn load(&self, checkpoint_id: &str) -> GraphResult<S> {
120        let checkpoints = self.checkpoints.lock().await;
121        checkpoints
122            .get(checkpoint_id)
123            .map(|d| d.state.clone())
124            .ok_or_else(|| {
125                GraphError::CheckpointError(format!("Checkpoint '{}' not found", checkpoint_id))
126            })
127    }
128
129    async fn list(&self) -> GraphResult<Vec<String>> {
130        Ok(self.checkpoints.lock().await.keys().cloned().collect())
131    }
132
133    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()> {
134        self.checkpoints.lock().await.remove(checkpoint_id);
135        Ok(())
136    }
137}
138
139/// File-based checkpointer for persistent storage
140pub struct FileCheckpointer<S: StateSchema> {
141    directory: std::path::PathBuf,
142    _phantom: std::marker::PhantomData<S>,
143}
144
145impl<S: StateSchema> FileCheckpointer<S> {
146    pub fn new(directory: impl Into<std::path::PathBuf>) -> GraphResult<Self> {
147        let dir = directory.into();
148        if !dir.exists() {
149            std::fs::create_dir_all(&dir).map_err(|e| {
150                GraphError::CheckpointError(format!(
151                    "Failed to create directory '{}': {}",
152                    dir.display(),
153                    e
154                ))
155            })?;
156        }
157        Ok(Self {
158            directory: dir,
159            _phantom: std::marker::PhantomData,
160        })
161    }
162
163    fn checkpoint_path(&self, id: &str) -> GraphResult<std::path::PathBuf> {
164        // Sanitize id to prevent path traversal: reject ".." and absolute paths
165        if id.contains("..") || id.contains('/') || id.contains('\\') {
166            return Err(GraphError::CheckpointError(format!(
167                "Invalid checkpoint id '{}': path traversal detected",
168                id
169            )));
170        }
171        if std::path::Path::new(id).is_absolute() {
172            return Err(GraphError::CheckpointError(format!(
173                "Invalid checkpoint id '{}': absolute path not allowed",
174                id
175            )));
176        }
177        Ok(self.directory.join(format!("{}.json", id)))
178    }
179}
180
181// NOTE: `Default` is intentionally NOT implemented for `FileCheckpointer` (Q1).
182// The default constructor would have to create the `.checkpoints` directory, which
183// is I/O that can fail (read-only cwd, disk full, permissions) — `Default` cannot
184// report that failure, so it would have to panic. Use `FileCheckpointer::new(...)`
185// which returns a `GraphResult` and surfaces the error instead.
186
187#[async_trait]
188impl<S: StateSchema> Checkpointer<S> for FileCheckpointer<S> {
189    async fn save(&self, state: &S) -> GraphResult<String> {
190        let data = CheckpointData::new(state.clone());
191        let id = data.id.clone();
192        let path = self.checkpoint_path(&id)?;
193
194        let json = serde_json::to_string_pretty(&data)
195            .map_err(|e| GraphError::CheckpointError(format!("Serialize error: {}", e)))?;
196
197        tokio::fs::write(&path, json)
198            .await
199            .map_err(|e| GraphError::CheckpointError(format!("Write error: {}", e)))?;
200
201        Ok(id)
202    }
203
204    async fn load(&self, checkpoint_id: &str) -> GraphResult<S> {
205        let path = self.checkpoint_path(checkpoint_id)?;
206
207        if !path.exists() {
208            return Err(GraphError::CheckpointError(format!(
209                "Checkpoint '{}' not found",
210                checkpoint_id
211            )));
212        }
213
214        let json = tokio::fs::read_to_string(&path)
215            .await
216            .map_err(|e| GraphError::CheckpointError(format!("Read error: {}", e)))?;
217
218        let data: CheckpointData<S> = serde_json::from_str(&json)
219            .map_err(|e| GraphError::CheckpointError(format!("Deserialize error: {}", e)))?;
220
221        Ok(data.state)
222    }
223
224    async fn list(&self) -> GraphResult<Vec<String>> {
225        let mut ids = Vec::new();
226
227        let entries = tokio::fs::read_dir(&self.directory)
228            .await
229            .map_err(|e| GraphError::CheckpointError(format!("Read dir error: {}", e)))?;
230
231        let mut entries = entries;
232        while let Some(entry) = entries
233            .next_entry()
234            .await
235            .map_err(|e| GraphError::CheckpointError(format!("Read dir entry error: {}", e)))?
236        {
237            let path = entry.path();
238            if path.extension().is_some_and(|ext| ext == "json") {
239                if let Some(id) = path.file_stem().and_then(|s| s.to_str()) {
240                    ids.push(id.to_string());
241                }
242            }
243        }
244
245        Ok(ids)
246    }
247
248    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()> {
249        let path = self.checkpoint_path(checkpoint_id)?;
250
251        if path.exists() {
252            tokio::fs::remove_file(&path)
253                .await
254                .map_err(|e| GraphError::CheckpointError(format!("Delete error: {}", e)))?;
255        }
256
257        Ok(())
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use crate::state::AgentState;
265
266    #[tokio::test]
267    async fn test_thread_safe_checkpointer() {
268        let checkpointer = ThreadSafeMemoryCheckpointer::<AgentState>::new();
269
270        let state = AgentState::new("test".to_string());
271        let id = checkpointer.save(&state).await.unwrap();
272
273        let loaded = checkpointer.load(&id).await.unwrap();
274        assert_eq!(loaded.input, "test");
275
276        let list = checkpointer.list().await.unwrap();
277        assert_eq!(list.len(), 1);
278
279        checkpointer.delete(&id).await.unwrap();
280        let list = checkpointer.list().await.unwrap();
281        assert!(list.is_empty());
282    }
283
284    #[tokio::test]
285    async fn test_file_checkpointer() {
286        let temp_dir = tempfile::tempdir().unwrap();
287        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
288
289        let state = AgentState::new("file_test".to_string());
290        let id = checkpointer.save(&state).await.unwrap();
291
292        let loaded = checkpointer.load(&id).await.unwrap();
293        assert_eq!(loaded.input, "file_test");
294
295        let list = checkpointer.list().await.unwrap();
296        assert_eq!(list.len(), 1);
297
298        checkpointer.delete(&id).await.unwrap();
299        let list = checkpointer.list().await.unwrap();
300        assert!(list.is_empty());
301    }
302
303    #[tokio::test]
304    async fn test_file_checkpointer_multiple() {
305        let temp_dir = tempfile::tempdir().unwrap();
306        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
307
308        let id1 = checkpointer
309            .save(&AgentState::new("state1".to_string()))
310            .await
311            .unwrap();
312        let id2 = checkpointer
313            .save(&AgentState::new("state2".to_string()))
314            .await
315            .unwrap();
316        let _id3 = checkpointer
317            .save(&AgentState::new("state3".to_string()))
318            .await
319            .unwrap();
320
321        let list = checkpointer.list().await.unwrap();
322        assert_eq!(list.len(), 3);
323
324        let loaded = checkpointer.load(&id2).await.unwrap();
325        assert_eq!(loaded.input, "state2");
326
327        checkpointer.delete(&id1).await.unwrap();
328        let list = checkpointer.list().await.unwrap();
329        assert_eq!(list.len(), 2);
330    }
331
332    #[tokio::test]
333    async fn test_file_checkpointer_path_traversal() {
334        let temp_dir = tempfile::tempdir().unwrap();
335        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
336
337        // Path traversal should be rejected
338        let result = checkpointer.load("..").await;
339        assert!(result.is_err());
340
341        let result = checkpointer.load("../etc/passwd").await;
342        assert!(result.is_err());
343    }
344}
lc_langgraph/checkpointer.rs

lc_langgraph/
checkpointer.rs
