diff --git a/codex-rs/ext/memories/src/backend.rs b/codex-rs/ext/memories/src/backend.rs index 929852f1b..fa635c830 100644 --- a/codex-rs/ext/memories/src/backend.rs +++ b/codex-rs/ext/memories/src/backend.rs @@ -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 diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs new file mode 100644 index 000000000..2b4955918 --- /dev/null +++ b/codex-rs/ext/memories/src/extension.rs @@ -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> + Send + 'a>> { + Box::pin(async move { + let Some(config) = thread_store.get::() 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 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> { + let Some(config) = thread_store.get::() 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) { + let extension = Arc::new(MemoriesExtension); + registry.thread_lifecycle_contributor(extension.clone()); + registry.prompt_contributor(extension.clone()); + registry.tool_contributor(extension); +} diff --git a/codex-rs/ext/memories/src/lib.rs b/codex-rs/ext/memories/src/lib.rs index 41143bef0..4b6f74c57 100644 --- a/codex-rs/ext/memories/src/lib.rs +++ b/codex-rs/ext/memories/src/lib.rs @@ -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> + Send + 'a>> { - Box::pin(async move { - let Some(config) = thread_store.get::() 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 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> { - let Some(config) = thread_store.get::() 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) { - 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::>(); - - 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; diff --git a/codex-rs/ext/memories/src/local.rs b/codex-rs/ext/memories/src/local.rs index 47c09b60d..8c663d73f 100644 --- a/codex-rs/ext/memories/src/local.rs +++ b/codex-rs/ext/memories/src/local.rs @@ -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) -> Self { + pub(crate) fn from_memory_root(root: impl Into) -> 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 { - 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::().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 { - 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 { - let queries = request - .queries - .iter() - .map(|query| query.trim().to_string()) - .collect::>(); - 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::().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, -) -> 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, -) -> 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::>(); - let line_matches = lines - .iter() - .map(|line| matcher.matched_query_flags(line)) - .collect::>(); - 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, -) -> 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, - prepared_queries: Vec, - comparison: SearchComparison, - match_mode: SearchMatchMode, -} - -impl SearchMatcher { - fn new( - queries: Vec, - match_mode: SearchMatchMode, - case_sensitive: bool, - normalized: bool, - ) -> Result { - let comparison = SearchComparison::new(case_sensitive, normalized); - let prepared_queries = queries - .iter() - .map(|query| comparison.prepare(query)) - .map(Cow::into_owned) - .collect::>(); - 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 { - 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 { - 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::(), - ) - } -} - -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 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::>() - .join("/") -} - -fn line_start_byte_offset( - content: &str, - line_offset: usize, -) -> Result { - 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 { - 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() -} diff --git a/codex-rs/ext/memories/src/local/list.rs b/codex-rs/ext/memories/src/local/list.rs new file mode 100644 index 000000000..60915aa1f --- /dev/null +++ b/codex-rs/ext/memories/src/local/list.rs @@ -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 { + 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::().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, + }) +} diff --git a/codex-rs/ext/memories/src/local/path.rs b/codex-rs/ext/memories/src/local/path.rs new file mode 100644 index 000000000..2052967f6 --- /dev/null +++ b/codex-rs/ext/memories/src/local/path.rs @@ -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, 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::>() + .join("/") +} diff --git a/codex-rs/ext/memories/src/local/read.rs b/codex-rs/ext/memories/src/local/read.rs new file mode 100644 index 000000000..d2e22bd1a --- /dev/null +++ b/codex-rs/ext/memories/src/local/read.rs @@ -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 { + 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 { + 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 { + 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() +} diff --git a/codex-rs/ext/memories/src/local/search.rs b/codex-rs/ext/memories/src/local/search.rs new file mode 100644 index 000000000..85551fd88 --- /dev/null +++ b/codex-rs/ext/memories/src/local/search.rs @@ -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 { + let queries = request + .queries + .iter() + .map(|query| query.trim().to_string()) + .collect::>(); + 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::().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, +) -> 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, +) -> 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::>(); + let line_matches = lines + .iter() + .map(|line| matcher.matched_query_flags(line)) + .collect::>(); + 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, +) -> 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, + prepared_queries: Vec, + comparison: SearchComparison, + match_mode: SearchMatchMode, +} + +impl SearchMatcher { + fn new( + queries: Vec, + match_mode: SearchMatchMode, + case_sensitive: bool, + normalized: bool, + ) -> Result { + let comparison = SearchComparison::new(case_sensitive, normalized); + let prepared_queries = queries + .iter() + .map(|query| comparison.prepare(query)) + .map(Cow::into_owned) + .collect::>(); + 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 { + 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 { + 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::(), + ) + } +} diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs new file mode 100644 index 000000000..3d9cd9a79 --- /dev/null +++ b/codex-rs/ext/memories/src/tests.rs @@ -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::>(); + + 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 { + 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")) +} diff --git a/codex-rs/ext/memories/src/tools/list.rs b/codex-rs/ext/memories/src/tools/list.rs index e9674f704..bd1599047 100644 --- a/codex-rs/ext/memories/src/tools/list.rs +++ b/codex-rs/ext/memories/src/tools/list.rs @@ -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 { + pub(super) backend: B, } -impl ExtensionToolExecutor for ListTool { +impl ExtensionToolExecutor for ListTool +where + B: MemoriesBackend, +{ fn tool_name(&self) -> ToolName { ToolName::plain(LIST_TOOL_NAME) } diff --git a/codex-rs/ext/memories/src/tools/mod.rs b/codex-rs/ext/memories/src/tools/mod.rs index 735126349..b488cc5e8 100644 --- a/codex-rs/ext/memories/src/tools/mod.rs +++ b/codex-rs/ext/memories/src/tools/mod.rs @@ -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> { +pub(crate) fn memory_tools(backend: B) -> Vec> +where + B: MemoriesBackend, +{ vec![ Arc::new(list::ListTool { backend: backend.clone(), diff --git a/codex-rs/ext/memories/src/tools/read.rs b/codex-rs/ext/memories/src/tools/read.rs index 96129ec0f..ec411a994 100644 --- a/codex-rs/ext/memories/src/tools/read.rs +++ b/codex-rs/ext/memories/src/tools/read.rs @@ -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 { + pub(super) backend: B, } -impl ExtensionToolExecutor for ReadTool { +impl ExtensionToolExecutor for ReadTool +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 - })) - ); - } -} diff --git a/codex-rs/ext/memories/src/tools/search.rs b/codex-rs/ext/memories/src/tools/search.rs index f437df004..38863bcad 100644 --- a/codex-rs/ext/memories/src/tools/search.rs +++ b/codex-rs/ext/memories/src/tools/search.rs @@ -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 { + pub(super) backend: B, } -impl ExtensionToolExecutor for SearchTool { +impl ExtensionToolExecutor for SearchTool +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::(json!({ - "query": "needle", - })) - .expect_err("legacy query field should be rejected"); - - assert!(err.to_string().contains("unknown field")); - assert!(err.to_string().contains("query")); - } -}