mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: memories mcp v1 (#20622)
Add an experimental MCP on memories This must never be used and is only here for testing purpose --------- Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
f48b777717
commit
d013155f40
Generated
+16
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
load("//:defs.bzl", "codex_rust_crate")
|
||||
|
||||
codex_rust_crate(
|
||||
name = "mcp",
|
||||
crate_name = "codex_memories_mcp",
|
||||
)
|
||||
@@ -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"] }
|
||||
@@ -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<Output = Result<ListMemoriesResponse, MemoriesBackendError>> + Send;
|
||||
|
||||
fn read(
|
||||
&self,
|
||||
request: ReadMemoryRequest,
|
||||
) -> impl Future<Output = Result<ReadMemoryResponse, MemoriesBackendError>> + Send;
|
||||
|
||||
fn search(
|
||||
&self,
|
||||
request: SearchMemoriesRequest,
|
||||
) -> impl Future<Output = Result<SearchMemoriesResponse, MemoriesBackendError>> + Send;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ListMemoriesRequest {
|
||||
pub path: Option<String>,
|
||||
pub max_results: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
pub struct ListMemoriesResponse {
|
||||
pub path: Option<String>,
|
||||
pub entries: Vec<MemoryEntry>,
|
||||
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<String>,
|
||||
pub max_results: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
pub struct SearchMemoriesResponse {
|
||||
pub query: String,
|
||||
pub path: Option<String>,
|
||||
pub matches: Vec<MemorySearchMatch>,
|
||||
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<String>, reason: impl Into<String>) -> Self {
|
||||
Self::InvalidPath {
|
||||
path: path.into(),
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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<PathBuf>) -> Self {
|
||||
Self { root: root.into() }
|
||||
}
|
||||
|
||||
pub fn root(&self) -> &Path {
|
||||
&self.root
|
||||
}
|
||||
|
||||
fn resolve_scoped_path(
|
||||
&self,
|
||||
relative_path: Option<&str>,
|
||||
) -> Result<PathBuf, MemoriesBackendError> {
|
||||
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<Option<std::fs::Metadata>, 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<ListMemoriesResponse, MemoriesBackendError> {
|
||||
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<ReadMemoryResponse, MemoriesBackendError> {
|
||||
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<SearchMemoriesResponse, MemoriesBackendError> {
|
||||
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<MemoryEntry>,
|
||||
max_results: usize,
|
||||
) -> Result<bool, MemoriesBackendError> {
|
||||
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<MemorySearchMatch>,
|
||||
max_results: usize,
|
||||
) -> Result<bool, MemoriesBackendError> {
|
||||
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<MemorySearchMatch>,
|
||||
max_results: usize,
|
||||
) -> Result<bool, MemoriesBackendError> {
|
||||
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<Vec<PathBuf>, 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::<Vec<_>>()
|
||||
.join("/")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "local_tests.rs"]
|
||||
mod tests;
|
||||
@@ -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<_>>(),
|
||||
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 { .. }));
|
||||
}
|
||||
@@ -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}"))
|
||||
}
|
||||
@@ -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<B> {
|
||||
backend: B,
|
||||
tools: Arc<Vec<Tool>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ListArgs {
|
||||
path: Option<String>,
|
||||
max_results: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ReadArgs {
|
||||
path: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SearchArgs {
|
||||
query: String,
|
||||
path: Option<String>,
|
||||
max_results: Option<usize>,
|
||||
}
|
||||
|
||||
impl<B: MemoriesBackend> MemoriesMcpServer<B> {
|
||||
pub fn new(backend: B) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
tools: Arc::new(vec![list_tool(), read_tool(), search_tool()]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: MemoriesBackend> ServerHandler for MemoriesMcpServer<B> {
|
||||
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<PaginatedRequestParams>,
|
||||
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
|
||||
) -> impl std::future::Future<Output = Result<ListToolsResult, McpError>> + 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<rmcp::service::RoleServer>,
|
||||
) -> Result<CallToolResult, McpError> {
|
||||
let value = serde_json::Value::Object(
|
||||
request
|
||||
.arguments
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.collect::<serde_json::Map<String, serde_json::Value>>(),
|
||||
);
|
||||
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<T: for<'de> Deserialize<'de>>(value: serde_json::Value) -> Result<T, McpError> {
|
||||
serde_json::from_value(value).map_err(|err| McpError::invalid_params(err.to_string(), None))
|
||||
}
|
||||
|
||||
fn clamp_max_results(requested: Option<usize>, 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),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user