vtcode_core/tools/
ast_grep_binary.rs

1use std::path::{Path, PathBuf};
2use std::sync::Mutex;
3
4use once_cell::sync::Lazy;
5
6pub const AST_GREP_BIN_ENV: &str = "VTCODE_AST_GREP_BIN";
7pub const AST_GREP_INSTALL_COMMAND: &str = "vtcode dependencies install ast-grep";
8
9static AST_GREP_OVERRIDE: Lazy<Mutex<AstGrepBinaryOverride>> =
10    Lazy::new(|| Mutex::new(AstGrepBinaryOverride::System));
11
12#[derive(Debug, Clone, Default)]
13enum AstGrepBinaryOverride {
14    #[default]
15    System,
16    Missing,
17    Path(PathBuf),
18}
19
20#[doc(hidden)]
21#[must_use]
22pub struct AstGrepBinaryOverrideGuard {
23    previous: AstGrepBinaryOverride,
24}
25
26impl Drop for AstGrepBinaryOverrideGuard {
27    fn drop(&mut self) {
28        *AST_GREP_OVERRIDE
29            .lock()
30            .expect("ast-grep override mutex must not be poisoned") = self.previous.clone();
31    }
32}
33
34#[doc(hidden)]
35pub fn set_ast_grep_binary_override_for_tests(path: Option<PathBuf>) -> AstGrepBinaryOverrideGuard {
36    let mut state = AST_GREP_OVERRIDE
37        .lock()
38        .expect("ast-grep override mutex must not be poisoned");
39    let previous = state.clone();
40    *state = match path {
41        Some(path) => AstGrepBinaryOverride::Path(path),
42        None => AstGrepBinaryOverride::Missing,
43    };
44    AstGrepBinaryOverrideGuard { previous }
45}
46
47pub fn managed_ast_grep_bin_dir() -> PathBuf {
48    dirs::home_dir()
49        .map(|home| managed_ast_grep_bin_dir_from_home(&home))
50        .unwrap_or_else(|| PathBuf::from(".vtcode/bin"))
51}
52
53pub fn managed_ast_grep_binary_path() -> PathBuf {
54    managed_ast_grep_bin_dir().join(canonical_ast_grep_binary_name())
55}
56
57pub fn managed_ast_grep_alias_path() -> Option<PathBuf> {
58    alias_ast_grep_binary_name().map(|name| managed_ast_grep_bin_dir().join(name))
59}
60
61pub fn managed_ast_grep_candidates() -> Vec<PathBuf> {
62    let mut candidates = vec![managed_ast_grep_binary_path()];
63    if let Some(alias) = managed_ast_grep_alias_path() {
64        candidates.push(alias);
65    }
66    candidates
67}
68
69pub fn resolve_ast_grep_binary_from_env_and_fs() -> Option<PathBuf> {
70    match AST_GREP_OVERRIDE
71        .lock()
72        .expect("ast-grep override mutex must not be poisoned")
73        .clone()
74    {
75        AstGrepBinaryOverride::System => {}
76        AstGrepBinaryOverride::Missing => return None,
77        AstGrepBinaryOverride::Path(path) => return Some(path),
78    }
79
80    let env_override = std::env::var_os(AST_GREP_BIN_ENV)
81        .filter(|value| !value.is_empty())
82        .map(PathBuf::from);
83
84    resolve_ast_grep_binary_with_sources(
85        env_override,
86        managed_ast_grep_candidates(),
87        resolve_ast_grep_binary_on_path(),
88    )
89}
90
91pub fn resolve_ast_grep_binary_on_path() -> Option<PathBuf> {
92    which::which(canonical_ast_grep_binary_name())
93        .ok()
94        .or_else(|| alias_ast_grep_binary_name().and_then(|alias| which::which(alias).ok()))
95}
96
97pub fn canonical_ast_grep_binary_name() -> &'static str {
98    if cfg!(target_os = "windows") {
99        "ast-grep.exe"
100    } else {
101        "ast-grep"
102    }
103}
104
105pub fn alias_ast_grep_binary_name() -> Option<&'static str> {
106    if cfg!(target_os = "linux") {
107        None
108    } else if cfg!(target_os = "windows") {
109        Some("sg.exe")
110    } else {
111        Some("sg")
112    }
113}
114
115#[cold]
116pub fn missing_ast_grep_message(suffix: &str) -> String {
117    let extra = if suffix.is_empty() {
118        String::new()
119    } else {
120        format!(" {suffix}")
121    };
122    format!(
123        "ast-grep is not available; run `{AST_GREP_INSTALL_COMMAND}` or install `ast-grep` manually.{extra}"
124    )
125}
126
127fn managed_ast_grep_bin_dir_from_home(home: &Path) -> PathBuf {
128    home.join(".vtcode").join("bin")
129}
130
131fn resolve_ast_grep_binary_with_sources(
132    env_override: Option<PathBuf>,
133    managed_candidates: Vec<PathBuf>,
134    path_candidate: Option<PathBuf>,
135) -> Option<PathBuf> {
136    env_override
137        .filter(|path| path.exists())
138        .or_else(|| {
139            managed_candidates
140                .into_iter()
141                .find(|candidate| candidate.exists())
142        })
143        .or(path_candidate)
144}
145
146#[cfg(test)]
147mod tests {
148    use super::{
149        alias_ast_grep_binary_name, canonical_ast_grep_binary_name,
150        managed_ast_grep_bin_dir_from_home, resolve_ast_grep_binary_from_env_and_fs,
151        resolve_ast_grep_binary_with_sources, set_ast_grep_binary_override_for_tests,
152    };
153    use std::path::{Path, PathBuf};
154    use tempfile::TempDir;
155
156    #[test]
157    fn managed_bin_dir_uses_vtcode_home_bin() {
158        let path = managed_ast_grep_bin_dir_from_home(Path::new("/tmp/example-home"));
159        assert_eq!(path, Path::new("/tmp/example-home/.vtcode/bin"));
160    }
161
162    #[test]
163    fn canonical_binary_name_matches_platform() {
164        if cfg!(target_os = "windows") {
165            assert_eq!(canonical_ast_grep_binary_name(), "ast-grep.exe");
166        } else {
167            assert_eq!(canonical_ast_grep_binary_name(), "ast-grep");
168        }
169    }
170
171    #[test]
172    fn alias_binary_name_skips_linux() {
173        if cfg!(target_os = "linux") {
174            assert_eq!(alias_ast_grep_binary_name(), None);
175        } else if cfg!(target_os = "windows") {
176            assert_eq!(alias_ast_grep_binary_name(), Some("sg.exe"));
177        } else {
178            assert_eq!(alias_ast_grep_binary_name(), Some("sg"));
179        }
180    }
181
182    #[test]
183    fn resolution_prefers_env_override() {
184        let temp_dir = TempDir::new().expect("temp dir");
185        let env_path = temp_dir.path().join("custom-ast-grep");
186        std::fs::write(&env_path, "binary").expect("env binary");
187
188        let managed_path = temp_dir.path().join("managed-ast-grep");
189        std::fs::write(&managed_path, "binary").expect("managed binary");
190
191        let path_fallback = temp_dir.path().join("path-ast-grep");
192        std::fs::write(&path_fallback, "binary").expect("path binary");
193
194        let resolved = resolve_ast_grep_binary_with_sources(
195            Some(env_path.clone()),
196            vec![managed_path],
197            Some(path_fallback),
198        );
199
200        assert_eq!(resolved, Some(env_path));
201    }
202
203    #[test]
204    fn resolution_prefers_managed_binary_before_path() {
205        let temp_dir = TempDir::new().expect("temp dir");
206        let managed_path = temp_dir.path().join("managed-ast-grep");
207        std::fs::write(&managed_path, "binary").expect("managed binary");
208
209        let path_fallback = temp_dir.path().join("path-ast-grep");
210        std::fs::write(&path_fallback, "binary").expect("path binary");
211
212        let resolved = resolve_ast_grep_binary_with_sources(
213            None,
214            vec![managed_path.clone()],
215            Some(path_fallback),
216        );
217
218        assert_eq!(resolved, Some(managed_path));
219    }
220
221    #[test]
222    fn resolution_uses_path_fallback_when_needed() {
223        let temp_dir = TempDir::new().expect("temp dir");
224        let path_fallback = temp_dir.path().join("path-ast-grep");
225        std::fs::write(&path_fallback, "binary").expect("path binary");
226
227        let resolved = resolve_ast_grep_binary_with_sources(
228            Some(PathBuf::from("/missing/env-ast-grep")),
229            vec![PathBuf::from("/missing/managed-ast-grep")],
230            Some(path_fallback.clone()),
231        );
232
233        assert_eq!(resolved, Some(path_fallback));
234    }
235
236    #[test]
237    fn resolution_uses_test_override_when_present() {
238        let temp_dir = TempDir::new().expect("temp dir");
239        let override_path = temp_dir.path().join("override-ast-grep");
240        std::fs::write(&override_path, "binary").expect("override binary");
241        let _guard = set_ast_grep_binary_override_for_tests(Some(override_path.clone()));
242
243        let resolved = resolve_ast_grep_binary_from_env_and_fs();
244
245        assert_eq!(resolved, Some(override_path));
246    }
247
248    #[test]
249    fn resolution_can_be_forced_missing_in_tests() {
250        let _guard = set_ast_grep_binary_override_for_tests(None);
251
252        let resolved = resolve_ast_grep_binary_from_env_and_fs();
253
254        assert_eq!(resolved, None);
255    }
256}
vtcode_core/tools/ast_grep_binary.rs

vtcode_core/tools/
ast_grep_binary.rs
