diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index f402a9b8a..929852f1b 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -73,6 +73,7 @@ pub struct SearchMemoriesRequest { pub cursor: Option, pub context_lines: usize, pub case_sensitive: bool, + pub normalized: bool, pub max_results: usize, } diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index f16dba195..97aacee11 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -250,7 +250,8 @@ impl MemoriesBackend for LocalMemoriesBackend { queries.clone(), request.match_mode.clone(), request.case_sensitive, - ); + request.normalized, + )?; let mut matches = Vec::new(); search_entries( &self.root, @@ -450,32 +451,38 @@ fn build_search_match( struct SearchMatcher { queries: Vec, - normalized_queries: Option>, + prepared_queries: Vec, + comparison: SearchComparison, match_mode: SearchMatchMode, } impl SearchMatcher { - fn new(queries: Vec, match_mode: SearchMatchMode, case_sensitive: bool) -> Self { - let normalized_queries = (!case_sensitive).then(|| { - queries - .iter() - .map(|query| query.to_lowercase()) - .collect::>() - }); - Self { - queries, - normalized_queries, - match_mode, + fn new( + queries: Vec, + match_mode: SearchMatchMode, + case_sensitive: bool, + normalized: bool, + ) -> Result { + let comparison = SearchComparison::new(case_sensitive, normalized); + let prepared_queries = queries + .iter() + .map(|query| comparison.prepare(query)) + .map(Cow::into_owned) + .collect::>(); + if prepared_queries.iter().any(std::string::String::is_empty) { + return Err(MemoriesBackendError::EmptyQuery); } + Ok(Self { + queries, + prepared_queries, + comparison, + match_mode, + }) } fn matched_query_flags(&self, line: &str) -> Vec { - let line = match self.normalized_queries.as_ref() { - Some(_) => Cow::Owned(line.to_lowercase()), - None => Cow::Borrowed(line), - }; - let queries = self.normalized_queries.as_deref().unwrap_or(&self.queries); - queries + let line = self.comparison.prepare(line); + self.prepared_queries .iter() .map(|query| line.as_ref().contains(query)) .collect() @@ -490,6 +497,43 @@ impl SearchMatcher { } } +#[derive(Clone, Copy)] +struct SearchComparison { + case_sensitive: bool, + normalized: bool, +} + +impl SearchComparison { + fn new(case_sensitive: bool, normalized: bool) -> Self { + Self { + case_sensitive, + normalized, + } + } + + fn prepare<'a>(self, value: &'a str) -> Cow<'a, str> { + if self.case_sensitive && !self.normalized { + return Cow::Borrowed(value); + } + + let value = if self.case_sensitive { + Cow::Borrowed(value) + } else { + Cow::Owned(value.to_lowercase()) + }; + if !self.normalized { + return value; + } + + Cow::Owned( + value + .chars() + .filter(|ch| ch.is_alphanumeric()) + .collect::(), + ) + } +} + async fn read_sorted_dir_paths(dir_path: &Path) -> Result, MemoriesBackendError> { let mut dir = match tokio::fs::read_dir(dir_path).await { Ok(dir) => dir, diff --git a/codex-rs/memories/mcp/src/local_tests.rs b/codex-rs/memories/mcp/src/local_tests.rs index 4edf06eef..a2dbbc047 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -16,6 +16,7 @@ fn search_request(queries: &[&str]) -> SearchMemoriesRequest { cursor: None, context_lines: 0, case_sensitive: true, + normalized: false, max_results: DEFAULT_SEARCH_MAX_RESULTS, } } @@ -751,6 +752,67 @@ async fn search_supports_case_insensitive_matching() { ); } +#[tokio::test] +async fn search_supports_normalized_matching() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write( + tempdir.path().join("MEMORY.md"), + "MultiAgentV2\ncold-resume\n", + ) + .await + .expect("write memory file"); + + let literal_response = backend(&tempdir) + .search(search_request(&["multi agent v2", "cold resume"])) + .await + .expect("search without normalization"); + assert_eq!(literal_response.matches, Vec::new()); + + let mut request = search_request(&["multi agent v2", "cold resume"]); + request.case_sensitive = false; + request.normalized = true; + let normalized_response = backend(&tempdir) + .search(request) + .await + .expect("search with normalization"); + assert_eq!( + normalized_response.matches, + vec![ + MemorySearchMatch { + path: "MEMORY.md".to_string(), + match_line_number: 1, + content_start_line_number: 1, + content: "MultiAgentV2".to_string(), + matched_queries: vec!["multi agent v2".to_string()], + }, + MemorySearchMatch { + path: "MEMORY.md".to_string(), + match_line_number: 2, + content_start_line_number: 2, + content: "cold-resume".to_string(), + matched_queries: vec!["cold resume".to_string()], + }, + ] + ); +} + +#[tokio::test] +async fn search_rejects_queries_that_normalize_to_empty_strings() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "needle\n") + .await + .expect("write memory file"); + + let mut request = search_request(&["-"]); + request.normalized = true; + let err = backend(&tempdir) + .search(request) + .await + .expect_err("separator-only normalized queries should be rejected"); + + assert!(matches!(err, MemoriesBackendError::EmptyQuery)); +} + #[tokio::test] async fn search_supports_any_and_all_on_same_line_match_modes() { let tempdir = TempDir::new().expect("tempdir"); diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index 88eeb8fa1..f25e75d37 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -74,6 +74,7 @@ struct SearchArgs { #[schemars(range(min = 0))] context_lines: Option, case_sensitive: Option, + normalized: Option, #[schemars(range(min = 1))] max_results: Option, } @@ -227,7 +228,7 @@ fn search_tool() -> Tool { let mut tool = Tool::new( Cow::Borrowed(SEARCH_TOOL_NAME), Cow::Borrowed( - "Search Codex memory files for substring matches, optionally requiring all query substrings on the same line or within a line window.", + "Search Codex memory files for substring matches, optionally normalizing separators or requiring all query substrings on the same line or within a line window.", ), Arc::new(schema::input_schema_for::()), ); @@ -251,6 +252,7 @@ impl SearchArgs { cursor: self.cursor, context_lines: self.context_lines.unwrap_or(0), case_sensitive: self.case_sensitive.unwrap_or(true), + normalized: self.normalized.unwrap_or(false), max_results: clamp_max_results( self.max_results, DEFAULT_SEARCH_MAX_RESULTS, @@ -306,6 +308,7 @@ mod tests { cursor: None, context_lines: 0, case_sensitive: false, + normalized: false, max_results: DEFAULT_SEARCH_MAX_RESULTS, } ); @@ -333,6 +336,33 @@ mod tests { cursor: None, context_lines: 0, case_sensitive: true, + normalized: false, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + } + ); + } + + #[test] + fn search_args_accept_normalized_matching() { + let args: SearchArgs = parse_args(json!({ + "queries": ["multi agent v2"], + "case_sensitive": false, + "normalized": true + })) + .expect("normalized args should parse"); + + let request = args.into_request(); + + assert_eq!( + request, + SearchMemoriesRequest { + queries: vec!["multi agent v2".to_string()], + match_mode: SearchMatchMode::Any, + path: None, + cursor: None, + context_lines: 0, + case_sensitive: false, + normalized: true, max_results: DEFAULT_SEARCH_MAX_RESULTS, } );