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

vtcode_core/tools/
outline_search.rs
