From 6b56ac102796ab0106b64b61937dfe8b89e0f55b Mon Sep 17 00:00:00 2001
From: starr-openai <starr@openai.com>
Date: Tue, 2 Jun 2026 22:23:38 -0700
Subject: [PATCH] refactor memories prompt source seam

---
 codex-rs/ext/memories/src/extension.rs     |  74 +++++++++------
 codex-rs/ext/memories/src/lib.rs           |   1 +
 codex-rs/ext/memories/src/local.rs         |  14 +++
 codex-rs/ext/memories/src/prompt_source.rs |  13 +++
 codex-rs/ext/memories/src/prompts.rs       |  20 ++--
 codex-rs/ext/memories/src/prompts_tests.rs |   4 +-
 codex-rs/ext/memories/src/tests.rs         | 102 +++++++++++++--------
 7 files changed, 150 insertions(+), 78 deletions(-)
 create mode 100644 codex-rs/ext/memories/src/prompt_source.rs

diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs
index 0be773b4c7..77a91aae26 100644
--- a/codex-rs/ext/memories/src/extension.rs
+++ b/codex-rs/ext/memories/src/extension.rs
@@ -1,3 +1,4 @@
+use std::marker::PhantomData;
 use std::sync::Arc;
 
 use codex_core::config::Config;
@@ -11,56 +12,74 @@ use codex_extension_api::ThreadStartInput;
 use codex_extension_api::ToolContributor;
 use codex_features::Feature;
 use codex_otel::MetricsClient;
-use codex_utils_absolute_path::AbsolutePathBuf;
 
