diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index 9e4fde40d..0e2048bc7 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -77,16 +77,66 @@ 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 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 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(ListMemoriesResponse { + path: request.path, + entries: Vec::new(), + next_cursor: None, + truncated: false, + }); + }; + reject_symlink(&display_relative_path(&self.root, &start), &metadata)?; + + let mut entries = if metadata.is_file() { + vec![MemoryEntry { + path: display_relative_path(&self.root, &start), + entry_type: MemoryEntryType::File, + }] + } else if metadata.is_dir() { + let mut entries = Vec::new(); + for path in read_sorted_dir_paths(&start).await? { + let Some(metadata) = Self::metadata_or_none(&path).await? else { + continue; + }; + if metadata.file_type().is_symlink() { + continue; + } + + let entry_type = if metadata.is_dir() { + MemoryEntryType::Directory + } else if metadata.is_file() { + MemoryEntryType::File + } else { + continue; + }; + entries.push(MemoryEntry { + path: display_relative_path(&self.root, &path), + entry_type, + }); + } + entries + } else { + Vec::new() + }; + if start_index > entries.len() { + return Err(MemoriesBackendError::invalid_cursor( + start_index.to_string(), + "exceeds result count", + )); + } + + let end_index = start_index.saturating_add(max_results).min(entries.len()); + let next_cursor = (end_index < entries.len()).then(|| end_index.to_string()); let truncated = next_cursor.is_some(); Ok(ListMemoriesResponse { path: request.path, - entries, + entries: entries.drain(start_index..end_index).collect(), next_cursor, truncated, }) @@ -159,87 +209,6 @@ impl MemoriesBackend for LocalMemoriesBackend { } } -async fn collect_entries_page( - root: &Path, - current: &Path, - start_index: usize, - stop_after: usize, - entries: &mut Vec, -) -> Result { - let Some(metadata) = LocalMemoriesBackend::metadata_or_none(current).await? else { - return Ok(0); - }; - reject_symlink(&display_relative_path(root, current), &metadata)?; - - let mut seen = 0usize; - if metadata.is_file() { - 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(0); - } - - let mut pending = vec![current.to_path_buf()]; - 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; - }; - push_list_entry( - entries, - &mut seen, - start_index, - stop_after, - MemoryEntry { - path: relative, - 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); - } - } - - if seen < start_index { - return Err(MemoriesBackendError::invalid_cursor( - start_index.to_string(), - "exceeds result count", - )); - } - - Ok(seen) -} - async fn search_entries( root: &Path, current: &Path, @@ -347,34 +316,6 @@ 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 d819e6887..f59abc283 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -9,7 +9,7 @@ fn backend(tempdir: &TempDir) -> LocalMemoriesBackend { } #[tokio::test] -async fn list_returns_recursive_memory_paths() { +async fn list_returns_shallow_memory_paths() { let tempdir = TempDir::new().expect("tempdir"); tokio::fs::create_dir_all(tempdir.path().join("skills/example")) .await @@ -41,14 +41,6 @@ async fn list_returns_recursive_memory_paths() { path: "skills".to_string(), entry_type: MemoryEntryType::Directory, }, - MemoryEntry { - path: "skills/example".to_string(), - entry_type: MemoryEntryType::Directory, - }, - MemoryEntry { - path: "skills/example/SKILL.md".to_string(), - entry_type: MemoryEntryType::File, - }, ] ); assert_eq!(response.next_cursor, None); @@ -58,15 +50,18 @@ async fn list_returns_recursive_memory_paths() { #[tokio::test] async fn list_supports_pagination() { let tempdir = TempDir::new().expect("tempdir"); - tokio::fs::create_dir_all(tempdir.path().join("skills/example")) + tokio::fs::create_dir_all(tempdir.path().join("skills")) .await .expect("create skills dir"); + tokio::fs::create_dir_all(tempdir.path().join("rollout_summaries")) + .await + .expect("create rollout 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") + tokio::fs::write(tempdir.path().join("memory_summary.md"), "summary") .await - .expect("write skill file"); + .expect("write memory summary"); let page1 = backend(&tempdir) .list(ListMemoriesRequest { @@ -84,8 +79,8 @@ async fn list_supports_pagination() { entry_type: MemoryEntryType::File, }, MemoryEntry { - path: "skills".to_string(), - entry_type: MemoryEntryType::Directory, + path: "memory_summary.md".to_string(), + entry_type: MemoryEntryType::File, }, ] ); @@ -104,12 +99,12 @@ async fn list_supports_pagination() { page2.entries, vec![ MemoryEntry { - path: "skills/example".to_string(), + path: "rollout_summaries".to_string(), entry_type: MemoryEntryType::Directory, }, MemoryEntry { - path: "skills/example/SKILL.md".to_string(), - entry_type: MemoryEntryType::File, + path: "skills".to_string(), + entry_type: MemoryEntryType::Directory, }, ] ); @@ -118,21 +113,15 @@ async fn list_supports_pagination() { } #[tokio::test] -async fn list_preserves_lexicographic_order_across_directories() { +async fn list_preserves_lexicographic_order_for_siblings() { let tempdir = TempDir::new().expect("tempdir"); - tokio::fs::create_dir_all(tempdir.path().join("a/nested")) + tokio::fs::create_dir_all(tempdir.path().join("a")) .await .expect("create a dir"); - tokio::fs::create_dir_all(tempdir.path().join("b")) + tokio::fs::write(tempdir.path().join("a.txt"), "a") .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") + .expect("write a.txt file"); + tokio::fs::write(tempdir.path().join("b.txt"), "b") .await .expect("write b file"); @@ -151,13 +140,43 @@ async fn list_preserves_lexicographic_order_across_directories() { .iter() .map(|entry| entry.path.as_str()) .collect::>(), + vec!["a", "a.txt", "b.txt"] + ); +} + +#[tokio::test] +async fn list_scoped_directory_is_shallow() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::create_dir_all(tempdir.path().join("skills/example")) + .await + .expect("create nested skills dir"); + tokio::fs::write(tempdir.path().join("skills/README.md"), "readme") + .await + .expect("write skills readme"); + tokio::fs::write(tempdir.path().join("skills/example/SKILL.md"), "skill") + .await + .expect("write nested skill file"); + + let response = backend(&tempdir) + .list(ListMemoriesRequest { + path: Some("skills".to_string()), + cursor: None, + max_results: DEFAULT_LIST_MAX_RESULTS, + }) + .await + .expect("list scoped directory"); + + assert_eq!( + response.entries, vec![ - "a", - "a/file.txt", - "a/nested", - "a/nested/inner.txt", - "b", - "b/file.txt", + MemoryEntry { + path: "skills/README.md".to_string(), + entry_type: MemoryEntryType::File, + }, + MemoryEntry { + path: "skills/example".to_string(), + entry_type: MemoryEntryType::Directory, + }, ] ); } diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index 36435fdb6..cd99a5e23 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -189,7 +189,9 @@ pub async fn run_stdio_server(codex_home: &AbsolutePathBuf) -> anyhow::Result<() fn list_tool() -> Tool { let mut tool = Tool::new( Cow::Borrowed(LIST_TOOL_NAME), - Cow::Borrowed("List files and directories under the Codex memories store."), + Cow::Borrowed( + "List immediate files and directories under a path in the Codex memories store.", + ), Arc::new(schema::list_input_schema()), ); tool.output_schema = Some(Arc::new(schema::list_output_schema()));