agent_harness_adapter_core/hook/
guard.rs

1//! The loop guard of a blocking hook: the hash of the last text the hook
2//! blocked on, one file per key (a repository, a session) in a directory the
3//! caller chooses. A hook that would block twice on the same text lets the
4//! agent stop instead.
5// @zen-component: KIT-LoopGuard
6
7use std::{fs, path::PathBuf};
8
9use crate::hash::{Fnv, hash_text};
10
11/// A loop guard whose files live in one directory.
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct LoopGuard {
14    dir: PathBuf,
15}
16
17impl LoopGuard {
18    /// A guard keeping its files in `dir`, created on the first write with a
19    /// `.gitignore` of `*` so a directory inside a repository stays out of
20    /// it.
21    pub fn new(dir: impl Into<PathBuf>) -> Self {
22        LoopGuard { dir: dir.into() }
23    }
24
25    /// The guard file of `key`: the key itself when it is short and made of
26    /// `[A-Za-z0-9_-]`, else its hash.
27    fn path(&self, key: &str) -> PathBuf {
28        let safe = !key.is_empty()
29            && key.len() <= 64
30            && key
31                .bytes()
32                .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_');
33        let name = if safe {
34            key.to_string()
35        } else {
36            let mut h = Fnv::new();
37            h.update(key.as_bytes());
38            h.render().replace(':', "-")
39        };
40        self.dir.join(format!("{name}.hash"))
41    }
42
43    /// The hash last recorded for `key`; `None` when there is none.
44    fn read(&self, key: &str) -> Option<String> {
45        let text = fs::read_to_string(self.path(key)).ok()?;
46        let text = text.trim().to_string();
47        (!text.is_empty()).then_some(text)
48    }
49
50    /// Record `hash` for `key`. Silent on failure: the guard is a cache.
51    // @zen-impl: KIT-12_AC-2
52    fn write(&self, key: &str, hash: &str) {
53        if fs::create_dir_all(&self.dir).is_err() {
54            return;
55        }
56        let ignore = self.dir.join(".gitignore");
57        if !ignore.exists() {
58            let _ = fs::write(&ignore, "*\n");
59        }
60        let _ = fs::write(self.path(key), format!("{hash}\n"));
61    }
62
63    /// Forget `key`, e.g. once the hook no longer blocks.
64    pub fn clear(&self, key: &str) {
65        let _ = fs::remove_file(self.path(key));
66    }
67
68    /// Whether the hook should block on `text` for `key`: `false` when it
69    /// blocked on the same text last time; otherwise `true`, and the text's
70    /// hash is recorded.
71    // @zen-impl: KIT-12_AC-1
72    pub fn should_block(&self, key: &str, text: &str) -> bool {
73        let hash = hash_text(text);
74        if self.read(key).as_deref() == Some(hash.as_str()) {
75            return false;
76        }
77        self.write(key, &hash);
78        true
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use proptest::prelude::*;
85
86    use super::*;
87    use crate::test_support::temp_dir;
88
89    // @zen-test: KIT-12_AC-2
90    #[test]
91    fn reads_what_it_wrote_per_key() {
92        let dir = temp_dir("guard");
93        let guard = LoopGuard::new(dir.join("cache"));
94        assert_eq!(guard.read("repo"), None);
95        guard.write("repo", "fnv1a64:01");
96        assert_eq!(guard.read("repo").as_deref(), Some("fnv1a64:01"));
97        assert_eq!(
98            fs::read_to_string(dir.join("cache/.gitignore")).unwrap(),
99            "*\n"
100        );
101        assert!(dir.join("cache/repo.hash").is_file());
102        // Keys are separate; an unsafe key is hashed into a file name.
103        assert_eq!(guard.read("other"), None);
104        guard.write("a/../b c", "fnv1a64:02");
105        assert_eq!(guard.read("a/../b c").as_deref(), Some("fnv1a64:02"));
106        assert!(!dir.join("b c.hash").exists());
107        guard.clear("repo");
108        assert_eq!(guard.read("repo"), None);
109        fs::remove_dir_all(&dir).unwrap();
110    }
111
112    // @zen-test: KIT-12_AC-1
113    #[test]
114    fn blocks_once_per_text() {
115        let dir = temp_dir("guard-block");
116        let guard = LoopGuard::new(&dir);
117        assert!(guard.should_block("s1", "2 errors"));
118        assert!(!guard.should_block("s1", "2 errors"));
119        assert!(!guard.should_block("s1", "2 errors\r\n".trim_end()));
120        assert!(guard.should_block("s2", "2 errors"));
121        assert!(guard.should_block("s1", "1 error"));
122        assert!(guard.should_block("s1", "2 errors"));
123        fs::remove_dir_all(&dir).unwrap();
124    }
125
126    // @zen-test: KIT-12_AC-2
127    #[test]
128    fn an_unwritable_directory_is_silent() {
129        let dir = temp_dir("guard-unwritable");
130        fs::write(dir.join("file"), "x").unwrap();
131        let guard = LoopGuard::new(dir.join("file"));
132        guard.write("k", "h");
133        assert_eq!(guard.read("k"), None);
134        // Never recorded, so it blocks every time.
135        assert!(guard.should_block("k", "t"));
136        assert!(guard.should_block("k", "t"));
137        fs::remove_dir_all(&dir).unwrap();
138    }
139
140    proptest! {
141        #![proptest_config(ProptestConfig::with_cases(64))]
142
143        // @zen-test: KIT_P-7
144        #[test]
145        fn the_guard_blocks_once_per_text(texts in prop::collection::vec(prop::sample::select(vec!["", "1 error", "2 errors"]), 1..10)) {
146            let dir = temp_dir("guard-prop");
147            let guard = LoopGuard::new(dir.join("cache"));
148            let mut last: Option<&str> = None;
149            for text in texts {
150                let blocked = guard.should_block("session", text);
151                prop_assert_eq!(blocked, last != Some(text));
152                if blocked {
153                    last = Some(text);
154                }
155            }
156            let _ = fs::remove_dir_all(&dir);
157        }
158    }
159}
agent_harness_adapter_core/hook/guard.rs

agent_harness_adapter_core/hook/
guard.rs
