ri_agent_graph/
checkpointer.rs

1use crate::checkpoint::Checkpoint;
2use crate::error::Result;
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8/// Trait for saving and loading checkpoints.
9#[async_trait]
10pub trait CheckpointSaver: Send + Sync {
11    /// Save a checkpoint
12    async fn save(&self, checkpoint: &Checkpoint) -> Result<()>;
13    /// Load the most recent checkpoint for a thread
14    async fn load(&self, thread_id: &str) -> Result<Option<Checkpoint>>;
15    /// Load all checkpoints for a thread (history)
16    async fn load_history(&self, thread_id: &str) -> Result<Vec<Checkpoint>>;
17    /// Clear all checkpoints for a thread
18    async fn clear(&self, thread_id: &str) -> Result<()>;
19}
20
21/// In-memory checkpoint storage (for testing and lightweight use).
22pub struct MemorySaver {
23    checkpoints: Arc<RwLock<HashMap<String, Vec<Checkpoint>>>>,
24}
25
26impl MemorySaver {
27    pub fn new() -> Self {
28        Self {
29            checkpoints: Arc::new(RwLock::new(HashMap::new())),
30        }
31    }
32}
33
34impl Default for MemorySaver {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40#[async_trait]
41impl CheckpointSaver for MemorySaver {
42    async fn save(&self, checkpoint: &Checkpoint) -> Result<()> {
43        let mut store = self.checkpoints.write().await;
44        store
45            .entry(checkpoint.execution_id.clone())
46            .or_default()
47            .push(checkpoint.clone());
48        Ok(())
49    }
50
51    async fn load(&self, thread_id: &str) -> Result<Option<Checkpoint>> {
52        let store = self.checkpoints.read().await;
53        Ok(store.get(thread_id).and_then(|v| v.last()).cloned())
54    }
55
56    async fn load_history(&self, thread_id: &str) -> Result<Vec<Checkpoint>> {
57        let store = self.checkpoints.read().await;
58        Ok(store.get(thread_id).cloned().unwrap_or_default())
59    }
60
61    async fn clear(&self, thread_id: &str) -> Result<()> {
62        let mut store = self.checkpoints.write().await;
63        store.remove(thread_id);
64        Ok(())
65    }
66}
67
68/// SQLite-based checkpoint storage (wraps existing CheckpointManager).
69#[cfg(feature = "checkpointing")]
70pub struct SqliteSaver {
71    manager: std::sync::Mutex<crate::checkpoint::CheckpointManager>,
72}
73
74#[cfg(feature = "checkpointing")]
75impl SqliteSaver {
76    pub fn new(db_path: &str) -> Result<Self> {
77        Ok(Self {
78            manager: std::sync::Mutex::new(crate::checkpoint::CheckpointManager::new(db_path)?),
79        })
80    }
81}
82
83#[cfg(feature = "checkpointing")]
84#[async_trait]
85impl CheckpointSaver for SqliteSaver {
86    async fn save(&self, checkpoint: &Checkpoint) -> Result<()> {
87        let mgr = self
88            .manager
89            .lock()
90            .map_err(|e| crate::error::AgentGraphError::CheckpointError(e.to_string()))?;
91        mgr.save(checkpoint)
92    }
93
94    async fn load(&self, thread_id: &str) -> Result<Option<Checkpoint>> {
95        let mgr = self
96            .manager
97            .lock()
98            .map_err(|e| crate::error::AgentGraphError::CheckpointError(e.to_string()))?;
99        mgr.load(thread_id)
100    }
101
102    async fn load_history(&self, thread_id: &str) -> Result<Vec<Checkpoint>> {
103        let mgr = self
104            .manager
105            .lock()
106            .map_err(|e| crate::error::AgentGraphError::CheckpointError(e.to_string()))?;
107        mgr.load_all(thread_id)
108    }
109
110    async fn clear(&self, thread_id: &str) -> Result<()> {
111        let mgr = self
112            .manager
113            .lock()
114            .map_err(|e| crate::error::AgentGraphError::CheckpointError(e.to_string()))?;
115        mgr.clear(thread_id)
116    }
117}
ri_agent_graph/checkpointer.rs

ri_agent_graph/
checkpointer.rs
