ri_agent_graph/
checkpoint.rs

1use crate::error::Result;
2use crate::state::StateSnapshot;
3use chrono::{DateTime, Utc};
4use rusqlite::{params, Connection, OptionalExtension};
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Checkpoint {
9    pub execution_id: String,
10    pub timestamp: DateTime<Utc>,
11    pub current_node: String,
12    pub iteration: usize,
13    pub state: StateSnapshot,
14    #[serde(default)]
15    pub step_number: usize,
16    #[serde(default)]
17    pub active_nodes: Vec<String>,
18}
19
20pub struct CheckpointManager {
21    conn: Connection,
22}
23
24impl CheckpointManager {
25    /// Create a new checkpoint manager
26    pub fn new(db_path: &str) -> Result<Self> {
27        let conn = Connection::open(db_path)?;
28
29        conn.execute(
30            "CREATE TABLE IF NOT EXISTS checkpoints (
31                execution_id TEXT NOT NULL,
32                timestamp TEXT NOT NULL,
33                current_node TEXT NOT NULL,
34                iteration INTEGER NOT NULL,
35                state_data TEXT NOT NULL,
36                PRIMARY KEY (execution_id, timestamp)
37            )",
38            [],
39        )?;
40
41        Ok(Self { conn })
42    }
43
44    /// Save a checkpoint
45    pub fn save(&self, checkpoint: &Checkpoint) -> Result<()> {
46        let state_json = serde_json::to_string(&checkpoint.state)?;
47
48        self.conn.execute(
49            "INSERT OR REPLACE INTO checkpoints (execution_id, timestamp, current_node, iteration, state_data)
50             VALUES (?1, ?2, ?3, ?4, ?5)",
51            params![
52                &checkpoint.execution_id,
53                checkpoint.timestamp.to_rfc3339(),
54                &checkpoint.current_node,
55                checkpoint.iteration as i64,
56                &state_json,
57            ],
58        )?;
59
60        Ok(())
61    }
62
63    /// Load the most recent checkpoint for an execution
64    pub fn load(&self, execution_id: &str) -> Result<Option<Checkpoint>> {
65        let mut stmt = self.conn.prepare(
66            "SELECT timestamp, current_node, iteration, state_data
67             FROM checkpoints
68             WHERE execution_id = ?1
69             ORDER BY timestamp DESC
70             LIMIT 1",
71        )?;
72
73        let checkpoint = stmt
74            .query_row(params![execution_id], |row| {
75                let timestamp_str: String = row.get(0)?;
76                let state_json: String = row.get(3)?;
77
78                let timestamp = DateTime::parse_from_rfc3339(&timestamp_str)
79                    .map_err(|e| {
80                        rusqlite::Error::FromSqlConversionFailure(
81                            0,
82                            rusqlite::types::Type::Text,
83                            Box::new(e),
84                        )
85                    })?
86                    .with_timezone(&Utc);
87                let state: StateSnapshot = serde_json::from_str(&state_json).map_err(|e| {
88                    rusqlite::Error::FromSqlConversionFailure(
89                        3,
90                        rusqlite::types::Type::Text,
91                        Box::new(e),
92                    )
93                })?;
94
95                Ok(Checkpoint {
96                    execution_id: execution_id.to_string(),
97                    timestamp,
98                    current_node: row.get(1)?,
99                    iteration: row.get::<_, i64>(2)? as usize,
100                    state,
101                    step_number: 0,
102                    active_nodes: Vec::new(),
103                })
104            })
105            .optional()?;
106
107        Ok(checkpoint)
108    }
109
110    /// Load all checkpoints for an execution (ordered by timestamp)
111    pub fn load_all(&self, execution_id: &str) -> Result<Vec<Checkpoint>> {
112        let mut stmt = self.conn.prepare(
113            "SELECT timestamp, current_node, iteration, state_data
114             FROM checkpoints
115             WHERE execution_id = ?1
116             ORDER BY timestamp ASC",
117        )?;
118
119        let checkpoints = stmt
120            .query_map(params![execution_id], |row| {
121                let timestamp_str: String = row.get(0)?;
122                let state_json: String = row.get(3)?;
123
124                let timestamp = DateTime::parse_from_rfc3339(&timestamp_str)
125                    .map_err(|e| {
126                        rusqlite::Error::FromSqlConversionFailure(
127                            0,
128                            rusqlite::types::Type::Text,
129                            Box::new(e),
130                        )
131                    })?
132                    .with_timezone(&Utc);
133                let state: StateSnapshot = serde_json::from_str(&state_json).map_err(|e| {
134                    rusqlite::Error::FromSqlConversionFailure(
135                        3,
136                        rusqlite::types::Type::Text,
137                        Box::new(e),
138                    )
139                })?;
140
141                Ok(Checkpoint {
142                    execution_id: execution_id.to_string(),
143                    timestamp,
144                    current_node: row.get(1)?,
145                    iteration: row.get::<_, i64>(2)? as usize,
146                    state,
147                    step_number: 0,
148                    active_nodes: Vec::new(),
149                })
150            })?
151            .collect::<std::result::Result<Vec<_>, _>>()?;
152
153        Ok(checkpoints)
154    }
155
156    /// Delete all checkpoints for an execution
157    pub fn clear(&self, execution_id: &str) -> Result<()> {
158        self.conn.execute(
159            "DELETE FROM checkpoints WHERE execution_id = ?1",
160            params![execution_id],
161        )?;
162        Ok(())
163    }
164}
ri_agent_graph/checkpoint.rs

ri_agent_graph/
checkpoint.rs
