diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1d7910fcb..c94785df6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3038,6 +3038,22 @@ dependencies = [ "wiremock", ] +[[package]] +name = "codex-memories-mcp" +version = "0.0.0" +dependencies = [ + "anyhow", + "codex-utils-absolute-path", + "codex-utils-output-truncation", + "pretty_assertions", + "rmcp", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", +] + [[package]] name = "codex-memories-read" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 44eade38d..ce3e91626 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -52,6 +52,7 @@ members = [ "login", "codex-mcp", "mcp-server", + "memories/mcp", "memories/read", "memories/write", "model-provider-info", @@ -168,6 +169,7 @@ codex-keyring-store = { path = "keyring-store" } codex-linux-sandbox = { path = "linux-sandbox" } codex-lmstudio = { path = "lmstudio" } codex-login = { path = "login" } +codex-memories-mcp = { path = "memories/mcp" } codex-memories-read = { path = "memories/read" } codex-memories-write = { path = "memories/write" } codex-mcp = { path = "codex-mcp" } @@ -461,6 +463,7 @@ unwrap_used = "deny" [workspace.metadata.cargo-shear] ignored = [ "codex-agent-graph-store", + "codex-memories-mcp", "icu_provider", "openssl-sys", "codex-utils-readiness", diff --git a/codex-rs/memories/mcp/BUILD.bazel b/codex-rs/memories/mcp/BUILD.bazel new file mode 100644 index 000000000..99048da38 --- /dev/null +++ b/codex-rs/memories/mcp/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "mcp", + crate_name = "codex_memories_mcp", +) diff --git a/codex-rs/memories/mcp/Cargo.toml b/codex-rs/memories/mcp/Cargo.toml new file mode 100644 index 000000000..da1d46579 --- /dev/null +++ b/codex-rs/memories/mcp/Cargo.toml @@ -0,0 +1,30 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-memories-mcp" +version.workspace = true + +[lib] +name = "codex_memories_mcp" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-utils-absolute-path = { workspace = true } +codex-utils-output-truncation = { workspace = true } +rmcp = { workspace = true, default-features = false, features = [ + "schemars", + "server", +] } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["fs", "io-std"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["fs", "macros"] } diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs new file mode 100644 index 000000000..71e4bafc4 --- /dev/null +++ b/codex-rs/memories/mcp/src/backend.rs @@ -0,0 +1,113 @@ +use serde::Serialize; +use std::future::Future; + +pub const DEFAULT_LIST_MAX_RESULTS: usize = 2_000; +pub const MAX_LIST_RESULTS: usize = 2_000; +pub const DEFAULT_SEARCH_MAX_RESULTS: usize = 200; +pub const MAX_SEARCH_RESULTS: usize = 200; +pub const DEFAULT_READ_MAX_TOKENS: usize = 20_000; + +/// Storage interface behind the memories MCP tools. +/// +/// Implementations should return paths relative to the memory store and enforce +/// their own storage-specific access rules. The local implementation uses the +/// filesystem today; a later implementation can satisfy the same contract from a +/// remote backend. +pub trait MemoriesBackend: Clone + Send + Sync + 'static { + fn list( + &self, + request: ListMemoriesRequest, + ) -> impl Future> + Send; + + fn read( + &self, + request: ReadMemoryRequest, + ) -> impl Future> + Send; + + fn search( + &self, + request: SearchMemoriesRequest, + ) -> impl Future> + Send; +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ListMemoriesRequest { + pub path: Option, + pub max_results: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct ListMemoriesResponse { + pub path: Option, + pub entries: Vec, + pub truncated: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReadMemoryRequest { + pub path: String, + pub max_tokens: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct ReadMemoryResponse { + pub path: String, + pub content: String, + pub truncated: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SearchMemoriesRequest { + pub query: String, + pub path: Option, + pub max_results: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct SearchMemoriesResponse { + pub query: String, + pub path: Option, + pub matches: Vec, + pub truncated: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct MemoryEntry { + pub path: String, + pub entry_type: MemoryEntryType, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum MemoryEntryType { + File, + Directory, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct MemorySearchMatch { + pub path: String, + pub line_number: usize, + pub line: String, +} + +#[derive(Debug, thiserror::Error)] +pub enum MemoriesBackendError { + #[error("path '{path}' {reason}")] + InvalidPath { path: String, reason: String }, + #[error("path '{path}' is not a file")] + NotFile { path: String }, + #[error("query must not be empty")] + EmptyQuery, + #[error("I/O error while reading memories: {0}")] + Io(#[from] std::io::Error), +} + +impl MemoriesBackendError { + pub fn invalid_path(path: impl Into, reason: impl Into) -> Self { + Self::InvalidPath { + path: path.into(), + reason: reason.into(), + } + } +} diff --git a/codex-rs/memories/mcp/src/lib.rs b/codex-rs/memories/mcp/src/lib.rs new file mode 100644 index 000000000..004d48ee7 --- /dev/null +++ b/codex-rs/memories/mcp/src/lib.rs @@ -0,0 +1,14 @@ +//! MCP access to Codex memories. +//! +//! This crate only exposes tools for discovering and reading memory files. The +//! policy that tells a model when to use those tools is injected elsewhere. + +pub mod backend; +pub mod local; + +mod schema; +mod server; + +pub use local::LocalMemoriesBackend; +pub use server::MemoriesMcpServer; +pub use server::run_stdio_server; diff --git a/codex-rs/memories/mcp/src/local.rs b/codex-rs/memories/mcp/src/local.rs new file mode 100644 index 000000000..1f444391b --- /dev/null +++ b/codex-rs/memories/mcp/src/local.rs @@ -0,0 +1,311 @@ +use crate::backend::DEFAULT_READ_MAX_TOKENS; +use crate::backend::ListMemoriesRequest; +use crate::backend::ListMemoriesResponse; +use crate::backend::MAX_LIST_RESULTS; +use crate::backend::MAX_SEARCH_RESULTS; +use crate::backend::MemoriesBackend; +use crate::backend::MemoriesBackendError; +use crate::backend::MemoryEntry; +use crate::backend::MemoryEntryType; +use crate::backend::MemorySearchMatch; +use crate::backend::ReadMemoryRequest; +use crate::backend::ReadMemoryResponse; +use crate::backend::SearchMemoriesRequest; +use crate::backend::SearchMemoriesResponse; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; +use std::path::Component; +use std::path::Path; +use std::path::PathBuf; + +#[derive(Debug, Clone)] +pub struct LocalMemoriesBackend { + root: PathBuf, +} + +impl LocalMemoriesBackend { + pub fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self { + Self::from_memory_root(codex_home.join("memories").to_path_buf()) + } + + pub fn from_memory_root(root: impl Into) -> Self { + Self { root: root.into() } + } + + pub fn root(&self) -> &Path { + &self.root + } + + fn resolve_scoped_path( + &self, + relative_path: Option<&str>, + ) -> Result { + let Some(relative_path) = relative_path else { + return Ok(self.root.clone()); + }; + let relative = Path::new(relative_path); + if relative.components().any(|component| { + matches!( + component, + Component::ParentDir | Component::RootDir | Component::Prefix(_) + ) + }) { + return Err(MemoriesBackendError::invalid_path( + relative_path, + "must stay within the memories root", + )); + } + Ok(self.root.join(relative)) + } + + async fn metadata_or_none( + path: &Path, + ) -> Result, MemoriesBackendError> { + match tokio::fs::symlink_metadata(path).await { + Ok(metadata) => Ok(Some(metadata)), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(err) => Err(err.into()), + } + } +} + +impl MemoriesBackend for LocalMemoriesBackend { + async fn list( + &self, + request: ListMemoriesRequest, + ) -> Result { + let max_results = request.max_results.min(MAX_LIST_RESULTS); + let start = self.resolve_scoped_path(request.path.as_deref())?; + 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)); + Ok(ListMemoriesResponse { + path: request.path, + entries, + truncated, + }) + } + + async fn read( + &self, + request: ReadMemoryRequest, + ) -> Result { + let path = self.resolve_scoped_path(Some(request.path.as_str()))?; + let Some(metadata) = Self::metadata_or_none(&path).await? else { + return Err(MemoriesBackendError::NotFile { path: request.path }); + }; + reject_symlink(&request.path, &metadata)?; + if !metadata.is_file() { + return Err(MemoriesBackendError::NotFile { path: request.path }); + } + + let original_content = tokio::fs::read_to_string(&path).await?; + let max_tokens = if request.max_tokens == 0 { + DEFAULT_READ_MAX_TOKENS + } else { + request.max_tokens + }; + let content = truncate_text(&original_content, TruncationPolicy::Tokens(max_tokens)); + let truncated = content != original_content; + Ok(ReadMemoryResponse { + path: request.path, + content, + truncated, + }) + } + + async fn search( + &self, + request: SearchMemoriesRequest, + ) -> Result { + let query = request.query.trim(); + if query.is_empty() { + return Err(MemoriesBackendError::EmptyQuery); + } + + let max_results = request.max_results.min(MAX_SEARCH_RESULTS); + let start = self.resolve_scoped_path(request.path.as_deref())?; + let mut matches = Vec::new(); + let truncated = + search_entries(&self.root, &start, query, &mut matches, max_results).await?; + matches.sort_by(|left, right| { + left.path + .cmp(&right.path) + .then(left.line_number.cmp(&right.line_number)) + }); + Ok(SearchMemoriesResponse { + query: request.query, + path: request.path, + matches, + truncated, + }) + } +} + +async fn collect_entries( + root: &Path, + current: &Path, + entries: &mut Vec, + max_results: usize, +) -> Result { + if max_results == 0 { + return Ok(false); + } + 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() { + entries.push(MemoryEntry { + path: display_relative_path(root, current), + entry_type: MemoryEntryType::File, + }); + return Ok(entries.len() >= max_results); + } + if !metadata.is_dir() { + return Ok(false); + } + + 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 { + continue; + }; + if metadata.file_type().is_symlink() { + continue; + } + + let relative = display_relative_path(root, &path); + if metadata.is_dir() { + entries.push(MemoryEntry { + path: relative, + entry_type: MemoryEntryType::Directory, + }); + pending.push(path); + } else if metadata.is_file() { + entries.push(MemoryEntry { + path: relative, + entry_type: MemoryEntryType::File, + }); + } + } + } + + Ok(false) +} + +async fn search_entries( + root: &Path, + current: &Path, + query: &str, + matches: &mut Vec, + max_results: usize, +) -> Result { + if max_results == 0 { + return Ok(false); + } + 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); + } + + 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; + }; + if metadata.file_type().is_symlink() { + continue; + } + if metadata.is_dir() { + pending.push(path); + } else if metadata.is_file() + && search_file(root, &path, query, matches, max_results).await? + { + return Ok(true); + } + } + } + + Ok(false) +} + +async fn search_file( + root: &Path, + path: &Path, + query: &str, + matches: &mut Vec, + max_results: usize, +) -> Result { + 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) => 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), + line_number: idx + 1, + line: line.to_string(), + }); + } + } + Ok(false) +} + +async fn read_sorted_dir_paths(dir_path: &Path) -> Result, MemoriesBackendError> { + let mut dir = match tokio::fs::read_dir(dir_path).await { + Ok(dir) => dir, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(err) => return Err(err.into()), + }; + let mut paths = Vec::new(); + while let Some(entry) = dir.next_entry().await? { + paths.push(entry.path()); + } + paths.sort(); + Ok(paths) +} + +fn reject_symlink(path: &str, metadata: &std::fs::Metadata) -> Result<(), MemoriesBackendError> { + if metadata.file_type().is_symlink() { + return Err(MemoriesBackendError::invalid_path( + path, + "must not be a symlink", + )); + } + Ok(()) +} + +fn display_relative_path(root: &Path, path: &Path) -> String { + path.strip_prefix(root) + .unwrap_or(path) + .components() + .map(|component| component.as_os_str().to_string_lossy()) + .filter(|component| !component.is_empty()) + .collect::>() + .join("/") +} + +#[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 new file mode 100644 index 000000000..c2cf16827 --- /dev/null +++ b/codex-rs/memories/mcp/src/local_tests.rs @@ -0,0 +1,169 @@ +use super::*; +use crate::backend::DEFAULT_LIST_MAX_RESULTS; +use crate::backend::DEFAULT_SEARCH_MAX_RESULTS; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +fn backend(tempdir: &TempDir) -> LocalMemoriesBackend { + LocalMemoriesBackend::from_memory_root(tempdir.path()) +} + +#[tokio::test] +async fn list_returns_recursive_memory_paths() { + 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 response = backend(&tempdir) + .list(ListMemoriesRequest { + path: None, + max_results: DEFAULT_LIST_MAX_RESULTS, + }) + .await + .expect("list memories"); + + assert_eq!( + response.entries, + vec![ + MemoryEntry { + path: "MEMORY.md".to_string(), + entry_type: MemoryEntryType::File, + }, + MemoryEntry { + 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.truncated, false); +} + +#[tokio::test] +async fn read_rejects_directory_and_returns_file_content() { + let tempdir = TempDir::new().expect("tempdir"); + tokio::fs::write(tempdir.path().join("MEMORY.md"), "remember this") + .await + .expect("write memory file"); + + let response = backend(&tempdir) + .read(ReadMemoryRequest { + path: "MEMORY.md".to_string(), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect("read memory"); + + assert_eq!( + response, + ReadMemoryResponse { + path: "MEMORY.md".to_string(), + content: "remember this".to_string(), + truncated: false, + } + ); + + let err = backend(&tempdir) + .read(ReadMemoryRequest { + path: ".".to_string(), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect_err("directory should not be readable as file"); + assert!(matches!(err, MemoriesBackendError::NotFile { .. })); +} + +#[tokio::test] +async fn search_supports_directory_and_file_scopes() { + 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"), "alpha\nneedle\n") + .await + .expect("write memory file"); + tokio::fs::write( + tempdir.path().join("rollout_summaries/a.jsonl"), + "needle again\n", + ) + .await + .expect("write rollout summary"); + + let response = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: None, + max_results: DEFAULT_SEARCH_MAX_RESULTS, + }) + .await + .expect("search all memories"); + assert_eq!( + response + .matches + .iter() + .map(|entry| (entry.path.as_str(), entry.line_number)) + .collect::>(), + vec![("MEMORY.md", 2), ("rollout_summaries/a.jsonl", 1)] + ); + + let file_response = backend(&tempdir) + .search(SearchMemoriesRequest { + query: "needle".to_string(), + path: Some("MEMORY.md".to_string()), + 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"); +} + +#[tokio::test] +async fn scoped_paths_reject_parent_segments() { + let tempdir = TempDir::new().expect("tempdir"); + let err = backend(&tempdir) + .read(ReadMemoryRequest { + path: "../secret".to_string(), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect_err("parent traversal should fail"); + + assert!(matches!(err, MemoriesBackendError::InvalidPath { .. })); +} + +#[cfg(unix)] +#[tokio::test] +async fn read_rejects_symlinked_files() { + let tempdir = TempDir::new().expect("tempdir"); + let outside = tempdir.path().join("outside.txt"); + tokio::fs::write(&outside, "outside") + .await + .expect("write outside file"); + std::os::unix::fs::symlink(&outside, tempdir.path().join("inside-link")) + .expect("create symlink"); + + let err = backend(&tempdir) + .read(ReadMemoryRequest { + path: "inside-link".to_string(), + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .expect_err("symlink should be rejected"); + + assert!(matches!(err, MemoriesBackendError::InvalidPath { .. })); +} diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs new file mode 100644 index 000000000..06fe4f0d9 --- /dev/null +++ b/codex-rs/memories/mcp/src/schema.rs @@ -0,0 +1,109 @@ +use rmcp::model::JsonObject; +use serde_json::json; + +pub(crate) fn list_input_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "path": { "type": "string" }, + "max_results": { "type": "integer", "minimum": 1 } + }, + "additionalProperties": false + })) +} + +pub(crate) fn list_output_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "path": { + "anyOf": [{ "type": "string" }, { "type": "null" }] + }, + "entries": { + "type": "array", + "items": { + "type": "object", + "properties": { + "path": { "type": "string" }, + "entry_type": { "type": "string", "enum": ["file", "directory"] } + }, + "required": ["path", "entry_type"], + "additionalProperties": false + } + }, + "truncated": { "type": "boolean" } + }, + "required": ["path", "entries", "truncated"], + "additionalProperties": false + })) +} + +pub(crate) fn read_input_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "path": { "type": "string" } + }, + "required": ["path"], + "additionalProperties": false + })) +} + +pub(crate) fn read_output_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "path": { "type": "string" }, + "content": { "type": "string" }, + "truncated": { "type": "boolean" } + }, + "required": ["path", "content", "truncated"], + "additionalProperties": false + })) +} + +pub(crate) fn search_input_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "path": { "type": "string" }, + "max_results": { "type": "integer", "minimum": 1 } + }, + "required": ["query"], + "additionalProperties": false + })) +} + +pub(crate) fn search_output_schema() -> JsonObject { + json_schema(json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "path": { + "anyOf": [{ "type": "string" }, { "type": "null" }] + }, + "matches": { + "type": "array", + "items": { + "type": "object", + "properties": { + "path": { "type": "string" }, + "line_number": { "type": "integer" }, + "line": { "type": "string" } + }, + "required": ["path", "line_number", "line"], + "additionalProperties": false + } + }, + "truncated": { "type": "boolean" } + }, + "required": ["query", "path", "matches", "truncated"], + "additionalProperties": false + })) +} + +fn json_schema(value: serde_json::Value) -> JsonObject { + serde_json::from_value(value) + .unwrap_or_else(|err| panic!("static tool schema should deserialize: {err}")) +} diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs new file mode 100644 index 000000000..bec15ae9f --- /dev/null +++ b/codex-rs/memories/mcp/src/server.rs @@ -0,0 +1,231 @@ +use crate::backend::DEFAULT_LIST_MAX_RESULTS; +use crate::backend::DEFAULT_READ_MAX_TOKENS; +use crate::backend::DEFAULT_SEARCH_MAX_RESULTS; +use crate::backend::ListMemoriesRequest; +use crate::backend::MAX_LIST_RESULTS; +use crate::backend::MAX_SEARCH_RESULTS; +use crate::backend::MemoriesBackend; +use crate::backend::MemoriesBackendError; +use crate::backend::ReadMemoryRequest; +use crate::backend::SearchMemoriesRequest; +use crate::local::LocalMemoriesBackend; +use crate::schema; +use anyhow::Context; +use codex_utils_absolute_path::AbsolutePathBuf; +use rmcp::ErrorData as McpError; +use rmcp::ServiceExt; +use rmcp::handler::server::ServerHandler; +use rmcp::model::CallToolRequestParams; +use rmcp::model::CallToolResult; +use rmcp::model::Content; +use rmcp::model::ListToolsResult; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerInfo; +use rmcp::model::Tool; +use rmcp::model::ToolAnnotations; +use serde::Deserialize; +use serde_json::json; +use std::borrow::Cow; +use std::sync::Arc; + +const LIST_TOOL_NAME: &str = "list"; +const READ_TOOL_NAME: &str = "read"; +const SEARCH_TOOL_NAME: &str = "search"; + +#[derive(Clone)] +pub struct MemoriesMcpServer { + backend: B, + tools: Arc>, +} + +#[derive(Deserialize)] +struct ListArgs { + path: Option, + max_results: Option, +} + +#[derive(Deserialize)] +struct ReadArgs { + path: String, +} + +#[derive(Deserialize)] +struct SearchArgs { + query: String, + path: Option, + max_results: Option, +} + +impl MemoriesMcpServer { + pub fn new(backend: B) -> Self { + Self { + backend, + tools: Arc::new(vec![list_tool(), read_tool(), search_tool()]), + } + } +} + +impl ServerHandler for MemoriesMcpServer { + fn get_info(&self) -> ServerInfo { + ServerInfo { + instructions: Some( + "Use these tools to list, read, and search Codex memory files.".to_string(), + ), + capabilities: ServerCapabilities::builder().enable_tools().build(), + ..ServerInfo::default() + } + } + + fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + Send + '_ { + let tools = Arc::clone(&self.tools); + async move { + Ok(ListToolsResult { + tools: (*tools).clone(), + next_cursor: None, + meta: None, + }) + } + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: rmcp::service::RequestContext, + ) -> Result { + let value = serde_json::Value::Object( + request + .arguments + .unwrap_or_default() + .into_iter() + .collect::>(), + ); + let structured_content = match request.name.as_ref() { + LIST_TOOL_NAME => { + let args: ListArgs = parse_args(value)?; + json!( + self.backend + .list(ListMemoriesRequest { + path: args.path, + max_results: clamp_max_results( + args.max_results, + DEFAULT_LIST_MAX_RESULTS, + MAX_LIST_RESULTS, + ), + }) + .await + .map_err(backend_error_to_mcp)? + ) + } + READ_TOOL_NAME => { + let args: ReadArgs = parse_args(value)?; + json!( + self.backend + .read(ReadMemoryRequest { + path: args.path, + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .map_err(backend_error_to_mcp)? + ) + } + SEARCH_TOOL_NAME => { + let args: SearchArgs = parse_args(value)?; + json!( + self.backend + .search(SearchMemoriesRequest { + query: args.query, + path: args.path, + max_results: clamp_max_results( + args.max_results, + DEFAULT_SEARCH_MAX_RESULTS, + MAX_SEARCH_RESULTS, + ), + }) + .await + .map_err(backend_error_to_mcp)? + ) + } + other => { + return Err(McpError::invalid_params( + format!("unknown tool: {other}"), + None, + )); + } + }; + + Ok(CallToolResult { + content: vec![Content::text(structured_content.to_string())], + structured_content: Some(structured_content), + is_error: Some(false), + meta: None, + }) + } +} + +pub async fn run_stdio_server(codex_home: &AbsolutePathBuf) -> anyhow::Result<()> { + let backend = LocalMemoriesBackend::from_codex_home(codex_home); + tokio::fs::create_dir_all(backend.root()) + .await + .with_context(|| format!("create memories root at {}", backend.root().display()))?; + MemoriesMcpServer::new(backend) + .serve((tokio::io::stdin(), tokio::io::stdout())) + .await? + .waiting() + .await?; + Ok(()) +} + +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."), + Arc::new(schema::list_input_schema()), + ); + tool.output_schema = Some(Arc::new(schema::list_output_schema())); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool +} + +fn read_tool() -> Tool { + let mut tool = Tool::new( + Cow::Borrowed(READ_TOOL_NAME), + Cow::Borrowed("Read a Codex memory file by relative path."), + Arc::new(schema::read_input_schema()), + ); + tool.output_schema = Some(Arc::new(schema::read_output_schema())); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + 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."), + Arc::new(schema::search_input_schema()), + ); + tool.output_schema = Some(Arc::new(schema::search_output_schema())); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool +} + +fn parse_args Deserialize<'de>>(value: serde_json::Value) -> Result { + serde_json::from_value(value).map_err(|err| McpError::invalid_params(err.to_string(), None)) +} + +fn clamp_max_results(requested: Option, default: usize, max: usize) -> usize { + requested.unwrap_or(default).clamp(1, max) +} + +fn backend_error_to_mcp(err: MemoriesBackendError) -> McpError { + match err { + MemoriesBackendError::InvalidPath { .. } + | MemoriesBackendError::NotFile { .. } + | MemoriesBackendError::EmptyQuery => McpError::invalid_params(err.to_string(), None), + MemoriesBackendError::Io(_) => McpError::internal_error(err.to_string(), None), + } +}