vtcode_safety/command_safety/
cache.rs

1//! Caching layer for command safety decisions.
2//!
3//! Caches safety decisions to avoid re-evaluating the same commands.
4//! Implements LRU eviction when cache exceeds size limit.
5
6use hashbrown::HashMap;
7use std::sync::Arc;
8use tokio::sync::Mutex;
9
10/// Cached safety decision
11#[derive(Clone, Debug)]
12pub struct CachedDecision {
13    /// True if command is safe
14    pub(crate) is_safe: bool,
15    /// Reason for decision
16    pub(crate) reason: String,
17    /// Access count (for LRU)
18    access_count: u64,
19}
20
21/// Thread-safe cache for command safety decisions
22pub struct SafetyDecisionCache {
23    cache: Arc<Mutex<HashMap<String, CachedDecision>>>,
24    max_size: usize,
25}
26
27impl SafetyDecisionCache {
28    /// Creates a new cache with given max size
29    pub(crate) fn new(max_size: usize) -> Self {
30        Self {
31            cache: Arc::new(Mutex::new(HashMap::new())),
32            max_size,
33        }
34    }
35
36    /// Creates a default cache (1000 entries)
37    pub fn default_cache() -> Self {
38        Self::new(1000)
39    }
40
41    /// Gets a cached decision
42    pub(crate) async fn get(&self, command: &str) -> Option<CachedDecision> {
43        let mut cache = self.cache.lock().await;
44        if let Some(decision) = cache.get_mut(command) {
45            decision.access_count += 1;
46            return Some(decision.clone());
47        }
48        None
49    }
50
51    /// Sets a cached decision
52    pub(crate) async fn put(&self, command: String, is_safe: bool, reason: String) {
53        let mut cache = self.cache.lock().await;
54
55        if cache.len() >= self.max_size
56            && !cache.contains_key(&command)
57            && let Some(least_used) = cache
58                .iter()
59                .min_by_key(|(_, decision)| decision.access_count)
60                .map(|(k, _)| k.clone())
61        {
62            cache.remove(&least_used);
63        }
64
65        cache.insert(command, CachedDecision { is_safe, reason, access_count: 1 });
66    }
67
68    /// Clears all cached entries
69    async fn clear(&self) {
70        let mut cache = self.cache.lock().await;
71        cache.clear();
72    }
73
74    /// Returns current cache size
75    async fn size(&self) -> usize {
76        let cache = self.cache.lock().await;
77        cache.len()
78    }
79
80    /// Returns cache hit rate statistics
81    async fn stats(&self) -> CacheStats {
82        let cache = self.cache.lock().await;
83        let total_accesses: u64 = cache.values().map(|d| d.access_count).sum();
84        let entry_count = cache.len();
85
86        CacheStats {
87            entry_count,
88            total_accesses,
89            avg_access_per_entry: if entry_count > 0 {
90                total_accesses / entry_count as u64
91            } else {
92                0
93            },
94        }
95    }
96}
97
98impl Clone for SafetyDecisionCache {
99    fn clone(&self) -> Self {
100        Self {
101            cache: Arc::clone(&self.cache),
102            max_size: self.max_size,
103        }
104    }
105}
106
107/// Cache statistics
108#[derive(Debug, Clone)]
109pub struct CacheStats {
110    entry_count: usize,
111    total_accesses: u64,
112    avg_access_per_entry: u64,
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[tokio::test]
120    async fn cache_stores_and_retrieves() {
121        let cache = SafetyDecisionCache::new(10);
122        cache
123            .put("git status".to_string(), true, "git status allowed".to_string())
124            .await;
125
126        let decision = cache.get("git status").await;
127        assert!(decision.is_some());
128        assert!(decision.unwrap().is_safe);
129    }
130
131    #[tokio::test]
132    async fn cache_returns_none_for_missing_key() {
133        let cache = SafetyDecisionCache::new(10);
134        let decision = cache.get("missing").await;
135        assert!(decision.is_none());
136    }
137
138    #[tokio::test]
139    async fn cache_tracks_access_count() {
140        let cache = SafetyDecisionCache::new(10);
141        cache.put("cmd".to_string(), true, "allowed".to_string()).await;
142
143        let d1 = cache.get("cmd").await.unwrap();
144        assert_eq!(d1.access_count, 2);
145
146        let d2 = cache.get("cmd").await.unwrap();
147        assert_eq!(d2.access_count, 3);
148    }
149
150    #[tokio::test]
151    async fn cache_respects_max_size() {
152        let cache = SafetyDecisionCache::new(3);
153
154        cache.put("cmd1".to_string(), true, "allowed".to_string()).await;
155        cache.put("cmd2".to_string(), true, "allowed".to_string()).await;
156        cache.put("cmd3".to_string(), true, "allowed".to_string()).await;
157
158        assert_eq!(cache.size().await, 3);
159
160        // Adding a 4th entry should evict the least-used
161        cache.put("cmd4".to_string(), true, "allowed".to_string()).await;
162        assert_eq!(cache.size().await, 3);
163    }
164
165    #[tokio::test]
166    async fn cache_clears() {
167        let cache = SafetyDecisionCache::new(10);
168        cache.put("cmd".to_string(), true, "allowed".to_string()).await;
169        assert_eq!(cache.size().await, 1);
170
171        cache.clear().await;
172        assert_eq!(cache.size().await, 0);
173    }
174
175    #[tokio::test]
176    async fn cache_stats() {
177        let cache = SafetyDecisionCache::new(10);
178        cache.put("cmd1".to_string(), true, "allowed".to_string()).await;
179        cache.put("cmd2".to_string(), true, "allowed".to_string()).await;
180
181        let _d1 = cache.get("cmd1").await;
182        let _d2 = cache.get("cmd2").await;
183        let _d3 = cache.get("cmd2").await;
184
185        let stats = cache.stats().await;
186        assert_eq!(stats.entry_count, 2);
187        assert_eq!(stats.total_accesses, 5); // 1+1 initial puts + 1+2 gets
188    }
189}
vtcode_safety/command_safety/cache.rs

vtcode_safety/command_safety/
cache.rs
