diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index 1c4b60791..2fe47997c 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -47,6 +47,7 @@ pub struct ListMemoriesResponse { pub struct ReadMemoryRequest { pub path: String, pub line_offset: usize, + pub max_lines: Option, pub max_tokens: usize, } @@ -99,6 +100,8 @@ pub enum MemoriesBackendError { InvalidPath { path: String, reason: String }, #[error("line_offset must be a 1-indexed line number")] InvalidLineOffset, + #[error("max_lines must be a positive integer")] + InvalidMaxLines, #[error("line_offset exceeds file length")] LineOffsetExceedsFileLength, #[error("path '{path}' is not a file")] diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs index 4d2ab17c3..f39180643 100644 --- a/codex-rs/memories/mcp/src/local.rs +++ b/codex-rs/memories/mcp/src/local.rs @@ -94,6 +94,9 @@ impl MemoriesBackend for LocalMemoriesBackend { if request.line_offset == 0 { return Err(MemoriesBackendError::InvalidLineOffset); } + if request.max_lines == Some(0) { + return Err(MemoriesBackendError::InvalidMaxLines); + } let path = self.resolve_scoped_path(Some(request.path.as_str()))?; let Some(metadata) = Self::metadata_or_none(&path).await? else { @@ -106,14 +109,15 @@ impl MemoriesBackend for LocalMemoriesBackend { let original_content = tokio::fs::read_to_string(&path).await?; let start_byte = line_start_byte_offset(&original_content, request.line_offset)?; - let content_from_offset = &original_content[start_byte..]; + let end_byte = line_end_byte_offset(&original_content, start_byte, request.max_lines); + let content_from_offset = &original_content[start_byte..end_byte]; let max_tokens = if request.max_tokens == 0 { DEFAULT_READ_MAX_TOKENS } else { request.max_tokens }; let content = truncate_text(content_from_offset, TruncationPolicy::Tokens(max_tokens)); - let truncated = content != content_from_offset; + let truncated = end_byte < original_content.len() || content != content_from_offset; Ok(ReadMemoryResponse { path: request.path, start_line_number: request.line_offset, @@ -334,6 +338,24 @@ fn line_start_byte_offset( Err(MemoriesBackendError::LineOffsetExceedsFileLength) } +fn line_end_byte_offset(content: &str, start_byte: usize, max_lines: Option) -> usize { + let Some(max_lines) = max_lines else { + return content.len(); + }; + + let mut lines_seen = 1; + for (relative_idx, ch) in content[start_byte..].char_indices() { + if ch == '\n' { + if lines_seen == max_lines { + return start_byte + relative_idx + 1; + } + lines_seen += 1; + } + } + + content.len() +} + #[cfg(test)] #[path = "local_tests.rs"] mod tests; diff --git a/codex-rs/memories/mcp/src/local_tests.rs b/codex-rs/memories/mcp/src/local_tests.rs index 64eabae96..7f2ffbdc3 100644 --- a/codex-rs/memories/mcp/src/local_tests.rs +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -64,6 +64,7 @@ async fn read_rejects_directory_and_returns_file_content() { .read(ReadMemoryRequest { path: "MEMORY.md".to_string(), line_offset: 1, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -83,6 +84,7 @@ async fn read_rejects_directory_and_returns_file_content() { .read(ReadMemoryRequest { path: ".".to_string(), line_offset: 1, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -101,6 +103,7 @@ async fn read_supports_line_offset() { .read(ReadMemoryRequest { path: "MEMORY.md".to_string(), line_offset: 2, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -118,7 +121,35 @@ async fn read_supports_line_offset() { } #[tokio::test] -async fn read_rejects_invalid_line_offsets() { +async fn read_supports_max_lines() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "alpha\nbeta\ngamma\n") + .await + .expect("write memory file"); + + let response = backend(&tempdir) + .read(ReadMemoryRequest { + path: "MEMORY.md".to_string(), + line_offset: 2, + max_lines: Some(1), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect("read memory with line limit"); + + assert_eq!( + response, + ReadMemoryResponse { + path: "MEMORY.md".to_string(), + start_line_number: 2, + content: "beta\n".to_string(), + truncated: true, + } + ); +} + +#[tokio::test] +async fn read_rejects_invalid_line_requests() { let tempdir = TempDir::new().expect("tempdir"); tokio::fs::write(tempdir.path().join("MEMORY.md"), "only\n") .await @@ -128,6 +159,7 @@ async fn read_rejects_invalid_line_offsets() { .read(ReadMemoryRequest { path: "MEMORY.md".to_string(), line_offset: 0, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -137,10 +169,25 @@ async fn read_rejects_invalid_line_offsets() { MemoriesBackendError::InvalidLineOffset )); + let zero_max_lines_err = backend(&tempdir) + .read(ReadMemoryRequest { + path: "MEMORY.md".to_string(), + line_offset: 1, + max_lines: Some(0), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect_err("zero max lines should fail"); + assert!(matches!( + zero_max_lines_err, + MemoriesBackendError::InvalidMaxLines + )); + let past_end_err = backend(&tempdir) .read(ReadMemoryRequest { path: "MEMORY.md".to_string(), line_offset: 3, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -203,6 +250,7 @@ async fn scoped_paths_reject_parent_segments() { .read(ReadMemoryRequest { path: "../secret".to_string(), line_offset: 1, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -226,6 +274,7 @@ async fn read_rejects_symlinked_files() { .read(ReadMemoryRequest { path: "inside-link".to_string(), line_offset: 1, + max_lines: None, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs index a39c40c60..2d138f4f3 100644 --- a/codex-rs/memories/mcp/src/schema.rs +++ b/codex-rs/memories/mcp/src/schema.rs @@ -43,7 +43,8 @@ pub(crate) fn read_input_schema() -> JsonObject { "type": "object", "properties": { "path": { "type": "string" }, - "line_offset": { "type": "integer", "minimum": 1 } + "line_offset": { "type": "integer", "minimum": 1 }, + "max_lines": { "type": "integer", "minimum": 1 } }, "required": ["path"], "additionalProperties": false diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index 8bacb7c0e..ed298c47a 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -49,6 +49,7 @@ struct ListArgs { struct ReadArgs { path: String, line_offset: Option, + max_lines: Option, } #[derive(Deserialize)] @@ -129,6 +130,7 @@ impl ServerHandler for MemoriesMcpServer { .read(ReadMemoryRequest { path: args.path, line_offset: args.line_offset.unwrap_or(1), + max_lines: args.max_lines, max_tokens: DEFAULT_READ_MAX_TOKENS, }) .await @@ -197,7 +199,7 @@ fn read_tool() -> Tool { let mut tool = Tool::new( Cow::Borrowed(READ_TOOL_NAME), Cow::Borrowed( - "Read a Codex memory file by relative path, optionally starting at a 1-indexed line offset.", + "Read a Codex memory file by relative path, optionally starting at a 1-indexed line offset and limiting the number of lines returned.", ), Arc::new(schema::read_input_schema()), ); @@ -229,6 +231,7 @@ fn backend_error_to_mcp(err: MemoriesBackendError) -> McpError { match err { MemoriesBackendError::InvalidPath { .. } | MemoriesBackendError::InvalidLineOffset + | MemoriesBackendError::InvalidMaxLines | MemoriesBackendError::LineOffsetExceedsFileLength | MemoriesBackendError::NotFile { .. } | MemoriesBackendError::EmptyQuery => McpError::invalid_params(err.to_string(), None),