vtcode_core/tools/
outline_search.rs

1//! Bounded, read-only declaration discovery for `code_search`.
2
3use anyhow::{Context, Result, bail};
4use serde::Deserialize;
5use std::path::{Path, PathBuf};
6use std::process::Stdio;
7use tokio::io::BufReader;
8use tokio::process::Command;
9
10use super::grep_backend::{BoundedRecordRead, read_bounded_record};
11use crate::tools::ast_grep_installer::AstGrepStatus;
12use crate::tools::ast_grep_language::AstGrepLanguage;
13use vtcode_commons::canonicalize;
14use vtcode_commons::exclusions::is_sensitive_file;
15use vtcode_commons::walk::{build_walker_single_threaded, is_excluded_dir};
16const CODE_SEARCH_OUTLINE_BYTE_CAP: usize = 1024 * 1024;
17const CODE_SEARCH_OUTLINE_PATH_BATCH_SIZE: usize = 64;
18const CODE_SEARCH_OUTLINE_ARG_BATCH_BYTE_CAP: usize = 16 * 1024;
19const CODE_SEARCH_OUTLINE_FIXED_ARGS: [&str; 4] = ["outline", "--json=stream", "--items", "all"];
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub(crate) struct DeclarationRange {
23    pub byte_start: usize,
24    pub byte_end: usize,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub(crate) struct DeclarationRecord {
29    pub name: String,
30    pub range: DeclarationRange,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub(crate) struct DeclarationFileRecord {
35    pub path: PathBuf,
36    pub language: AstGrepLanguage,
37    pub declarations: Vec<DeclarationRecord>,
38    pub complete: bool,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
42pub(crate) struct DeclarationSearchOutcome {
43    pub files: Vec<DeclarationFileRecord>,
44    pub stream_complete: bool,
45    pub truncated: bool,
46}
47
48async fn kill_and_reap_declaration_child(child: &mut tokio::process::Child) {
49    let _ = child.start_kill();
50    let _ = child.wait().await;
51}
52
53fn smart_case_eq(left: &str, query: &str) -> bool {
54    if query.chars().any(char::is_uppercase) {
55        left == query
56    } else {
57        left.to_lowercase() == query.to_lowercase()
58    }
59}
60
61fn matching_declarations(file: &OutlineFile, query: &str, candidate_cap: usize) -> (Vec<DeclarationRecord>, bool) {
62    let matching = file
63        .items
64        .iter()
65        .filter(|item| !item.is_import)
66        .flat_map(|item| {
67            std::iter::once((&item.name, &item.range))
68                .chain(item.members.iter().map(|member| (&member.name, &member.range)))
69        })
70        .filter(|(name, _)| smart_case_eq(name, query))
71        .collect::<Vec<_>>();
72    let matching_count = matching.len();
73    let records = matching
74        .into_iter()
75        .filter_map(|(name, range)| {
76            let range = range.as_ref()?;
77            Some(DeclarationRecord {
78                name: name.clone(),
79                range: DeclarationRange {
80                    byte_start: usize::try_from(range.byte_offset.start).ok()?,
81                    byte_end: usize::try_from(range.byte_offset.end).ok()?,
82                },
83            })
84        })
85        .collect::<Vec<_>>();
86    let complete = records.len() == matching_count && records.len() <= candidate_cap;
87    (records.into_iter().take(candidate_cap).collect(), complete)
88}
89
90/// Stream recognised declarations using only an already installed outline
91/// executable. No installation or cache mutation is attempted.
92pub(crate) async fn search_declarations_bounded(
93    workspace_root: &Path,
94    resolved_path: &Path,
95    query: &str,
96    languages: &[AstGrepLanguage],
97    candidate_cap: usize,
98) -> Result<DeclarationSearchOutcome> {
99    let binary = match AstGrepStatus::check() {
100        AstGrepStatus::Available { binary, .. } => binary,
101        AstGrepStatus::NotFound => bail!("definition search is unavailable"),
102        AstGrepStatus::Error { .. } => bail!("definition search is unavailable"),
103    };
104    let mut command_paths = outline_paths(workspace_root, resolved_path, languages).peekable();
105    let mut bytes_read = 0usize;
106    let mut retained = 0usize;
107    let mut files = Vec::new();
108    let mut truncated = false;
109    while command_paths.peek().is_some() && !truncated {
110        let command_args = next_outline_path_batch(&binary, &mut command_paths)?;
111        let mut command = Command::new(&binary);
112        command
113            .current_dir(workspace_root)
114            .env("RAYON_NUM_THREADS", "1")
115            .args(CODE_SEARCH_OUTLINE_FIXED_ARGS)
116            .args(command_args)
117            .stdout(Stdio::piped())
118            .stderr(Stdio::null());
119
120        let mut child = command.spawn().context("failed to run definition search")?;
121        let Some(stdout) = child.stdout.take() else {
122            kill_and_reap_declaration_child(&mut child).await;
123            bail!("failed to capture definition search output");
124        };
125        let mut reader = BufReader::new(stdout);
126        let mut line_buf = Vec::new();
127
128        loop {
129            line_buf.clear();
130            let read =
131                match read_bounded_record(&mut reader, &mut line_buf, &mut bytes_read, CODE_SEARCH_OUTLINE_BYTE_CAP)
132                    .await
133                {
134                    Ok(read) => read,
135                    Err(error) => {
136                        drop(reader);
137                        kill_and_reap_declaration_child(&mut child).await;
138                        return Err(error).context("failed to read definition stream");
139                    }
140                };
141            match read {
142                BoundedRecordRead::Record => {}
143                BoundedRecordRead::Eof => break,
144                BoundedRecordRead::Exhausted => {
145                    truncated = true;
146                    break;
147                }
148            }
149            let file = match serde_json::from_slice::<OutlineFile>(&line_buf) {
150                Ok(file) => file,
151                Err(error) => {
152                    drop(reader);
153                    kill_and_reap_declaration_child(&mut child).await;
154                    return Err(error).context("failed to parse definition stream record");
155                }
156            };
157            let path = PathBuf::from(&file.path);
158            let Some(language) =
159                AstGrepLanguage::from_path(&path).or_else(|| AstGrepLanguage::from_user_value(&file.lang))
160            else {
161                continue;
162            };
163            if !languages.is_empty() && !languages.contains(&language) {
164                continue;
165            }
166            let remaining = candidate_cap.saturating_sub(retained);
167            let (declarations, complete) = matching_declarations(&file, query, remaining);
168            retained = retained.saturating_add(declarations.len());
169            files.push(DeclarationFileRecord { path, language, declarations, complete });
170            if !complete || retained >= candidate_cap {
171                truncated = true;
172                break;
173            }
174        }
175
176        drop(reader);
177        if truncated {
178            let _ = child.start_kill();
179        }
180        let status = child.wait().await.context("failed to reap definition search process")?;
181        if !truncated && !status.success() {
182            bail!("definition search failed");
183        }
184    }
185
186    Ok(DeclarationSearchOutcome { files, stream_complete: !truncated, truncated })
187}
188
189fn outline_paths<'a>(
190    workspace_root: &'a Path,
191    resolved_path: &'a Path,
192    languages: &'a [AstGrepLanguage],
193) -> Box<dyn Iterator<Item = String> + Send + 'a> {
194    if resolved_path.is_file() {
195        return Box::new(std::iter::once(command_path_arg(workspace_root, resolved_path)));
196    }
197    let mut builder = build_walker_single_threaded(resolved_path);
198    builder.filter_entry(|entry| !is_excluded_dir(entry));
199    builder.sort_by_file_path(|left, right| left.cmp(right));
200    let paths = builder
201        .build()
202        .filter_map(std::result::Result::ok)
203        .filter(|entry| entry.file_type().is_some_and(|kind| kind.is_file()))
204        .map(|entry| entry.into_path())
205        .filter(|path| {
206            !path.file_name().and_then(|name| name.to_str()).is_some_and(is_sensitive_file)
207                && (languages.is_empty()
208                    || AstGrepLanguage::from_path(path).is_some_and(|language| languages.contains(&language)))
209        })
210        .map(|path| command_path_arg(workspace_root, &path));
211    Box::new(paths)
212}
213
214fn next_outline_path_batch(
215    executable: &Path,
216    paths: &mut std::iter::Peekable<impl Iterator<Item = String>>,
217) -> Result<Vec<String>> {
218    next_outline_path_batch_with_cap(executable, paths, CODE_SEARCH_OUTLINE_ARG_BATCH_BYTE_CAP)
219}
220
221fn next_outline_path_batch_with_cap(
222    executable: &Path,
223    paths: &mut std::iter::Peekable<impl Iterator<Item = String>>,
224    arg_byte_cap: usize,
225) -> Result<Vec<String>> {
226    let mut batch = Vec::with_capacity(CODE_SEARCH_OUTLINE_PATH_BATCH_SIZE);
227    let mut total_arg_bytes = arg_os_bytes(executable.as_os_str());
228    total_arg_bytes = CODE_SEARCH_OUTLINE_FIXED_ARGS
229        .iter()
230        .fold(total_arg_bytes, |total, arg| total.saturating_add(arg_bytes(arg)));
231    while batch.len() < CODE_SEARCH_OUTLINE_PATH_BATCH_SIZE {
232        let Some(next) = paths.peek() else {
233            break;
234        };
235        let next_bytes = arg_bytes(next);
236        if total_arg_bytes.saturating_add(next_bytes) > arg_byte_cap {
237            if batch.is_empty() {
238                bail!("definition search path exceeds the command argument byte limit");
239            }
240            break;
241        }
242        let Some(next) = paths.next() else {
243            break;
244        };
245        total_arg_bytes = total_arg_bytes.saturating_add(next_bytes);
246        batch.push(next);
247    }
248    Ok(batch)
249}
250
251fn arg_bytes(arg: &str) -> usize {
252    arg.len().saturating_add(1)
253}
254
255fn arg_os_bytes(arg: &std::ffi::OsStr) -> usize {
256    arg.as_encoded_bytes().len().saturating_add(1)
257}
258
259/// Build the path argument passed to ast-grep. Use the workspace-relative form
260/// when possible so the emitted `path` field is relative and readable.
261fn command_path_arg(workspace_root: &Path, resolved: &Path) -> String {
262    let workspace_canonical = canonicalize(workspace_root).unwrap_or_else(|_| workspace_root.to_path_buf());
263    if let Ok(relative) = resolved.strip_prefix(&workspace_canonical) {
264        if relative.as_os_str().is_empty() {
265            ".".to_string()
266        } else {
267            relative.to_string_lossy().replace('\\', "/")
268        }
269    } else {
270        resolved.to_string_lossy().to_string()
271    }
272}
273
274// ---------------------------------------------------------------------------
275// Tolerant deserialization of the ast-grep outline JSON stream.
276//
277// The outline JSON is an alpha preview surface in ast-grep 0.44.0. Unknown
278// keys may appear in future versions, so every struct uses `#[serde(default)]`
279// and none set `deny_unknown_fields`. Missing fields degrade gracefully
280// instead of failing the whole call.
281// ---------------------------------------------------------------------------
282
283#[derive(Debug, Default, Clone, Deserialize)]
284struct OutlineFile {
285    #[serde(default)]
286    path: String,
287    #[serde(default, rename = "language")]
288    lang: String,
289    #[serde(default)]
290    items: Vec<OutlineItem>,
291}
292
293/// Source range as reported by ast-grep outline. All line/column values are
294/// zero-based in the raw stream. We expose the raw range plus a derived
295/// 1-based `lineRange` in the `full` view so callers can feed the lines
296/// straight to `file_operation` `read` (`offset_lines` is 1-based, inclusive).
297///
298/// Forward-compat tolerant: `#[serde(default)]` and no `deny_unknown_fields`,
299/// so unknown keys from future ast-grep versions are ignored instead of
300/// failing the whole call.
301#[derive(Debug, Default, Clone, Deserialize)]
302struct OutlineRange {
303    #[serde(default, rename = "byteOffset")]
304    byte_offset: OutlineByteOffset,
305}
306
307#[derive(Debug, Default, Clone, Deserialize)]
308struct OutlineByteOffset {
309    #[serde(default)]
310    start: u64,
311    #[serde(default)]
312    end: u64,
313}
314
315#[derive(Debug, Default, Clone, Deserialize)]
316struct OutlineItem {
317    #[serde(default)]
318    name: String,
319    #[serde(default, rename = "isImport")]
320    is_import: bool,
321    #[serde(default)]
322    range: Option<OutlineRange>,
323    #[serde(default)]
324    members: Vec<OutlineMember>,
325}
326
327#[derive(Debug, Default, Clone, Deserialize)]
328struct OutlineMember {
329    #[serde(default)]
330    name: String,
331    #[serde(default)]
332    range: Option<OutlineRange>,
333}
334
335#[cfg(test)]
336mod tests;
vtcode_core/tools/outline_search.rs

vtcode_core/tools/
outline_search.rs
