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

code_kb_core/edit_check/
session.rs
