diff --git a/src/mcp/tools/handlers/ast_grep_search.rs b/src/mcp/tools/handlers/ast_grep_search.rs index 49888e838..530905753 100644 --- a/src/mcp/tools/handlers/ast_grep_search.rs +++ b/src/mcp/tools/handlers/ast_grep_search.rs @@ -18,21 +18,13 @@ use crate::tracedecay::TraceDecay; use super::super::ToolResult; use super::super::render::{self, Md}; -use super::support::unique_file_paths; +use super::support::{CancelSearchOnDrop, unique_file_paths}; /// Hard cap on `max_results` regardless of what the caller requests. const MAX_RESULTS_CAP: usize = 200; /// Default `max_results` when the caller omits it. const DEFAULT_MAX_RESULTS: usize = 50; -struct CancelSearchOnDrop(Arc); - -impl Drop for CancelSearchOnDrop { - fn drop(&mut self) { - self.0.store(true, Ordering::Release); - } -} - async fn search_tree_off_thread( project_root: std::path::PathBuf, pattern: String, @@ -43,7 +35,7 @@ async fn search_tree_off_thread( ) -> Result { let query = pattern.clone(); let cancelled = Arc::new(AtomicBool::new(false)); - let cancel_on_drop = CancelSearchOnDrop(cancelled.clone()); + let cancel_on_drop = CancelSearchOnDrop::new(cancelled.clone()); let result = tokio::task::spawn_blocking(move || { search_tree_scoped_with_cancel( &project_root, @@ -224,7 +216,7 @@ mod tests { fn cancellation_guard_signals_worker_on_drop() { let cancelled = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); { - let _guard = CancelSearchOnDrop(cancelled.clone()); + let _guard = CancelSearchOnDrop::new(cancelled.clone()); } assert!(cancelled.load(std::sync::atomic::Ordering::Acquire)); } diff --git a/src/mcp/tools/handlers/grep.rs b/src/mcp/tools/handlers/grep.rs index 7ded9fc7b..d221b4d1d 100644 --- a/src/mcp/tools/handlers/grep.rs +++ b/src/mcp/tools/handlers/grep.rs @@ -7,19 +7,23 @@ //! matches symbol *names*, not file *content*. use std::fmt::Write as _; -use std::path::Path; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; use ignore::WalkBuilder; use ignore::overrides::{Override, OverrideBuilder}; use regex::{Regex, RegexBuilder}; use serde_json::{Value, json}; +use tokio::sync::Semaphore; use crate::errors::{Result, TraceDecayError}; use crate::tracedecay::TraceDecay; use super::super::ToolResult; use super::super::render::{self, Md}; -use super::support::{filter_by_scope, unique_file_paths}; +use super::support::{CancelSearchOnDrop, filter_by_scope, unique_file_paths}; /// Hard cap on `max_results` regardless of what the caller requests. const MAX_RESULTS_CAP: usize = 200; @@ -33,6 +37,12 @@ const MAX_HITS_PER_FILE: usize = 20; const BINARY_SNIFF_BYTES: usize = 8_192; /// Skip individual lines longer than this (minified bundles, embedded blobs). const MAX_LINE_BYTES: usize = 4_096; +/// Skip files too large for a bounded interactive content search. +const MAX_FILE_BYTES: u64 = 2_000_000; +/// Bound each grep request, including time spent waiting for a worker permit. +const GREP_SCAN_TIMEOUT: Duration = Duration::from_secs(10); +/// Keep concurrent blocking scans from monopolizing the daemon's worker pool. +static GREP_SCAN_SEMAPHORE: Semaphore = Semaphore::const_new(2); /// A single content-search hit, enriched with the enclosing graph symbol. struct GrepHit { @@ -72,7 +82,10 @@ pub(super) async fn handle_grep( .get("case_sensitive") .and_then(Value::as_bool) .unwrap_or(false); - let path_glob = args.get("path_glob").and_then(Value::as_str); + let path_glob = args + .get("path_glob") + .and_then(Value::as_str) + .map(str::to_owned); let max_results = args .get("max_results") .and_then(Value::as_u64) @@ -90,7 +103,7 @@ pub(super) async fn handle_grep( // Optional path filter. A caller-supplied glob whitelists candidate files // via the `ignore` crate's override mechanism (same glob semantics as a // `.gitignore` line), so it prunes at the walker level. - let overrides = match path_glob { + let overrides = match path_glob.as_deref() { Some(raw) if !raw.trim().is_empty() => { let mut builder = OverrideBuilder::new(&project_root); builder.add(raw).map_err(|err| TraceDecayError::Config { @@ -104,13 +117,15 @@ pub(super) async fn handle_grep( }; // Collect one extra past the cap so we can honestly report truncation. - let scan = scan_tree( - &project_root, - &matcher, + let scan = scan_tree_off_thread( + project_root, + matcher, overrides, + path_glob, context_lines, max_results, - ); + ) + .await?; // Scope filtering mirrors `tracedecay_search`: when the client pins a // subtree, only hits under it are returned. @@ -155,22 +170,148 @@ fn build_matcher(pattern: &str, fixed_strings: bool, case_sensitive: bool) -> Re }) } +async fn scan_tree_off_thread( + project_root: PathBuf, + matcher: Regex, + overrides: Option, + path_glob: Option, + context_lines: usize, + max_results: usize, +) -> Result { + let query = matcher.as_str().to_string(); + let cancelled = Arc::new(AtomicBool::new(false)); + let cancel_on_drop = CancelSearchOnDrop::new(cancelled.clone()); + let worker_cancelled = cancelled.clone(); + let worker_query = query.clone(); + + let result = tokio::time::timeout(GREP_SCAN_TIMEOUT, async move { + let permit = + GREP_SCAN_SEMAPHORE + .acquire() + .await + .map_err(|err| TraceDecayError::Search { + message: format!("grep scan concurrency gate closed: {err}"), + query: worker_query.clone(), + })?; + tokio::task::spawn_blocking(move || { + let _permit = permit; + scan_tree( + &project_root, + &matcher, + overrides, + path_glob.as_deref(), + context_lines, + max_results, + || worker_cancelled.load(Ordering::Acquire), + ) + }) + .await + .map_err(|err| TraceDecayError::Search { + message: format!("grep scan worker failed: {err}"), + query: worker_query, + }) + }) + .await; + + match result { + Ok(scan) => { + drop(cancel_on_drop); + scan + } + Err(_) => Err(TraceDecayError::Search { + message: format!( + "grep scan timed out after {} seconds; narrow the search with path_glob", + GREP_SCAN_TIMEOUT.as_secs() + ), + query, + }), + } +} + struct ScanResult { hits: Vec, files_scanned: usize, truncated: bool, } +struct GeneratedDirScope { + literal_prefix: PathBuf, + may_match_descendants: bool, +} + +impl GeneratedDirScope { + fn from_path_glob(path_glob: &str) -> Option { + let path_glob = path_glob.trim(); + if path_glob.is_empty() || path_glob.starts_with('!') { + return None; + } + let segments: Vec<&str> = path_glob + .trim_start_matches('/') + .split('/') + .filter(|segment| !segment.is_empty()) + .collect(); + let wildcard_start = segments + .iter() + .position(|segment| { + segment.contains('*') + || segment.contains('?') + || segment.contains('[') + || segment.contains('{') + }) + .unwrap_or(segments.len()); + let literal_prefix = + segments[..wildcard_start] + .iter() + .fold(PathBuf::new(), |mut prefix, segment| { + prefix.push(segment); + prefix + }); + let wildcard_suffix = &segments[wildcard_start..]; + let may_match_descendants = wildcard_suffix + .iter() + .enumerate() + .any(|(index, segment)| index > 0 || *segment == "**"); + + Some(Self { + literal_prefix, + may_match_descendants, + }) + } + + fn allows(&self, project_root: &Path, path: &Path) -> bool { + let Ok(relative) = path.strip_prefix(project_root) else { + return false; + }; + if self.literal_prefix.as_os_str().is_empty() { + return self.may_match_descendants; + } + self.literal_prefix.starts_with(relative) + || relative == self.literal_prefix + || (self.may_match_descendants && relative.starts_with(&self.literal_prefix)) + } +} + /// Walks the working tree respecting `.gitignore`, skipping binary files, and /// collects matching lines. Stops early once `max_results` + 1 hits are found /// so the caller can report truncation without scanning the whole tree. -fn scan_tree( +fn scan_tree( project_root: &Path, matcher: &Regex, overrides: Option, + path_glob: Option<&str>, context_lines: usize, max_results: usize, -) -> ScanResult { + is_cancelled: F, +) -> ScanResult +where + F: Fn() -> bool, +{ + let has_positive_override = overrides + .as_ref() + .is_some_and(|overrides| overrides.num_whitelists() > 0); + let generated_dir_overrides = overrides.clone(); + let generated_dir_scope = path_glob.and_then(GeneratedDirScope::from_path_glob); + let filter_root = project_root.to_path_buf(); let mut builder = WalkBuilder::new(project_root); builder .follow_links(false) @@ -178,7 +319,26 @@ fn scan_tree( .git_ignore(true) .git_global(true) .git_exclude(true) - .add_custom_ignore_filename(".gitignore"); + .add_custom_ignore_filename(".gitignore") + .filter_entry(move |entry| { + if entry.depth() == 0 { + return true; + } + let segment = entry.file_name().to_string_lossy(); + if segment == ".git" || segment == ".tracedecay" { + return false; + } + let requested_generated_dir = has_positive_override + && (generated_dir_overrides + .as_ref() + .is_some_and(|overrides| overrides.matched(entry.path(), true).is_whitelist()) + || generated_dir_scope + .as_ref() + .is_some_and(|scope| scope.allows(&filter_root, entry.path()))); + !entry.file_type().is_some_and(|kind| kind.is_dir()) + || requested_generated_dir + || !crate::config::is_generated_dir_segment(&segment) + }); if let Some(overrides) = overrides { builder.overrides(overrides); } @@ -189,6 +349,9 @@ fn scan_tree( let mut truncated = false; for entry in walker { + if is_cancelled() { + break; + } let Ok(entry) = entry else { continue }; let Some(ft) = entry.file_type() else { continue; @@ -202,6 +365,18 @@ fn scan_tree( }; let rel_str = rel.to_string_lossy().replace('\\', "/"); + if is_cancelled() { + break; + } + let Ok(metadata) = entry.metadata() else { + continue; + }; + if metadata.len() > MAX_FILE_BYTES { + continue; + } + if is_cancelled() { + break; + } let Ok(bytes) = std::fs::read(path) else { continue; }; @@ -216,6 +391,13 @@ fn scan_tree( let lines: Vec<&str> = content.lines().collect(); let mut file_hits = 0usize; for (idx, line) in lines.iter().enumerate() { + if is_cancelled() { + return ScanResult { + hits, + files_scanned, + truncated, + }; + } if line.len() > MAX_LINE_BYTES { continue; } @@ -353,3 +535,161 @@ fn render_grep_md(hits: &[GrepHit], truncated: bool, files_scanned: usize) -> St md.line(&summary); md.render() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scan_tree_prunes_generated_dependency_directories_without_gitignore() { + let project = tempfile::tempdir().expect("temp project"); + let generated = project.path().join(".venv/lib/python/site-packages/pkg"); + std::fs::create_dir_all(&generated).expect("generated fixture directory"); + std::fs::create_dir_all(project.path().join("src")).expect("source fixture directory"); + std::fs::create_dir_all(project.path().join(".tracedecay")) + .expect("metadata fixture directory"); + std::fs::write( + generated.join("generated.py"), + "UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .expect("generated fixture"); + std::fs::write( + project.path().join("src/tracked.rs"), + "// UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .expect("source fixture"); + std::fs::write( + project.path().join(".git"), + "gitdir: UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .expect("linked-worktree git file fixture"); + std::fs::write( + project.path().join(".tracedecay/internal.txt"), + "UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .expect("metadata fixture"); + + let matcher = Regex::new("UNIQUE_GENERATED_DIR_TOKEN").expect("matcher"); + let scan = scan_tree(project.path(), &matcher, None, None, 0, 10, || false); + let files = scan + .hits + .iter() + .map(|hit| hit.file.as_str()) + .collect::>(); + + assert!(files.contains(&"src/tracked.rs"), "{files:?}"); + assert!( + !files.iter().any(|file| file.starts_with(".venv/")), + "generated dependency trees must be pruned: {files:?}" + ); + assert!( + !files.contains(&".git"), + "git metadata must be pruned: {files:?}" + ); + assert!( + !files.iter().any(|file| file.starts_with(".tracedecay/")), + "TraceDecay metadata must be pruned: {files:?}" + ); + } + + #[test] + fn scan_tree_path_glob_prunes_unrelated_generated_directories() { + let project = tempfile::tempdir().expect("temp project"); + std::fs::create_dir_all(project.path().join("src")).expect("source fixture directory"); + std::fs::write( + project.path().join("src/selected.rs"), + "NORMAL_PATH_GLOB_TOKEN\n", + ) + .expect("source fixture"); + + let overrides = |root: &Path| { + let mut builder = OverrideBuilder::new(root); + builder.add("src/**").expect("path glob"); + builder.build().expect("overrides") + }; + let matcher = Regex::new("NORMAL_PATH_GLOB_TOKEN").expect("matcher"); + let baseline_checks = std::sync::atomic::AtomicUsize::new(0); + let baseline = scan_tree( + project.path(), + &matcher, + Some(overrides(project.path())), + Some("src/**"), + 0, + 10, + || { + baseline_checks.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + false + }, + ); + + std::fs::create_dir_all(project.path().join("target/generated")) + .expect("generated fixture directory"); + std::fs::write( + project.path().join("target/generated/unrelated.rs"), + "NORMAL_PATH_GLOB_TOKEN\n", + ) + .expect("generated fixture"); + + let checks = std::sync::atomic::AtomicUsize::new(0); + let scan = scan_tree( + project.path(), + &matcher, + Some(overrides(project.path())), + Some("src/**"), + 0, + 10, + || { + checks.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + false + }, + ); + + assert_eq!(scan.hits.len(), baseline.hits.len()); + assert_eq!( + checks.load(std::sync::atomic::Ordering::Relaxed), + baseline_checks.load(std::sync::atomic::Ordering::Relaxed), + "unrelated generated directories must not reach the scan loop" + ); + } + + #[test] + fn scan_tree_stops_when_cancelled_during_line_matching() { + let project = tempfile::tempdir().expect("temp project"); + let source = "CANCELLATION_TOKEN\n".repeat(100); + std::fs::write(project.path().join("fixture.txt"), source).expect("fixture"); + + let matcher = Regex::new("CANCELLATION_TOKEN").expect("matcher"); + let checks = std::sync::atomic::AtomicUsize::new(0); + let scan = scan_tree(project.path(), &matcher, None, None, 0, 200, || { + checks.fetch_add(1, Ordering::Relaxed) >= 10 + }); + + assert!( + scan.hits.len() < MAX_HITS_PER_FILE, + "cancelled scan should stop before the per-file cap: {}", + scan.hits.len() + ); + assert!(checks.load(Ordering::Relaxed) > 10); + } + + #[test] + fn scan_tree_skips_files_larger_than_two_megabytes() { + let project = tempfile::tempdir().expect("temp project"); + let mut oversized = b"OVERSIZED_FILE_TOKEN\n".to_vec(); + oversized.resize((MAX_FILE_BYTES as usize) + 1, b'x'); + std::fs::write(project.path().join("oversized.txt"), oversized).expect("oversized fixture"); + std::fs::write(project.path().join("tracked.txt"), "OVERSIZED_FILE_TOKEN\n") + .expect("source fixture"); + + let matcher = Regex::new("OVERSIZED_FILE_TOKEN").expect("matcher"); + let scan = scan_tree(project.path(), &matcher, None, None, 0, 10, || false); + let files = scan + .hits + .iter() + .map(|hit| hit.file.as_str()) + .collect::>(); + + assert!(files.contains(&"tracked.txt"), "{files:?}"); + assert!(!files.contains(&"oversized.txt"), "{files:?}"); + } +} diff --git a/src/mcp/tools/handlers/support.rs b/src/mcp/tools/handlers/support.rs index 98af06f5b..ad836305b 100644 --- a/src/mcp/tools/handlers/support.rs +++ b/src/mcp/tools/handlers/support.rs @@ -5,6 +5,8 @@ use std::collections::HashSet; use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use serde_json::{Value, json}; @@ -13,6 +15,20 @@ use super::super::render; use crate::errors::{Result, TraceDecayError}; use crate::global_db::{CodeProjectRecord, GlobalDb, ProjectRegistryContext}; +pub(super) struct CancelSearchOnDrop(Arc); + +impl CancelSearchOnDrop { + pub(super) fn new(cancelled: Arc) -> Self { + Self(cancelled) + } +} + +impl Drop for CancelSearchOnDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} + /// Trimmed, non-empty string argument by key, or `None` when absent, non-string, /// or blank after trimming. pub(super) fn string_arg<'a>(args: &'a Value, key: &str) -> Option<&'a str> { diff --git a/tests/mcp_suite/mcp_handler_test.rs b/tests/mcp_suite/mcp_handler_test.rs index 4ee73c794..b3b41f2d2 100644 --- a/tests/mcp_suite/mcp_handler_test.rs +++ b/tests/mcp_suite/mcp_handler_test.rs @@ -2006,6 +2006,84 @@ async fn test_grep_respects_gitignore() { ); } +#[tokio::test] +async fn test_grep_prunes_generated_dependency_directories_without_gitignore() { + let (cg, _dir) = setup_project().await; + let root = cg.project_root().to_path_buf(); + fs::create_dir_all(root.join(".venv/lib/python/site-packages/pkg")).unwrap(); + fs::write( + root.join(".venv/lib/python/site-packages/pkg/generated.py"), + "UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .unwrap(); + fs::write( + root.join("src/tracked.rs"), + "// UNIQUE_GENERATED_DIR_TOKEN\n", + ) + .unwrap(); + + let result = handle_tool_call( + &cg, + "tracedecay_grep", + json!({"pattern": "UNIQUE_GENERATED_DIR_TOKEN"}), + None, + None, + ) + .await + .unwrap(); + let payload = extract_json(&result.value); + let files: Vec<&str> = payload["results"] + .as_array() + .unwrap() + .iter() + .map(|hit| hit["file"].as_str().unwrap()) + .collect(); + assert!( + files.contains(&"src/tracked.rs"), + "source file should match: {payload}" + ); + assert!( + !files.iter().any(|file| file.starts_with(".venv/")), + "generated dependency trees must be pruned even without .gitignore coverage: {payload}" + ); +} + +#[tokio::test] +async fn test_grep_path_glob_includes_explicit_generated_directory() { + let (cg, _dir) = setup_project().await; + let root = cg.project_root().to_path_buf(); + fs::create_dir_all(root.join("dist")).unwrap(); + fs::write( + root.join("dist/generated.js"), + "UNIQUE_GENERATED_WHITELIST_TOKEN\n", + ) + .unwrap(); + + let result = handle_tool_call( + &cg, + "tracedecay_grep", + json!({ + "pattern": "UNIQUE_GENERATED_WHITELIST_TOKEN", + "path_glob": "dist/**" + }), + None, + None, + ) + .await + .unwrap(); + let payload = extract_json(&result.value); + let files: Vec<&str> = payload["results"] + .as_array() + .unwrap() + .iter() + .map(|hit| hit["file"].as_str().unwrap()) + .collect(); + assert!( + files.contains(&"dist/generated.js"), + "explicit path_glob should include generated directory: {payload}" + ); +} + #[tokio::test] async fn test_grep_skips_binary_files() { let (cg, _dir) = setup_project().await;