agent_engine/tools/
confirm_latch.rs

1//! Per-session, per-extension "Allow all" latch for extension `Confirm`
2//! results on `before_tool_call`.
3//!
4//! The latch is keyed by (extension id, load generation). Both are stamped
5//! by `HookBus::emit`, never by the extension itself. Answering "Allow all"
6//! on a prompt from extension `A` auto-approves later confirms from that
7//! same loaded instance of `A` only: confirms from any other extension still
8//! prompt, and so does `A` after it is unloaded, reloaded or reinstalled
9//! (the generation changes). A confirm with no known source is never
10//! auto-approved by the latch.
11//!
12//! Clones of the session runtime share the `Arc`, so a revoke applies to
13//! them live. Once [`SessionConfirmLatch::revoke`] has run (conversation
14//! boundary) the latch can never be armed again, so a stale clone of an old
15//! conversation cannot re-arm an orphaned set nobody can see or reset. It
16//! never affects `Block`/`Modify` and grants no `activate_tools` authority.
17
18use std::collections::BTreeSet;
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::sync::Mutex;
21
22/// (extension id, load generation).
23pub type LatchKey = (String, u64);
24
25/// Pair `sources` with `generations` (missing generations are 0).
26pub fn latch_keys(sources: &[String], generations: &[u64]) -> Vec<LatchKey> {
27    sources
28        .iter()
29        .enumerate()
30        .map(|(i, s)| (s.clone(), generations.get(i).copied().unwrap_or(0)))
31        .collect()
32}
33
34#[derive(Debug, Default)]
35pub struct SessionConfirmLatch {
36    allowed: Mutex<BTreeSet<LatchKey>>,
37    revoked: AtomicBool,
38}
39
40impl SessionConfirmLatch {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    fn set(&self) -> std::sync::MutexGuard<'_, BTreeSet<LatchKey>> {
46        self.allowed.lock().unwrap_or_else(|e| e.into_inner())
47    }
48
49    /// True iff `keys` is non-empty and every key is latched.
50    pub fn allows_all(&self, keys: &[LatchKey]) -> bool {
51        if keys.is_empty() || self.revoked.load(Ordering::Acquire) {
52            return false;
53        }
54        let set = self.set();
55        keys.iter().all(|k| set.contains(k))
56    }
57
58    /// Latch every key. Returns false (and latches nothing) when `keys` is
59    /// empty — an unattributed confirm cannot be scoped — or the latch has
60    /// been revoked for good.
61    pub fn allow(&self, keys: &[LatchKey]) -> bool {
62        if keys.is_empty() || self.revoked.load(Ordering::Acquire) {
63            return false;
64        }
65        let mut set = self.set();
66        for k in keys {
67            set.insert(k.clone());
68        }
69        true
70    }
71
72    /// Revoke every latched extension; returns the revoked ids (sorted,
73    /// de-duplicated). The latch stays usable (`/confirms reset`).
74    pub fn clear(&self) -> Vec<String> {
75        let taken = std::mem::take(&mut *self.set());
76        let mut ids: Vec<String> = taken.into_iter().map(|(id, _)| id).collect();
77        ids.dedup();
78        ids
79    }
80
81    /// Clear and permanently disable (conversation boundary).
82    pub fn revoke(&self) {
83        self.revoked.store(true, Ordering::Release);
84        self.clear();
85    }
86
87    /// Latched extension ids (sorted, de-duplicated).
88    pub fn extensions(&self) -> Vec<String> {
89        let mut ids: Vec<String> = self.set().iter().map(|(id, _)| id.clone()).collect();
90        ids.dedup();
91        ids
92    }
93
94    pub fn is_empty(&self) -> bool {
95        self.set().is_empty()
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    fn k(v: &[&str]) -> Vec<LatchKey> {
104        v.iter().map(|x| (x.to_string(), 0)).collect()
105    }
106
107    #[test]
108    fn latch_is_scoped_to_the_answering_extension() {
109        let latch = SessionConfirmLatch::new();
110        assert!(!latch.allows_all(&k(&["jev"])));
111        assert!(latch.allow(&k(&["jev"])));
112        assert!(latch.allows_all(&k(&["jev"])));
113        assert!(
114            !latch.allows_all(&k(&["other"])),
115            "other extensions still prompt"
116        );
117        assert!(
118            !latch.allows_all(&k(&["jev", "other"])),
119            "a multi-source confirm needs every source latched"
120        );
121    }
122
123    #[test]
124    fn reinstalled_extension_under_same_id_is_not_latched() {
125        let latch = SessionConfirmLatch::new();
126        latch.allow(&[("jev".to_string(), 0)]);
127        assert!(latch.allows_all(&[("jev".to_string(), 0)]));
128        assert!(
129            !latch.allows_all(&[("jev".to_string(), 1)]),
130            "a new load generation of the same id must prompt again"
131        );
132        assert_eq!(
133            latch_keys(&["a".into(), "b".into()], &[4]),
134            vec![("a".to_string(), 4), ("b".to_string(), 0)]
135        );
136    }
137
138    #[test]
139    fn unattributed_confirm_is_never_latched() {
140        let latch = SessionConfirmLatch::new();
141        assert!(!latch.allow(&[]));
142        assert!(latch.is_empty());
143        assert!(!latch.allows_all(&[]));
144    }
145
146    #[test]
147    fn clear_revokes_everything() {
148        let latch = SessionConfirmLatch::new();
149        latch.allow(&k(&["b", "a"]));
150        assert_eq!(latch.extensions(), vec!["a".to_string(), "b".to_string()]);
151        assert_eq!(latch.clear(), vec!["a".to_string(), "b".to_string()]);
152        assert!(latch.is_empty());
153        assert!(!latch.allows_all(&k(&["a"])));
154        assert!(latch.allow(&k(&["a"])), "reset latch can be armed again");
155    }
156
157    #[test]
158    fn revoked_latch_can_never_be_rearmed() {
159        let latch = SessionConfirmLatch::new();
160        latch.allow(&k(&["a"]));
161        latch.revoke();
162        assert!(latch.is_empty());
163        assert!(!latch.allow(&k(&["a"])), "stale clone cannot re-arm");
164        assert!(!latch.allows_all(&k(&["a"])));
165    }
166}
agent_engine/tools/confirm_latch.rs

agent_engine/tools/
confirm_latch.rs
