vtcode_core/tools/
grep_search.rs1use anyhow::Result;
15use serde_json;
16use std::num::NonZeroUsize;
17use std::path::PathBuf;
18use std::sync::Arc;
19use std::sync::Mutex;
20use std::sync::atomic::AtomicBool;
21use std::sync::atomic::Ordering;
22use std::thread;
23use std::time::Duration;
24
25const MAX_SEARCH_RESULTS: NonZeroUsize = NonZeroUsize::new(100).unwrap();
27
28const NUM_SEARCH_THREADS: NonZeroUsize = NonZeroUsize::new(2).unwrap();
30
31const SEARCH_DEBOUNCE: Duration = Duration::from_millis(150);
34
35const ACTIVE_SEARCH_COMPLETE_POLL_INTERVAL: Duration = Duration::from_millis(20);
37
38#[derive(Debug, Clone)]
40pub struct GrepSearchInput {
41 pub pattern: String,
42 pub path: String,
43 pub case_sensitive: Option<bool>,
44 pub literal: Option<bool>,
45 pub glob_pattern: Option<String>,
46 pub context_lines: Option<usize>,
47 pub include_hidden: Option<bool>,
48 pub max_results: Option<usize>,
49}
50
51#[derive(Debug, Clone)]
53pub struct GrepSearchResult {
54 pub query: String,
55 pub matches: Vec<serde_json::Value>,
56}
57
58pub struct GrepSearchManager {
60 state: Arc<Mutex<SearchState>>,
62
63 search_dir: PathBuf,
64}
65
66struct SearchState {
67 latest_query: String,
69
70 is_search_scheduled: bool,
72
73 active_search: Option<ActiveSearch>,
75 last_result: Option<GrepSearchResult>,
76}
77
78struct ActiveSearch {
79 query: String,
80 cancellation_token: Arc<AtomicBool>,
81}
82
83impl GrepSearchManager {
84 pub fn new(search_dir: PathBuf) -> Self {
85 Self {
86 state: Arc::new(Mutex::new(SearchState {
87 latest_query: String::new(),
88 is_search_scheduled: false,
89 active_search: None,
90 last_result: None,
91 })),
92 search_dir,
93 }
94 }
95
96 pub fn on_user_query(&self, query: String) {
98 {
99 #[expect(clippy::unwrap_used)]
100 let mut st = self.state.lock().unwrap();
101 if query == st.latest_query {
102 return;
104 }
105
106 st.latest_query.clear();
108 st.latest_query.push_str(&query);
109
110 if let Some(active_search) = &st.active_search
113 && !query.starts_with(&active_search.query)
114 {
115 active_search
116 .cancellation_token
117 .store(true, Ordering::Relaxed);
118 st.active_search = None;
119 }
120
121 if !st.is_search_scheduled {
123 st.is_search_scheduled = true;
124 } else {
125 return;
126 }
127 }
128
129 let state = self.state.clone();
133 let search_dir = self.search_dir.clone();
134 thread::spawn(move || {
135 thread::sleep(SEARCH_DEBOUNCE);
138 loop {
139 #[expect(clippy::unwrap_used)]
140 if state.lock().unwrap().active_search.is_none() {
141 break;
142 }
143 thread::sleep(ACTIVE_SEARCH_COMPLETE_POLL_INTERVAL);
144 }
145
146 let cancellation_token = Arc::new(AtomicBool::new(false));
149 let token = cancellation_token.clone();
150 let query = {
151 #[expect(clippy::unwrap_used)]
152 let mut st = state.lock().unwrap();
153 let query = st.latest_query.clone();
154 st.is_search_scheduled = false;
155 st.active_search = Some(ActiveSearch {
156 query: query.clone(),
157 cancellation_token: token,
158 });
159 query
160 };
161
162 GrepSearchManager::spawn_rp_search(query, search_dir, cancellation_token, state);
163 });
164 }
165
166 pub fn last_result(&self) -> Option<GrepSearchResult> {
168 #[expect(clippy::unwrap_used)]
169 let st = self.state.lock().unwrap();
170 st.last_result.clone()
171 }
172
173 fn spawn_rp_search(
174 query: String,
175 search_dir: PathBuf,
176 cancellation_token: Arc<AtomicBool>,
177 search_state: Arc<Mutex<SearchState>>,
178 ) {
179 use std::process::Command;
180
181 thread::spawn(move || {
182 if cancellation_token.load(Ordering::Relaxed) {
184 {
186 #[expect(clippy::unwrap_used)]
187 let mut st = search_state.lock().unwrap();
188 if let Some(active_search) = &st.active_search
189 && Arc::ptr_eq(&active_search.cancellation_token, &cancellation_token)
190 {
191 st.active_search = None;
192 }
193 }
194 return;
195 }
196
197 let mut cmd = Command::new("rg");
199 cmd.arg("-j").arg(NUM_SEARCH_THREADS.get().to_string());
200
201 cmd.arg(&query);
203
204 cmd.arg(search_dir.to_string_lossy().as_ref());
206
207 cmd.arg("--json");
209
210 cmd.arg("--max-count")
212 .arg(MAX_SEARCH_RESULTS.get().to_string());
213
214 let output = cmd.output();
216
217 let is_cancelled = cancellation_token.load(Ordering::Relaxed);
218 if !is_cancelled {
219 if let Ok(output) = output {
220 if output.status.success() {
221 let output_str = String::from_utf8_lossy(&output.stdout);
222 let mut matches = Vec::new();
223 for line in output_str.lines() {
224 if let Ok(val) = serde_json::from_str::<serde_json::Value>(line) {
225 matches.push(val);
226 }
227 }
228 let result = GrepSearchResult {
229 query: query.clone(),
230 matches,
231 };
232 #[expect(clippy::unwrap_used)]
233 let mut st = search_state.lock().unwrap();
234 st.last_result = Some(result);
235 }
236 }
237 }
238
239 {
241 #[expect(clippy::unwrap_used)]
242 let mut st = search_state.lock().unwrap();
243 if let Some(active_search) = &st.active_search
244 && Arc::ptr_eq(&active_search.cancellation_token, &cancellation_token)
245 {
246 st.active_search = None;
247 }
248 }
249 });
250 }
251
252 pub async fn perform_search(&self, input: GrepSearchInput) -> Result<GrepSearchResult> {
254 use std::process::Command;
256
257 let mut cmd = Command::new("rg");
259
260 cmd.arg(&input.pattern);
262
263 cmd.arg(&input.path);
265
266 if let Some(case_sensitive) = input.case_sensitive {
268 if case_sensitive {
269 cmd.arg("--case-sensitive");
270 } else {
271 cmd.arg("--ignore-case");
272 }
273 }
274
275 if let Some(literal) = input.literal {
276 if literal {
277 cmd.arg("--fixed-strings");
278 }
279 }
280
281 if let Some(glob_pattern) = &input.glob_pattern {
282 cmd.arg("--glob").arg(glob_pattern);
283 }
284
285 if let Some(context_lines) = input.context_lines {
286 cmd.arg("--context").arg(context_lines.to_string());
287 }
288
289 if let Some(include_hidden) = input.include_hidden {
290 if include_hidden {
291 cmd.arg("--hidden");
292 }
293 }
294
295 let max_results = input.max_results.unwrap_or(MAX_SEARCH_RESULTS.get());
297 cmd.arg("--max-count").arg(max_results.to_string());
298
299 cmd.arg("--json");
301
302 let output = cmd.output()?;
304
305 if !output.status.success() {
306 if String::from_utf8_lossy(&output.stderr).contains("not found") {
308 return Err(anyhow::anyhow!(
309 "ripgrep (rg) command not found. Please install ripgrep to use search functionality."
310 ));
311 }
312
313 }
315
316 let output_str = String::from_utf8_lossy(&output.stdout);
318 let mut matches = Vec::new();
319
320 for line in output_str.lines() {
321 if !line.trim().is_empty() {
322 if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(line) {
323 matches.push(json_value);
324 }
325 }
326 }
327
328 Ok(GrepSearchResult {
329 query: input.pattern,
330 matches,
331 })
332 }
333}