ghidra_mcp/
poison.rs

1//! Bounded poison-set (spec §4.3): remembers (tool, selector) pairs that reproducibly crashed the
2//! worker, so an identical later request is rejected up-front without a third/fourth respawn. Bounded
3//! LRU (evicts oldest) so a stubborn agent spamming distinct payloads can't grow it without limit.
4
5use lru::LruCache;
6use std::num::NonZeroUsize;
7
8/// Default LRU bound (spec §9.5).
9pub const DEFAULT_POISON_CAPACITY: usize = 64;
10
11pub struct PoisonSet {
12    cache: LruCache<String, ()>,
13}
14
15impl PoisonSet {
16    pub fn new(capacity: usize) -> Self {
17        Self {
18            cache: LruCache::new(
19                NonZeroUsize::new(capacity.max(1)).expect("capacity.max(1) is always non-zero"),
20            ),
21        }
22    }
23
24    /// The key for a tool call. `\u{1}` cannot appear in a tool name and is vanishingly unlikely in a
25    /// selector, so it is a safe separator (avoids `tool="a", sel="b"` colliding with `tool="a b"`).
26    pub fn key(tool: &str, selector: &str) -> String {
27        format!("{tool}\u{1}{selector}")
28    }
29
30    /// Record a poison key (marks it terminal). Touches LRU recency.
31    pub fn insert(&mut self, tool: &str, selector: &str) {
32        self.cache.put(Self::key(tool, selector), ());
33    }
34
35    /// True if this exact (tool, selector) is known-poison. Touches recency so hot poisons stay.
36    pub fn contains(&mut self, tool: &str, selector: &str) -> bool {
37        self.cache.get(&Self::key(tool, selector)).is_some()
38    }
39
40    pub fn len(&self) -> usize {
41        self.cache.len()
42    }
43
44    pub fn is_empty(&self) -> bool {
45        self.cache.is_empty()
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn remembers_then_short_circuits_same_selector() {
55        let mut p = PoisonSet::new(64);
56        assert!(!p.contains("inspect_function", "boom"));
57        p.insert("inspect_function", "boom");
58        assert!(p.contains("inspect_function", "boom"));
59        // A different selector or tool is not poisoned.
60        assert!(!p.contains("inspect_function", "safe"));
61        assert!(!p.contains("attach_program", "boom"));
62    }
63
64    #[test]
65    fn eviction_bounds_the_set() {
66        let mut p = PoisonSet::new(2);
67        p.insert("t", "a");
68        p.insert("t", "b");
69        p.insert("t", "c"); // evicts the least-recently-used ("a")
70        assert_eq!(p.len(), 2);
71        assert!(!p.contains("t", "a"));
72        assert!(p.contains("t", "c"));
73    }
74
75    #[test]
76    fn key_separator_avoids_collision() {
77        assert_ne!(PoisonSet::key("a", "b"), PoisonSet::key("a b", ""));
78    }
79}
ghidra_mcp/poison.rs

ghidra_mcp/
poison.rs
