Skip to main content

vtcode_core/tools/
search_runtime.rs

1use std::path::{Path, PathBuf};
2use std::sync::{Mutex, OnceLock};
3use tracing::warn;
4
5use crate::command_safety::shell_parser::prewarm_bash_parser;
6use crate::tools::tree_sitter_runtime::prewarm_workspace_languages;
7use crate::tools::{AstGrepStatus, RipgrepStatus};
8use crate::utils::common::detect_workspace_languages;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum SearchToolReadiness {
12    Ready,
13    Missing,
14    Error,
15}
16
17impl SearchToolReadiness {
18    #[must_use]
19    pub fn is_ready(self) -> bool {
20        matches!(self, Self::Ready)
21    }
22
23    #[must_use]
24    pub fn label(self) -> &'static str {
25        match self {
26            Self::Ready => "ready",
27            Self::Missing => "missing",
28            Self::Error => "error",
29        }
30    }
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct SearchToolBundleStatus {
35    pub ripgrep: SearchToolReadiness,
36    pub ast_grep: SearchToolReadiness,
37}
38
39impl SearchToolBundleStatus {
40    #[must_use]
41    pub fn all_ready(self) -> bool {
42        self.ripgrep.is_ready() && self.ast_grep.is_ready()
43    }
44
45    #[must_use]
46    pub fn all_unavailable(self) -> bool {
47        !self.ripgrep.is_ready() && !self.ast_grep.is_ready()
48    }
49
50    #[must_use]
51    pub fn has_errors(self) -> bool {
52        matches!(self.ripgrep, SearchToolReadiness::Error) || matches!(self.ast_grep, SearchToolReadiness::Error)
53    }
54
55    #[must_use]
56    pub fn header_summary(self) -> String {
57        if self.all_ready() {
58            return "Search: ripgrep \u{00b7} ast-grep".to_string();
59        }
60        format!("Search: ripgrep {} \u{00b7} ast-grep {}", self.ripgrep.label(), self.ast_grep.label())
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub(crate) struct SearchRuntimeSnapshot {
66    pub(crate) workspace_languages: Vec<String>,
67    pub(crate) search_tools: SearchToolBundleStatus,
68    pub(crate) ripgrep_ready: bool,
69    pub(crate) ast_grep_ready: bool,
70    pub(crate) code_tree_sitter_languages: Vec<String>,
71    pub(crate) bash_tree_sitter_ready: bool,
72}
73
74/// Bounded per-workspace search-runtime snapshot cache. Snapshots are cheap to
75/// recompute relative to the value of caching (detect/prewarm only happens on
76/// miss), and keying by workspace path means an unbounded map would leak one
77/// entry per distinct workspace ever seen in a long-lived process. A small LRU
78/// cap bounds memory while keeping the hot workspace cached.
79const SEARCH_RUNTIME_CACHE_CAP: usize = 16;
80
81static SEARCH_RUNTIME_CACHE: OnceLock<Mutex<lru::LruCache<PathBuf, SearchRuntimeSnapshot>>> = OnceLock::new();
82
83pub(crate) fn snapshot_for_workspace(workspace_root: &Path) -> SearchRuntimeSnapshot {
84    let workspace_root = workspace_root.to_path_buf();
85    let cache = SEARCH_RUNTIME_CACHE.get_or_init(|| {
86        let capacity = std::num::NonZeroUsize::new(SEARCH_RUNTIME_CACHE_CAP)
87            .expect("SEARCH_RUNTIME_CACHE_CAP is a non-zero const");
88        Mutex::new(lru::LruCache::new(capacity))
89    });
90
91    // Check cache first; recover from poison by clearing and returning None to recompute
92    {
93        let mut cache_guard = cache.lock().unwrap_or_else(|poisoned| {
94            warn!("search runtime cache mutex poisoned; recovering");
95            poisoned.into_inner()
96        });
97        if let Some(snapshot) = cache_guard.get(&workspace_root).cloned() {
98            return snapshot;
99        }
100    }
101
102    let workspace_languages = detect_workspace_languages(&workspace_root);
103    let search_tools = SearchToolBundleStatus {
104        ripgrep: SearchToolReadiness::from_ripgrep_status(RipgrepStatus::check()),
105        ast_grep: SearchToolReadiness::from_ast_grep_status(AstGrepStatus::check()),
106    };
107    let snapshot = SearchRuntimeSnapshot {
108        code_tree_sitter_languages: prewarm_workspace_languages(workspace_languages.iter().map(String::as_str)),
109        workspace_languages,
110        search_tools,
111        ripgrep_ready: search_tools.ripgrep.is_ready(),
112        ast_grep_ready: search_tools.ast_grep.is_ready(),
113        bash_tree_sitter_ready: prewarm_bash_parser().is_ok(),
114    };
115
116    let mut guard = cache.lock().unwrap_or_else(|poisoned| {
117        warn!("search runtime cache mutex poisoned; recovering");
118        poisoned.into_inner()
119    });
120    guard.get_or_insert(workspace_root, || snapshot.clone()).clone()
121}
122
123pub fn dominant_workspace_language(workspace_root: &Path) -> Option<String> {
124    snapshot_for_workspace(workspace_root).workspace_languages.into_iter().next()
125}
126
127pub fn search_tool_bundle_status(workspace_root: &Path) -> SearchToolBundleStatus {
128    snapshot_for_workspace(workspace_root).search_tools
129}
130
131impl SearchToolReadiness {
132    fn from_ripgrep_status(status: RipgrepStatus) -> Self {
133        match status {
134            RipgrepStatus::Available { .. } => Self::Ready,
135            RipgrepStatus::NotFound => Self::Missing,
136            RipgrepStatus::Error { .. } => Self::Error,
137        }
138    }
139
140    fn from_ast_grep_status(status: AstGrepStatus) -> Self {
141        match status {
142            AstGrepStatus::Available { .. } => Self::Ready,
143            AstGrepStatus::NotFound => Self::Missing,
144            AstGrepStatus::Error { .. } => Self::Error,
145        }
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::{SearchToolBundleStatus, SearchToolReadiness, dominant_workspace_language, snapshot_for_workspace};
152    use std::fs;
153    use tempfile::TempDir;
154
155    #[test]
156    fn snapshot_for_workspace_captures_languages_and_bash_parser_state() {
157        let workspace = TempDir::new().expect("workspace tempdir");
158        fs::create_dir_all(workspace.path().join("src")).expect("create src");
159        fs::create_dir_all(workspace.path().join("web")).expect("create web");
160        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
161        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
162
163        let snapshot = snapshot_for_workspace(workspace.path());
164
165        assert_eq!(snapshot.workspace_languages, vec!["Rust".to_string(), "TypeScript".to_string()]);
166        assert_eq!(snapshot.ripgrep_ready, snapshot.search_tools.ripgrep.is_ready());
167        assert_eq!(snapshot.ast_grep_ready, snapshot.search_tools.ast_grep.is_ready());
168        assert_eq!(snapshot.code_tree_sitter_languages, vec!["Rust".to_string(), "TypeScript".to_string()]);
169        assert!(snapshot.bash_tree_sitter_ready);
170    }
171
172    #[test]
173    fn snapshot_for_workspace_reuses_cached_languages() {
174        let workspace = TempDir::new().expect("workspace tempdir");
175        fs::create_dir_all(workspace.path().join("src")).expect("create src");
176        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
177
178        let initial = snapshot_for_workspace(workspace.path());
179
180        fs::create_dir_all(workspace.path().join("web")).expect("create web");
181        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
182
183        let cached = snapshot_for_workspace(workspace.path());
184
185        assert_eq!(initial.workspace_languages, cached.workspace_languages);
186    }
187
188    #[test]
189    fn dominant_workspace_language_returns_first_detected_language() {
190        let workspace = TempDir::new().expect("workspace tempdir");
191        fs::create_dir_all(workspace.path().join("src")).expect("create src");
192        fs::create_dir_all(workspace.path().join("web")).expect("create web");
193        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
194        fs::write(workspace.path().join("src/main.rs"), "fn beta() {}\n").expect("write rust");
195        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
196
197        assert_eq!(dominant_workspace_language(workspace.path()).as_deref(), Some("Rust"));
198    }
199
200    #[test]
201    fn snapshot_for_workspace_ignores_languages_without_shared_parser_support() {
202        let workspace = TempDir::new().expect("workspace tempdir");
203        fs::create_dir_all(workspace.path().join("ios")).expect("create ios");
204        fs::write(workspace.path().join("ios/App.swift"), "struct App {}\n").expect("write swift");
205
206        let snapshot = snapshot_for_workspace(workspace.path());
207
208        assert_eq!(snapshot.workspace_languages, vec!["Swift".to_string()]);
209        assert!(snapshot.code_tree_sitter_languages.is_empty());
210    }
211
212    #[test]
213    fn search_tool_bundle_header_summary_reflects_readiness() {
214        let status = SearchToolBundleStatus {
215            ripgrep: SearchToolReadiness::Ready,
216            ast_grep: SearchToolReadiness::Missing,
217        };
218
219        assert_eq!(status.header_summary(), "Search: ripgrep ready \u{00b7} ast-grep missing");
220        assert!(!status.all_ready());
221        assert!(!status.all_unavailable());
222        assert!(!status.has_errors());
223    }
224
225    #[test]
226    fn search_tool_bundle_header_summary_all_ready() {
227        let status = SearchToolBundleStatus {
228            ripgrep: SearchToolReadiness::Ready,
229            ast_grep: SearchToolReadiness::Ready,
230        };
231
232        assert_eq!(status.header_summary(), "Search: ripgrep \u{00b7} ast-grep");
233        assert!(status.all_ready());
234    }
235}