Skip to main content

vtcode_core/tools/
search_runtime.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::{Mutex, OnceLock};
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)
53            || matches!(self.ast_grep, SearchToolReadiness::Error)
54    }
55
56    #[must_use]
57    pub fn header_summary(self) -> String {
58        if self.all_ready() {
59            return "Search: ripgrep · ast-grep".to_string();
60        }
61        format!(
62            "Search: ripgrep {} · ast-grep {}",
63            self.ripgrep.label(),
64            self.ast_grep.label()
65        )
66    }
67}
68
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub(crate) struct SearchRuntimeSnapshot {
71    pub(crate) workspace_languages: Vec<String>,
72    pub(crate) search_tools: SearchToolBundleStatus,
73    pub(crate) ripgrep_ready: bool,
74    pub(crate) ast_grep_ready: bool,
75    pub(crate) code_tree_sitter_languages: Vec<String>,
76    pub(crate) bash_tree_sitter_ready: bool,
77}
78
79static SEARCH_RUNTIME_CACHE: OnceLock<Mutex<HashMap<PathBuf, SearchRuntimeSnapshot>>> =
80    OnceLock::new();
81
82pub(crate) fn snapshot_for_workspace(workspace_root: &Path) -> SearchRuntimeSnapshot {
83    let workspace_root = workspace_root.to_path_buf();
84    let cache = SEARCH_RUNTIME_CACHE.get_or_init(|| Mutex::new(HashMap::new()));
85
86    if let Some(snapshot) = cache
87        .lock()
88        .expect("search runtime cache mutex must not be poisoned")
89        .get(&workspace_root)
90        .cloned()
91    {
92        return snapshot;
93    }
94
95    let workspace_languages = detect_workspace_languages(&workspace_root);
96    let search_tools = SearchToolBundleStatus {
97        ripgrep: SearchToolReadiness::from_ripgrep_status(RipgrepStatus::check()),
98        ast_grep: SearchToolReadiness::from_ast_grep_status(AstGrepStatus::check()),
99    };
100    let snapshot = SearchRuntimeSnapshot {
101        code_tree_sitter_languages: prewarm_workspace_languages(
102            workspace_languages.iter().map(String::as_str),
103        ),
104        workspace_languages,
105        search_tools,
106        ripgrep_ready: search_tools.ripgrep.is_ready(),
107        ast_grep_ready: search_tools.ast_grep.is_ready(),
108        bash_tree_sitter_ready: prewarm_bash_parser().is_ok(),
109    };
110
111    let mut guard = cache
112        .lock()
113        .expect("search runtime cache mutex must not be poisoned");
114    guard
115        .entry(workspace_root)
116        .or_insert_with(|| snapshot.clone())
117        .clone()
118}
119
120pub fn dominant_workspace_language(workspace_root: &Path) -> Option<String> {
121    snapshot_for_workspace(workspace_root)
122        .workspace_languages
123        .into_iter()
124        .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::{
152        SearchToolBundleStatus, SearchToolReadiness, dominant_workspace_language,
153        snapshot_for_workspace,
154    };
155    use std::fs;
156    use tempfile::TempDir;
157
158    #[test]
159    fn snapshot_for_workspace_captures_languages_and_bash_parser_state() {
160        let workspace = TempDir::new().expect("workspace tempdir");
161        fs::create_dir_all(workspace.path().join("src")).expect("create src");
162        fs::create_dir_all(workspace.path().join("web")).expect("create web");
163        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
164        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
165
166        let snapshot = snapshot_for_workspace(workspace.path());
167
168        assert_eq!(
169            snapshot.workspace_languages,
170            vec!["Rust".to_string(), "TypeScript".to_string()]
171        );
172        assert_eq!(
173            snapshot.ripgrep_ready,
174            snapshot.search_tools.ripgrep.is_ready()
175        );
176        assert_eq!(
177            snapshot.ast_grep_ready,
178            snapshot.search_tools.ast_grep.is_ready()
179        );
180        assert_eq!(
181            snapshot.code_tree_sitter_languages,
182            vec!["Rust".to_string(), "TypeScript".to_string()]
183        );
184        assert!(snapshot.bash_tree_sitter_ready);
185    }
186
187    #[test]
188    fn snapshot_for_workspace_reuses_cached_languages() {
189        let workspace = TempDir::new().expect("workspace tempdir");
190        fs::create_dir_all(workspace.path().join("src")).expect("create src");
191        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
192
193        let initial = snapshot_for_workspace(workspace.path());
194
195        fs::create_dir_all(workspace.path().join("web")).expect("create web");
196        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
197
198        let cached = snapshot_for_workspace(workspace.path());
199
200        assert_eq!(initial.workspace_languages, cached.workspace_languages);
201    }
202
203    #[test]
204    fn dominant_workspace_language_returns_first_detected_language() {
205        let workspace = TempDir::new().expect("workspace tempdir");
206        fs::create_dir_all(workspace.path().join("src")).expect("create src");
207        fs::create_dir_all(workspace.path().join("web")).expect("create web");
208        fs::write(workspace.path().join("src/lib.rs"), "fn alpha() {}\n").expect("write rust");
209        fs::write(workspace.path().join("src/main.rs"), "fn beta() {}\n").expect("write rust");
210        fs::write(workspace.path().join("web/app.ts"), "const app = 1;\n").expect("write ts");
211
212        assert_eq!(
213            dominant_workspace_language(workspace.path()).as_deref(),
214            Some("Rust")
215        );
216    }
217
218    #[test]
219    fn snapshot_for_workspace_ignores_languages_without_shared_parser_support() {
220        let workspace = TempDir::new().expect("workspace tempdir");
221        fs::create_dir_all(workspace.path().join("ios")).expect("create ios");
222        fs::write(workspace.path().join("ios/App.swift"), "struct App {}\n").expect("write swift");
223
224        let snapshot = snapshot_for_workspace(workspace.path());
225
226        assert_eq!(snapshot.workspace_languages, vec!["Swift".to_string()]);
227        assert!(snapshot.code_tree_sitter_languages.is_empty());
228    }
229
230    #[test]
231    fn search_tool_bundle_header_summary_reflects_readiness() {
232        let status = SearchToolBundleStatus {
233            ripgrep: SearchToolReadiness::Ready,
234            ast_grep: SearchToolReadiness::Missing,
235        };
236
237        assert_eq!(
238            status.header_summary(),
239            "Search: ripgrep ready · ast-grep missing"
240        );
241        assert!(!status.all_ready());
242        assert!(!status.all_unavailable());
243        assert!(!status.has_errors());
244    }
245
246    #[test]
247    fn search_tool_bundle_header_summary_all_ready() {
248        let status = SearchToolBundleStatus {
249            ripgrep: SearchToolReadiness::Ready,
250            ast_grep: SearchToolReadiness::Ready,
251        };
252
253        assert_eq!(status.header_summary(), "Search: ripgrep · ast-grep");
254        assert!(status.all_ready());
255    }
256}