From affbdaadd640387a3e4976faaba75c1c88efe923 Mon Sep 17 00:00:00 2001 From: Kh05ifr4nD Date: Sun, 2 Aug 2026 20:28:30 +0800 Subject: [PATCH] fix(mcp): preserve multi_grep constraints and OR results --- Cargo.lock | 1 + crates/fff-mcp/Cargo.toml | 3 + crates/fff-mcp/src/cursor.rs | 92 +++++- crates/fff-mcp/src/server.rs | 387 ++++++++++++++++++++++---- crates/fff-query-parser/src/lib.rs | 23 ++ crates/fff-query-parser/src/parser.rs | 19 +- 6 files changed, 457 insertions(+), 68 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bfe636b18..2d1be8b4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -662,6 +662,7 @@ dependencies = [ "schemars", "serde", "serde_json", + "tempfile", "tokio", "tracing", ] diff --git a/crates/fff-mcp/Cargo.toml b/crates/fff-mcp/Cargo.toml index 636f94bfc..c73f3e303 100644 --- a/crates/fff-mcp/Cargo.toml +++ b/crates/fff-mcp/Cargo.toml @@ -30,3 +30,6 @@ tokio = { version = "1", features = ["full"] } tracing = { workspace = true } git2 = { workspace = true } clap = { version = "4", features = ["derive", "env"] } + +[dev-dependencies] +tempfile = "3.8" diff --git a/crates/fff-mcp/src/cursor.rs b/crates/fff-mcp/src/cursor.rs index fc5ad651c..5e488bf96 100644 --- a/crates/fff-mcp/src/cursor.rs +++ b/crates/fff-mcp/src/cursor.rs @@ -3,15 +3,79 @@ //! Maintains an in-memory map of opaque cursor IDs to file offsets. //! Cursors are evicted LRU-style when the store exceeds capacity. -use std::collections::{HashMap, VecDeque}; +use fff::grep::GrepMatch; +use std::collections::{HashMap, HashSet, VecDeque}; const MAX_CURSORS: usize = 20; +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct MatchKey { + pub(crate) file_path: String, + pub(crate) line_number: u64, + pub(crate) byte_offset: u64, +} + +impl MatchKey { + pub(crate) fn new(file_path: String, line_number: u64, byte_offset: u64) -> Self { + Self { + file_path, + line_number, + byte_offset, + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct PendingMatch { + pub(crate) file_path: String, + pub(crate) match_data: GrepMatch, +} + +#[derive(Clone, Debug)] +pub(crate) struct MultiGrepCursor { + pub(crate) patterns: Vec, + pub(crate) constraints: String, + pub(crate) next_offsets: Vec>, + pub(crate) pending: Vec, + pub(crate) seen: HashSet, + pub(crate) seen_files: HashSet, +} + +impl MultiGrepCursor { + pub(crate) fn new(patterns: Vec, constraints: String) -> Self { + Self { + next_offsets: vec![Some(0); patterns.len()], + patterns, + constraints, + pending: Vec::new(), + seen: HashSet::new(), + seen_files: HashSet::new(), + } + } + + pub(crate) fn has_more(&self) -> bool { + !self.pending.is_empty() || self.next_offsets.iter().any(Option::is_some) + } + + pub(crate) fn remember_match(&mut self, key: MatchKey) -> bool { + self.seen.insert(key) + } + + pub(crate) fn remember_file(&mut self, path: String) -> bool { + self.seen_files.insert(path) + } +} + +enum CursorState { + FileOffset(usize), + MultiGrep(MultiGrepCursor), +} + /// Stores cursor state for paginated grep results. pub struct CursorStore { counter: u64, - /// Map from cursor ID string → file offset for next page. - cursors: HashMap, + /// Map from cursor ID string to the state required for the next page. + cursors: HashMap, /// Insertion order for LRU eviction. insertion_order: VecDeque, } @@ -27,10 +91,25 @@ impl CursorStore { /// Store a cursor and return its opaque ID string. pub fn store(&mut self, file_offset: usize) -> String { + self.store_state(CursorState::FileOffset(file_offset)) + } + + pub(crate) fn store_multi_grep(&mut self, cursor: MultiGrepCursor) -> String { + self.store_state(CursorState::MultiGrep(cursor)) + } + + pub(crate) fn get_multi_grep(&self, id: &str) -> Option { + match self.cursors.get(id) { + Some(CursorState::MultiGrep(cursor)) => Some(cursor.clone()), + Some(CursorState::FileOffset(_)) | None => None, + } + } + + fn store_state(&mut self, state: CursorState) -> String { self.counter = self.counter.wrapping_add(1); let id = self.counter.to_string(); - self.cursors.insert(id.clone(), file_offset); + self.cursors.insert(id.clone(), state); self.insertion_order.push_back(id.clone()); // Evict oldest cursors @@ -47,6 +126,9 @@ impl CursorStore { /// Retrieve the file offset for a cursor ID. pub fn get(&self, id: &str) -> Option { - self.cursors.get(id).copied() + match self.cursors.get(id) { + Some(CursorState::FileOffset(offset)) => Some(*offset), + Some(CursorState::MultiGrep(_)) | None => None, + } } } diff --git a/crates/fff-mcp/src/server.rs b/crates/fff-mcp/src/server.rs index c4affa334..06933ea0e 100644 --- a/crates/fff-mcp/src/server.rs +++ b/crates/fff-mcp/src/server.rs @@ -1,4 +1,4 @@ -use crate::cursor::CursorStore; +use crate::cursor::{CursorStore, MatchKey, MultiGrepCursor, PendingMatch}; use crate::output::{GrepFormatter, OutputMode, file_suffix}; use fff::grep::{GrepMode, GrepSearchOptions, has_regex_metacharacters}; use fff::types::{FileItem, PaginationArgs}; @@ -71,6 +71,25 @@ fn make_grep_options( ) } +fn file_index_for<'a>(files: &mut Vec<&'a FileItem>, file: &'a FileItem) -> usize { + if let Some(index) = files + .iter() + .position(|candidate| std::ptr::eq(*candidate, file)) + { + index + } else { + files.push(file); + files.len() - 1 + } +} + +fn append_multi_grep_cursor(text: &mut String, cursors: &mut CursorStore, cursor: MultiGrepCursor) { + if cursor.has_more() { + let cursor_id = cursors.store_multi_grep(cursor); + text.push_str(&format!("\ncursor: {cursor_id}")); + } +} + #[derive(Debug, serde::Deserialize, schemars::JsonSchema)] pub struct FindFilesParams { /// Fuzzy search query. Supports path prefixes and glob constraints. @@ -594,100 +613,174 @@ impl FffServer { let context = params.context.map(|v| v.round() as usize); let output_mode = OutputMode::new(params.output_mode.as_deref()); - let file_offset = params - .cursor - .as_deref() - .and_then(|id| self.cursor_store.lock().ok()?.get(id)) - .unwrap_or(0); + let (multi_cursor, file_offset) = { + let cursors = self.lock_cursors()?; + let multi_cursor = params + .cursor + .as_deref() + .and_then(|id| cursors.get_multi_grep(id)); + let file_offset = params + .cursor + .as_deref() + .and_then(|id| cursors.get(id)) + .unwrap_or(0); + (multi_cursor, file_offset) + }; + let is_multi_cursor = multi_cursor.is_some(); + let (patterns, constraint_query) = match multi_cursor.as_ref() { + Some(cursor) => (cursor.patterns.clone(), cursor.constraints.clone()), + None => (params.patterns, params.constraints.unwrap_or_default()), + }; let (options, auto_expand) = make_grep_options(output_mode, GrepMode::PlainText, file_offset, context); let ctx_lines = options.before_context; - let constraint_query = params.constraints.as_deref().unwrap_or(""); let guard = self.picker.read().map_err(|e| { ErrorData::internal_error(format!("Failed to acquire picker lock: {e}"), None) })?; let picker = guard .as_ref() .ok_or_else(|| ErrorData::internal_error("File picker not initialized", None))?; - let patterns_refs: Vec<&str> = params.patterns.iter().map(|s| s.as_str()).collect(); let parser = fff_query_parser::QueryParser::new(fff_query_parser::AiGrepConfig); - let parsed_constraints = parser.parse(constraint_query); - let constraints = parsed_constraints.constraints.as_slice(); + if !is_multi_cursor { + let parsed_constraints = parser.parse_constraints(&constraint_query); + let constraints = parsed_constraints.constraints.as_slice(); + let patterns_refs: Vec<&str> = patterns.iter().map(String::as_str).collect(); + let result = picker.multi_grep(&patterns_refs, constraints, &options); + + if !result.matches.is_empty() || file_offset > 0 { + if result.matches.is_empty() { + return Ok(CallToolResult::success(vec![Content::text( + "0 matches.".to_string(), + )])); + } - let result = picker.multi_grep(&patterns_refs, constraints, &options); - let file_refs: Vec<&FileItem> = result.files.to_vec(); + let file_refs: Vec<&FileItem> = result.files.to_vec(); + let mut cs = self.lock_cursors()?; + let text = &GrepFormatter { + matches: &result.matches, + files: &file_refs, + total_matched: result.matches.len(), + next_file_offset: result.next_file_offset, + output_mode, + max_results, + show_context: ctx_lines > 0, + auto_expand_defs: auto_expand, + picker, + } + .format(&mut cs); + return Ok(CallToolResult::success(vec![Content::text(text)])); + } + } - if result.matches.is_empty() && file_offset == 0 { - // Fallback: try individual patterns with plain grep - let (fallback_options, _) = - make_grep_options(output_mode, GrepMode::PlainText, 0, context); - - let fallback_options = GrepSearchOptions { - time_budget_ms: 3000, - before_context: 0, - ..fallback_options + let mut cursor = multi_cursor + .unwrap_or_else(|| MultiGrepCursor::new(patterns.clone(), constraint_query.clone())); + let (fallback_options, _) = make_grep_options(output_mode, GrepMode::PlainText, 0, context); + let fallback_options = GrepSearchOptions { + time_budget_ms: 3000, + before_context: 0, + page_limit: max_results, + ..fallback_options + }; + let fallback_patterns = cursor.patterns.clone(); + let mut fallback_matches = Vec::new(); + let mut fallback_files: Vec<&FileItem> = Vec::new(); + + for pending in std::mem::take(&mut cursor.pending) { + let Some(file) = picker + .get_files() + .iter() + .find(|file| file.relative_path(picker) == pending.file_path) + else { + continue; }; + let mut match_data = pending.match_data; + match_data.file_index = file_index_for(&mut fallback_files, file); + fallback_matches.push(match_data); + } + + while fallback_matches.len() <= max_results { + let mut advanced = false; + for (pattern_index, pat) in fallback_patterns.iter().enumerate() { + if fallback_matches.len() > max_results { + break; + } + let Some(file_offset) = cursor.next_offsets[pattern_index] else { + continue; + }; + advanced = true; - for pat in ¶ms.patterns { - let full_query: Cow = if !constraint_query.is_empty() { - Cow::Owned(format!("{} {}", constraint_query, pat)) + let full_query: Cow = if !cursor.constraints.is_empty() { + Cow::Owned(format!("{} {}", cursor.constraints, pat)) } else { Cow::Borrowed(pat) }; - let parsed = parser.parse(&full_query); - let fb_result = picker.grep(&parsed, &fallback_options); - - if !fb_result.matches.is_empty() { - let fb_file_refs: Vec<&FileItem> = fb_result.files.to_vec(); - let mut cs = self.lock_cursors()?; - let text = &GrepFormatter { - matches: &fb_result.matches, - files: &fb_file_refs, - total_matched: fb_result.matches.len(), - next_file_offset: fb_result.next_file_offset, - output_mode, - max_results, - show_context: false, - auto_expand_defs: auto_expand, - picker, + let mut options = fallback_options.clone(); + options.file_offset = file_offset; + let result = picker.grep(&parsed, &options); + cursor.next_offsets[pattern_index] = + (result.next_file_offset > 0).then_some(result.next_file_offset); + + for mut match_data in result.matches { + let file = result.files[match_data.file_index]; + let file_path = file.relative_path(picker).to_string(); + if output_mode == OutputMode::FilesWithMatches + && !cursor.remember_file(file_path.clone()) + { + continue; + } + let key = + MatchKey::new(file_path, match_data.line_number, match_data.byte_offset); + if !cursor.remember_match(key) { + continue; } - .format(&mut cs); - return Ok(CallToolResult::success(vec![Content::text(format!( - "0 multi-pattern matches. Plain grep fallback for \"{}\":\n{}", - pat, text - ))])); + match_data.file_index = file_index_for(&mut fallback_files, file); + fallback_matches.push(match_data); } } - return Ok(CallToolResult::success(vec![Content::text( - "0 matches.".to_string(), - )])); + let has_offsets = cursor.next_offsets.iter().any(Option::is_some); + if !advanced || !has_offsets || fallback_matches.len() > max_results { + break; + } } - if result.matches.is_empty() { - return Ok(CallToolResult::success(vec![Content::text( - "0 matches.".to_string(), - )])); + if fallback_matches.is_empty() { + let mut text = "0 matches.".to_string(); + let mut cs = self.lock_cursors()?; + append_multi_grep_cursor(&mut text, &mut cs, cursor); + return Ok(CallToolResult::success(vec![Content::text(text)])); } + cursor.pending = fallback_matches + .iter() + .skip(max_results) + .map(|match_data| PendingMatch { + file_path: fallback_files[match_data.file_index] + .relative_path(picker) + .to_string(), + match_data: match_data.clone(), + }) + .collect(); + let mut cs = self.lock_cursors()?; - let text = &GrepFormatter { - matches: &result.matches, - files: &file_refs, - total_matched: result.matches.len(), - next_file_offset: result.next_file_offset, + let text = GrepFormatter { + matches: &fallback_matches, + files: &fallback_files, + total_matched: fallback_matches.len(), + next_file_offset: 0, output_mode, max_results, - show_context: ctx_lines > 0, + show_context: false, auto_expand_defs: auto_expand, picker, } .format(&mut cs); - + let mut text = format!("0 multi-pattern matches. Plain grep fallback:\n{text}"); + append_multi_grep_cursor(&mut text, &mut cs, cursor); Ok(CallToolResult::success(vec![Content::text(text)])) } } @@ -711,6 +804,47 @@ impl ServerHandler for FffServer { #[cfg(test)] mod tests { use super::*; + use std::fs; + use std::sync::atomic::Ordering; + use tempfile::TempDir; + + fn test_server(tmp: &TempDir, files: &[(&str, &str)]) -> FffServer { + for (relative_path, contents) in files { + let path = tmp.path().join(relative_path); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).expect("create fixture directory"); + } + fs::write(path, contents).expect("write fixture file"); + } + + let mut picker = fff::FilePicker::new(fff::FilePickerOptions { + base_path: tmp.path().to_string_lossy().into_owned(), + watch: false, + ..Default::default() + }) + .expect("create file picker"); + picker.collect_files().expect("collect fixture files"); + + let shared_picker = SharedFilePicker::default(); + *shared_picker.write().expect("acquire picker lock") = Some(picker); + let server = FffServer::new(shared_picker); + server.scan_ready.store(true, Ordering::Relaxed); + server + } + + fn response_text(result: CallToolResult) -> String { + serde_json::to_string(&result).expect("serialize MCP response") + } + + fn cursor_id(text: &str) -> String { + text.split("cursor: ") + .nth(1) + .expect("response cursor") + .split('"') + .next() + .expect("cursor id") + .to_string() + } #[test] fn normalize_max_results_none_uses_default() { @@ -765,4 +899,139 @@ mod tests { serde_json::from_str(r#"{"pattern":"foo"}"#).expect("pattern alias"); assert_eq!(via_pattern.query, "foo"); } + + #[test] + fn multi_grep_preserves_positive_constraints_and_or_results() { + let tmp = TempDir::new().expect("create fixture directory"); + let server = test_server( + &tmp, + &[ + ("scope-a/one.txt", "alpha\n"), + ("scope-a/two.txt", "beta\n"), + ("scope-a/three.rs", "alpha\n"), + ("scope-b/three.txt", "alpha\n"), + ("scope-b/four.rs", "alpha\n"), + ], + ); + let cases = [ + ( + vec!["alpha"], + Some("scope-a/"), + vec!["scope-a/one.txt", "scope-a/three.rs"], + vec!["scope-b/three.txt", "scope-b/four.rs"], + ), + ( + vec!["alpha"], + Some("scope-a/*.txt"), + vec!["scope-a/one.txt"], + vec!["scope-a/three.rs", "scope-b/three.txt"], + ), + ( + vec!["alpha"], + Some("scope-a/one.txt"), + vec!["scope-a/one.txt"], + vec!["scope-a/two.txt", "scope-a/three.rs"], + ), + ( + vec!["alpha"], + Some("!scope-b/"), + vec!["scope-a/one.txt", "scope-a/three.rs"], + vec!["scope-b/three.txt", "scope-b/four.rs"], + ), + ( + vec!["alpha", "missing"], + Some("scope-a/"), + vec!["scope-a/one.txt", "scope-a/three.rs"], + vec!["scope-a/two.txt", "scope-b/three.txt"], + ), + ( + vec!["alpha", "beta", "missing"], + Some("scope-a/"), + vec!["scope-a/one.txt", "scope-a/two.txt", "scope-a/three.rs"], + vec!["scope-b/three.txt", "scope-b/four.rs"], + ), + ]; + + for (patterns, constraints, expected, excluded) in cases { + let text = response_text( + server + .multi_grep_inner(MultiGrepParams { + patterns: patterns.into_iter().map(str::to_string).collect(), + constraints: constraints.map(str::to_string), + max_results: None, + cursor: None, + output_mode: None, + context: None, + }) + .expect("multi_grep response"), + ); + for path in expected { + assert!(text.contains(path), "missing {path} in {text}"); + } + for path in excluded { + assert!(!text.contains(path), "unexpected {path} in {text}"); + } + } + } + + #[test] + fn multi_grep_fallback_returns_union_and_continues() { + let tmp = TempDir::new().expect("create fixture directory"); + let server = test_server( + &tmp, + &[("one.txt", "alpha only\n"), ("two.txt", "beta only\n")], + ); + let params = |patterns: &[&str], cursor, max_results| MultiGrepParams { + patterns: patterns.iter().map(ToString::to_string).collect(), + constraints: None, + max_results, + cursor, + output_mode: None, + context: None, + }; + + let full = response_text( + server + .multi_grep_inner(params( + &["*.txt alpha", "*.txt beta", "*.txt missing"], + None, + None, + )) + .expect("fallback response"), + ); + assert!(full.contains("one.txt")); + assert!(full.contains("two.txt")); + + let one_match = response_text( + server + .multi_grep_inner(params(&["*.txt alpha", "*.txt missing"], None, None)) + .expect("single fallback match response"), + ); + assert!(one_match.contains("one.txt")); + assert!(!one_match.contains("two.txt")); + + let first = response_text( + server + .multi_grep_inner(params( + &["*.txt alpha", "*.txt beta", "*.txt missing"], + None, + Some(1.0), + )) + .expect("first fallback page"), + ); + let cursor = cursor_id(&first); + let second = response_text( + server + .multi_grep_inner(params( + &["*.txt alpha", "*.txt beta", "*.txt missing"], + Some(cursor), + Some(1.0), + )) + .expect("second fallback page"), + ); + assert!(first.contains("one.txt")); + assert!(!first.contains("two.txt")); + assert!(second.contains("two.txt")); + assert!(!second.contains("cursor: ")); + } } diff --git a/crates/fff-query-parser/src/lib.rs b/crates/fff-query-parser/src/lib.rs index 8dfcd4ddb..e44f847da 100644 --- a/crates/fff-query-parser/src/lib.rs +++ b/crates/fff-query-parser/src/lib.rs @@ -84,6 +84,29 @@ mod tests { assert_eq!(result.fuzzy_query, FuzzyQuery::Text("hello")); } + #[test] + fn parse_constraints_keeps_lone_path_and_filename_filters() { + let parser = QueryParser::new(AiGrepConfig); + + let path = parser.parse_constraints("scope-a/"); + assert!(matches!( + path.constraints.as_slice(), + [Constraint::PathSegment("scope-a")] + )); + assert_eq!(path.fuzzy_query, FuzzyQuery::Empty); + + let file = parser.parse_constraints("scope-a/one.txt"); + assert!(matches!( + file.constraints.as_slice(), + [Constraint::FilePath("scope-a/one.txt")] + )); + assert_eq!(file.fuzzy_query, FuzzyQuery::Empty); + + let ordinary = parser.parse("scope-a/"); + assert!(ordinary.constraints.is_empty()); + assert_eq!(ordinary.fuzzy_query, FuzzyQuery::Text("scope-a/")); + } + #[test] fn test_simple_text() { let parser = QueryParser::default(); diff --git a/crates/fff-query-parser/src/parser.rs b/crates/fff-query-parser/src/parser.rs index 4fd72116d..609429315 100644 --- a/crates/fff-query-parser/src/parser.rs +++ b/crates/fff-query-parser/src/parser.rs @@ -50,6 +50,17 @@ impl QueryParser { } pub fn parse<'a>(&self, query: &'a str) -> FFFQuery<'a> { + self.parse_inner(query, true) + } + + /// Parse a query supplied as an independent set of search constraints. + /// Unlike ordinary single-token search, a lone path or filename remains a + /// constraint instead of being promoted to fuzzy text. + pub fn parse_constraints<'a>(&self, query: &'a str) -> FFFQuery<'a> { + self.parse_inner(query, false) + } + + fn parse_inner<'a>(&self, query: &'a str, demote_lone_constraints: bool) -> FFFQuery<'a> { let raw_query = query; let config: &C = &self.config; let mut constraints = ConstraintVec::new(); @@ -79,12 +90,12 @@ impl QueryParser { // for grep we don't want to treat a part of path like pathname let treat_as_text = matches!(constraint, Constraint::PathSegment(_)) + && demote_lone_constraints && config.treat_lone_path_as_text(); + let file_path_as_text = + matches!(constraint, Constraint::FilePath(_)) && demote_lone_constraints; - if !matches!(constraint, Constraint::FilePath(_)) - && !has_location_suffix - && !treat_as_text - { + if !file_path_as_text && !has_location_suffix && !treat_as_text { constraints.push(constraint); return FFFQuery { raw_query,