mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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`
This commit is contained in:
committed by
GitHub
Unverified
parent
6b6581ac59
commit
5730615e75
@@ -33,6 +33,7 @@ pub trait MemoriesBackend: Clone + Send + Sync + 'static {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ListMemoriesRequest {
|
||||
pub path: Option<String>,
|
||||
pub cursor: Option<String>,
|
||||
pub max_results: usize,
|
||||
}
|
||||
|
||||
@@ -40,6 +41,7 @@ pub struct ListMemoriesRequest {
|
||||
pub struct ListMemoriesResponse {
|
||||
pub path: Option<String>,
|
||||
pub entries: Vec<MemoryEntry>,
|
||||
pub next_cursor: Option<String>,
|
||||
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<String>, reason: impl Into<String>) -> Self {
|
||||
Self::InvalidCursor {
|
||||
cursor: cursor.into(),
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,12 +77,17 @@ impl MemoriesBackend for LocalMemoriesBackend {
|
||||
) -> Result<ListMemoriesResponse, MemoriesBackendError> {
|
||||
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<MemoryEntry>,
|
||||
max_results: usize,
|
||||
) -> Result<bool, MemoriesBackendError> {
|
||||
if max_results == 0 {
|
||||
return Ok(false);
|
||||
}
|
||||
) -> Result<usize, MemoriesBackendError> {
|
||||
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<usize, MemoriesBackendError> {
|
||||
let Some(cursor) = cursor else {
|
||||
return Ok(0);
|
||||
};
|
||||
|
||||
let start_index = cursor.parse::<usize>().map_err(|_| {
|
||||
MemoriesBackendError::invalid_cursor(cursor, "must be a non-negative integer")
|
||||
})?;
|
||||
|
||||
Ok(start_index)
|
||||
}
|
||||
|
||||
fn push_list_entry(
|
||||
entries: &mut Vec<MemoryEntry>,
|
||||
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,
|
||||
|
||||
@@ -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<_>>(),
|
||||
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");
|
||||
|
||||
@@ -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
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -42,6 +42,7 @@ pub struct MemoriesMcpServer<B> {
|
||||
#[derive(Deserialize)]
|
||||
struct ListArgs {
|
||||
path: Option<String>,
|
||||
cursor: Option<String>,
|
||||
max_results: Option<usize>,
|
||||
}
|
||||
|
||||
@@ -113,6 +114,7 @@ impl<B: MemoriesBackend> ServerHandler for MemoriesMcpServer<B> {
|
||||
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<usize>, 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
|
||||
|
||||
Reference in New Issue
Block a user