From 67a76fc49549578827a1ac0f836b896103cbf21d Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 3 Jun 2026 10:49:43 -0700 Subject: [PATCH] finish memories deps config split --- codex-rs/ext/memories/src/extension.rs | 65 ++++++++++++++++++-------- codex-rs/ext/memories/src/tests.rs | 61 +++++++++++++++--------- 2 files changed, 83 insertions(+), 43 deletions(-) diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs index 77a91aae26..e7fce0f785 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/codex-rs/ext/memories/src/extension.rs @@ -1,6 +1,11 @@ use std::marker::PhantomData; use std::sync::Arc; +use crate::backend::MemoriesBackend; +use crate::local::LocalMemoriesBackend; +use crate::prompt_source::MemoryPromptSource; +use crate::prompts::build_memory_tool_developer_instructions; +use crate::tools; use codex_core::config::Config; use codex_extension_api::ConfigContributor; use codex_extension_api::ContextContributor; @@ -13,12 +18,6 @@ use codex_extension_api::ToolContributor; use codex_features::Feature; use codex_otel::MetricsClient; -use crate::backend::MemoriesBackend; -use crate::local::LocalMemoriesBackend; -use crate::prompt_source::MemoryPromptSource; -use crate::prompts::build_memory_tool_developer_instructions; -use crate::tools; - /// Contributes Codex memory read-path prompt context and memory read tools. #[derive(Clone)] pub(crate) struct MemoriesExtension { @@ -39,24 +38,46 @@ impl MemoriesExtension { storage: PhantomData, } } + + fn store_thread_state(thread_store: &ExtensionData, config: &Config) { + thread_store.insert(MemoriesExtensionConfig::from_config(config)); + thread_store.insert(MemoriesExtensionStorageDeps::local(config)); + } } -#[derive(Clone)] -pub(crate) struct MemoriesExtensionConfig { +#[derive(Clone, Debug)] +pub(crate) struct MemoriesExtensionConfig { pub(crate) enabled: bool, pub(crate) dedicated_tools: bool, - pub(crate) backend: B, - pub(crate) prompt_source: S, } impl MemoriesExtensionConfig { fn from_config(config: &Config) -> Self { - let backend = LocalMemoriesBackend::from_codex_home(&config.codex_home); Self { enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories, dedicated_tools: config.memories.dedicated_tools, - prompt_source: backend.clone(), + } + } +} + +#[derive(Clone)] +pub(crate) struct MemoriesExtensionStorageDeps { + pub(crate) backend: B, + pub(crate) prompt_source: S, +} + +impl MemoriesExtensionStorageDeps { + fn local(config: &Config) -> Self { + let backend = LocalMemoriesBackend::from_codex_home(&config.codex_home); + Self::new(backend.clone(), backend) + } +} + +impl MemoriesExtensionStorageDeps { + pub(crate) fn new(backend: B, prompt_source: S) -> Self { + Self { backend, + prompt_source, } } } @@ -72,14 +93,17 @@ where thread_store: &'a ExtensionData, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { - let Some(config) = thread_store.get::>() else { + let Some(config) = thread_store.get::() else { return Vec::new(); }; if !config.enabled { return Vec::new(); } + let Some(deps) = thread_store.get::>() else { + return Vec::new(); + }; - build_memory_tool_developer_instructions(&config.prompt_source) + build_memory_tool_developer_instructions(&deps.prompt_source) .await .map(PromptFragment::developer_policy) .into_iter() @@ -91,9 +115,7 @@ where #[async_trait::async_trait] impl ThreadLifecycleContributor for MemoriesExtension { async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - input - .thread_store - .insert(MemoriesExtensionConfig::from_config(input.config)); + Self::store_thread_state(input.thread_store, input.config); } } @@ -105,7 +127,7 @@ impl ConfigContributor for MemoriesExtension { _previous_config: &Config, new_config: &Config, ) { - thread_store.insert(MemoriesExtensionConfig::from_config(new_config)); + Self::store_thread_state(thread_store, new_config); } } @@ -119,14 +141,17 @@ where _session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec>> { - let Some(config) = thread_store.get::>() else { + let Some(config) = thread_store.get::() else { return Vec::new(); }; if !config.enabled || !config.dedicated_tools { return Vec::new(); } + let Some(deps) = thread_store.get::>() else { + return Vec::new(); + }; - tools::memory_tools(config.backend.clone(), self.metrics_client.clone()) + tools::memory_tools(deps.backend.clone(), self.metrics_client.clone()) } } diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index 9d3da96bf6..2d9ee6c5ed 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -18,6 +18,7 @@ use serde_json::json; use crate::extension::MemoriesExtension; use crate::extension::MemoriesExtensionConfig; +use crate::extension::MemoriesExtensionStorageDeps; use crate::local::LocalMemoriesBackend; use crate::prompt_source::MemoryPromptSource; use crate::prompt_source::MemoryPromptSummary; @@ -33,20 +34,30 @@ impl MemoryPromptSource for FakeMemoryPromptSource { } } -fn memories_extension_config( - enabled: bool, - dedicated_tools: bool, - memory_root: &Path, -) -> MemoriesExtensionConfig { - let backend = LocalMemoriesBackend::from_memory_root(memory_root); +fn memories_extension_config(enabled: bool, dedicated_tools: bool) -> MemoriesExtensionConfig { MemoriesExtensionConfig { enabled, dedicated_tools, - prompt_source: backend.clone(), - backend, } } +fn memories_extension_storage_deps( + memory_root: &Path, +) -> MemoriesExtensionStorageDeps { + let backend = LocalMemoriesBackend::from_memory_root(memory_root); + MemoriesExtensionStorageDeps::new(backend.clone(), backend) +} + +fn insert_memories_extension_state( + thread_store: &ExtensionData, + enabled: bool, + dedicated_tools: bool, + memory_root: &Path, +) { + thread_store.insert(memories_extension_config(enabled, dedicated_tools)); + thread_store.insert(memories_extension_storage_deps(memory_root)); +} + #[test] fn memory_tool_namespace_matches_responses_api_identifier() { assert!(!crate::MEMORY_TOOLS_NAMESPACE.is_empty()); @@ -75,11 +86,12 @@ fn tools_are_not_contributed_without_thread_config() { fn tools_are_not_contributed_when_disabled() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); - thread_store.insert(memories_extension_config( + insert_memories_extension_state( + &thread_store, /*enabled*/ false, /*dedicated_tools*/ true, Path::new("/tmp/codex-home/memories"), - )); + ); assert!( extension @@ -92,11 +104,12 @@ fn tools_are_not_contributed_when_disabled() { fn tools_are_not_contributed_when_dedicated_tools_disabled() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); - thread_store.insert(memories_extension_config( + insert_memories_extension_state( + &thread_store, /*enabled*/ true, /*dedicated_tools*/ false, Path::new("/tmp/codex-home/memories"), - )); + ); assert!( extension @@ -109,11 +122,12 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() { fn tools_are_contributed_when_enabled_with_dedicated_tools() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); - thread_store.insert(memories_extension_config( + insert_memories_extension_state( + &thread_store, /*enabled*/ true, /*dedicated_tools*/ true, Path::new("/tmp/codex-home/memories"), - )); + ); let tool_names = extension .tools(&ExtensionData::new("session"), &thread_store) @@ -138,11 +152,12 @@ fn install_registers_dedicated_tool_contributor() { crate::install(&mut builder, /*metrics_client*/ None); let registry = builder.build(); let thread_store = ExtensionData::new("thread"); - thread_store.insert(memories_extension_config( + insert_memories_extension_state( + &thread_store, /*enabled*/ true, /*dedicated_tools*/ true, Path::new("/tmp/codex-home/memories"), - )); + ); let tool_names = registry .tool_contributors() @@ -188,18 +203,18 @@ async fn prompt_contribution_uses_injected_prompt_source() { /*metrics_client*/ None, ); let thread_store = ExtensionData::new("thread"); - let backend = LocalMemoriesBackend::from_memory_root("/unused/memories"); - thread_store.insert(MemoriesExtensionConfig { - enabled: true, - dedicated_tools: false, - prompt_source: FakeMemoryPromptSource { + thread_store.insert(memories_extension_config( + /*enabled*/ true, /*dedicated_tools*/ false, + )); + thread_store.insert(MemoriesExtensionStorageDeps::new( + LocalMemoriesBackend::from_memory_root("/unused/memories"), + FakeMemoryPromptSource { summary: MemoryPromptSummary { base_path: Path::new("/fake/memories").to_path_buf(), content: "Remember repository-specific implementation preferences.".to_string(), }, }, - backend, - }); + )); let fragments = extension .contribute(&ExtensionData::new("session"), &thread_store)