finish memories deps config split

This commit is contained in:
starr-openai
2026-06-03 10:49:43 -07:00
parent 02f75796cd
commit 67a76fc495
2 changed files with 83 additions and 43 deletions
+45 -20
View File
@@ -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<B = LocalMemoriesBackend, S = LocalMemoriesBackend> {
@@ -39,24 +38,46 @@ impl<B, S> MemoriesExtension<B, S> {
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<B = LocalMemoriesBackend, S = LocalMemoriesBackend> {
#[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<B = LocalMemoriesBackend, S = LocalMemoriesBackend> {
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<B, S> MemoriesExtensionStorageDeps<B, S> {
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<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
Box::pin(async move {
let Some(config) = thread_store.get::<MemoriesExtensionConfig<B, S>>() else {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled {
return Vec::new();
}
let Some(deps) = thread_store.get::<MemoriesExtensionStorageDeps<B, S>>() 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<Config> 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<Config> 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<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
let Some(config) = thread_store.get::<MemoriesExtensionConfig<B, S>>() else {
let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
return Vec::new();
};
if !config.enabled || !config.dedicated_tools {
return Vec::new();
}
let Some(deps) = thread_store.get::<MemoriesExtensionStorageDeps<B, S>>() else {
return Vec::new();
};
tools::memory_tools(config.backend.clone(), self.metrics_client.clone())
tools::memory_tools(deps.backend.clone(), self.metrics_client.clone())
}
}
+38 -23
View File
@@ -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<LocalMemoriesBackend, LocalMemoriesBackend> {
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)