From 5730615e7591cc5dd76338c799a0ba16a4a542d1 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 4 May 2026 14:59:56 +0200 Subject: [PATCH] feat: paginate MCP memories list (#20993) ## Why Large memories trees do not fit well into a single MCP `list` response. This change makes the memories MCP server page `list` results so callers can continue walking the tree without overfetching or relying on ambiguous truncation. ## What changed - add an optional `cursor` input to the memories MCP `list` API and return `next_cursor` alongside `truncated` in the response - paginate recursive local-memory traversal while preserving lexicographic path order across directories - reject malformed and out-of-range cursors as invalid MCP requests - update the server/schema wiring and add coverage for pagination, ordering, and cursor validation in `memories/mcp/src/local_tests.rs` ## Testing - `cargo test -p codex-memories-mcp` --- codex-rs/memories/mcp/src/backend.rs | 11 ++ codex-rs/memories/mcp/src/local.rs | 132 ++++++++++++++------ codex-rs/memories/mcp/src/local_tests.rs | 147 +++++++++++++++++++++++ codex-rs/memories/mcp/src/schema.rs | 6 +- codex-rs/memories/mcp/src/server.rs | 3 + 5 files changed, 261 insertions(+), 38 deletions(-) diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index 2fe47997c..1886a78ba 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -33,6 +33,7 @@ pub trait MemoriesBackend: Clone + Send + Sync + 'static { #[derive(Debug, Clone, PartialEq, Eq)] pub struct ListMemoriesRequest { pub path: Option, + pub cursor: Option, pub max_results: usize, } @@ -40,6 +41,7 @@ pub struct ListMemoriesRequest { pub struct ListMemoriesResponse { pub path: Option, pub entries: Vec, + pub next_cursor: Option, pub truncated: bool, } @@ -98,6 +100,8 @@ pub struct MemorySearchMatch { pub enum MemoriesBackendError { #[error("path '{path}' {reason}")] InvalidPath { path: String, reason: String }, + #[error("cursor '{cursor}' {reason}")] + InvalidCursor { cursor: String, reason: String }, #[error("line_offset must be a 1-indexed line number")] InvalidLineOffset, #[error("max_lines must be a positive integer")] @@ -119,4 +123,11 @@ impl MemoriesBackendError { reason: reason.into(), } } + + pub fn invalid_cursor(cursor: impl Into, reason: impl Into) -> Self { + Self::InvalidCursor { + cursor: cursor.into(), + reason: reason.into(), + } + } } diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index f39180643..9e4fde40d 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -77,12 +77,17 @@ impl MemoriesBackend for LocalMemoriesBackend { ) -> Result { let max_results = request.max_results.min(MAX_LIST_RESULTS); let start = self.resolve_scoped_path(request.path.as_deref())?; + let start_index = parse_list_cursor(request.cursor.as_deref())?; + let stop_after = start_index.saturating_add(max_results); let mut entries = Vec::new(); - let truncated = collect_entries(&self.root, &start, &mut entries, max_results).await?; - entries.sort_by(|left, right| left.path.cmp(&right.path)); + let listed_count = + collect_entries_page(&self.root, &start, start_index, stop_after, &mut entries).await?; + let next_cursor = (listed_count > stop_after).then(|| stop_after.to_string()); + let truncated = next_cursor.is_some(); Ok(ListMemoriesResponse { path: request.path, entries, + next_cursor, truncated, }) } @@ -154,60 +159,85 @@ impl MemoriesBackend for LocalMemoriesBackend { } } -async fn collect_entries( +async fn collect_entries_page( root: &Path, current: &Path, + start_index: usize, + stop_after: usize, entries: &mut Vec, - max_results: usize, -) -> Result { - if max_results == 0 { - return Ok(false); - } +) -> Result { let Some(metadata) = LocalMemoriesBackend::metadata_or_none(current).await? else { - return Ok(false); + return Ok(0); }; reject_symlink(&display_relative_path(root, current), &metadata)?; + + let mut seen = 0usize; if metadata.is_file() { - entries.push(MemoryEntry { - path: display_relative_path(root, current), - entry_type: MemoryEntryType::File, - }); - return Ok(entries.len() >= max_results); + push_list_entry( + entries, + &mut seen, + start_index, + stop_after, + MemoryEntry { + path: display_relative_path(root, current), + entry_type: MemoryEntryType::File, + }, + ); + return Ok(seen); } if !metadata.is_dir() { - return Ok(false); + return Ok(0); } 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 entries.len() >= max_results { - return Ok(true); - } - let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&path).await? else { + while let Some(path) = pending.pop() { + let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&path).await? else { + continue; + }; + if metadata.file_type().is_symlink() { + continue; + } + if path != current { + let relative = display_relative_path(root, &path); + let entry_type = if metadata.is_dir() { + MemoryEntryType::Directory + } else if metadata.is_file() { + MemoryEntryType::File + } else { continue; }; - if metadata.file_type().is_symlink() { - continue; - } - - let relative = display_relative_path(root, &path); - if metadata.is_dir() { - entries.push(MemoryEntry { + push_list_entry( + entries, + &mut seen, + start_index, + stop_after, + MemoryEntry { path: relative, - entry_type: MemoryEntryType::Directory, - }); - pending.push(path); - } else if metadata.is_file() { - entries.push(MemoryEntry { - path: relative, - entry_type: MemoryEntryType::File, - }); + entry_type, + }, + ); + if seen > stop_after { + return Ok(seen); } } + if metadata.is_dir() { + let mut children = read_sorted_dir_paths(&path).await?; + children.reverse(); + pending.extend(children); + } + if seen > stop_after { + return Ok(seen); + } } - Ok(false) + if seen < start_index { + return Err(MemoriesBackendError::invalid_cursor( + start_index.to_string(), + "exceeds result count", + )); + } + + Ok(seen) } async fn search_entries( @@ -317,6 +347,34 @@ fn display_relative_path(root: &Path, path: &Path) -> String { .join("/") } +fn parse_list_cursor(cursor: Option<&str>) -> Result { + let Some(cursor) = cursor else { + return Ok(0); + }; + + let start_index = cursor.parse::().map_err(|_| { + MemoriesBackendError::invalid_cursor(cursor, "must be a non-negative integer") + })?; + + Ok(start_index) +} + +fn push_list_entry( + entries: &mut Vec, + seen: &mut usize, + start_index: usize, + stop_after: usize, + entry: MemoryEntry, +) { + *seen += 1; + if *seen <= start_index { + return; + } + if *seen <= stop_after { + entries.push(entry); + } +} + fn line_start_byte_offset( content: &str, line_offset: usize, diff --git a/codex-rs/memories/mcp/src/local_tests.rs b/codex-rs/memories/mcp/src/local_tests.rs index 7f2ffbdc3..d819e6887 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -24,6 +24,7 @@ async fn list_returns_recursive_memory_paths() { let response = backend(&tempdir) .list(ListMemoriesRequest { path: None, + cursor: None, max_results: DEFAULT_LIST_MAX_RESULTS, }) .await @@ -50,9 +51,155 @@ async fn list_returns_recursive_memory_paths() { }, ] ); + assert_eq!(response.next_cursor, None); assert_eq!(response.truncated, false); } +#[tokio::test] +async fn list_supports_pagination() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::create_dir_all(tempdir.path().join("skills/example")) + .await + .expect("create skills dir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "summary") + .await + .expect("write memory file"); + tokio::fs::write(tempdir.path().join("skills/example/SKILL.md"), "skill") + .await + .expect("write skill file"); + + let page1 = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + cursor: None, + max_results: 2, + }) + .await + .expect("list first page"); + assert_eq!( + page1.entries, + vec![ + MemoryEntry { + path: "MEMORY.md".to_string(), + entry_type: MemoryEntryType::File, + }, + MemoryEntry { + path: "skills".to_string(), + entry_type: MemoryEntryType::Directory, + }, + ] + ); + assert_eq!(page1.next_cursor.as_deref(), Some("2")); + assert_eq!(page1.truncated, true); + + let page2 = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + cursor: page1.next_cursor, + max_results: 2, + }) + .await + .expect("list second page"); + assert_eq!( + page2.entries, + vec![ + MemoryEntry { + path: "skills/example".to_string(), + entry_type: MemoryEntryType::Directory, + }, + MemoryEntry { + path: "skills/example/SKILL.md".to_string(), + entry_type: MemoryEntryType::File, + }, + ] + ); + assert_eq!(page2.next_cursor, None); + assert_eq!(page2.truncated, false); +} + +#[tokio::test] +async fn list_preserves_lexicographic_order_across_directories() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::create_dir_all(tempdir.path().join("a/nested")) + .await + .expect("create a dir"); + tokio::fs::create_dir_all(tempdir.path().join("b")) + .await + .expect("create b dir"); + tokio::fs::write(tempdir.path().join("a/file.txt"), "a") + .await + .expect("write a file"); + tokio::fs::write(tempdir.path().join("a/nested/inner.txt"), "inner") + .await + .expect("write nested file"); + tokio::fs::write(tempdir.path().join("b/file.txt"), "b") + .await + .expect("write b file"); + + let response = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + cursor: None, + max_results: DEFAULT_LIST_MAX_RESULTS, + }) + .await + .expect("list memories"); + + assert_eq!( + response + .entries + .iter() + .map(|entry| entry.path.as_str()) + .collect::>(), + vec![ + "a", + "a/file.txt", + "a/nested", + "a/nested/inner.txt", + "b", + "b/file.txt", + ] + ); +} + +#[tokio::test] +async fn list_rejects_invalid_cursor() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "summary") + .await + .expect("write memory file"); + + let err = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + cursor: Some("bogus".to_string()), + max_results: DEFAULT_LIST_MAX_RESULTS, + }) + .await + .expect_err("cursor should be rejected"); + + assert!(matches!(err, MemoriesBackendError::InvalidCursor { .. })); +} + +#[tokio::test] +async fn list_rejects_cursor_past_end() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "summary") + .await + .expect("write memory file"); + + let err = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + cursor: Some("2".to_string()), + max_results: DEFAULT_LIST_MAX_RESULTS, + }) + .await + .expect_err("cursor past end should be rejected"); + + assert!(matches!(err, MemoriesBackendError::InvalidCursor { .. })); +} + #[tokio::test] async fn read_rejects_directory_and_returns_file_content() { let tempdir = TempDir::new().expect("tempdir"); diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs index 2d138f4f3..3674b234a 100644 --- a/codex-rs/memories/mcp/src/schema.rs +++ b/codex-rs/memories/mcp/src/schema.rs @@ -6,6 +6,7 @@ pub(crate) fn list_input_schema() -> JsonObject { "type": "object", "properties": { "path": { "type": "string" }, + "cursor": { "type": "string" }, "max_results": { "type": "integer", "minimum": 1 } }, "additionalProperties": false @@ -19,6 +20,9 @@ pub(crate) fn list_output_schema() -> JsonObject { "path": { "anyOf": [{ "type": "string" }, { "type": "null" }] }, + "next_cursor": { + "anyOf": [{ "type": "string" }, { "type": "null" }] + }, "entries": { "type": "array", "items": { @@ -33,7 +37,7 @@ pub(crate) fn list_output_schema() -> JsonObject { }, "truncated": { "type": "boolean" } }, - "required": ["path", "entries", "truncated"], + "required": ["path", "entries", "next_cursor", "truncated"], "additionalProperties": false })) } diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index ed298c47a..36435fdb6 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -42,6 +42,7 @@ pub struct MemoriesMcpServer { #[derive(Deserialize)] struct ListArgs { path: Option, + cursor: Option, max_results: Option, } @@ -113,6 +114,7 @@ impl ServerHandler for MemoriesMcpServer { self.backend .list(ListMemoriesRequest { path: args.path, + cursor: args.cursor, max_results: clamp_max_results( args.max_results, DEFAULT_LIST_MAX_RESULTS, @@ -230,6 +232,7 @@ fn clamp_max_results(requested: Option, default: usize, max: usize) -> us fn backend_error_to_mcp(err: MemoriesBackendError) -> McpError { match err { MemoriesBackendError::InvalidPath { .. } + | MemoriesBackendError::InvalidCursor { .. } | MemoriesBackendError::InvalidLineOffset | MemoriesBackendError::InvalidMaxLines | MemoriesBackendError::LineOffsetExceedsFileLength