agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: split memories extension crate modules (#22500)

Browse Source 
## Why

The memories extension has several distinct responsibilities:
registering its prompt and tool contributors, enforcing local-memory
filesystem boundaries, implementing list/read/search behavior, and
wrapping that backend as extension tools. Those responsibilities were
concentrated in `lib.rs`, `local.rs`, and the tool modules, which made
follow-up work harder to review and risked growing files through
unrelated edits.

This PR reorganizes the crate so each responsibility has a narrower
owner while preserving the same extension entrypoint and memory tool
behavior.

## What Changed

- Moved extension lifecycle, prompt, and tool registration into
`src/extension.rs`, leaving `src/lib.rs` as the small crate entrypoint.
- Split `LocalMemoriesBackend` helpers into `local/list.rs`,
`local/path.rs`, `local/read.rs`, and `local/search.rs`.
- Centralized tool names and limits at the crate level, and kept the
backend and extension implementation crate-private.
- Made `memory_list`, `memory_read`, and `memory_search` tool executors
generic over `MemoriesBackend`, so tests can exercise the full executor
path without depending on tool internals.
- Consolidated and expanded memory extension tests in `src/tests.rs`,
including read/search tool output coverage, multi-query search, windowed
`all_within_lines`, and legacy `query` rejection.

## Testing

- Not run locally.
This commit is contained in:
jif-oai
2026-05-13 17:39:50 +02:00
committed by GitHub
Unverified
parent 3d517fbd00
commit 2edae8d858
13 changed files with 1029 additions and 866 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/memories/src/backend.rs
-6
View File
@@ -3,12 +3,6 @@ use serde::Deserialize;
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
codex-rs/ext/memories/src/extension.rs
+84
View File
@@ -0,0 +1,84 @@
use std::sync::Arc;
use codex_core::config::Config;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PromptFragment;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ToolContributor;
use codex_features::Feature;
use codex_memories_read::build_memory_tool_developer_instructions;
use codex_utils_absolute_path::AbsolutePathBuf;
use crate::local::LocalMemoriesBackend;
use crate::tools;
/// Contributes Codex memory read-path prompt context and memory read tools.
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct MemoriesExtension;
#[derive(Clone, Debug)]
pub(crate) struct MemoriesExtensionConfig {
pub(crate) enabled: bool,
pub(crate) codex_home: AbsolutePathBuf,
}
impl ContextContributor for MemoriesExtension {
fn contribute<'a>(
&'a self,
_session_store: &'a ExtensionData,
thread_store: &'a ExtensionData,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
Box::pin(async move {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled {
return Vec::new();
}
build_memory_tool_developer_instructions(&config.codex_home)
.await
.map(PromptFragment::developer_policy)
.into_iter()
.collect()
})
}
}
impl ThreadLifecycleContributor<Config> for MemoriesExtension {
fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
input.thread_store.insert(MemoriesExtensionConfig {
enabled: input.config.features.enabled(Feature::MemoryTool)
&& input.config.memories.use_memories,
codex_home: input.config.codex_home.clone(),
});
}
}
impl ToolContributor for MemoriesExtension {
fn tools(
&self,
_session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<Arc<dyn codex_extension_api::ExtensionToolExecutor>> {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled {
return Vec::new();
}
tools::memory_tools(LocalMemoriesBackend::from_codex_home(&config.codex_home))
}
}
/// Installs the memories extension contributors into the extension registry.
pub fn install(registry: &mut ExtensionRegistryBuilder<Config>) {
let extension = Arc::new(MemoriesExtension);
registry.thread_lifecycle_contributor(extension.clone());
registry.prompt_contributor(extension.clone());
registry.tool_contributor(extension);
}
codex-rs/ext/memories/src/lib.rs
+11 -183
View File
@@ -1,192 +1,20 @@
use std::sync::Arc;
use codex_core::config::Config;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PromptFragment;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ToolContributor;
use codex_features::Feature;
use codex_memories_read::build_memory_tool_developer_instructions;
use codex_utils_absolute_path::AbsolutePathBuf;
mod backend;
mod extension;
mod local;
mod schema;
mod tools;
use local::LocalMemoriesBackend;
pub use extension::install;
/// Contributes Codex memory read-path prompt context and memory read tools.
#[derive(Clone, Copy, Debug, Default)]
pub struct MemoriesExtension;
pub(crate) const DEFAULT_LIST_MAX_RESULTS: usize = 2_000;
pub(crate) const MAX_LIST_RESULTS: usize = 2_000;
pub(crate) const DEFAULT_SEARCH_MAX_RESULTS: usize = 200;
pub(crate) const MAX_SEARCH_RESULTS: usize = 200;
pub(crate) const DEFAULT_READ_MAX_TOKENS: usize = 20_000;
#[derive(Clone, Debug)]
struct MemoriesExtensionConfig {
enabled: bool,
codex_home: AbsolutePathBuf,
}
impl ContextContributor for MemoriesExtension {
fn contribute<'a>(
&'a self,
_session_store: &'a ExtensionData,
thread_store: &'a ExtensionData,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
Box::pin(async move {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled {
return Vec::new();
}
build_memory_tool_developer_instructions(&config.codex_home)
.await
.map(PromptFragment::developer_policy)
.into_iter()
.collect()
})
}
}
impl ThreadLifecycleContributor<Config> for MemoriesExtension {
fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
input.thread_store.insert(MemoriesExtensionConfig {
enabled: input.config.features.enabled(Feature::MemoryTool)
&& input.config.memories.use_memories,
codex_home: input.config.codex_home.clone(),
});
}
}
impl ToolContributor for MemoriesExtension {
fn tools(
&self,
_session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<Arc<dyn codex_extension_api::ExtensionToolExecutor>> {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled {
return Vec::new();
}
tools::memory_tools(LocalMemoriesBackend::from_codex_home(&config.codex_home))
}
}
/// Installs the memories extension contributors into the extension registry.
pub fn install(registry: &mut ExtensionRegistryBuilder<Config>) {
let extension = Arc::new(MemoriesExtension);
registry.thread_lifecycle_contributor(extension.clone());
registry.prompt_contributor(extension.clone());
registry.tool_contributor(extension);
}
pub(crate) const LIST_TOOL_NAME: &str = "memory_list";
pub(crate) const READ_TOOL_NAME: &str = "memory_read";
pub(crate) const SEARCH_TOOL_NAME: &str = "memory_search";
#[cfg(test)]
mod tests {
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::PromptSlot;
use codex_extension_api::ToolContributor;
use codex_utils_absolute_path::test_support::PathBufExt;
use codex_utils_absolute_path::test_support::PathExt;
use codex_utils_absolute_path::test_support::test_path_buf;
use pretty_assertions::assert_eq;
use super::MemoriesExtension;
use super::MemoriesExtensionConfig;
#[test]
fn tools_are_not_contributed_without_thread_config() {
let extension = MemoriesExtension;
assert!(
extension
.tools(
&ExtensionData::new("session"),
&ExtensionData::new("thread")
)
.is_empty()
);
}
#[test]
fn tools_are_not_contributed_when_disabled() {
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: false,
codex_home: test_path_buf("/tmp/codex-home").abs(),
});
assert!(
extension
.tools(&ExtensionData::new("session"), &thread_store)
.is_empty()
);
}
#[test]
fn tools_are_contributed_when_enabled() {
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: true,
codex_home: test_path_buf("/tmp/codex-home").abs(),
});
let tool_names = extension
.tools(&ExtensionData::new("session"), &thread_store)
.into_iter()
.map(|tool| tool.tool_name().to_string())
.collect::<Vec<_>>();
assert_eq!(
tool_names,
vec![
"memory_list".to_string(),
"memory_read".to_string(),
"memory_search".to_string()
]
);
}
#[tokio::test]
async fn prompt_contribution_uses_memory_summary_when_enabled() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memories_dir = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memories_dir)
.await
.expect("create memories dir");
tokio::fs::write(
memories_dir.join("memory_summary.md"),
"Remember repository-specific implementation preferences.",
)
.await
.expect("write memory summary");
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: true,
codex_home: tempdir.path().abs(),
});
let fragments = extension
.contribute(&ExtensionData::new("session"), &thread_store)
.await;
assert_eq!(fragments.len(), 1);
assert_eq!(fragments[0].slot(), PromptSlot::DeveloperPolicy);
assert!(
fragments[0]
.text()
.contains("Remember repository-specific implementation preferences.")
);
}
}
mod tests;
codex-rs/ext/memories/src/local.rs
+27 -524
View File
@@ -1,37 +1,34 @@
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::SearchMatchMode;
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::borrow::Cow;
use std::path::Component;
use std::path::Path;
use std::path::PathBuf;
use codex_utils_absolute_path::AbsolutePathBuf;
use crate::backend::ListMemoriesRequest;
use crate::backend::ListMemoriesResponse;
use crate::backend::MemoriesBackend;
use crate::backend::MemoriesBackendError;
use crate::backend::ReadMemoryRequest;
use crate::backend::ReadMemoryResponse;
use crate::backend::SearchMemoriesRequest;
use crate::backend::SearchMemoriesResponse;
mod list;
mod path;
mod read;
mod search;
#[derive(Debug, Clone)]
pub struct LocalMemoriesBackend {
pub(crate) struct LocalMemoriesBackend {
root: PathBuf,
}
impl LocalMemoriesBackend {
pub fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self {
pub(crate) 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 {
pub(crate) fn from_memory_root(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
@@ -54,7 +51,7 @@ impl LocalMemoriesBackend {
"must stay within the memories root",
));
}
if relative.components().any(is_hidden_component) {
if relative.components().any(path::is_hidden_component) {
return Err(MemoriesBackendError::NotFound {
path: relative_path.to_string(),
});
@@ -72,7 +69,10 @@ impl LocalMemoriesBackend {
return Ok(scoped_path);
};
reject_symlink(&display_relative_path(&self.root, &scoped_path), &metadata)?;
path::reject_symlink(
&path::display_relative_path(&self.root, &scoped_path),
&metadata,
)?;
if idx + 1 < components.len() && !metadata.is_dir() {
return Err(MemoriesBackendError::invalid_path(
relative_path,
@@ -100,517 +100,20 @@ impl MemoriesBackend for LocalMemoriesBackend {
&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()).await?;
let start_index = match request.cursor.as_deref() {
Some(cursor) => cursor.parse::<usize>().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 Err(MemoriesBackendError::NotFound {
path: request.path.unwrap_or_default(),
});
};
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? {
if is_hidden_path(&path) {
continue;
}
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.drain(start_index..end_index).collect(),
next_cursor,
truncated,
})
list::list(self, request).await
}
async fn read(
&self,
request: ReadMemoryRequest,
) -> Result<ReadMemoryResponse, MemoriesBackendError> {
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()))
.await?;
let Some(metadata) = Self::metadata_or_none(&path).await? else {
return Err(MemoriesBackendError::NotFound { 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 start_byte = line_start_byte_offset(&original_content, request.line_offset)?;
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 = end_byte < original_content.len() || content != content_from_offset;
Ok(ReadMemoryResponse {
path: request.path,
start_line_number: request.line_offset,
content,
truncated,
})
read::read(self, request).await
}
async fn search(
&self,
request: SearchMemoriesRequest,
) -> Result<SearchMemoriesResponse, MemoriesBackendError> {
let queries = request
.queries
.iter()
.map(|query| query.trim().to_string())
.collect::<Vec<_>>();
if queries.is_empty() || queries.iter().any(std::string::String::is_empty) {
return Err(MemoriesBackendError::EmptyQuery);
}
if matches!(
request.match_mode,
SearchMatchMode::AllWithinLines { line_count: 0 }
) {
return Err(MemoriesBackendError::InvalidMatchWindow);
}
let max_results = request.max_results.min(MAX_SEARCH_RESULTS);
let start = self.resolve_scoped_path(request.path.as_deref()).await?;
let start_index = match request.cursor.as_deref() {
Some(cursor) => cursor.parse::<usize>().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 Err(MemoriesBackendError::NotFound {
path: request.path.unwrap_or_default(),
});
};
reject_symlink(&display_relative_path(&self.root, &start), &metadata)?;
let matcher = SearchMatcher::new(
queries.clone(),
request.match_mode.clone(),
request.case_sensitive,
request.normalized,
)?;
let mut matches = Vec::new();
search_entries(
&self.root,
&start,
&metadata,
&matcher,
request.context_lines,
&mut matches,
)
.await?;
matches.sort_by(|left, right| {
left.path
.cmp(&right.path)
.then(left.match_line_number.cmp(&right.match_line_number))
});
if start_index > matches.len() {
return Err(MemoriesBackendError::invalid_cursor(
start_index.to_string(),
"exceeds result count",
));
}
let end_index = start_index.saturating_add(max_results).min(matches.len());
let next_cursor = (end_index < matches.len()).then(|| end_index.to_string());
let truncated = next_cursor.is_some();
Ok(SearchMemoriesResponse {
queries,
match_mode: request.match_mode,
path: request.path,
matches: matches.drain(start_index..end_index).collect(),
next_cursor,
truncated,
})
search::search(self, request).await
}
}
async fn search_entries(
root: &Path,
current: &Path,
current_metadata: &std::fs::Metadata,
matcher: &SearchMatcher,
context_lines: usize,
matches: &mut Vec<MemorySearchMatch>,
) -> Result<(), MemoriesBackendError> {
if current_metadata.is_file() {
search_file(root, current, matcher, context_lines, matches).await?;
return Ok(());
}
if !current_metadata.is_dir() {
return Ok(());
}
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 is_hidden_path(&path) {
continue;
}
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, matcher, context_lines, matches).await?;
}
}
}
Ok(())
}
async fn search_file(
root: &Path,
path: &Path,
matcher: &SearchMatcher,
context_lines: usize,
matches: &mut Vec<MemorySearchMatch>,
) -> Result<(), 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(()),
Err(err) => return Err(err.into()),
};
let lines = content.lines().collect::<Vec<_>>();
let line_matches = lines
.iter()
.map(|line| matcher.matched_query_flags(line))
.collect::<Vec<_>>();
match &matcher.match_mode {
SearchMatchMode::Any => {
for (idx, matched_query_flags) in line_matches.iter().enumerate() {
if matched_query_flags.iter().any(|matched| *matched) {
matches.push(build_search_match(
root,
path,
&lines,
idx,
idx,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
SearchMatchMode::AllOnSameLine => {
for (idx, matched_query_flags) in line_matches.iter().enumerate() {
if matched_query_flags.iter().all(|matched| *matched) {
matches.push(build_search_match(
root,
path,
&lines,
idx,
idx,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
SearchMatchMode::AllWithinLines { line_count } => {
let mut windows = Vec::new();
for start_index in 0..lines.len() {
if !line_matches[start_index].iter().any(|matched| *matched) {
continue;
}
let last_allowed_index = start_index
.saturating_add(line_count.saturating_sub(1))
.min(lines.len().saturating_sub(1));
let mut matched_query_flags = vec![false; matcher.queries.len()];
for (end_index, line_match_flags) in line_matches
.iter()
.enumerate()
.take(last_allowed_index + 1)
.skip(start_index)
{
for (idx, matched) in line_match_flags.iter().enumerate() {
matched_query_flags[idx] |= matched;
}
if matched_query_flags.iter().all(|matched| *matched) {
windows.push((start_index, end_index, matched_query_flags));
break;
}
}
}
for (idx, (start_index, end_index, matched_query_flags)) in windows.iter().enumerate() {
let strictly_contains_another_window = windows.iter().enumerate().any(
|(other_idx, (other_start_index, other_end_index, _))| {
idx != other_idx
&& start_index <= other_start_index
&& end_index >= other_end_index
&& (start_index != other_start_index || end_index != other_end_index)
},
);
if strictly_contains_another_window {
continue;
}
matches.push(build_search_match(
root,
path,
&lines,
*start_index,
*end_index,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
Ok(())
}
fn build_search_match(
root: &Path,
path: &Path,
lines: &[&str],
match_start_index: usize,
match_end_index: usize,
context_lines: usize,
matched_queries: Vec<String>,
) -> MemorySearchMatch {
let content_start_index = match_start_index.saturating_sub(context_lines);
let content_end_index = match_end_index
.saturating_add(context_lines)
.saturating_add(1)
.min(lines.len());
MemorySearchMatch {
path: display_relative_path(root, path),
match_line_number: match_start_index + 1,
content_start_line_number: content_start_index + 1,
content: lines[content_start_index..content_end_index].join("\n"),
matched_queries,
}
}
struct SearchMatcher {
queries: Vec<String>,
prepared_queries: Vec<String>,
comparison: SearchComparison,
match_mode: SearchMatchMode,
}
impl SearchMatcher {
fn new(
queries: Vec<String>,
match_mode: SearchMatchMode,
case_sensitive: bool,
normalized: bool,
) -> Result<Self, MemoriesBackendError> {
let comparison = SearchComparison::new(case_sensitive, normalized);
let prepared_queries = queries
.iter()
.map(|query| comparison.prepare(query))
.map(Cow::into_owned)
.collect::<Vec<_>>();
if prepared_queries.iter().any(std::string::String::is_empty) {
return Err(MemoriesBackendError::EmptyQuery);
}
Ok(Self {
queries,
prepared_queries,
comparison,
match_mode,
})
}
fn matched_query_flags(&self, line: &str) -> Vec<bool> {
let line = self.comparison.prepare(line);
self.prepared_queries
.iter()
.map(|query| line.as_ref().contains(query))
.collect()
}
fn matched_queries(&self, matched_query_flags: &[bool]) -> Vec<String> {
self.queries
.iter()
.zip(matched_query_flags)
.filter_map(|(query, matched)| matched.then_some(query.clone()))
.collect()
}
}
#[derive(Clone, Copy)]
struct SearchComparison {
case_sensitive: bool,
normalized: bool,
}
impl SearchComparison {
fn new(case_sensitive: bool, normalized: bool) -> Self {
Self {
case_sensitive,
normalized,
}
}
fn prepare<'a>(self, value: &'a str) -> Cow<'a, str> {
if self.case_sensitive && !self.normalized {
return Cow::Borrowed(value);
}
let value = if self.case_sensitive {
Cow::Borrowed(value)
} else {
Cow::Owned(value.to_lowercase())
};
if !self.normalized {
return value;
}
Cow::Owned(
value
.chars()
.filter(|ch| ch.is_alphanumeric())
.collect::<String>(),
)
}
}
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 is_hidden_component(component: Component<'_>) -> bool {
matches!(
component,
Component::Normal(name) if name.to_string_lossy().starts_with('.')
)
}
fn is_hidden_path(path: &Path) -> bool {
path.file_name()
.is_some_and(|name| name.to_string_lossy().starts_with('.'))
}
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("/")
}
fn line_start_byte_offset(
content: &str,
line_offset: usize,
) -> Result<usize, MemoriesBackendError> {
if line_offset == 1 {
return Ok(0);
}
let mut current_line = 1;
for (idx, ch) in content.char_indices() {
if ch == '\n' {
current_line += 1;
if current_line == line_offset {
return Ok(idx + 1);
}
}
}
Err(MemoriesBackendError::LineOffsetExceedsFileLength)
}
fn line_end_byte_offset(content: &str, start_byte: usize, max_lines: Option<usize>) -> 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()
}
codex-rs/ext/memories/src/local/list.rs
+83
View File
@@ -0,0 +1,83 @@
use crate::MAX_LIST_RESULTS;
use crate::backend::ListMemoriesRequest;
use crate::backend::ListMemoriesResponse;
use crate::backend::MemoriesBackendError;
use crate::backend::MemoryEntry;
use crate::backend::MemoryEntryType;
use super::LocalMemoriesBackend;
use super::path::display_relative_path;
use super::path::is_hidden_path;
use super::path::read_sorted_dir_paths;
use super::path::reject_symlink;
pub(super) async fn list(
backend: &LocalMemoriesBackend,
request: ListMemoriesRequest,
) -> Result<ListMemoriesResponse, MemoriesBackendError> {
let max_results = request.max_results.min(MAX_LIST_RESULTS);
let start = backend.resolve_scoped_path(request.path.as_deref()).await?;
let start_index = match request.cursor.as_deref() {
Some(cursor) => cursor.parse::<usize>().map_err(|_| {
MemoriesBackendError::invalid_cursor(cursor, "must be a non-negative integer")
})?,
None => 0,
};
let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&start).await? else {
return Err(MemoriesBackendError::NotFound {
path: request.path.unwrap_or_default(),
});
};
reject_symlink(&display_relative_path(&backend.root, &start), &metadata)?;
let mut entries = if metadata.is_file() {
vec![MemoryEntry {
path: display_relative_path(&backend.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? {
if is_hidden_path(&path) {
continue;
}
let Some(metadata) = LocalMemoriesBackend::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(&backend.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.drain(start_index..end_index).collect(),
next_cursor,
truncated,
})
}
codex-rs/ext/memories/src/local/path.rs
+56
View File
@@ -0,0 +1,56 @@
use std::path::Component;
use std::path::Path;
use std::path::PathBuf;
use crate::backend::MemoriesBackendError;
pub(super) 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)
}
pub(super) 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(())
}
pub(super) fn is_hidden_component(component: Component<'_>) -> bool {
matches!(
component,
Component::Normal(name) if name.to_string_lossy().starts_with('.')
)
}
pub(super) fn is_hidden_path(path: &Path) -> bool {
path.file_name()
.is_some_and(|name| name.to_string_lossy().starts_with('.'))
}
pub(super) 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("/")
}
codex-rs/ext/memories/src/local/read.rs
+90
View File
@@ -0,0 +1,90 @@
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::truncate_text;
use crate::DEFAULT_READ_MAX_TOKENS;
use crate::backend::MemoriesBackendError;
use crate::backend::ReadMemoryRequest;
use crate::backend::ReadMemoryResponse;
use super::LocalMemoriesBackend;
use super::path::reject_symlink;
pub(super) async fn read(
backend: &LocalMemoriesBackend,
request: ReadMemoryRequest,
) -> Result<ReadMemoryResponse, MemoriesBackendError> {
if request.line_offset == 0 {
return Err(MemoriesBackendError::InvalidLineOffset);
}
if request.max_lines == Some(0) {
return Err(MemoriesBackendError::InvalidMaxLines);
}
let path = backend
.resolve_scoped_path(Some(request.path.as_str()))
.await?;
let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&path).await? else {
return Err(MemoriesBackendError::NotFound { 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 start_byte = line_start_byte_offset(&original_content, request.line_offset)?;
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 = end_byte < original_content.len() || content != content_from_offset;
Ok(ReadMemoryResponse {
path: request.path,
start_line_number: request.line_offset,
content,
truncated,
})
}
fn line_start_byte_offset(
content: &str,
line_offset: usize,
) -> Result<usize, MemoriesBackendError> {
if line_offset == 1 {
return Ok(0);
}
let mut current_line = 1;
for (idx, ch) in content.char_indices() {
if ch == '\n' {
current_line += 1;
if current_line == line_offset {
return Ok(idx + 1);
}
}
}
Err(MemoriesBackendError::LineOffsetExceedsFileLength)
}
fn line_end_byte_offset(content: &str, start_byte: usize, max_lines: Option<usize>) -> 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()
}
codex-rs/ext/memories/src/local/search.rs
+336
View File
@@ -0,0 +1,336 @@
use std::borrow::Cow;
use std::path::Path;
use crate::MAX_SEARCH_RESULTS;
use crate::backend::MemoriesBackendError;
use crate::backend::MemorySearchMatch;
use crate::backend::SearchMatchMode;
use crate::backend::SearchMemoriesRequest;
use crate::backend::SearchMemoriesResponse;
use super::LocalMemoriesBackend;
use super::path::display_relative_path;
use super::path::is_hidden_path;
use super::path::read_sorted_dir_paths;
use super::path::reject_symlink;
pub(super) async fn search(
backend: &LocalMemoriesBackend,
request: SearchMemoriesRequest,
) -> Result<SearchMemoriesResponse, MemoriesBackendError> {
let queries = request
.queries
.iter()
.map(|query| query.trim().to_string())
.collect::<Vec<_>>();
if queries.is_empty() || queries.iter().any(std::string::String::is_empty) {
return Err(MemoriesBackendError::EmptyQuery);
}
if matches!(
request.match_mode,
SearchMatchMode::AllWithinLines { line_count: 0 }
) {
return Err(MemoriesBackendError::InvalidMatchWindow);
}
let max_results = request.max_results.min(MAX_SEARCH_RESULTS);
let start = backend.resolve_scoped_path(request.path.as_deref()).await?;
let start_index = match request.cursor.as_deref() {
Some(cursor) => cursor.parse::<usize>().map_err(|_| {
MemoriesBackendError::invalid_cursor(cursor, "must be a non-negative integer")
})?,
None => 0,
};
let Some(metadata) = LocalMemoriesBackend::metadata_or_none(&start).await? else {
return Err(MemoriesBackendError::NotFound {
path: request.path.unwrap_or_default(),
});
};
reject_symlink(&display_relative_path(&backend.root, &start), &metadata)?;
let matcher = SearchMatcher::new(
queries.clone(),
request.match_mode.clone(),
request.case_sensitive,
request.normalized,
)?;
let mut matches = Vec::new();
search_entries(
&backend.root,
&start,
&metadata,
&matcher,
request.context_lines,
&mut matches,
)
.await?;
matches.sort_by(|left, right| {
left.path
.cmp(&right.path)
.then(left.match_line_number.cmp(&right.match_line_number))
});
if start_index > matches.len() {
return Err(MemoriesBackendError::invalid_cursor(
start_index.to_string(),
"exceeds result count",
));
}
let end_index = start_index.saturating_add(max_results).min(matches.len());
let next_cursor = (end_index < matches.len()).then(|| end_index.to_string());
let truncated = next_cursor.is_some();
Ok(SearchMemoriesResponse {
queries,
match_mode: request.match_mode,
path: request.path,
matches: matches.drain(start_index..end_index).collect(),
next_cursor,
truncated,
})
}
async fn search_entries(
root: &Path,
current: &Path,
current_metadata: &std::fs::Metadata,
matcher: &SearchMatcher,
context_lines: usize,
matches: &mut Vec<MemorySearchMatch>,
) -> Result<(), MemoriesBackendError> {
if current_metadata.is_file() {
search_file(root, current, matcher, context_lines, matches).await?;
return Ok(());
}
if !current_metadata.is_dir() {
return Ok(());
}
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 is_hidden_path(&path) {
continue;
}
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, matcher, context_lines, matches).await?;
}
}
}
Ok(())
}
async fn search_file(
root: &Path,
path: &Path,
matcher: &SearchMatcher,
context_lines: usize,
matches: &mut Vec<MemorySearchMatch>,
) -> Result<(), 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(()),
Err(err) => return Err(err.into()),
};
let lines = content.lines().collect::<Vec<_>>();
let line_matches = lines
.iter()
.map(|line| matcher.matched_query_flags(line))
.collect::<Vec<_>>();
match &matcher.match_mode {
SearchMatchMode::Any => {
for (idx, matched_query_flags) in line_matches.iter().enumerate() {
if matched_query_flags.iter().any(|matched| *matched) {
matches.push(build_search_match(
root,
path,
&lines,
idx,
idx,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
SearchMatchMode::AllOnSameLine => {
for (idx, matched_query_flags) in line_matches.iter().enumerate() {
if matched_query_flags.iter().all(|matched| *matched) {
matches.push(build_search_match(
root,
path,
&lines,
idx,
idx,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
SearchMatchMode::AllWithinLines { line_count } => {
let mut windows = Vec::new();
for start_index in 0..lines.len() {
if !line_matches[start_index].iter().any(|matched| *matched) {
continue;
}
let last_allowed_index = start_index
.saturating_add(line_count.saturating_sub(1))
.min(lines.len().saturating_sub(1));
let mut matched_query_flags = vec![false; matcher.queries.len()];
for (end_index, line_match_flags) in line_matches
.iter()
.enumerate()
.take(last_allowed_index + 1)
.skip(start_index)
{
for (idx, matched) in line_match_flags.iter().enumerate() {
matched_query_flags[idx] |= matched;
}
if matched_query_flags.iter().all(|matched| *matched) {
windows.push((start_index, end_index, matched_query_flags));
break;
}
}
}
for (idx, (start_index, end_index, matched_query_flags)) in windows.iter().enumerate() {
let strictly_contains_another_window = windows.iter().enumerate().any(
|(other_idx, (other_start_index, other_end_index, _))| {
idx != other_idx
&& start_index <= other_start_index
&& end_index >= other_end_index
&& (start_index != other_start_index || end_index != other_end_index)
},
);
if strictly_contains_another_window {
continue;
}
matches.push(build_search_match(
root,
path,
&lines,
*start_index,
*end_index,
context_lines,
matcher.matched_queries(matched_query_flags),
));
}
}
}
Ok(())
}
fn build_search_match(
root: &Path,
path: &Path,
lines: &[&str],
match_start_index: usize,
match_end_index: usize,
context_lines: usize,
matched_queries: Vec<String>,
) -> MemorySearchMatch {
let content_start_index = match_start_index.saturating_sub(context_lines);
let content_end_index = match_end_index
.saturating_add(context_lines)
.saturating_add(1)
.min(lines.len());
MemorySearchMatch {
path: display_relative_path(root, path),
match_line_number: match_start_index + 1,
content_start_line_number: content_start_index + 1,
content: lines[content_start_index..content_end_index].join("\n"),
matched_queries,
}
}
struct SearchMatcher {
queries: Vec<String>,
prepared_queries: Vec<String>,
comparison: SearchComparison,
match_mode: SearchMatchMode,
}
impl SearchMatcher {
fn new(
queries: Vec<String>,
match_mode: SearchMatchMode,
case_sensitive: bool,
normalized: bool,
) -> Result<Self, MemoriesBackendError> {
let comparison = SearchComparison::new(case_sensitive, normalized);
let prepared_queries = queries
.iter()
.map(|query| comparison.prepare(query))
.map(Cow::into_owned)
.collect::<Vec<_>>();
if prepared_queries.iter().any(std::string::String::is_empty) {
return Err(MemoriesBackendError::EmptyQuery);
}
Ok(Self {
queries,
prepared_queries,
comparison,
match_mode,
})
}
fn matched_query_flags(&self, line: &str) -> Vec<bool> {
let line = self.comparison.prepare(line);
self.prepared_queries
.iter()
.map(|query| line.as_ref().contains(query))
.collect()
}
fn matched_queries(&self, matched_query_flags: &[bool]) -> Vec<String> {
self.queries
.iter()
.zip(matched_query_flags)
.filter_map(|(query, matched)| matched.then_some(query.clone()))
.collect()
}
}
#[derive(Clone, Copy)]
struct SearchComparison {
case_sensitive: bool,
normalized: bool,
}
impl SearchComparison {
fn new(case_sensitive: bool, normalized: bool) -> Self {
Self {
case_sensitive,
normalized,
}
}
fn prepare<'a>(self, value: &'a str) -> Cow<'a, str> {
if self.case_sensitive && !self.normalized {
return Cow::Borrowed(value);
}
let value = if self.case_sensitive {
Cow::Borrowed(value)
} else {
Cow::Owned(value.to_lowercase())
};
if !self.normalized {
return value;
}
Cow::Owned(
value
.chars()
.filter(|ch| ch.is_alphanumeric())
.collect::<String>(),
)
}
}
codex-rs/ext/memories/src/tests.rs
+311
View File
@@ -0,0 +1,311 @@
use std::path::Path;
use std::sync::Arc;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionToolExecutor;
use codex_extension_api::PromptSlot;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolContributor;
use codex_extension_api::ToolName;
use codex_extension_api::ToolPayload;
use codex_tools::ToolOutput;
use codex_utils_absolute_path::test_support::PathBufExt;
use codex_utils_absolute_path::test_support::PathExt;
use codex_utils_absolute_path::test_support::test_path_buf;
use pretty_assertions::assert_eq;
use serde_json::json;
use crate::extension::MemoriesExtension;
use crate::extension::MemoriesExtensionConfig;
use crate::local::LocalMemoriesBackend;
#[test]
fn tools_are_not_contributed_without_thread_config() {
let extension = MemoriesExtension;
assert!(
extension
.tools(
&ExtensionData::new("session"),
&ExtensionData::new("thread")
)
.is_empty()
);
}
#[test]
fn tools_are_not_contributed_when_disabled() {
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: false,
codex_home: test_path_buf("/tmp/codex-home").abs(),
});
assert!(
extension
.tools(&ExtensionData::new("session"), &thread_store)
.is_empty()
);
}
#[test]
fn tools_are_contributed_when_enabled() {
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: true,
codex_home: test_path_buf("/tmp/codex-home").abs(),
});
let tool_names = extension
.tools(&ExtensionData::new("session"), &thread_store)
.into_iter()
.map(|tool| tool.tool_name().to_string())
.collect::<Vec<_>>();
assert_eq!(
tool_names,
vec![
"memory_list".to_string(),
"memory_read".to_string(),
"memory_search".to_string()
]
);
}
#[tokio::test]
async fn prompt_contribution_uses_memory_summary_when_enabled() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memories_dir = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memories_dir)
.await
.expect("create memories dir");
tokio::fs::write(
memories_dir.join("memory_summary.md"),
"Remember repository-specific implementation preferences.",
)
.await
.expect("write memory summary");
let extension = MemoriesExtension;
let thread_store = ExtensionData::new("thread");
thread_store.insert(MemoriesExtensionConfig {
enabled: true,
codex_home: tempdir.path().abs(),
});
let fragments = extension
.contribute(&ExtensionData::new("session"), &thread_store)
.await;
assert_eq!(fragments.len(), 1);
assert_eq!(fragments[0].slot(), PromptSlot::DeveloperPolicy);
assert!(
fragments[0]
.text()
.contains("Remember repository-specific implementation preferences.")
);
}
#[tokio::test]
async fn read_tool_reads_memory_file() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memory_root = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memory_root)
.await
.expect("create memories dir");
tokio::fs::write(
memory_root.join("MEMORY.md"),
"first line\nsecond needle line\nthird line\n",
)
.await
.expect("write memory");
let tool = memory_tool(&memory_root, crate::READ_TOOL_NAME);
let payload = ToolPayload::Function {
arguments: json!({
"path": "MEMORY.md",
"line_offset": 2,
"max_lines": 1
})
.to_string(),
};
let output = tool
.handle(ToolCall {
call_id: "call-1".to_string(),
tool_name: ToolName::plain(crate::READ_TOOL_NAME),
payload: payload.clone(),
})
.await
.expect("read should succeed");
assert_eq!(
output.post_tool_use_response("call-1", &payload),
Some(json!({
"path": "MEMORY.md",
"content": "second needle line\n",
"start_line_number": 2,
"truncated": true
}))
);
}
#[tokio::test]
async fn search_tool_accepts_multiple_queries() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memory_root = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memory_root)
.await
.expect("create memories dir");
tokio::fs::write(
memory_root.join("MEMORY.md"),
"alpha only\nneedle only\nalpha needle\n",
)
.await
.expect("write memory");
let tool = memory_tool(&memory_root, crate::SEARCH_TOOL_NAME);
let payload = ToolPayload::Function {
arguments: json!({
"queries": ["alpha", "needle"],
"case_sensitive": false
})
.to_string(),
};
let output = tool
.handle(ToolCall {
call_id: "call-1".to_string(),
tool_name: ToolName::plain(crate::SEARCH_TOOL_NAME),
payload: payload.clone(),
})
.await
.expect("search should succeed");
assert_eq!(
output.post_tool_use_response("call-1", &payload),
Some(json!({
"queries": ["alpha", "needle"],
"match_mode": {
"type": "any"
},
"path": null,
"matches": [
{
"path": "MEMORY.md",
"match_line_number": 1,
"content_start_line_number": 1,
"content": "alpha only",
"matched_queries": ["alpha"]
},
{
"path": "MEMORY.md",
"match_line_number": 2,
"content_start_line_number": 2,
"content": "needle only",
"matched_queries": ["needle"]
},
{
"path": "MEMORY.md",
"match_line_number": 3,
"content_start_line_number": 3,
"content": "alpha needle",
"matched_queries": ["alpha", "needle"]
}
],
"next_cursor": null,
"truncated": false
}))
);
}
#[tokio::test]
async fn search_tool_accepts_windowed_all_match_mode() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memory_root = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memory_root)
.await
.expect("create memories dir");
tokio::fs::write(memory_root.join("MEMORY.md"), "alpha\nmiddle\nneedle\n")
.await
.expect("write memory");
let tool = memory_tool(&memory_root, crate::SEARCH_TOOL_NAME);
let payload = ToolPayload::Function {
arguments: json!({
"queries": ["alpha", "needle"],
"match_mode": {
"type": "all_within_lines",
"line_count": 3
}
})
.to_string(),
};
let output = tool
.handle(ToolCall {
call_id: "call-1".to_string(),
tool_name: ToolName::plain(crate::SEARCH_TOOL_NAME),
payload: payload.clone(),
})
.await
.expect("search should succeed");
assert_eq!(
output.post_tool_use_response("call-1", &payload),
Some(json!({
"queries": ["alpha", "needle"],
"match_mode": {
"type": "all_within_lines",
"line_count": 3
},
"path": null,
"matches": [
{
"path": "MEMORY.md",
"match_line_number": 1,
"content_start_line_number": 1,
"content": "alpha\nmiddle\nneedle",
"matched_queries": ["alpha", "needle"]
}
],
"next_cursor": null,
"truncated": false
}))
);
}
#[tokio::test]
async fn search_tool_rejects_legacy_single_query() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memory_root = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memory_root)
.await
.expect("create memories dir");
let tool = memory_tool(&memory_root, crate::SEARCH_TOOL_NAME);
let payload = ToolPayload::Function {
arguments: json!({
"query": "needle",
})
.to_string(),
};
let err = tool
.handle(ToolCall {
call_id: "call-1".to_string(),
tool_name: ToolName::plain(crate::SEARCH_TOOL_NAME),
payload,
})
.await
.expect_err("legacy query field should be rejected");
assert!(err.to_string().contains("unknown field"));
assert!(err.to_string().contains("query"));
}
fn memory_tool(memory_root: &Path, tool_name: &str) -> Arc<dyn ExtensionToolExecutor> {
crate::tools::memory_tools(LocalMemoriesBackend::from_memory_root(memory_root))
.into_iter()
.find(|tool| tool.tool_name().to_string() == tool_name)
.unwrap_or_else(|| panic!("{tool_name} tool should be registered"))
}
codex-rs/ext/memories/src/tools/list.rs
+9 -7
View File
@@ -8,14 +8,13 @@ use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use crate::backend::DEFAULT_LIST_MAX_RESULTS;
use crate::DEFAULT_LIST_MAX_RESULTS;
use crate::LIST_TOOL_NAME;
use crate::MAX_LIST_RESULTS;
use crate::backend::ListMemoriesRequest;
use crate::backend::ListMemoriesResponse;
use crate::backend::MAX_LIST_RESULTS;
use crate::backend::MemoriesBackend;
use crate::local::LocalMemoriesBackend;
use super::LIST_TOOL_NAME;
use super::backend_error_to_function_call;
use super::clamp_max_results;
use super::function_tool;
@@ -31,11 +30,14 @@ struct ListArgs {
}
#[derive(Clone)]
pub(super) struct ListTool {
pub(super) backend: LocalMemoriesBackend,
pub(super) struct ListTool<B> {
pub(super) backend: B,
}
impl ExtensionToolExecutor for ListTool {
impl<B> ExtensionToolExecutor for ListTool<B>
where
B: MemoriesBackend,
{
fn tool_name(&self) -> ToolName {
ToolName::plain(LIST_TOOL_NAME)
}
codex-rs/ext/memories/src/tools/mod.rs
+5 -6
View File
@@ -10,19 +10,18 @@ use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::Value;
use crate::backend::MemoriesBackend;
use crate::backend::MemoriesBackendError;
use crate::local::LocalMemoriesBackend;
use crate::schema;
mod list;
mod read;
mod search;
const LIST_TOOL_NAME: &str = "memory_list";
const READ_TOOL_NAME: &str = "memory_read";
const SEARCH_TOOL_NAME: &str = "memory_search";
pub(crate) fn memory_tools(backend: LocalMemoriesBackend) -> Vec<Arc<dyn ExtensionToolExecutor>> {
pub(crate) fn memory_tools<B>(backend: B) -> Vec<Arc<dyn ExtensionToolExecutor>>
where
B: MemoriesBackend,
{
vec![
Arc::new(list::ListTool {
backend: backend.clone(),
codex-rs/ext/memories/src/tools/read.rs
+8 -61
View File
@@ -8,13 +8,12 @@ use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use crate::backend::DEFAULT_READ_MAX_TOKENS;
use crate::DEFAULT_READ_MAX_TOKENS;
use crate::READ_TOOL_NAME;
use crate::backend::MemoriesBackend;
use crate::backend::ReadMemoryRequest;
use crate::backend::ReadMemoryResponse;
use crate::local::LocalMemoriesBackend;
use super::READ_TOOL_NAME;
use super::backend_error_to_function_call;
use super::function_tool;
use super::parse_args;
@@ -30,11 +29,14 @@ struct ReadArgs {
}
#[derive(Clone)]
pub(super) struct ReadTool {
pub(super) backend: LocalMemoriesBackend,
pub(super) struct ReadTool<B> {
pub(super) backend: B,
}
impl ExtensionToolExecutor for ReadTool {
impl<B> ExtensionToolExecutor for ReadTool<B>
where
B: MemoriesBackend,
{
fn tool_name(&self) -> ToolName {
ToolName::plain(READ_TOOL_NAME)
}
@@ -63,58 +65,3 @@ impl ExtensionToolExecutor for ReadTool {
})
}
}
#[cfg(test)]
mod tests {
use codex_extension_api::ToolPayload;
use codex_tools::ToolOutput;
use pretty_assertions::assert_eq;
use serde_json::json;
use super::*;
#[tokio::test]
async fn read_tool_reads_memory_file() {
let tempdir = tempfile::tempdir().expect("tempdir");
let memory_root = tempdir.path().join("memories");
tokio::fs::create_dir_all(&memory_root)
.await
.expect("create memories dir");
tokio::fs::write(
memory_root.join("MEMORY.md"),
"first line\nsecond needle line\nthird line\n",
)
.await
.expect("write memory");
let tool = ReadTool {
backend: LocalMemoriesBackend::from_memory_root(&memory_root),
};
let payload = ToolPayload::Function {
arguments: json!({
"path": "MEMORY.md",
"line_offset": 2,
"max_lines": 1
})
.to_string(),
};
let output = tool
.handle(ToolCall {
call_id: "call-1".to_string(),
tool_name: ToolName::plain(READ_TOOL_NAME),
payload: payload.clone(),
})
.await
.expect("read should succeed");
assert_eq!(
output.post_tool_use_response("call-1", &payload),
Some(json!({
"path": "MEMORY.md",
"content": "second needle line\n",
"start_line_number": 2,
"truncated": true
}))
);
}
}
codex-rs/ext/memories/src/tools/search.rs
+9 -79
View File
@@ -8,15 +8,14 @@ use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use crate::backend::DEFAULT_SEARCH_MAX_RESULTS;
use crate::backend::MAX_SEARCH_RESULTS;
use crate::DEFAULT_SEARCH_MAX_RESULTS;
use crate::MAX_SEARCH_RESULTS;
use crate::SEARCH_TOOL_NAME;
use crate::backend::MemoriesBackend;
use crate::backend::SearchMatchMode;
use crate::backend::SearchMemoriesRequest;
use crate::backend::SearchMemoriesResponse;
use crate::local::LocalMemoriesBackend;
use super::SEARCH_TOOL_NAME;
use super::backend_error_to_function_call;
use super::clamp_max_results;
use super::function_tool;
@@ -39,11 +38,14 @@ struct SearchArgs {
}
#[derive(Clone)]
pub(super) struct SearchTool {
pub(super) backend: LocalMemoriesBackend,
pub(super) struct SearchTool<B> {
pub(super) backend: B,
}
impl ExtensionToolExecutor for SearchTool {
impl<B> ExtensionToolExecutor for SearchTool<B>
where
B: MemoriesBackend,
{
fn tool_name(&self) -> ToolName {
ToolName::plain(SEARCH_TOOL_NAME)
}
@@ -86,75 +88,3 @@ impl SearchArgs {
}
}
}
#[cfg(test)]
mod tests {
use pretty_assertions::assert_eq;
use serde_json::json;
use super::*;
#[test]
fn search_args_accept_multiple_queries() {
let args: SearchArgs = serde_json::from_value(json!({
"queries": ["alpha", "needle"],
"case_sensitive": false
}))
.expect("multi-query args should parse");
let request = args.into_request();
assert_eq!(
request,
SearchMemoriesRequest {
queries: vec!["alpha".to_string(), "needle".to_string()],
match_mode: SearchMatchMode::Any,
path: None,
cursor: None,
context_lines: 0,
case_sensitive: false,
normalized: false,
max_results: DEFAULT_SEARCH_MAX_RESULTS,
}
);
}
#[test]
fn search_args_accept_windowed_all_match_mode() {
let args: SearchArgs = serde_json::from_value(json!({
"queries": ["alpha", "needle"],
"match_mode": {
"type": "all_within_lines",
"line_count": 3
}
}))
.expect("windowed all args should parse");
let request = args.into_request();
assert_eq!(
request,
SearchMemoriesRequest {
queries: vec!["alpha".to_string(), "needle".to_string()],
match_mode: SearchMatchMode::AllWithinLines { line_count: 3 },
path: None,
cursor: None,
context_lines: 0,
case_sensitive: true,
normalized: false,
max_results: DEFAULT_SEARCH_MAX_RESULTS,
}
);
}
#[test]
fn search_args_reject_legacy_single_query() {
let err = serde_json::from_value::<SearchArgs>(json!({
"query": "needle",
}))
.expect_err("legacy query field should be rejected");
assert!(err.to_string().contains("unknown field"));
assert!(err.to_string().contains("query"));
}
}