refactor memories prompt source seam

This commit is contained in:
starr-openai
2026-06-02 22:53:05 -07:00
parent d36a3ead3c
commit 6b56ac1027
7 changed files with 150 additions and 78 deletions
+47 -27
View File
@@ -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())
}
}
+1
View File
@@ -2,6 +2,7 @@ mod backend;
mod extension;
mod local;
mod metrics;
mod prompt_source;
mod prompts;
mod schema;
mod tools;
+14
View File
@@ -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,
})
}
}
@@ -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,
}
+8 -12
View File
@@ -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),
+3 -1
View File
@@ -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();
+64 -38
View File
@@ -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]