vtcode_core/tools/
command_cache.rs

1use crate::acp::permission_cache::PermissionGrant;
2pub use crate::acp::permission_cache::ToolPermissionCache as PermissionCache;
3use crate::cache::{CacheKey, EvictionPolicy, UnifiedCache};
4use crate::tools::shell::ShellOutput;
5use hashbrown::HashMap;
6use once_cell::sync::Lazy;
7use parking_lot::Mutex;
8use std::path::{Path, PathBuf};
9use tokio::sync::{Mutex as TokioMutex, oneshot};
10use vtcode_config::CommandCacheConfig;
11
12#[derive(Debug, Clone, Hash, Eq, PartialEq)]
13struct CommandCacheKey {
14    command: String,
15    cwd: PathBuf,
16}
17
18impl CacheKey for CommandCacheKey {
19    fn to_cache_key(&self) -> String {
20        format!("{}::{}", self.command, self.cwd.display())
21    }
22}
23
24fn command_cache_key(command: &str, cwd: &Path) -> CommandCacheKey {
25    CommandCacheKey {
26        command: command.to_string(),
27        cwd: cwd.to_path_buf(),
28    }
29}
30
31struct CommandCache {
32    inner: Mutex<CommandCacheInner>,
33}
34
35struct CommandCacheInner {
36    config: CommandCacheConfig,
37    cache: UnifiedCache<CommandCacheKey, ShellOutput>,
38}
39
40static COMMAND_CACHE: Lazy<CommandCache> =
41    Lazy::new(|| CommandCache::new(CommandCacheConfig::default()));
42
43pub type InFlightResult = Result<ShellOutput, String>;
44
45pub struct InFlightToken(CommandCacheKey);
46
47pub enum InFlightState {
48    Owner(InFlightToken),
49    Wait(oneshot::Receiver<InFlightResult>),
50}
51
52static IN_FLIGHT: Lazy<TokioMutex<HashMap<CommandCacheKey, Vec<oneshot::Sender<InFlightResult>>>>> =
53    Lazy::new(|| TokioMutex::new(HashMap::new()));
54
55impl PermissionCache {
56    pub fn get(&mut self, key: &str) -> Option<bool> {
57        match self.get_permission(key) {
58            Some(PermissionGrant::Denied) => Some(false),
59            Some(_) => Some(true),
60            None => None,
61        }
62    }
63
64    pub fn put(&mut self, key: &str, allowed: bool, _reason: &str) {
65        let grant = if allowed {
66            PermissionGrant::Session
67        } else {
68            PermissionGrant::Denied
69        };
70        self.cache_grant(key.to_string(), grant);
71    }
72}
73
74impl CommandCache {
75    fn new(config: CommandCacheConfig) -> Self {
76        let cache = Self::build_cache(&config);
77        Self {
78            inner: Mutex::new(CommandCacheInner { config, cache }),
79        }
80    }
81
82    fn build_cache(config: &CommandCacheConfig) -> UnifiedCache<CommandCacheKey, ShellOutput> {
83        UnifiedCache::new(
84            config.max_entries.max(1),
85            std::time::Duration::from_millis(config.ttl_ms),
86            EvictionPolicy::Lru,
87        )
88    }
89
90    fn configure(&self, config: &CommandCacheConfig) {
91        let mut inner = self.inner.lock();
92        inner.config = config.clone();
93        inner.cache = Self::build_cache(config);
94    }
95
96    fn allowlisted_with_config(cfg: &CommandCacheConfig, command: &str) -> bool {
97        if !cfg.enabled {
98            return false;
99        }
100        let trimmed = command.trim();
101        cfg.allowlist.iter().any(|entry| {
102            let entry = entry.trim();
103            trimmed == entry || trimmed.starts_with(&format!("{entry} "))
104        })
105    }
106
107    fn allowlisted(&self, command: &str) -> bool {
108        let inner = self.inner.lock();
109        Self::allowlisted_with_config(&inner.config, command)
110    }
111
112    fn get(&self, command: &str, cwd: &Path) -> Option<ShellOutput> {
113        let inner = self.inner.lock();
114        if !Self::allowlisted_with_config(&inner.config, command) {
115            return None;
116        }
117        let key = command_cache_key(command, cwd);
118        inner.cache.get_owned(&key)
119    }
120
121    fn put(&self, command: &str, cwd: &Path, output: ShellOutput) {
122        let inner = self.inner.lock();
123        if !Self::allowlisted_with_config(&inner.config, command) || output.exit_code != 0 {
124            return;
125        }
126        let key = command_cache_key(command, cwd);
127        let size = (output.stdout.len() + output.stderr.len()) as u64;
128        inner.cache.insert(key, output, size);
129    }
130}
131
132pub fn configure_command_cache(config: &CommandCacheConfig) {
133    COMMAND_CACHE.configure(config);
134}
135
136pub fn get_cached_output(command: &str, cwd: &Path) -> Option<ShellOutput> {
137    COMMAND_CACHE.get(command, cwd)
138}
139
140pub fn cache_output(command: &str, cwd: &Path, output: ShellOutput) {
141    COMMAND_CACHE.put(command, cwd, output);
142}
143
144pub async fn enter_inflight(command: &str, cwd: &Path) -> Option<InFlightState> {
145    if !COMMAND_CACHE.allowlisted(command) {
146        return None;
147    }
148
149    let key = command_cache_key(command, cwd);
150
151    let mut inflight = IN_FLIGHT.lock().await;
152    if let Some(waiters) = inflight.get_mut(&key) {
153        let (tx, rx) = oneshot::channel();
154        waiters.push(tx);
155        return Some(InFlightState::Wait(rx));
156    }
157
158    inflight.insert(key.clone(), Vec::new());
159    Some(InFlightState::Owner(InFlightToken(key)))
160}
161
162pub async fn finish_inflight(token: InFlightToken, result: InFlightResult) {
163    let key = token.0;
164    let mut inflight = IN_FLIGHT.lock().await;
165    if let Some(waiters) = inflight.remove(&key) {
166        for waiter in waiters {
167            let _ = waiter.send(result.clone());
168        }
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    fn test_config() -> CommandCacheConfig {
177        CommandCacheConfig {
178            enabled: true,
179            ttl_ms: 10_000,
180            max_entries: 8,
181            allowlist: vec!["git status".to_string(), "echo".to_string()],
182        }
183    }
184
185    #[test]
186    fn allowlist_matches_prefix() {
187        let cache = CommandCache::new(test_config());
188        assert!(cache.allowlisted("git status"));
189        assert!(cache.allowlisted("git status -s"));
190        assert!(!cache.allowlisted("git diff"));
191    }
192
193    #[test]
194    fn cache_stores_only_successes() {
195        let cache = CommandCache::new(test_config());
196        let cwd = Path::new("/tmp");
197
198        let failed = ShellOutput {
199            stdout: "nope".to_string(),
200            stderr: "err".to_string(),
201            exit_code: 1,
202        };
203        cache.put("echo bad", cwd, failed);
204        assert!(cache.get("echo bad", cwd).is_none());
205
206        let ok = ShellOutput {
207            stdout: "ok".to_string(),
208            stderr: String::new(),
209            exit_code: 0,
210        };
211        cache.put("echo ok", cwd, ok.clone());
212        let cached = cache.get("echo ok", cwd).expect("cached output");
213        assert_eq!(cached.stdout, ok.stdout);
214    }
215}
vtcode_core/tools/command_cache.rs

vtcode_core/tools/
command_cache.rs
