mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
refactor memories prompt source seam
This commit is contained in:
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ mod backend;
|
||||
mod extension;
|
||||
mod local;
|
||||
mod metrics;
|
||||
mod prompt_source;
|
||||
mod prompts;
|
||||
mod schema;
|
||||
mod tools;
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user