Skip to main content

code_kb_core/edit_check/
session.rs

1//! Per-session state of the post-edit check, one file per session at
2//! `<root>/.code-kb/sessions/{session_id}.json`: the commit the session started at, the diff the
3//! last hook check ran on, and the diff a Stop hook blocked on. A `{session_id}.lock` file next to
4//! it serializes the hooks of one session.
5
6use std::fs::{self, File};
7use std::io;
8use std::path::{Path, PathBuf};
9use std::time::{Duration, Instant, SystemTime};
10
11use serde::{Deserialize, Serialize};
12
13const LATEST_MAX_AGE: Duration = Duration::from_secs(24 * 3600);
14const PRUNE_AGE: Duration = Duration::from_secs(7 * 24 * 3600);
15const MAX_ID_LEN: usize = 128;
16const LOCK_STALE_AGE: Duration = Duration::from_secs(60);
17const LOCK_POLL: Duration = Duration::from_millis(25);
18
19/// What the hooks remember about one session.
20#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
21pub struct SessionState {
22    #[serde(default, skip_serializing_if = "Option::is_none")]
23    pub base: Option<String>,
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    pub blocked_diff: Option<String>,
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub checked_diff: Option<String>,
28}
29
30impl SessionState {
31    /// True when a hook of this session already checked the diff `hash` or blocked on it.
32    pub fn has_seen(&self, hash: &str) -> bool {
33        [&self.checked_diff, &self.blocked_diff]
34            .into_iter()
35            .any(|seen| seen.as_deref() == Some(hash))
36    }
37}
38
39/// True for a session id that is safe in a file name: letters, digits, `_` and `-` only.
40pub fn is_session_id(id: &str) -> bool {
41    !id.is_empty()
42        && id.len() <= MAX_ID_LEN
43        && id
44            .bytes()
45            .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-')
46}
47
48fn sessions_dir(root: &Path) -> PathBuf {
49    root.join(".code-kb").join("sessions")
50}
51
52/// Creates `<root>/.code-kb/sessions` and, like an index does, the `.gitignore` that hides
53/// `.code-kb`, so a project with no index yet stays out of its own commits.
54fn create_sessions_dir(root: &Path) -> io::Result<()> {
55    crate::sync::ensure_index_dir(&root.join(".code-kb").join("artifact.db"))?;
56    fs::create_dir_all(sessions_dir(root))
57}
58
59fn session_file(root: &Path, id: &str) -> io::Result<PathBuf> {
60    if is_session_id(id) {
61        Ok(sessions_dir(root).join(format!("{id}.json")))
62    } else {
63        Err(io::Error::new(
64            io::ErrorKind::InvalidInput,
65            "a session id holds only letters, digits, `_` and `-`",
66        ))
67    }
68}
69
70fn parse(path: &Path) -> Option<SessionState> {
71    serde_json::from_str(&fs::read_to_string(path).ok()?).ok()
72}
73
74pub fn read_session(root: &Path, id: &str) -> Option<SessionState> {
75    parse(&session_file(root, id).ok()?)
76}
77
78pub fn write_session(root: &Path, id: &str, state: &SessionState) -> io::Result<()> {
79    let path = session_file(root, id)?;
80    create_sessions_dir(root)?;
81    let temp = path.with_extension(format!("json.{}.tmp", std::process::id()));
82    fs::write(&temp, serde_json::to_string(state)?)?;
83    fs::rename(&temp, &path).inspect_err(|_| {
84        let _ = fs::remove_file(&temp);
85    })
86}
87
88/// Records `head` as the base of a session whose file has none, and keeps the other fields of the
89/// file. A session that has a base keeps it, so a resume or a compaction does not move it. The
90/// caller holds the session lock, because a `Stop` hook rewrites the whole file.
91pub fn record_session_base(root: &Path, id: &str, head: &str) -> io::Result<()> {
92    let state = read_session(root, id).unwrap_or_default();
93    if state.base.is_some() {
94        return Ok(());
95    }
96    write_session(
97        root,
98        id,
99        &SessionState {
100            base: Some(head.to_string()),
101            ..state
102        },
103    )
104}
105
106/// The lock of one session. Dropping it removes the lock file.
107#[derive(Debug)]
108pub struct SessionLock {
109    path: PathBuf,
110}
111
112impl Drop for SessionLock {
113    fn drop(&mut self) {
114        let _ = fs::remove_file(&self.path);
115    }
116}
117
118/// Takes the lock of a session by creating `{session_id}.lock` exclusively, polling for up to
119/// `wait`. `None` when the lock stays held, or the id or the folder is unusable. A lock file
120/// older than the hook timeout belongs to a hook that died, and the next caller takes it over.
121pub fn lock_session(root: &Path, id: &str, wait: Duration) -> Option<SessionLock> {
122    let path = session_file(root, id).ok()?.with_extension("lock");
123    create_sessions_dir(root).ok()?;
124    let deadline = Instant::now() + wait;
125    loop {
126        match File::options().write(true).create_new(true).open(&path) {
127            Ok(_) => return Some(SessionLock { path }),
128            Err(error)
129                if matches!(
130                    error.kind(),
131                    io::ErrorKind::AlreadyExists | io::ErrorKind::PermissionDenied
132                ) =>
133            {
134                if is_stale(&path) && fs::remove_file(&path).is_ok() {
135                    continue;
136                }
137                if Instant::now() >= deadline {
138                    return None;
139                }
140                std::thread::sleep(LOCK_POLL);
141            }
142            Err(_) => return None,
143        }
144    }
145}
146
147fn is_stale(path: &Path) -> bool {
148    fs::metadata(path)
149        .and_then(|meta| meta.modified())
150        .is_ok_and(|modified| age_of(modified) > LOCK_STALE_AGE)
151}
152
153fn age_of(modified: SystemTime) -> Duration {
154    SystemTime::now()
155        .duration_since(modified)
156        .unwrap_or(Duration::ZERO)
157}
158
159/// The base of the session file written last within the past day, for a caller that has no
160/// session id.
161pub fn latest_session_base(root: &Path) -> Option<String> {
162    let mut files: Vec<(SystemTime, PathBuf)> = fs::read_dir(sessions_dir(root))
163        .ok()?
164        .filter_map(Result::ok)
165        .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "json"))
166        .filter_map(|entry| Some((entry.metadata().ok()?.modified().ok()?, entry.path())))
167        .filter(|(modified, _)| age_of(*modified) <= LATEST_MAX_AGE)
168        .collect();
169    files.sort_by_key(|(modified, _)| std::cmp::Reverse(*modified));
170    files.iter().find_map(|(_, path)| parse(path)?.base)
171}
172
173/// Removes the session files that no session can still use.
174pub fn prune_sessions(root: &Path) {
175    let Ok(entries) = fs::read_dir(sessions_dir(root)) else {
176        return;
177    };
178    for entry in entries.filter_map(Result::ok) {
179        let stale = entry
180            .metadata()
181            .and_then(|meta| meta.modified())
182            .is_ok_and(|modified| age_of(modified) > PRUNE_AGE);
183        if stale {
184            let _ = fs::remove_file(entry.path());
185        }
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    fn root() -> tempfile::TempDir {
194        tempfile::tempdir().unwrap()
195    }
196
197    fn age(root: &Path, id: &str, by: Duration) {
198        let file = File::options()
199            .write(true)
200            .open(session_file(root, id).unwrap())
201            .unwrap();
202        file.set_modified(SystemTime::now() - by).unwrap();
203    }
204
205    fn based(base: &str) -> SessionState {
206        SessionState {
207            base: Some(base.to_string()),
208            ..SessionState::default()
209        }
210    }
211
212    #[test]
213    fn only_letters_digits_underscore_and_hyphen_make_a_session_id() {
214        for id in ["abc", "A-1_b", "0123456789abcdef-0123"] {
215            assert!(is_session_id(id), "{id}");
216        }
217        for id in [
218            "", ".", "..", "a.b", "a/b", "a\\b", "a b", "../x", "é", "a\0b",
219        ] {
220            assert!(!is_session_id(id), "{id:?}");
221        }
222        assert!(is_session_id(&"a".repeat(MAX_ID_LEN)));
223        assert!(!is_session_id(&"a".repeat(MAX_ID_LEN + 1)));
224    }
225
226    #[test]
227    fn a_state_round_trips_and_leaves_out_what_it_lacks() {
228        let root = root();
229        write_session(root.path(), "s", &based("abc")).unwrap();
230        let text = fs::read_to_string(session_file(root.path(), "s").unwrap()).unwrap();
231        assert_eq!(text, "{\"base\":\"abc\"}");
232        assert_eq!(read_session(root.path(), "s"), Some(based("abc")));
233
234        let all = SessionState {
235            base: Some("abc".to_string()),
236            blocked_diff: Some("h".to_string()),
237            checked_diff: Some("k".to_string()),
238        };
239        write_session(root.path(), "s", &all).unwrap();
240        assert_eq!(read_session(root.path(), "s"), Some(all));
241    }
242
243    #[test]
244    fn a_missing_or_unreadable_file_reads_as_no_session() {
245        let root = root();
246        assert_eq!(read_session(root.path(), "none"), None);
247        fs::create_dir_all(sessions_dir(root.path())).unwrap();
248        fs::write(sessions_dir(root.path()).join("bad.json"), "not json").unwrap();
249        assert_eq!(read_session(root.path(), "bad"), None);
250        assert_eq!(read_session(root.path(), "../bad"), None);
251    }
252
253    #[test]
254    fn a_write_replaces_the_file_and_leaves_no_temporary_file() {
255        let root = root();
256        write_session(root.path(), "s", &based("one")).unwrap();
257        write_session(root.path(), "s", &based("two")).unwrap();
258        assert_eq!(read_session(root.path(), "s"), Some(based("two")));
259        let names: Vec<_> = fs::read_dir(sessions_dir(root.path()))
260            .unwrap()
261            .map(|entry| entry.unwrap().file_name())
262            .collect();
263        assert_eq!(names, ["s.json"]);
264    }
265
266    #[test]
267    fn a_write_or_record_with_a_bad_id_fails_and_creates_nothing() {
268        let root = root();
269        assert!(write_session(root.path(), "../x", &based("a")).is_err());
270        assert!(record_session_base(root.path(), "a/b", "a").is_err());
271        assert!(!root.path().join(".code-kb").exists());
272    }
273
274    #[test]
275    fn the_first_session_file_of_a_root_comes_with_the_gitignore_of_the_index_folder() {
276        let root = root();
277        record_session_base(root.path(), "s", "abc").unwrap();
278        let gitignore = root.path().join(".code-kb").join(".gitignore");
279        assert_eq!(fs::read_to_string(gitignore).unwrap(), "*\n");
280    }
281
282    #[test]
283    fn a_recorded_base_is_kept_by_a_later_record() {
284        let root = root();
285        record_session_base(root.path(), "s", "first").unwrap();
286        record_session_base(root.path(), "s", "second").unwrap();
287        assert_eq!(read_session(root.path(), "s"), Some(based("first")));
288    }
289
290    #[test]
291    fn a_record_adds_the_base_to_a_file_that_has_none_and_keeps_its_other_fields() {
292        let root = root();
293        let seen = SessionState {
294            blocked_diff: Some("b".to_string()),
295            checked_diff: Some("c".to_string()),
296            ..SessionState::default()
297        };
298        write_session(root.path(), "s", &seen).unwrap();
299        record_session_base(root.path(), "s", "head").unwrap();
300        let expected = SessionState {
301            base: Some("head".to_string()),
302            ..seen
303        };
304        assert_eq!(read_session(root.path(), "s"), Some(expected));
305    }
306
307    #[test]
308    fn the_latest_base_is_that_of_the_newest_file() {
309        let root = root();
310        assert_eq!(latest_session_base(root.path()), None);
311        record_session_base(root.path(), "old", "older").unwrap();
312        record_session_base(root.path(), "new", "newer").unwrap();
313        age(root.path(), "old", Duration::from_secs(3600));
314        assert_eq!(latest_session_base(root.path()).as_deref(), Some("newer"));
315    }
316
317    #[test]
318    fn the_latest_base_skips_a_file_with_no_base_and_a_file_older_than_a_day() {
319        let root = root();
320        record_session_base(root.path(), "usable", "kept").unwrap();
321        age(root.path(), "usable", Duration::from_secs(3600));
322        write_session(
323            root.path(),
324            "blocked-only",
325            &SessionState {
326                blocked_diff: Some("h".to_string()),
327                ..SessionState::default()
328            },
329        )
330        .unwrap();
331        assert_eq!(latest_session_base(root.path()).as_deref(), Some("kept"));
332
333        age(root.path(), "usable", Duration::from_secs(25 * 3600));
334        assert_eq!(latest_session_base(root.path()), None);
335    }
336
337    #[test]
338    fn pruning_removes_only_files_older_than_a_week() {
339        let root = root();
340        record_session_base(root.path(), "stale", "a").unwrap();
341        record_session_base(root.path(), "recent", "b").unwrap();
342        age(root.path(), "stale", Duration::from_secs(8 * 24 * 3600));
343        age(root.path(), "recent", Duration::from_secs(2 * 24 * 3600));
344        prune_sessions(root.path());
345        assert_eq!(read_session(root.path(), "stale"), None);
346        assert_eq!(read_session(root.path(), "recent"), Some(based("b")));
347    }
348
349    #[test]
350    fn a_held_lock_keeps_a_second_taker_out_until_it_is_dropped() {
351        let root = root();
352        let first = lock_session(root.path(), "s", Duration::ZERO).unwrap();
353        assert!(lock_session(root.path(), "s", Duration::from_millis(60)).is_none());
354        assert!(lock_session(root.path(), "other", Duration::ZERO).is_some());
355        drop(first);
356        assert!(lock_session(root.path(), "s", Duration::ZERO).is_some());
357    }
358
359    #[test]
360    fn a_lock_file_older_than_the_stale_age_is_taken_over() {
361        let root = root();
362        fs::create_dir_all(sessions_dir(root.path())).unwrap();
363        let lock = sessions_dir(root.path()).join("s.lock");
364        let file = File::create(&lock).unwrap();
365        file.set_modified(SystemTime::now() - LOCK_STALE_AGE - Duration::from_secs(1))
366            .unwrap();
367        drop(file);
368        let taken = lock_session(root.path(), "s", Duration::ZERO);
369        assert!(taken.is_some());
370        drop(taken);
371        assert!(!lock.exists());
372    }
373
374    #[test]
375    fn a_lock_for_a_bad_session_id_is_refused_and_creates_nothing() {
376        let root = root();
377        assert!(lock_session(root.path(), "../x", Duration::ZERO).is_none());
378        assert!(!root.path().join(".code-kb").exists());
379    }
380
381    #[test]
382    fn a_lock_file_is_no_session_file() {
383        let root = root();
384        let _lock = lock_session(root.path(), "s", Duration::ZERO).unwrap();
385        assert_eq!(latest_session_base(root.path()), None);
386        assert_eq!(read_session(root.path(), "s"), None);
387    }
388
389    #[test]
390    fn a_state_has_seen_the_diffs_it_checked_or_blocked_on() {
391        let state = SessionState {
392            blocked_diff: Some("b".to_string()),
393            checked_diff: Some("c".to_string()),
394            ..SessionState::default()
395        };
396        assert!(state.has_seen("b") && state.has_seen("c"));
397        assert!(!state.has_seen("d") && !SessionState::default().has_seen(""));
398    }
399
400    #[test]
401    fn pruning_a_root_with_no_sessions_does_nothing() {
402        let root = root();
403        prune_sessions(root.path());
404        assert!(!root.path().join(".code-kb").exists());
405    }
406}