vtcode_core/tools/
ast_grep_binary.rs

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

vtcode_core/tools/
ast_grep_binary.rs
