From 554223ab8010ecebab57cb4e6c5bf107c626eb4c Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 4 May 2026 15:23:10 +0200 Subject: [PATCH] feat: paginate memories MCP search results (#20996) ## Why The memories MCP `search` tool previously stopped once it hit `max_results`, so callers could tell there were more matches via `truncated` but had no way to fetch the rest of the result set. That made large searches awkward for clients that need to keep paging through a stable, deterministic view of the matches. ## What changed - add an optional `cursor` field to `SearchMemoriesRequest` / tool input and return `next_cursor` in `SearchMemoriesResponse` - update the MCP schemas and tool wiring so clients can request subsequent pages explicitly - change the local memories backend to collect and sort the full scoped match list, then slice the requested page and reject invalid cursors - add unit coverage for paginated search results and invalid cursor handling in `memories/mcp/src/local_tests.rs` ## Testing - Added targeted unit coverage in `memories/mcp/src/local_tests.rs` - GitHub Actions are running for the branch --- codex-rs/memories/mcp/src/backend.rs | 2 + codex-rs/memories/mcp/src/local.rs | 73 ++++++++++++--------- codex-rs/memories/mcp/src/local_tests.rs | 83 ++++++++++++++++++++++++ codex-rs/memories/mcp/src/schema.rs | 6 +- codex-rs/memories/mcp/src/server.rs | 4 +- 5 files changed, 135 insertions(+), 33 deletions(-) diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index 1886a78ba..b287ce5df 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -65,6 +65,7 @@ pub struct ReadMemoryResponse { pub struct SearchMemoriesRequest { pub query: String, pub path: Option, + pub cursor: Option, pub max_results: usize, } @@ -73,6 +74,7 @@ pub struct SearchMemoriesResponse { pub query: String, pub path: Option, pub matches: Vec, + pub next_cursor: Option, pub truncated: bool, } diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index 0e2048bc7..a4b8889a1 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -192,18 +192,44 @@ impl MemoriesBackend for LocalMemoriesBackend { let max_results = request.max_results.min(MAX_SEARCH_RESULTS); let start = self.resolve_scoped_path(request.path.as_deref())?; + let start_index = match request.cursor.as_deref() { + Some(cursor) => cursor.parse::().map_err(|_| { + MemoriesBackendError::invalid_cursor(cursor, "must be a non-negative integer") + })?, + None => 0, + }; + let Some(metadata) = Self::metadata_or_none(&start).await? else { + return Ok(SearchMemoriesResponse { + query: request.query, + path: request.path, + matches: Vec::new(), + next_cursor: None, + truncated: false, + }); + }; + reject_symlink(&display_relative_path(&self.root, &start), &metadata)?; + let mut matches = Vec::new(); - let truncated = - search_entries(&self.root, &start, query, &mut matches, max_results).await?; + search_entries(&self.root, &start, &metadata, query, &mut matches).await?; matches.sort_by(|left, right| { left.path .cmp(&right.path) .then(left.line_number.cmp(&right.line_number)) }); + if start_index > matches.len() { + return Err(MemoriesBackendError::invalid_cursor( + start_index.to_string(), + "exceeds result count", + )); + } + let end_index = start_index.saturating_add(max_results).min(matches.len()); + let next_cursor = (end_index < matches.len()).then(|| end_index.to_string()); + let truncated = next_cursor.is_some(); Ok(SearchMemoriesResponse { query: request.query, path: request.path, - matches, + matches: matches.drain(start_index..end_index).collect(), + next_cursor, truncated, }) } @@ -212,30 +238,21 @@ impl MemoriesBackend for LocalMemoriesBackend { async fn search_entries( root: &Path, current: &Path, + current_metadata: &std::fs::Metadata, query: &str, matches: &mut Vec, - max_results: usize, -) -> Result { - if max_results == 0 { - return Ok(false); +) -> Result<(), MemoriesBackendError> { + if current_metadata.is_file() { + search_file(root, current, query, matches).await?; + return Ok(()); } - let Some(metadata) = LocalMemoriesBackend::metadata_or_none(current).await? else { - return Ok(false); - }; - reject_symlink(&display_relative_path(root, current), &metadata)?; - if metadata.is_file() { - return search_file(root, current, query, matches, max_results).await; - } - if !metadata.is_dir() { - return Ok(false); + if !current_metadata.is_dir() { + return Ok(()); } let mut pending = vec![current.to_path_buf()]; while let Some(dir_path) = pending.pop() { for path in read_sorted_dir_paths(&dir_path).await? { - if matches.len() >= max_results { - return Ok(true); - } let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&path).await? else { continue; }; @@ -244,15 +261,13 @@ async fn search_entries( } if metadata.is_dir() { pending.push(path); - } else if metadata.is_file() - && search_file(root, &path, query, matches, max_results).await? - { - return Ok(true); + } else if metadata.is_file() { + search_file(root, &path, query, matches).await?; } } } - Ok(false) + Ok(()) } async fn search_file( @@ -260,17 +275,13 @@ async fn search_file( path: &Path, query: &str, matches: &mut Vec, - max_results: usize, -) -> Result { +) -> Result<(), MemoriesBackendError> { let content = match tokio::fs::read_to_string(path).await { Ok(content) => content, - Err(err) if err.kind() == std::io::ErrorKind::InvalidData => return Ok(false), + 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 matches.len() >= max_results { - return Ok(true); - } if line.contains(query) { matches.push(MemorySearchMatch { path: display_relative_path(root, path), @@ -279,7 +290,7 @@ async fn search_file( }); } } - Ok(false) + Ok(()) } async fn read_sorted_dir_paths(dir_path: &Path) -> Result, MemoriesBackendError> { diff --git a/codex-rs/memories/mcp/src/local_tests.rs b/codex-rs/memories/mcp/src/local_tests.rs index f59abc283..d1444e473 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -384,6 +384,7 @@ async fn search_supports_directory_and_file_scopes() { .search(SearchMemoriesRequest { query: "needle".to_string(), path: None, + cursor: None, max_results: DEFAULT_SEARCH_MAX_RESULTS, }) .await @@ -396,17 +397,99 @@ async fn search_supports_directory_and_file_scopes() { .collect::>(), vec![("MEMORY.md", 2), ("rollout_summaries/a.jsonl", 1)] ); + assert_eq!(response.next_cursor, None); + assert_eq!(response.truncated, false); let file_response = backend(&tempdir) .search(SearchMemoriesRequest { query: "needle".to_string(), path: Some("MEMORY.md".to_string()), + cursor: None, 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.next_cursor, None); + assert_eq!(file_response.truncated, false); +} + +#[tokio::test] +async fn search_supports_pagination() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::create_dir_all(tempdir.path().join("rollout_summaries")) + .await + .expect("create rollout summaries dir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "needle one\nneedle two\n") + .await + .expect("write memory file"); + tokio::fs::write( + tempdir.path().join("rollout_summaries/a.jsonl"), + "needle three\n", + ) + .await + .expect("write rollout summary"); + + let page1 = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: None, + 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)] + ); + assert_eq!(page1.next_cursor.as_deref(), Some("2")); + assert_eq!(page1.truncated, true); + + let page2 = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: page1.next_cursor, + 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)] + ); + assert_eq!(page2.next_cursor, None); + assert_eq!(page2.truncated, false); +} + +#[tokio::test] +async fn search_rejects_invalid_cursor() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "needle\n") + .await + .expect("write memory file"); + + let err = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + cursor: Some("bogus".to_string()), + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect_err("cursor should be rejected"); + + assert!(matches!(err, MemoriesBackendError::InvalidCursor { .. })); } #[tokio::test] diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs index 3674b234a..c382aba51 100644 --- a/codex-rs/memories/mcp/src/schema.rs +++ b/codex-rs/memories/mcp/src/schema.rs @@ -75,6 +75,7 @@ pub(crate) fn search_input_schema() -> JsonObject { "properties": { "query": { "type": "string" }, "path": { "type": "string" }, + "cursor": { "type": "string" }, "max_results": { "type": "integer", "minimum": 1 } }, "required": ["query"], @@ -90,6 +91,9 @@ pub(crate) fn search_output_schema() -> JsonObject { "path": { "anyOf": [{ "type": "string" }, { "type": "null" }] }, + "next_cursor": { + "anyOf": [{ "type": "string" }, { "type": "null" }] + }, "matches": { "type": "array", "items": { @@ -105,7 +109,7 @@ pub(crate) fn search_output_schema() -> JsonObject { }, "truncated": { "type": "boolean" } }, - "required": ["query", "path", "matches", "truncated"], + "required": ["query", "path", "matches", "next_cursor", "truncated"], "additionalProperties": false })) } diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index cd99a5e23..a2fe12baf 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -57,6 +57,7 @@ struct ReadArgs { struct SearchArgs { query: String, path: Option, + cursor: Option, max_results: Option, } @@ -146,6 +147,7 @@ impl ServerHandler for MemoriesMcpServer { .search(SearchMemoriesRequest { query: args.query, path: args.path, + cursor: args.cursor, max_results: clamp_max_results( args.max_results, DEFAULT_SEARCH_MAX_RESULTS, @@ -215,7 +217,7 @@ 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."), + Cow::Borrowed("Search Codex memory files for exact text matches, with pagination."), Arc::new(schema::search_input_schema()), ); tool.output_schema = Some(Arc::new(schema::search_output_schema()));