From 0269a46ab17ef832ab763b178c5ea74d3db248ef Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 4 May 2026 15:32:57 +0200 Subject: [PATCH] feat: add context lines to memories MCP search (#20997) ## Why The paginated memories MCP `search` tool still returned only the matching line text, which made it harder for clients to present useful search results or decide whether they needed to follow up with a separate `read` call. Adding a small amount of surrounding context makes individual hits much more usable while keeping the search response deterministic and line-addressable. ## What changed - add an optional `context_lines` search argument and thread it through the MCP server into the local memories backend - change search matches to return the matched `line_number` plus a `start_line_number` and multi-line `content` block for the requested context window - update the search tool schema and description to document the new request/response shape - extend the local backend tests to cover zero-context matches, contextual results, pagination, and invalid cursors that point past the end of the result set ## Testing - Added targeted unit coverage in `memories/mcp/src/local_tests.rs` - GitHub Actions are running for the branch --------- Co-authored-by: Codex --- codex-rs/memories/mcp/src/backend.rs | 5 +- codex-rs/memories/mcp/src/local.rs | 37 ++++- codex-rs/memories/mcp/src/local_tests.rs | 198 ++++++++++++++++++++--- codex-rs/memories/mcp/src/schema.rs | 7 +- codex-rs/memories/mcp/src/server.rs | 8 +- 5 files changed, 225 insertions(+), 30 deletions(-) diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index b287ce5df..96a004827 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -66,6 +66,8 @@ pub struct SearchMemoriesRequest { pub query: String, pub path: Option, pub cursor: Option, + pub context_lines: usize, + pub case_sensitive: bool, pub max_results: usize, } @@ -95,7 +97,8 @@ pub enum MemoryEntryType { pub struct MemorySearchMatch { pub path: String, pub line_number: usize, - pub line: String, + pub start_line_number: usize, + pub content: String, } #[derive(Debug, thiserror::Error)] diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index a4b8889a1..7fc342f1b 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -210,7 +210,16 @@ impl MemoriesBackend for LocalMemoriesBackend { reject_symlink(&display_relative_path(&self.root, &start), &metadata)?; let mut matches = Vec::new(); - search_entries(&self.root, &start, &metadata, query, &mut matches).await?; + search_entries( + &self.root, + &start, + &metadata, + query, + request.context_lines, + request.case_sensitive, + &mut matches, + ) + .await?; matches.sort_by(|left, right| { left.path .cmp(&right.path) @@ -240,10 +249,12 @@ async fn search_entries( current: &Path, current_metadata: &std::fs::Metadata, query: &str, + context_lines: usize, + case_sensitive: bool, matches: &mut Vec, ) -> Result<(), MemoriesBackendError> { if current_metadata.is_file() { - search_file(root, current, query, matches).await?; + search_file(root, current, query, context_lines, case_sensitive, matches).await?; return Ok(()); } if !current_metadata.is_dir() { @@ -262,7 +273,7 @@ async fn search_entries( if metadata.is_dir() { pending.push(path); } else if metadata.is_file() { - search_file(root, &path, query, matches).await?; + search_file(root, &path, query, context_lines, case_sensitive, matches).await?; } } } @@ -274,6 +285,8 @@ async fn search_file( root: &Path, path: &Path, query: &str, + context_lines: usize, + case_sensitive: bool, matches: &mut Vec, ) -> Result<(), MemoriesBackendError> { let content = match tokio::fs::read_to_string(path).await { @@ -281,12 +294,24 @@ async fn search_file( Err(err) if err.kind() == std::io::ErrorKind::InvalidData => return Ok(()), Err(err) => return Err(err.into()), }; - for (idx, line) in content.lines().enumerate() { - if line.contains(query) { + let lines = content.lines().collect::>(); + let normalized_query = (!case_sensitive).then(|| query.to_lowercase()); + for (idx, line) in lines.iter().enumerate() { + let is_match = match normalized_query.as_deref() { + Some(query) => line.to_lowercase().contains(query), + None => line.contains(query), + }; + if is_match { + let start_index = idx.saturating_sub(context_lines); + let end_index = idx + .saturating_add(context_lines) + .saturating_add(1) + .min(lines.len()); matches.push(MemorySearchMatch { path: display_relative_path(root, path), line_number: idx + 1, - line: line.to_string(), + start_line_number: start_index + 1, + content: lines[start_index..end_index].join("\n"), }); } } diff --git a/codex-rs/memories/mcp/src/local_tests.rs b/codex-rs/memories/mcp/src/local_tests.rs index d1444e473..361dbe049 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -385,17 +385,28 @@ async fn search_supports_directory_and_file_scopes() { query: "needle".to_string(), path: None, cursor: None, + context_lines: 0, + case_sensitive: true, max_results: DEFAULT_SEARCH_MAX_RESULTS, }) .await .expect("search all memories"); assert_eq!( - response - .matches - .iter() - .map(|entry| (entry.path.as_str(), entry.line_number)) - .collect::>(), - vec![("MEMORY.md", 2), ("rollout_summaries/a.jsonl", 1)] + response.matches, + vec![ + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 2, + content: "needle".to_string(), + }, + MemorySearchMatch { + path: "rollout_summaries/a.jsonl".to_string(), + line_number: 1, + start_line_number: 1, + content: "needle again".to_string(), + }, + ] ); assert_eq!(response.next_cursor, None); assert_eq!(response.truncated, false); @@ -405,12 +416,21 @@ async fn search_supports_directory_and_file_scopes() { query: "needle".to_string(), path: Some("MEMORY.md".to_string()), cursor: None, + context_lines: 0, + case_sensitive: true, max_results: DEFAULT_SEARCH_MAX_RESULTS, }) .await .expect("search one memory file"); - assert_eq!(file_response.matches.len(), 1); - assert_eq!(file_response.matches[0].path, "MEMORY.md"); + assert_eq!( + file_response.matches, + vec![MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 2, + content: "needle".to_string(), + }] + ); assert_eq!(file_response.next_cursor, None); assert_eq!(file_response.truncated, false); } @@ -436,17 +456,28 @@ async fn search_supports_pagination() { query: "needle".to_string(), path: None, cursor: None, + context_lines: 0, + case_sensitive: true, max_results: 2, }) .await .expect("search first page"); assert_eq!( - page1 - .matches - .iter() - .map(|entry| (entry.path.as_str(), entry.line_number)) - .collect::>(), - vec![("MEMORY.md", 1), ("MEMORY.md", 2)] + page1.matches, + vec![ + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 1, + start_line_number: 1, + content: "needle one".to_string(), + }, + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 2, + content: "needle two".to_string(), + }, + ] ); assert_eq!(page1.next_cursor.as_deref(), Some("2")); assert_eq!(page1.truncated, true); @@ -456,22 +487,130 @@ async fn search_supports_pagination() { query: "needle".to_string(), path: None, cursor: page1.next_cursor, + context_lines: 0, + case_sensitive: true, max_results: 2, }) .await .expect("search second page"); assert_eq!( - page2 - .matches - .iter() - .map(|entry| (entry.path.as_str(), entry.line_number)) - .collect::>(), - vec![("rollout_summaries/a.jsonl", 1)] + page2.matches, + vec![MemorySearchMatch { + path: "rollout_summaries/a.jsonl".to_string(), + line_number: 1, + start_line_number: 1, + content: "needle three".to_string(), + }] ); assert_eq!(page2.next_cursor, None); assert_eq!(page2.truncated, false); } +#[tokio::test] +async fn search_supports_context_lines() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write( + tempdir.path().join("MEMORY.md"), + "alpha\nneedle\nomega\nneedle again\n", + ) + .await + .expect("write memory file"); + + let response = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: None, + context_lines: 1, + case_sensitive: true, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect("search with context"); + + assert_eq!( + response.matches, + vec![ + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 1, + content: "alpha\nneedle\nomega".to_string(), + }, + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 4, + start_line_number: 3, + content: "omega\nneedle again".to_string(), + }, + ] + ); +} + +#[tokio::test] +async fn search_supports_case_insensitive_matching() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "Needle\nneedle\nNEEDLE\n") + .await + .expect("write memory file"); + + let sensitive_response = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: None, + context_lines: 0, + case_sensitive: true, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect("search with case-sensitive matching"); + assert_eq!( + sensitive_response.matches, + vec![MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 2, + content: "needle".to_string(), + }] + ); + + let insensitive_response = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: None, + context_lines: 0, + case_sensitive: false, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect("search with case-insensitive matching"); + assert_eq!( + insensitive_response.matches, + vec![ + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 1, + start_line_number: 1, + content: "Needle".to_string(), + }, + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 2, + start_line_number: 2, + content: "needle".to_string(), + }, + MemorySearchMatch { + path: "MEMORY.md".to_string(), + line_number: 3, + start_line_number: 3, + content: "NEEDLE".to_string(), + }, + ] + ); +} + #[tokio::test] async fn search_rejects_invalid_cursor() { let tempdir = TempDir::new().expect("tempdir"); @@ -484,12 +623,31 @@ async fn search_rejects_invalid_cursor() { query: "needle".to_string(), path: None, cursor: Some("bogus".to_string()), + context_lines: 0, + case_sensitive: true, max_results: DEFAULT_SEARCH_MAX_RESULTS, }) .await .expect_err("cursor should be rejected"); assert!(matches!(err, MemoriesBackendError::InvalidCursor { .. })); + + let past_end_err = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: Some("2".to_string()), + context_lines: 0, + case_sensitive: true, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect_err("cursor past end should be rejected"); + + assert!(matches!( + past_end_err, + MemoriesBackendError::InvalidCursor { .. } + )); } #[tokio::test] diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs index c382aba51..c1d7347a7 100644 --- a/codex-rs/memories/mcp/src/schema.rs +++ b/codex-rs/memories/mcp/src/schema.rs @@ -76,6 +76,8 @@ pub(crate) fn search_input_schema() -> JsonObject { "query": { "type": "string" }, "path": { "type": "string" }, "cursor": { "type": "string" }, + "context_lines": { "type": "integer", "minimum": 0 }, + "case_sensitive": { "type": "boolean" }, "max_results": { "type": "integer", "minimum": 1 } }, "required": ["query"], @@ -101,9 +103,10 @@ pub(crate) fn search_output_schema() -> JsonObject { "properties": { "path": { "type": "string" }, "line_number": { "type": "integer" }, - "line": { "type": "string" } + "start_line_number": { "type": "integer" }, + "content": { "type": "string" } }, - "required": ["path", "line_number", "line"], + "required": ["path", "line_number", "start_line_number", "content"], "additionalProperties": false } }, diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index a2fe12baf..f16919d7e 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -58,6 +58,8 @@ struct SearchArgs { query: String, path: Option, cursor: Option, + context_lines: Option, + case_sensitive: Option, max_results: Option, } @@ -148,6 +150,8 @@ impl ServerHandler for MemoriesMcpServer { query: args.query, path: args.path, cursor: args.cursor, + context_lines: args.context_lines.unwrap_or(0), + case_sensitive: args.case_sensitive.unwrap_or(true), max_results: clamp_max_results( args.max_results, DEFAULT_SEARCH_MAX_RESULTS, @@ -217,7 +221,9 @@ fn read_tool() -> Tool { fn search_tool() -> Tool { let mut tool = Tool::new( Cow::Borrowed(SEARCH_TOOL_NAME), - Cow::Borrowed("Search Codex memory files for exact text matches, with pagination."), + Cow::Borrowed( + "Search Codex memory files for exact text matches, with pagination, optional surrounding context lines, and optional case-insensitive matching.", + ), Arc::new(schema::search_input_schema()), ); tool.output_schema = Some(Arc::new(schema::search_output_schema()));