+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, Default)]
-pub(crate) struct MemoriesExtension {
+#[derive(Clone)]
+pub(crate) struct MemoriesExtension<B = LocalMemoriesBackend, S = LocalMemoriesBackend> {
     metrics_client: Option<MetricsClient>,
+    storage: PhantomData<fn() -> (B, S)>,
 }
 
-impl MemoriesExtension {
-    fn new(metrics_client: Option<MetricsClient>) -> Self {
-        Self { metrics_client }
+impl Default for MemoriesExtension {
+    fn default() -> Self {
+        Self::new(/*metrics_client*/ None)
     }
 }
 
-#[derive(Clone, Debug)]
-pub(crate) struct MemoriesExtensionConfig {
-    pub(crate) enabled: bool,
-    pub(crate) dedicated_tools: bool,
-    pub(crate) codex_home: AbsolutePathBuf,
-}
-
-impl MemoriesExtensionConfig {
-    fn from_config(config: &Config) -> Self {
+impl<B, S> MemoriesExtension<B, S> {
+    pub(crate) fn new(metrics_client: Option<MetricsClient>) -> Self {
         Self {
-            enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories,
-            dedicated_tools: config.memories.dedicated_tools,
-            codex_home: config.codex_home.clone(),
+            metrics_client,
+            storage: PhantomData,
         }
     }
 }
 
-impl ContextContributor for MemoriesExtension {
+#[derive(Clone)]
+pub(crate) struct MemoriesExtensionConfig<B = LocalMemoriesBackend, S = LocalMemoriesBackend> {
+    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(),
+            backend,
+        }
+    }
+}
+
+impl<B, S> ContextContributor for MemoriesExtension<B, S>
+where
+    B: MemoriesBackend,
+    S: MemoryPromptSource,
+{
     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 {
+            let Some(config) = thread_store.get::<MemoriesExtensionConfig<B, S>>() else {
                 return Vec::new();
             };
             if !config.enabled {
                 return Vec::new();
             }
 
-            build_memory_tool_developer_instructions(&config.codex_home)
+            build_memory_tool_developer_instructions(&config.prompt_source)
                 .await
                 .map(PromptFragment::developer_policy)
                 .into_iter()
@@ -90,23 +109,24 @@ impl ConfigContributor<Config> for MemoriesExtension {
     }
 }
 
-impl ToolContributor for MemoriesExtension {
+impl<B, S> ToolContributor for MemoriesExtension<B, S>
+where
+    B: MemoriesBackend,
+    S: MemoryPromptSource,
+{
     fn tools(
         &self,
         _session_store: &ExtensionData,
         thread_store: &ExtensionData,
     ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
-        let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
+        let Some(config) = thread_store.get::<MemoriesExtensionConfig<B, S>>() else {
             return Vec::new();
         };
         if !config.enabled || !config.dedicated_tools {
             return Vec::new();
         }
 
-        tools::memory_tools(
-            LocalMemoriesBackend::from_codex_home(&config.codex_home),
-            self.metrics_client.clone(),
-        )
+        tools::memory_tools(config.backend.clone(), self.metrics_client.clone())
     }
 }
 
diff --git a/codex-rs/ext/memories/src/lib.rs b/codex-rs/ext/memories/src/lib.rs
index 25120c7a3a..43fe88f929 100644
--- a/codex-rs/ext/memories/src/lib.rs
+++ b/codex-rs/ext/memories/src/lib.rs
@@ -2,6 +2,7 @@ mod backend;
 mod extension;
 mod local;
 mod metrics;
+mod prompt_source;
 mod prompts;
 mod schema;
 mod tools;
diff --git a/codex-rs/ext/memories/src/local.rs b/codex-rs/ext/memories/src/local.rs
index f4b6cbf1c4..e8cba3eaf6 100644
--- a/codex-rs/ext/memories/src/local.rs
+++ b/codex-rs/ext/memories/src/local.rs
@@ -14,6 +14,8 @@ use crate::backend::ReadMemoryRequest;
 use crate::backend::ReadMemoryResponse;
 use crate::backend::SearchMemoriesRequest;
 use crate::backend::SearchMemoriesResponse;
+use crate::prompt_source::MemoryPromptSource;
+use crate::prompt_source::MemoryPromptSummary;
 
 mod ad_hoc_note;
 mod list;
@@ -127,3 +129,15 @@ impl MemoriesBackend for LocalMemoriesBackend {
         search::search(self, request).await
     }
 }
+
+impl MemoryPromptSource for LocalMemoriesBackend {
+    async fn read_summary(&self) -> Option<MemoryPromptSummary> {
+        let content = tokio::fs::read_to_string(self.root.join("memory_summary.md"))
+            .await
+            .ok()?;
+        Some(MemoryPromptSummary {
+            base_path: self.root.clone(),
+            content,
+        })
+    }
+}
diff --git a/codex-rs/ext/memories/src/prompt_source.rs b/codex-rs/ext/memories/src/prompt_source.rs
new file mode 100644
index 0000000000..a34ea86691
--- /dev/null
+++ b/codex-rs/ext/memories/src/prompt_source.rs
@@ -0,0 +1,13 @@
+use std::future::Future;
+use std::path::PathBuf;
+
+/// Supplies the bounded memory summary used by prompt injection.
+pub(crate) trait MemoryPromptSource: Clone + Send + Sync + 'static {
+    fn read_summary(&self) -> impl Future<Output = Option<MemoryPromptSummary>> + Send;
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) struct MemoryPromptSummary {
+    pub(crate) base_path: PathBuf,
+    pub(crate) content: String,
+}
diff --git a/codex-rs/ext/memories/src/prompts.rs b/codex-rs/ext/memories/src/prompts.rs
index 6546206b1c..5bd47ea7b4 100644
--- a/codex-rs/ext/memories/src/prompts.rs
+++ b/codex-rs/ext/memories/src/prompts.rs
@@ -1,10 +1,10 @@
 use crate::MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT;
-use codex_utils_absolute_path::AbsolutePathBuf;
+use crate::prompt_source::MemoryPromptSource;
+use crate::prompt_source::MemoryPromptSummary;
 use codex_utils_output_truncation::TruncationPolicy;
 use codex_utils_output_truncation::truncate_text;
 use codex_utils_template::Template;
 use std::sync::LazyLock;
-use tokio::fs;
 
 static MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
     parse_embedded_template(
@@ -24,16 +24,12 @@ fn parse_embedded_template(source: &'static str, template_name: &str) -> Templat
 ///
 /// Large `memory_summary.md` files are truncated at
 /// [MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT].
-pub(crate) async fn build_memory_tool_developer_instructions(
-    codex_home: &AbsolutePathBuf,
-) -> Option<String> {
-    let base_path = codex_home.join("memories");
-    let memory_summary_path = base_path.join("memory_summary.md");
-    let memory_summary = fs::read_to_string(&memory_summary_path)
-        .await
-        .ok()?
-        .trim()
-        .to_string();
+pub(crate) async fn build_memory_tool_developer_instructions<S>(source: &S) -> Option<String>
+where
+    S: MemoryPromptSource,
+{
+    let MemoryPromptSummary { base_path, content } = source.read_summary().await?;
+    let memory_summary = content.trim().to_string();
     let memory_summary = truncate_text(
         &memory_summary,
         TruncationPolicy::Tokens(MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT),
diff --git a/codex-rs/ext/memories/src/prompts_tests.rs b/codex-rs/ext/memories/src/prompts_tests.rs
index b4a8cef134..396fe17d12 100644
--- a/codex-rs/ext/memories/src/prompts_tests.rs
+++ b/codex-rs/ext/memories/src/prompts_tests.rs
@@ -1,4 +1,5 @@
 use super::*;
+use crate::local::LocalMemoriesBackend;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use pretty_assertions::assert_eq;
 use tempfile::tempdir;
@@ -17,7 +18,8 @@ async fn build_memory_tool_developer_instructions_renders_embedded_template() {
     .await
     .unwrap();
 
-    let instructions = build_memory_tool_developer_instructions(&codex_home)
+    let source = LocalMemoriesBackend::from_codex_home(&codex_home);
+    let instructions = build_memory_tool_developer_instructions(&source)
         .await
         .unwrap();
 
diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs
index ce03d74965..9d3da96bf6 100644
--- a/codex-rs/ext/memories/src/tests.rs
+++ b/codex-rs/ext/memories/src/tests.rs
@@ -12,9 +12,6 @@ use codex_extension_api::ToolExecutor;
 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 codex_utils_output_truncation::TruncationPolicy;
 use pretty_assertions::assert_eq;
 use serde_json::json;
@@ -22,6 +19,33 @@ use serde_json::json;
 use crate::extension::MemoriesExtension;
 use crate::extension::MemoriesExtensionConfig;
 use crate::local::LocalMemoriesBackend;
+use crate::prompt_source::MemoryPromptSource;
+use crate::prompt_source::MemoryPromptSummary;
+
+#[derive(Clone)]
+struct FakeMemoryPromptSource {
+    summary: MemoryPromptSummary,
+}
+
+impl MemoryPromptSource for FakeMemoryPromptSource {
+    async fn read_summary(&self) -> Option<MemoryPromptSummary> {
+        Some(self.summary.clone())
+    }
+}
+
+fn memories_extension_config(
+    enabled: bool,
+    dedicated_tools: bool,
+    memory_root: &Path,
+) -> MemoriesExtensionConfig {
+    let backend = LocalMemoriesBackend::from_memory_root(memory_root);
+    MemoriesExtensionConfig {
+        enabled,
+        dedicated_tools,
+        prompt_source: backend.clone(),
+        backend,
+    }
+}
 
 #[test]
 fn memory_tool_namespace_matches_responses_api_identifier() {
@@ -51,11 +75,11 @@ 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(MemoriesExtensionConfig {
-        enabled: false,
-        dedicated_tools: true,
-        codex_home: test_path_buf("/tmp/codex-home").abs(),
-    });
+    thread_store.insert(memories_extension_config(
+        /*enabled*/ false,
+        /*dedicated_tools*/ true,
+        Path::new("/tmp/codex-home/memories"),
+    ));
 
     assert!(
         extension
@@ -68,11 +92,11 @@ 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(MemoriesExtensionConfig {
-        enabled: true,
-        dedicated_tools: false,
-        codex_home: test_path_buf("/tmp/codex-home").abs(),
-    });
+    thread_store.insert(memories_extension_config(
+        /*enabled*/ true,
+        /*dedicated_tools*/ false,
+        Path::new("/tmp/codex-home/memories"),
+    ));
 
     assert!(
         extension
@@ -85,11 +109,11 @@ 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(MemoriesExtensionConfig {
-        enabled: true,
-        dedicated_tools: true,
-        codex_home: test_path_buf("/tmp/codex-home").abs(),
-    });
+    thread_store.insert(memories_extension_config(
+        /*enabled*/ true,
+        /*dedicated_tools*/ true,
+        Path::new("/tmp/codex-home/memories"),
+    ));
 
     let tool_names = extension
         .tools(&ExtensionData::new("session"), &thread_store)
@@ -114,11 +138,11 @@ 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(MemoriesExtensionConfig {
-        enabled: true,
-        dedicated_tools: true,
-        codex_home: test_path_buf("/tmp/codex-home").abs(),
-    });
+    thread_store.insert(memories_extension_config(
+        /*enabled*/ true,
+        /*dedicated_tools*/ true,
+        Path::new("/tmp/codex-home/memories"),
+    ));
 
     let tool_names = registry
         .tool_contributors()
@@ -159,25 +183,22 @@ fn ad_hoc_tool_definition_includes_filename_contract() {
 }
 
 #[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::default();
+async fn prompt_contribution_uses_injected_prompt_source() {
+    let extension = MemoriesExtension::<LocalMemoriesBackend, FakeMemoryPromptSource>::new(
+        /*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,
-        codex_home: tempdir.path().abs(),
+        prompt_source: FakeMemoryPromptSource {
+            summary: MemoryPromptSummary {
+                base_path: Path::new("/fake/memories").to_path_buf(),
+                content: "Remember repository-specific implementation preferences.".to_string(),
+            },
+        },
+        backend,
     });
 
     let fragments = extension
@@ -191,6 +212,11 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() {
             .text()
             .contains("Remember repository-specific implementation preferences.")
     );
+    assert!(
+        fragments[0]
+            .text()
+            .contains("- /fake/memories/memory_summary.md")
+    );
 }
 
 #[tokio::test]