1use super::file_search_bridge::{self, FileSearchConfig};
15use super::grep_cache::GrepSearchCache;
16use crate::cache::estimate_json_size;
17use anyhow::Result;
18use serde_json::Value;
19use std::path::PathBuf;
20use std::sync::Arc;
21use std::sync::Mutex;
22use std::sync::atomic::AtomicBool;
23use std::sync::atomic::Ordering;
24use std::thread;
25use std::time::Duration;
26use tracing::warn;
27
28pub(crate) use super::grep_backend::*;
31
32const SEARCH_DEBOUNCE: Duration = Duration::from_millis(150);
35
36const ACTIVE_SEARCH_COMPLETE_POLL_INTERVAL: Duration = Duration::from_millis(20);
38
39pub struct GrepSearchManager {
41 state: Arc<Mutex<SearchState>>,
43
44 search_dir: PathBuf,
45
46 cache: Arc<GrepSearchCache>,
48}
49
50struct SearchState {
51 latest_query: String,
53
54 is_search_scheduled: bool,
56
57 active_search: Option<ActiveSearch>,
59 last_result: Option<GrepSearchResult>,
60}
61
62struct ActiveSearch {
63 query: String,
64 cancellation_token: Arc<AtomicBool>,
65}
66
67impl GrepSearchManager {
68 pub fn new(search_dir: PathBuf) -> Self {
69 Self {
70 state: Arc::new(Mutex::new(SearchState {
71 latest_query: String::new(),
72 is_search_scheduled: false,
73 active_search: None,
74 last_result: None,
75 })),
76 search_dir,
77 cache: Arc::new(GrepSearchCache::new(100)), }
79 }
80
81 fn cached_result(cache: &GrepSearchCache, input: &GrepSearchInput) -> Option<GrepSearchResult> {
82 cache.get(input).map(|cached| GrepSearchResult {
83 query: cached.query.clone(),
84 matches: cached.matches.clone(),
85 truncated: cached.truncated,
86 total_matches: cached.total_matches,
87 })
88 }
89
90 pub fn on_user_query(&self, query: &str) {
92 {
93 let mut st = match self.state.lock() {
94 Ok(state) => state,
95 Err(err) => {
96 warn!("grep search state lock poisoned while handling query update: {err}");
97 return;
98 }
99 };
100 if query != st.latest_query {
101 st.latest_query.clear();
102 st.latest_query.push_str(query);
103 } else {
104 return;
105 }
106
107 if let Some(active_search) = &st.active_search
110 && !query.starts_with(&active_search.query)
111 {
112 active_search.cancellation_token.store(true, Ordering::Relaxed);
113 st.active_search = None;
114 }
115
116 if !st.is_search_scheduled {
118 st.is_search_scheduled = true;
119 } else {
120 return;
121 }
122 }
123
124 let state = self.state.clone();
128 let search_dir = self.search_dir.clone();
129 let cache = self.cache.clone();
130 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
131 warn!("grep search query ignored: no active Tokio runtime");
132 if let Ok(mut st) = self.state.lock() {
133 st.is_search_scheduled = false;
134 st.latest_query.clear();
136 }
137 return;
138 };
139 let search_runtime = runtime.clone();
142 drop(runtime.spawn_blocking(move || {
143 thread::sleep(SEARCH_DEBOUNCE);
146 loop {
147 let active_is_none = match state.lock() {
148 Ok(st) => st.active_search.is_none(),
149 Err(err) => {
150 warn!("grep search state lock poisoned while waiting for active search: {err}");
151 return;
152 }
153 };
154 if active_is_none {
155 break;
156 }
157 thread::sleep(ACTIVE_SEARCH_COMPLETE_POLL_INTERVAL);
158 }
159
160 let cancellation_token = Arc::new(AtomicBool::new(false));
163 let token = cancellation_token.clone();
164 let query = {
165 let mut st = match state.lock() {
166 Ok(state) => state,
167 Err(err) => {
168 warn!("grep search state lock poisoned while preparing debounced search: {err}");
169 return;
170 }
171 };
172 let query = st.latest_query.clone();
173 st.is_search_scheduled = false;
174 st.active_search = Some(ActiveSearch { query: query.clone(), cancellation_token: token });
175 query
176 };
177
178 drop(search_runtime.spawn(GrepSearchManager::spawn_grep_file(
182 query,
183 search_dir,
184 cancellation_token,
185 state,
186 Some(cache),
187 )));
188 }));
189 }
190
191 pub fn last_result(&self) -> Option<GrepSearchResult> {
193 match self.state.lock() {
194 Ok(st) => st.last_result.clone(),
195 Err(err) => {
196 warn!("grep search state lock poisoned while reading last result: {err}");
197 None
198 }
199 }
200 }
201
202 async fn execute_with_backends(
203 input: &GrepSearchInput,
204 cancel_flag: Option<&AtomicBool>,
205 ) -> Result<(Vec<Value>, bool, usize)> {
206 let matches = run_ripgrep_backend_async(input, cancel_flag).await?;
207 Ok(Self::finalize_matches(matches, input))
208 }
209
210 fn finalize_matches(mut matches: Vec<Value>, input: &GrepSearchInput) -> (Vec<Value>, bool, usize) {
211 let mut truncated = false;
212 let max_results = input.max_results.unwrap_or(MAX_SEARCH_RESULTS.get());
213
214 if max_results == 0 {
215 return (Vec::new(), !matches.is_empty(), 0);
216 }
217
218 let total_match_count = matches
220 .iter()
221 .filter(|e| e.get("type").and_then(Value::as_str) == Some("match"))
222 .count();
223
224 let mut match_count = 0usize;
227 let mut cut_index = matches.len();
228 for (i, entry) in matches.iter().enumerate() {
229 let is_match = entry.get("type").and_then(Value::as_str).is_some_and(|t| t == "match");
230 if is_match {
231 match_count += 1;
232 if match_count >= max_results {
233 cut_index = i + 1;
236 for rest in matches.iter().skip(i + 1) {
238 let tp = rest.get("type").and_then(Value::as_str);
239 if tp == Some("context") {
240 cut_index += 1;
241 } else {
242 break;
243 }
244 }
245 break;
246 }
247 }
248 }
249 if matches[cut_index..]
251 .iter()
252 .any(|e| e.get("type").and_then(Value::as_str) == Some("match"))
253 {
254 truncated = true;
255 }
256 if cut_index < matches.len() {
257 matches.truncate(cut_index);
258 }
259
260 if let Some(limit) = input.max_result_bytes {
261 let mut total = 0usize;
262 let mut kept_count = 0;
263 for entry in &matches {
264 let entry_bytes = estimate_json_size(entry) as usize;
265 if total + entry_bytes > limit {
266 truncated = true;
267 break;
268 }
269 total += entry_bytes;
270 kept_count += 1;
271 }
272 matches.truncate(kept_count);
273 }
274
275 (matches, truncated, total_match_count)
276 }
277
278 async fn spawn_grep_file(
279 query: String,
280 search_dir: PathBuf,
281 cancellation_token: Arc<AtomicBool>,
282 search_state: Arc<Mutex<SearchState>>,
283 cache: Option<Arc<GrepSearchCache>>,
284 ) {
285 if cancellation_token.load(Ordering::Relaxed) {
287 {
289 let mut st = match search_state.lock() {
290 Ok(state) => state,
291 Err(err) => {
292 warn!("grep search state lock poisoned while cancelling search: {err}");
293 return;
294 }
295 };
296 if let Some(active_search) = &st.active_search
297 && Arc::ptr_eq(&active_search.cancellation_token, &cancellation_token)
298 {
299 st.active_search = None;
300 }
301 }
302 return;
303 }
304
305 let input = GrepSearchInput::with_defaults(query.clone(), search_dir.to_string_lossy().into_owned());
306
307 if let Some(ref cache) = cache
309 && let Some(cached_result) = Self::cached_result(cache, &input)
310 {
311 let mut st = match search_state.lock() {
312 Ok(state) => state,
313 Err(err) => {
314 warn!("grep search state lock poisoned while loading cached result: {err}");
315 return;
316 }
317 };
318 st.last_result = Some(cached_result);
319 return;
320 }
321
322 let timeout = input.timeout.unwrap_or(DEFAULT_SEARCH_TIMEOUT);
326 let search_result = match tokio::time::timeout(
327 timeout,
328 GrepSearchManager::execute_with_backends(&input, Some(&cancellation_token)),
329 )
330 .await
331 {
332 Ok(result) => result,
333 Err(_elapsed) => Err(anyhow::anyhow!("ripgrep search timed out after {}s", timeout.as_secs())),
334 };
335
336 let is_cancelled = cancellation_token.load(Ordering::Relaxed);
337 if !is_cancelled
338 && let Ok((matches, truncated, total_match_count)) = search_result
339 && !matches.is_empty()
340 {
341 let result = GrepSearchResult {
342 query,
343 matches,
344 truncated,
345 total_matches: if truncated { Some(total_match_count) } else { None },
346 };
347
348 if let Some(ref cache) = cache
350 && GrepSearchCache::should_cache(&result)
351 {
352 cache.put(&input, result.clone());
353 }
354
355 let mut st = match search_state.lock() {
356 Ok(state) => state,
357 Err(err) => {
358 warn!("grep search state lock poisoned while storing search result: {err}");
359 return;
360 }
361 };
362 st.last_result = Some(result);
363 }
364
365 {
367 let mut st = match search_state.lock() {
368 Ok(state) => state,
369 Err(err) => {
370 warn!("grep search state lock poisoned while clearing active search: {err}");
371 return;
372 }
373 };
374 if let Some(active_search) = &st.active_search
375 && Arc::ptr_eq(&active_search.cancellation_token, &cancellation_token)
376 {
377 st.active_search = None;
378 }
379 }
380 }
381
382 pub async fn perform_search(&self, input: GrepSearchInput) -> Result<GrepSearchResult> {
384 if let Some(cached_result) = Self::cached_result(&self.cache, &input) {
386 return Ok(cached_result);
387 }
388
389 let query = input.pattern.clone();
390 let input_clone = input.clone();
391
392 let timeout = input.timeout.unwrap_or(DEFAULT_SEARCH_TIMEOUT);
393 let outcome = tokio::time::timeout(timeout, Self::execute_with_backends(&input_clone, None)).await;
399 let (matches, truncated, total_match_count) = match outcome {
400 Ok(Ok(result)) => result,
401 Ok(Err(worker_err)) => {
402 return Err(worker_err.context("ripgrep search worker failed"));
403 }
404 Err(_elapsed) => {
405 return Err(anyhow::anyhow!(
406 "ripgrep search timed out after {}s; ripgrep was killed",
407 timeout.as_secs()
408 ));
409 }
410 };
411
412 let result = GrepSearchResult {
413 query,
414 matches,
415 truncated,
416 total_matches: if truncated { Some(total_match_count) } else { None },
417 };
418
419 if GrepSearchCache::should_cache(&result) {
421 self.cache.put(&input, result.clone());
422 }
423
424 Ok(result)
425 }
426
427 pub fn enumerate_files_with_pattern(
446 &self,
447 pattern: String,
448 max_results: usize,
449 cancel_flag: Option<Arc<AtomicBool>>,
450 ) -> Result<Vec<String>> {
451 let config = FileSearchConfig::new(pattern, self.search_dir.clone())
452 .with_limit(max_results)
453 .respect_gitignore(true);
454
455 let results = file_search_bridge::search_files(config, cancel_flag)?;
456
457 Ok(file_search_bridge::file_matches_only(results.matches)
458 .into_iter()
459 .map(|m| m.path)
460 .collect())
461 }
462
463 pub fn list_all_files(&self, max_results: usize, exclude_patterns: Vec<String>) -> Result<Vec<String>> {
477 let mut config = FileSearchConfig::new("".to_string(), self.search_dir.clone())
478 .with_limit(max_results)
479 .respect_gitignore(true);
480
481 for pattern in exclude_patterns {
482 config = config.exclude(pattern);
483 }
484
485 let results = file_search_bridge::search_files(config, None)?;
486
487 Ok(file_search_bridge::file_matches_only(results.matches)
488 .into_iter()
489 .map(|m| m.path)
490 .collect())
491 }
492}
493
494#[cfg(test)]
495mod tests {
496 use super::*;
497 use crate::tools::ast_grep_language::AstGrepLanguage;
498 use serde_json::json;
499 use tempfile::TempDir;
500 use tokio::io::BufReader;
501
502 #[tokio::test]
503 async fn code_search_literal_stream_reaps_at_candidate_cap() {
504 let workspace = TempDir::new().expect("workspace");
505 std::fs::write(workspace.path().join("matches.txt"), "Widget\nWidget\nWidget\n").expect("fixture");
506
507 let outcome = search_literal_bounded("Widget", workspace.path(), &[], 2)
508 .await
509 .expect("bounded literal search");
510
511 assert_eq!(outcome.candidates.len(), 2);
512 assert!(outcome.truncated);
513 }
514
515 #[tokio::test]
516 async fn code_search_literal_candidate_cap_selects_stable_path_prefix() {
517 let workspace = TempDir::new().expect("workspace");
518 for name in ["z.txt", "a.txt", "m.txt"] {
519 std::fs::write(workspace.path().join(name), "Widget\nWidget\n").expect("fixture");
520 }
521
522 let first = search_literal_bounded("Widget", workspace.path(), &[], 2)
523 .await
524 .expect("first bounded search");
525 let second = search_literal_bounded("Widget", workspace.path(), &[], 2)
526 .await
527 .expect("second bounded search");
528
529 assert_eq!(first, second);
530 assert!(first.truncated);
531 assert!(
532 first.candidates.iter().all(|candidate| candidate.path.ends_with("a.txt")),
533 "sorted cap prefix should come from a.txt: {first:?}"
534 );
535 }
536
537 #[tokio::test]
538 async fn code_search_literal_stream_reaps_at_byte_cap() {
539 let workspace = TempDir::new().expect("workspace");
540 std::fs::write(
541 workspace.path().join("large.txt"),
542 format!("Widget{}\n", "x".repeat(CODE_SEARCH_STREAM_BYTE_CAP)),
543 )
544 .expect("fixture");
545
546 let outcome = search_literal_bounded("Widget", workspace.path(), &[], 20)
547 .await
548 .expect("byte-bounded literal search");
549
550 assert!(outcome.candidates.is_empty());
551 assert!(outcome.truncated);
552 }
553
554 #[tokio::test]
555 async fn code_search_bounded_record_distinguishes_exact_eof_from_exhaustion() {
556 let mut exact_reader = BufReader::new(&b"{}\n"[..]);
557 let mut exact_record = Vec::with_capacity(3);
558 let mut exact_bytes_read = 0;
559 assert_eq!(
560 read_bounded_record(&mut exact_reader, &mut exact_record, &mut exact_bytes_read, 3,)
561 .await
562 .expect("exact record"),
563 BoundedRecordRead::Record
564 );
565 exact_record.clear();
566 assert_eq!(
567 read_bounded_record(&mut exact_reader, &mut exact_record, &mut exact_bytes_read, 3,)
568 .await
569 .expect("exact EOF probe"),
570 BoundedRecordRead::Eof
571 );
572
573 let mut oversized_reader = BufReader::new(&b"xxxx"[..]);
574 let mut oversized_record = Vec::with_capacity(3);
575 let mut oversized_bytes_read = 0;
576 assert_eq!(
577 read_bounded_record(&mut oversized_reader, &mut oversized_record, &mut oversized_bytes_read, 3,)
578 .await
579 .expect("oversized record"),
580 BoundedRecordRead::Exhausted
581 );
582 assert_eq!(oversized_record.len(), 3);
583 }
584
585 #[tokio::test]
586 async fn code_search_literal_language_prefilters_are_case_insensitive() {
587 let workspace = TempDir::new().expect("workspace");
588 std::fs::write(workspace.path().join("UPPER.RS"), "Widget\n").expect("Rust fixture");
589 std::fs::write(workspace.path().join("DOCKERFILE"), "Widget\n").expect("Dockerfile fixture");
590
591 let rust = search_literal_bounded("Widget", workspace.path(), &[AstGrepLanguage::Rust], 20)
592 .await
593 .expect("uppercase Rust extension");
594 assert!(rust.candidates.iter().any(|candidate| {
595 candidate.path.ends_with("UPPER.RS")
596 && AstGrepLanguage::from_path(&candidate.path) == Some(AstGrepLanguage::Rust)
597 }));
598
599 let dockerfile = search_literal_bounded("Widget", workspace.path(), &[AstGrepLanguage::Dockerfile], 20)
600 .await
601 .expect("uppercase Dockerfile name");
602 assert!(dockerfile.candidates.iter().any(|candidate| {
603 candidate.path.ends_with("DOCKERFILE")
604 && AstGrepLanguage::from_path(&candidate.path) == Some(AstGrepLanguage::Dockerfile)
605 }));
606 }
607
608 #[tokio::test]
609 async fn code_search_literal_alternation_matches_each_term() {
610 let workspace = TempDir::new().expect("workspace");
614 std::fs::write(workspace.path().join("a.rs"), "tokio\n").expect("fixture a");
615 std::fs::write(workspace.path().join("b.rs"), "async-std\n").expect("fixture b");
616 std::fs::write(workspace.path().join("c.rs"), "runtime\n").expect("fixture c");
617 std::fs::write(workspace.path().join("d.rs"), "unrelated\n").expect("fixture d");
618
619 let outcome = search_literal_bounded("tokio|async-std|runtime", workspace.path(), &[], 20)
620 .await
621 .expect("alternation search");
622
623 let matched_paths: Vec<String> = outcome
624 .candidates
625 .iter()
626 .map(|candidate| {
627 candidate
628 .path
629 .file_name()
630 .map(|name| name.to_string_lossy().into_owned())
631 .unwrap_or_default()
632 })
633 .collect();
634 assert!(matched_paths.contains(&"a.rs".to_string()), "tokio term should match: {matched_paths:?}");
635 assert!(matched_paths.contains(&"b.rs".to_string()), "async-std term should match: {matched_paths:?}");
636 assert!(matched_paths.contains(&"c.rs".to_string()), "runtime term should match: {matched_paths:?}");
637 assert!(
638 !matched_paths.contains(&"d.rs".to_string()),
639 "unrelated fixture should not match: {matched_paths:?}"
640 );
641 assert!(!outcome.truncated);
642 }
643
644 #[tokio::test]
645 async fn code_search_literal_single_term_still_uses_fixed_strings() {
646 let workspace = TempDir::new().expect("workspace");
649 std::fs::write(workspace.path().join("src.rs"), "fn main() {}\n").expect("fixture");
650
651 let outcome = search_literal_bounded("fn main()", workspace.path(), &[], 20)
652 .await
653 .expect("single-term literal search");
654
655 assert_eq!(outcome.candidates.len(), 1);
656 assert!(outcome.candidates[0].path.ends_with("src.rs"));
657 }
658
659 #[test]
660 fn split_alternation_handles_edge_cases() {
661 assert_eq!(split_alternation(""), Vec::<String>::new());
662 assert_eq!(split_alternation(" "), Vec::<String>::new());
663 assert_eq!(split_alternation("Widget"), vec!["Widget".to_string()]);
664 assert_eq!(split_alternation(" Widget "), vec!["Widget".to_string()]);
665 assert_eq!(
666 split_alternation("tokio|async-std|runtime"),
667 vec!["tokio".to_string(), "async-std".to_string(), "runtime".to_string()]
668 );
669 assert_eq!(split_alternation("|foo|"), vec!["foo".to_string()]);
671 assert_eq!(split_alternation("a||b"), vec!["a".to_string(), "b".to_string()]);
672 assert_eq!(split_alternation("||"), Vec::<String>::new());
674 }
675
676 #[test]
677 fn finalize_matches_respects_max_bytes() {
678 let mut input = GrepSearchInput::with_defaults("pat".into(), ".".into());
679 input.max_result_bytes = Some(100);
680 input.max_results = Some(5);
681
682 let matches = vec![json!({"text": "12345"}), json!({"text": "6789"})];
683
684 let (kept, truncated, _total) = GrepSearchManager::finalize_matches(matches, &input);
685 assert!(!truncated);
686 assert_eq!(kept.len(), 2);
687
688 input.max_result_bytes = Some(20);
690 let matches = vec![json!({"text": "12345"}), json!({"text": "6789"})];
691 let (kept, truncated, _total) = GrepSearchManager::finalize_matches(matches, &input);
692 assert!(truncated);
693 assert_eq!(kept.len(), 1); }
695
696 #[test]
697 fn finalize_matches_counts_only_match_type_entries() {
698 let mut input = GrepSearchInput::with_defaults("pat".into(), ".".into());
699 input.max_results = Some(2);
700
701 let matches = vec![
703 json!({"type": "begin", "data": {"path": {"text": "Cargo.lock"}}}),
704 json!({"type": "context", "data": {"line_number": 538, "lines": {"text": "ctx1"}}}),
705 json!({"type": "context", "data": {"line_number": 539, "lines": {"text": "ctx2"}}}),
706 json!({"type": "match", "data": {"line_number": 553, "lines": {"text": "match1"}}}),
707 json!({"type": "context", "data": {"line_number": 554, "lines": {"text": "ctx3"}}}),
708 json!({"type": "context", "data": {"line_number": 555, "lines": {"text": "ctx4"}}}),
709 json!({"type": "context", "data": {"line_number": 560, "lines": {"text": "ctx5"}}}),
710 json!({"type": "match", "data": {"line_number": 563, "lines": {"text": "match2"}}}),
711 json!({"type": "context", "data": {"line_number": 564, "lines": {"text": "ctx6"}}}),
712 json!({"type": "end", "data": {"path": {"text": "Cargo.lock"}}}),
713 ];
714
715 let (kept, truncated, total) = GrepSearchManager::finalize_matches(matches, &input);
716 assert!(!truncated);
719 assert_eq!(kept.len(), 9);
720 assert_eq!(kept[3]["type"], "match");
721 assert_eq!(kept[7]["type"], "match");
722 assert_eq!(total, 2);
723 }
724
725 #[test]
726 fn finalize_matches_truncates_when_more_match_types_than_limit() {
727 let mut input = GrepSearchInput::with_defaults("pat".into(), ".".into());
728 input.max_results = Some(1);
729
730 let matches = vec![
731 json!({"type": "begin", "data": {"path": {"text": "f.txt"}}}),
732 json!({"type": "match", "data": {"line_number": 1, "lines": {"text": "m1"}}}),
733 json!({"type": "context", "data": {"line_number": 2, "lines": {"text": "c1"}}}),
734 json!({"type": "match", "data": {"line_number": 10, "lines": {"text": "m2"}}}),
735 json!({"type": "context", "data": {"line_number": 11, "lines": {"text": "c2"}}}),
736 ];
737
738 let (kept, truncated, total) = GrepSearchManager::finalize_matches(matches, &input);
739 assert!(truncated);
740 assert_eq!(kept.len(), 3);
742 assert_eq!(kept[1]["type"], "match");
743 assert_eq!(kept[2]["type"], "context");
744 assert_eq!(total, 2); }
746
747 #[test]
748 fn test_grep_search_manager_creation() {
749 let manager = GrepSearchManager::new(PathBuf::from("."));
750 assert_eq!(manager.search_dir, PathBuf::from("."));
751 }
752
753 #[test]
754 fn on_user_query_without_runtime_does_not_panic_and_allows_retry() {
755 let manager = GrepSearchManager::new(PathBuf::from("."));
756 manager.on_user_query("first");
757 manager.on_user_query("first");
758 let state = manager.state.lock().expect("state lock");
759 assert!(state.latest_query.is_empty());
760 assert!(!state.is_search_scheduled);
761 }
762
763 #[test]
764 fn test_grep_search_input_new() {
765 let input = GrepSearchInput::new("pattern".to_string(), "/path/to/search".to_string());
766 assert_eq!(input.pattern, "pattern");
767 assert_eq!(input.path, "/path/to/search");
768 assert!(input.case_sensitive.is_none());
769 }
770
771 #[test]
772 fn test_grep_search_input_with_defaults() {
773 let input = GrepSearchInput::with_defaults("pattern".to_string(), "/path".to_string());
774 assert_eq!(input.pattern, "pattern");
775 assert_eq!(input.path, "/path");
776 assert_eq!(input.case_sensitive, Some(true));
777 assert_eq!(input.include_hidden, Some(false));
778 assert_eq!(input.max_results, Some(MAX_SEARCH_RESULTS.get()));
779 }
780}