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: std::sync::Mutex<HashMap<String, CheckpointData<S>>>,
94}
95
96impl<S: StateSchema> ThreadSafeMemoryCheckpointer<S> {
97    pub fn new() -> Self {
98        Self {
99            checkpoints: std::sync::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
116            .lock()
117            .unwrap_or_else(|e| e.into_inner())
118            .insert(id.clone(), data);
119        Ok(id)
120    }
121
122    async fn load(&self, checkpoint_id: &str) -> GraphResult<S> {
123        let checkpoints = self.checkpoints.lock().unwrap_or_else(|e| e.into_inner());
124        checkpoints
125            .get(checkpoint_id)
126            .map(|d| d.state.clone())
127            .ok_or_else(|| {
128                GraphError::CheckpointError(format!("Checkpoint '{}' not found", checkpoint_id))
129            })
130    }
131
132    async fn list(&self) -> GraphResult<Vec<String>> {
133        Ok(self
134            .checkpoints
135            .lock()
136            .unwrap_or_else(|e| e.into_inner())
137            .keys()
138            .cloned()
139            .collect())
140    }
141
142    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()> {
143        self.checkpoints
144            .lock()
145            .unwrap_or_else(|e| e.into_inner())
146            .remove(checkpoint_id);
147        Ok(())
148    }
149}
150
151/// File-based checkpointer for persistent storage
152pub struct FileCheckpointer<S: StateSchema> {
153    directory: std::path::PathBuf,
154    _phantom: std::marker::PhantomData<S>,
155}
156
157impl<S: StateSchema> FileCheckpointer<S> {
158    pub fn new(directory: impl Into<std::path::PathBuf>) -> GraphResult<Self> {
159        let dir = directory.into();
160        if !dir.exists() {
161            std::fs::create_dir_all(&dir).map_err(|e| {
162                GraphError::CheckpointError(format!(
163                    "Failed to create directory '{}': {}",
164                    dir.display(),
165                    e
166                ))
167            })?;
168        }
169        Ok(Self {
170            directory: dir,
171            _phantom: std::marker::PhantomData,
172        })
173    }
174
175    fn checkpoint_path(&self, id: &str) -> GraphResult<std::path::PathBuf> {
176        // Sanitize id to prevent path traversal: reject ".." and absolute paths
177        if id.contains("..") || id.contains('/') || id.contains('\\') {
178            return Err(GraphError::CheckpointError(format!(
179                "Invalid checkpoint id '{}': path traversal detected",
180                id
181            )));
182        }
183        if std::path::Path::new(id).is_absolute() {
184            return Err(GraphError::CheckpointError(format!(
185                "Invalid checkpoint id '{}': absolute path not allowed",
186                id
187            )));
188        }
189        Ok(self.directory.join(format!("{}.json", id)))
190    }
191}
192
193impl<S: StateSchema> Default for FileCheckpointer<S> {
194    fn default() -> Self {
195        Self::new(".checkpoints").expect("Failed to create default checkpoint directory")
196    }
197}
198
199#[async_trait]
200impl<S: StateSchema> Checkpointer<S> for FileCheckpointer<S> {
201    async fn save(&self, state: &S) -> GraphResult<String> {
202        let data = CheckpointData::new(state.clone());
203        let id = data.id.clone();
204        let path = self.checkpoint_path(&id)?;
205
206        let json = serde_json::to_string_pretty(&data)
207            .map_err(|e| GraphError::CheckpointError(format!("Serialize error: {}", e)))?;
208
209        tokio::fs::write(&path, json)
210            .await
211            .map_err(|e| GraphError::CheckpointError(format!("Write error: {}", e)))?;
212
213        Ok(id)
214    }
215
216    async fn load(&self, checkpoint_id: &str) -> GraphResult<S> {
217        let path = self.checkpoint_path(checkpoint_id)?;
218
219        if !path.exists() {
220            return Err(GraphError::CheckpointError(format!(
221                "Checkpoint '{}' not found",
222                checkpoint_id
223            )));
224        }
225
226        let json = tokio::fs::read_to_string(&path)
227            .await
228            .map_err(|e| GraphError::CheckpointError(format!("Read error: {}", e)))?;
229
230        let data: CheckpointData<S> = serde_json::from_str(&json)
231            .map_err(|e| GraphError::CheckpointError(format!("Deserialize error: {}", e)))?;
232
233        Ok(data.state)
234    }
235
236    async fn list(&self) -> GraphResult<Vec<String>> {
237        let mut ids = Vec::new();
238
239        let entries = tokio::fs::read_dir(&self.directory)
240            .await
241            .map_err(|e| GraphError::CheckpointError(format!("Read dir error: {}", e)))?;
242
243        let mut entries = entries;
244        while let Some(entry) = entries
245            .next_entry()
246            .await
247            .map_err(|e| GraphError::CheckpointError(format!("Read dir entry error: {}", e)))?
248        {
249            let path = entry.path();
250            if path.extension().is_some_and(|ext| ext == "json") {
251                if let Some(id) = path.file_stem().and_then(|s| s.to_str()) {
252                    ids.push(id.to_string());
253                }
254            }
255        }
256
257        Ok(ids)
258    }
259
260    async fn delete(&self, checkpoint_id: &str) -> GraphResult<()> {
261        let path = self.checkpoint_path(checkpoint_id)?;
262
263        if path.exists() {
264            tokio::fs::remove_file(&path)
265                .await
266                .map_err(|e| GraphError::CheckpointError(format!("Delete error: {}", e)))?;
267        }
268
269        Ok(())
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276    use crate::state::AgentState;
277
278    #[tokio::test]
279    async fn test_thread_safe_checkpointer() {
280        let checkpointer = ThreadSafeMemoryCheckpointer::<AgentState>::new();
281
282        let state = AgentState::new("test".to_string());
283        let id = checkpointer.save(&state).await.unwrap();
284
285        let loaded = checkpointer.load(&id).await.unwrap();
286        assert_eq!(loaded.input, "test");
287
288        let list = checkpointer.list().await.unwrap();
289        assert_eq!(list.len(), 1);
290
291        checkpointer.delete(&id).await.unwrap();
292        let list = checkpointer.list().await.unwrap();
293        assert!(list.is_empty());
294    }
295
296    #[tokio::test]
297    async fn test_file_checkpointer() {
298        let temp_dir = tempfile::tempdir().unwrap();
299        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
300
301        let state = AgentState::new("file_test".to_string());
302        let id = checkpointer.save(&state).await.unwrap();
303
304        let loaded = checkpointer.load(&id).await.unwrap();
305        assert_eq!(loaded.input, "file_test");
306
307        let list = checkpointer.list().await.unwrap();
308        assert_eq!(list.len(), 1);
309
310        checkpointer.delete(&id).await.unwrap();
311        let list = checkpointer.list().await.unwrap();
312        assert!(list.is_empty());
313    }
314
315    #[tokio::test]
316    async fn test_file_checkpointer_multiple() {
317        let temp_dir = tempfile::tempdir().unwrap();
318        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
319
320        let id1 = checkpointer
321            .save(&AgentState::new("state1".to_string()))
322            .await
323            .unwrap();
324        let id2 = checkpointer
325            .save(&AgentState::new("state2".to_string()))
326            .await
327            .unwrap();
328        let _id3 = checkpointer
329            .save(&AgentState::new("state3".to_string()))
330            .await
331            .unwrap();
332
333        let list = checkpointer.list().await.unwrap();
334        assert_eq!(list.len(), 3);
335
336        let loaded = checkpointer.load(&id2).await.unwrap();
337        assert_eq!(loaded.input, "state2");
338
339        checkpointer.delete(&id1).await.unwrap();
340        let list = checkpointer.list().await.unwrap();
341        assert_eq!(list.len(), 2);
342    }
343
344    #[tokio::test]
345    async fn test_file_checkpointer_path_traversal() {
346        let temp_dir = tempfile::tempdir().unwrap();
347        let checkpointer = FileCheckpointer::<AgentState>::new(temp_dir.path()).unwrap();
348
349        // Path traversal should be rejected
350        let result = checkpointer.load("..").await;
351        assert!(result.is_err());
352
353        let result = checkpointer.load("../etc/passwd").await;
354        assert!(result.is_err());
355    }
356}
lc_langgraph/checkpointer.rs

lc_langgraph/
checkpointer.rs
