oino_core/session/
checkpoints.rs

1//! File checkpoints for per-turn undo and rewind (`_oino/checkpoints`, `_oino/rewind`; Oino
2//! addition, runtime §13).
3//!
4//! Before a built-in `write` or `edit` call changes a file, the session records the file's
5//! previous content (or its absence) under the checkpoint of the user message that started
6//! the run, once per file and checkpoint. Rewinding to a user message restores every file
7//! touched at or after that checkpoint to the state it had before it and, optionally, forks
8//! the conversation before that message. Shell commands and extension tools are
9//! not tracked.
10//!
11//! Persisted sessions keep the checkpoints in `<agentDir>/checkpoints/<sessionId>/`
12//! (`index.json` plus one blob per snapshot); in-memory sessions keep them in memory.
13
14use std::collections::BTreeMap;
15use std::path::{Path, PathBuf};
16
17use serde::{Deserialize, Serialize};
18
19/// Most checkpoints kept per session (oldest dropped first).
20pub const MAX_CHECKPOINTS: usize = 100;
21/// Files larger than this are not snapshotted.
22pub const MAX_SNAPSHOT_BYTES: u64 = 16 * 1024 * 1024;
23
24/// One file of a checkpoint as frontends see it.
25#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "camelCase")]
27pub struct CheckpointFile {
28    /// Absolute path.
29    pub path: String,
30    /// The file existed before the change (`false`: rewind deletes it).
31    pub existed: bool,
32}
33
34/// One checkpoint (`_oino/checkpoints`).
35#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
36#[serde(rename_all = "camelCase")]
37pub struct Checkpoint {
38    /// User message entry that started the run.
39    pub entry_id: String,
40    /// Creation time (epoch ms).
41    pub created_at: i64,
42    /// Files changed after this checkpoint started, in first-change order.
43    pub files: Vec<CheckpointFile>,
44}
45
46/// Outcome of `_oino/rewind`.
47#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
48#[serde(rename_all = "camelCase")]
49pub struct RewindResult {
50    /// Files written back.
51    pub restored: Vec<String>,
52    /// Files deleted (created after the checkpoint).
53    pub deleted: Vec<String>,
54    /// Files that could not be restored, with the error.
55    #[serde(default, skip_serializing_if = "Vec::is_empty")]
56    pub failed: Vec<RewindFailure>,
57    /// Text of the user message (composer prefill) when the conversation was forked.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub text: Option<String>,
60    /// A `session_before_fork` handler cancelled the conversation fork.
61    pub cancelled: bool,
62}
63
64/// A file `_oino/rewind` could not restore.
65#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
66pub struct RewindFailure {
67    /// Path.
68    pub path: String,
69    /// Error text.
70    pub error: String,
71}
72
73/// Where a snapshot's bytes live.
74#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
75#[serde(rename_all = "camelCase", tag = "kind")]
76enum Blob {
77    /// The file did not exist.
78    Absent,
79    /// Blob file name under `blobs/`.
80    File {
81        /// Name.
82        name: String,
83    },
84    /// Kept in memory (sessions without a file).
85    #[serde(skip)]
86    Memory(Vec<u8>),
87}
88
89#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
90#[serde(rename_all = "camelCase")]
91struct Stored {
92    entry_id: String,
93    created_at: i64,
94    files: Vec<(String, Blob)>,
95}
96
97/// The checkpoints of one session.
98#[derive(Debug, Default)]
99pub(crate) struct CheckpointStore {
100    dir: Option<PathBuf>,
101    loaded: bool,
102    list: Vec<Stored>,
103    next_blob: u64,
104}
105
106fn now_ms() -> i64 {
107    std::time::SystemTime::now()
108        .duration_since(std::time::UNIX_EPOCH)
109        .ok()
110        .and_then(|d| i64::try_from(d.as_millis()).ok())
111        .unwrap_or_default()
112}
113
114impl CheckpointStore {
115    /// Points the store at `dir` (persisted sessions) and loads its index once.
116    pub(crate) fn bind(&mut self, dir: Option<PathBuf>) {
117        if self.loaded && self.dir == dir {
118            return;
119        }
120        self.dir = dir;
121        self.loaded = true;
122        self.list.clear();
123        let Some(dir) = &self.dir else {
124            return;
125        };
126        let index = std::fs::read(dir.join("index.json"))
127            .ok()
128            .and_then(|bytes| serde_json::from_slice::<Vec<Stored>>(&bytes).ok());
129        if let Some(index) = index {
130            self.next_blob = index
131                .iter()
132                .flat_map(|c| &c.files)
133                .filter_map(|(_, blob)| match blob {
134                    Blob::File { name } => name.parse::<u64>().ok(),
135                    _ => None,
136                })
137                .max()
138                .map_or(0, |n| n + 1);
139            self.list = index;
140        }
141    }
142
143    fn persist(&self) {
144        let Some(dir) = &self.dir else {
145            return;
146        };
147        if let Ok(bytes) = serde_json::to_vec(&self.list) {
148            let _ = std::fs::create_dir_all(dir);
149            let tmp = dir.join("index.json.tmp");
150            if std::fs::write(&tmp, bytes).is_ok() {
151                let _ = std::fs::rename(&tmp, dir.join("index.json"));
152            }
153        }
154    }
155
156    /// Records `path`'s current content under the checkpoint of `entry_id` unless that
157    /// checkpoint already has it. Returns whether a snapshot was taken.
158    pub(crate) fn record(&mut self, entry_id: &str, path: &Path) -> bool {
159        let key = path.display().to_string();
160        let index = if let Some(index) = self.list.iter().position(|c| c.entry_id == entry_id) {
161            index
162        } else {
163            self.list.push(Stored {
164                entry_id: entry_id.to_owned(),
165                created_at: now_ms(),
166                files: Vec::new(),
167            });
168            self.list.len() - 1
169        };
170        if self.list[index].files.iter().any(|(p, _)| *p == key) {
171            return false;
172        }
173        let blob = match std::fs::metadata(path) {
174            Err(_) => Blob::Absent,
175            Ok(meta) if !meta.is_file() || meta.len() > MAX_SNAPSHOT_BYTES => return false,
176            Ok(_) => match std::fs::read(path) {
177                Err(_) => return false,
178                Ok(bytes) => self.store_blob(bytes),
179            },
180        };
181        self.list[index].files.push((key, blob));
182        self.evict();
183        self.persist();
184        true
185    }
186
187    fn store_blob(&mut self, bytes: Vec<u8>) -> Blob {
188        let Some(dir) = &self.dir else {
189            return Blob::Memory(bytes);
190        };
191        let name = self.next_blob.to_string();
192        self.next_blob += 1;
193        let blobs = dir.join("blobs");
194        let _ = std::fs::create_dir_all(&blobs);
195        match std::fs::write(blobs.join(&name), &bytes) {
196            Ok(()) => Blob::File { name },
197            Err(_) => Blob::Memory(bytes),
198        }
199    }
200
201    fn evict(&mut self) {
202        while self.list.len() > MAX_CHECKPOINTS {
203            let old = self.list.remove(0);
204            self.remove_blobs(&old);
205        }
206    }
207
208    fn remove_blobs(&self, stored: &Stored) {
209        let Some(dir) = &self.dir else {
210            return;
211        };
212        for (_, blob) in &stored.files {
213            if let Blob::File { name } = blob {
214                let _ = std::fs::remove_file(dir.join("blobs").join(name));
215            }
216        }
217    }
218
219    /// Every checkpoint, oldest first.
220    pub(crate) fn list(&self) -> Vec<Checkpoint> {
221        self.list
222            .iter()
223            .map(|c| Checkpoint {
224                entry_id: c.entry_id.clone(),
225                created_at: c.created_at,
226                files: c
227                    .files
228                    .iter()
229                    .map(|(path, blob)| CheckpointFile {
230                        path: path.clone(),
231                        existed: *blob != Blob::Absent,
232                    })
233                    .collect(),
234            })
235            .collect()
236    }
237
238    /// Whether a checkpoint for `entry_id` exists.
239    pub(crate) fn contains(&self, entry_id: &str) -> bool {
240        self.list.iter().any(|c| c.entry_id == entry_id)
241    }
242
243    /// Restores every file changed at or after the checkpoint of `entry_id` and drops those
244    /// checkpoints. `None` when there is no such checkpoint.
245    pub(crate) fn rewind(&mut self, entry_id: &str) -> Option<RewindResult> {
246        let start = self.list.iter().position(|c| c.entry_id == entry_id)?;
247        // The earliest snapshot of each path at or after `start` holds its state before.
248        let mut earliest: BTreeMap<String, Blob> = BTreeMap::new();
249        for stored in &self.list[start..] {
250            for (path, blob) in &stored.files {
251                earliest.entry(path.clone()).or_insert_with(|| blob.clone());
252            }
253        }
254        let mut result = RewindResult::default();
255        for (path, blob) in &earliest {
256            let outcome = match blob {
257                Blob::Absent => match std::fs::remove_file(path) {
258                    Ok(()) => {
259                        result.deleted.push(path.clone());
260                        Ok(())
261                    }
262                    Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
263                    Err(e) => Err(e.to_string()),
264                },
265                Blob::Memory(bytes) => write_back(path, bytes).map(|()| {
266                    result.restored.push(path.clone());
267                }),
268                Blob::File { name } => self
269                    .dir
270                    .as_ref()
271                    .ok_or_else(|| "checkpoint blob missing".to_owned())
272                    .and_then(|dir| {
273                        std::fs::read(dir.join("blobs").join(name)).map_err(|e| e.to_string())
274                    })
275                    .and_then(|bytes| write_back(path, &bytes))
276                    .map(|()| result.restored.push(path.clone())),
277            };
278            if let Err(error) = outcome {
279                result.failed.push(RewindFailure {
280                    path: path.clone(),
281                    error,
282                });
283            }
284        }
285        let dropped: Vec<Stored> = self.list.drain(start..).collect();
286        for stored in &dropped {
287            self.remove_blobs(stored);
288        }
289        self.persist();
290        Some(result)
291    }
292
293    /// Copies the checkpoints whose entry `keep` accepts into `target` (the session a fork
294    /// created keeps the file history of the entries it kept).
295    pub(crate) fn copy_matching(&self, keep: impl Fn(&str) -> bool, target: &mut Self) {
296        for stored in self.list.iter().filter(|c| keep(&c.entry_id)) {
297            let mut copy = Stored {
298                entry_id: stored.entry_id.clone(),
299                created_at: stored.created_at,
300                files: Vec::new(),
301            };
302            for (path, blob) in &stored.files {
303                let blob = match blob {
304                    Blob::Absent => Blob::Absent,
305                    Blob::Memory(bytes) => target.store_blob(bytes.clone()),
306                    Blob::File { name } => {
307                        let bytes = self
308                            .dir
309                            .as_ref()
310                            .and_then(|dir| std::fs::read(dir.join("blobs").join(name)).ok());
311                        match bytes {
312                            Some(bytes) => target.store_blob(bytes),
313                            None => continue,
314                        }
315                    }
316                };
317                copy.files.push((path.clone(), blob));
318            }
319            target.list.push(copy);
320        }
321        target.persist();
322    }
323}
324
325fn write_back(path: &str, bytes: &[u8]) -> Result<(), String> {
326    let path = Path::new(path);
327    if let Some(parent) = path.parent() {
328        std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
329    }
330    std::fs::write(path, bytes).map_err(|e| e.to_string())
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336
337    fn roundtrip(persisted: bool) {
338        let dir = tempfile::tempdir().unwrap();
339        let work = dir.path().join("work");
340        std::fs::create_dir_all(&work).unwrap();
341        let a = work.join("a.txt");
342        let b = work.join("b.txt");
343        std::fs::write(&a, "a0").unwrap();
344        let mut store = CheckpointStore::default();
345        store.bind(persisted.then(|| dir.path().join("cp")));
346        assert!(store.record("u1", &a));
347        assert!(!store.record("u1", &a), "once per checkpoint");
348        std::fs::write(&a, "a1").unwrap();
349        assert!(store.record("u2", &a));
350        assert!(store.record("u2", &b));
351        std::fs::write(&a, "a2").unwrap();
352        std::fs::write(&b, "b2").unwrap();
353        assert_eq!(store.list().len(), 2);
354        assert!(!store.list()[1].files[1].existed);
355        if persisted {
356            let mut reloaded = CheckpointStore::default();
357            reloaded.bind(Some(dir.path().join("cp")));
358            assert_eq!(reloaded.list(), store.list());
359            store = reloaded;
360        }
361        let result = store.rewind("u2").unwrap();
362        assert_eq!(std::fs::read_to_string(&a).unwrap(), "a1");
363        assert!(!b.exists());
364        assert_eq!(result.restored.len(), 1);
365        assert_eq!(result.deleted.len(), 1);
366        assert_eq!(store.list().len(), 1);
367        let result = store.rewind("u1").unwrap();
368        assert_eq!(std::fs::read_to_string(&a).unwrap(), "a0");
369        assert_eq!(result.restored, vec![a.display().to_string()]);
370        assert!(store.rewind("u1").is_none());
371    }
372
373    #[test]
374    fn rewind_in_memory() {
375        roundtrip(false);
376    }
377
378    #[test]
379    fn rewind_persisted() {
380        roundtrip(true);
381    }
382
383    #[test]
384    fn copy_kept_entries() {
385        let dir = tempfile::tempdir().unwrap();
386        let a = dir.path().join("a");
387        std::fs::write(&a, "x").unwrap();
388        let mut store = CheckpointStore::default();
389        store.bind(None);
390        store.record("u1", &a);
391        store.record("u2", &a);
392        let mut target = CheckpointStore::default();
393        target.bind(None);
394        store.copy_matching(|id| id == "u1", &mut target);
395        assert_eq!(target.list().len(), 1);
396        assert!(target.contains("u1"));
397    }
398}
oino_core/session/checkpoints.rs

oino_core/session/
checkpoints.rs
