Better memory extension architecture

This commit is contained in:
Kevin Liu
2026-04-07 14:55:21 -07:00
parent 73472c8967
commit 8fbdc4df04
6 changed files with 68 additions and 103 deletions

View File

@@ -13,9 +13,10 @@ Memory prompt templates live under `codex-rs/core/templates/memories/`.
- `read_path.md`
- In `codex`, edit those undated template files in place.
- The dated snapshot-copy workflow is used in the separate `openai/project/agent_memory/write` harness repo, not here.
- At runtime, Phase 2 also appends any `*.md` files from
`~/.codex/memories/consolidation/` to the end of the consolidation prompt, in
sorted path order. Missing directories are skipped.
- The Phase 2 prompt includes a generic optional-extension rule. Extension-specific
instructions live outside the memory root, under
`~/.codex/memories_extensions/<extension-name>/instruction.md`, with optional
`resources/` files used only as that instruction directs.
## When it runs

View File

@@ -25,7 +25,7 @@ pub(crate) use control::clear_memory_root_contents;
pub(crate) use start::start_memories_startup_task;
mod artifacts {
pub(super) const CONSOLIDATION_SUBDIR: &str = "consolidation";
pub(super) const EXTENSIONS_SUBDIR: &str = "memories_extensions";
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
pub(super) const RAW_MEMORIES_FILENAME: &str = "raw_memories.md";
}
@@ -103,12 +103,12 @@ pub fn memory_root(codex_home: &Path) -> PathBuf {
codex_home.join("memories")
}
fn rollout_summaries_dir(root: &Path) -> PathBuf {
root.join(artifacts::ROLLOUT_SUMMARIES_SUBDIR)
fn memory_extensions_root(codex_home: &Path) -> PathBuf {
codex_home.join(artifacts::EXTENSIONS_SUBDIR)
}
fn consolidation_prompt_modules_dir(root: &Path) -> PathBuf {
root.join(artifacts::CONSOLIDATION_SUBDIR)
fn rollout_summaries_dir(root: &Path) -> PathBuf {
root.join(artifacts::ROLLOUT_SUMMARIES_SUBDIR)
}
fn raw_memories_file(root: &Path) -> PathBuf {

View File

@@ -324,8 +324,7 @@ mod agent {
config: Arc<Config>,
selection: &codex_state::Phase2InputSelection,
) -> Vec<UserInput> {
let root = memory_root(&config.codex_home);
let prompt = build_consolidation_prompt(&root, selection).await;
let prompt = build_consolidation_prompt(&config.codex_home, selection);
vec![UserInput::Text {
text: prompt,
text_elements: vec![],

View File

@@ -1,4 +1,4 @@
use crate::memories::consolidation_prompt_modules_dir;
use crate::memories::memory_extensions_root;
use crate::memories::memory_root;
use crate::memories::phase_one;
use crate::memories::storage::rollout_summary_file_stem_from_parts;
@@ -40,18 +40,20 @@ fn parse_embedded_template(source: &'static str, template_name: &str) -> Templat
}
}
/// Builds the consolidation subagent prompt for a specific memory root.
pub(super) async fn build_consolidation_prompt(
memory_root: &Path,
/// Builds the consolidation subagent prompt for a specific Codex home.
pub(super) fn build_consolidation_prompt(
codex_home: &Path,
selection: &Phase2InputSelection,
) -> String {
// additional prompt modules can be added here
let consolidation_prompt_modules = read_consolidation_prompt_modules(memory_root).await;
let memory_root = memory_root(codex_home);
let memory_extensions_root = memory_extensions_root(codex_home);
let memory_root = memory_root.display().to_string();
let memory_extensions_root = memory_extensions_root.display().to_string();
let phase2_input_selection = render_phase2_input_selection(selection);
let mut prompt = CONSOLIDATION_PROMPT_TEMPLATE
CONSOLIDATION_PROMPT_TEMPLATE
.render([
("memory_root", memory_root.as_str()),
("memory_extensions_root", memory_extensions_root.as_str()),
("phase2_input_selection", phase2_input_selection.as_str()),
])
.unwrap_or_else(|err| {
@@ -59,71 +61,7 @@ pub(super) async fn build_consolidation_prompt(
format!(
"## Memory Phase 2 (Consolidation)\nConsolidate Codex memories in: {memory_root}\n\n{phase2_input_selection}"
)
});
if !consolidation_prompt_modules.is_empty() {
prompt.push_str("\n\n");
prompt.push_str(&consolidation_prompt_modules);
}
prompt
}
async fn read_consolidation_prompt_modules(memory_root: &Path) -> String {
let modules_dir = consolidation_prompt_modules_dir(memory_root);
let mut dir = match fs::read_dir(&modules_dir).await {
Ok(dir) => dir,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return String::new(),
Err(err) => {
warn!(
"failed to read memories consolidation prompt modules dir {}: {err}",
modules_dir.display()
);
return String::new();
}
};
let mut module_paths = Vec::new();
loop {
match dir.next_entry().await {
Ok(Some(entry)) => {
let path = entry.path();
if path.extension().and_then(|extension| extension.to_str()) == Some("md") {
module_paths.push(path);
}
}
Ok(None) => break,
Err(err) => {
warn!(
"failed to list memories consolidation prompt modules dir {}: {err}",
modules_dir.display()
);
return String::new();
}
}
}
module_paths.sort();
let mut modules_concat = String::new();
for path in module_paths {
let module = match fs::read_to_string(&path).await {
Ok(module) => module,
Err(err) => {
warn!(
"failed to read memories consolidation prompt module {}: {err}",
path.display()
);
continue;
}
};
let module = module.trim();
if module.is_empty() {
continue;
}
modules_concat.push_str("\n\n");
modules_concat.push_str(module);
}
modules_concat
})
}
fn render_phase2_input_selection(selection: &Phase2InputSelection) -> String {

View File

@@ -53,39 +53,48 @@ fn build_stage_one_input_message_uses_default_limit_when_model_context_window_mi
assert!(message.contains(&expected_truncated));
}
#[tokio::test]
async fn build_consolidation_prompt_renders_embedded_template_without_modules() {
#[test]
fn build_consolidation_prompt_renders_embedded_template() {
let prompt =
build_consolidation_prompt(Path::new("/tmp/memories"), &Phase2InputSelection::default())
.await;
build_consolidation_prompt(Path::new("/tmp/codex"), &Phase2InputSelection::default());
assert!(prompt.contains("Folder structure (under /tmp/memories/):"));
assert!(prompt.contains("Folder structure (under /tmp/codex/memories/):"));
assert!(prompt.contains("Optional memory extensions (under /tmp/codex/memories_extensions/)"));
assert!(prompt.contains("**Diff since last consolidation:**"));
assert!(prompt.contains("- selected inputs this run: 0"));
}
#[tokio::test]
async fn build_consolidation_prompt_appends_modules_in_sorted_order() {
async fn build_consolidation_prompt_points_to_extensions_without_inlining_them() {
let temp = tempdir().unwrap();
let memories_dir = temp.path();
let modules_dir = memories_dir.join("consolidation");
tokio_fs::create_dir_all(&modules_dir).await.unwrap();
tokio_fs::write(modules_dir.join("02-second.md"), "second module\n")
.await
.unwrap();
tokio_fs::write(modules_dir.join("01-first.md"), "first module\n")
.await
.unwrap();
tokio_fs::write(modules_dir.join("ignored.txt"), "ignored module\n")
let codex_home = temp.path();
let extension_dir = codex_home.join("memories_extensions/tape_recorder");
tokio_fs::create_dir_all(extension_dir.join("resources"))
.await
.unwrap();
tokio_fs::write(
extension_dir.join("instruction.md"),
"source-specific instructions\n",
)
.await
.unwrap();
tokio_fs::write(
extension_dir.join("resources/notes.md"),
"source-specific resource\n",
)
.await
.unwrap();
let prompt = build_consolidation_prompt(memories_dir, &Phase2InputSelection::default()).await;
let prompt = build_consolidation_prompt(codex_home, &Phase2InputSelection::default());
let first_index = prompt.find("first module").unwrap();
let second_index = prompt.find("second module").unwrap();
assert!(first_index < second_index);
assert!(!prompt.contains("ignored module"));
assert!(prompt.contains(&format!(
"Optional memory extensions (under {}/memories_extensions/)",
codex_home.display()
)));
assert!(prompt.contains("<extension_name>/instruction.md"));
assert!(prompt.contains("<extension_name>/resources/"));
assert!(!prompt.contains("source-specific instructions"));
assert!(!prompt.contains("source-specific resource"));
}
#[tokio::test]

View File

@@ -34,6 +34,17 @@ Folder structure (under {{ memory_root }}/):
pointers/references, and pruned raw evidence snippets. Distilled version of
everything valuable from the raw rollout.
Memory extensions (under {{ memory_extensions_root }}/):
- <extension_name>/instruction.md
- Source-specific guidance for interpreting additional memory signals. If an
extension folder exists, you must read its instruction.md to determine how to use this memory
source.
If the user has any memory extensions, you MUST read the instructions for each extension to
determine how to use the memory source. If `{{ memory_extensions_root }}` does not exist or has no
extension folders, continue with the standard memory inputs only.
============================================================
GLOBAL SAFETY, HYGIENE, AND NO-FILLER RULES (STRICT)
============================================================
@@ -136,6 +147,13 @@ Under `{{ memory_root }}/`:
- `skills/*`
- read existing skills so updates are incremental and non-duplicative
Optional source-specific inputs:
Under `{{ memory_extensions_root }}/`:
- `<extension_name>/instruction.md`
- If extension folders exist, read each instruction.md first and follow it when interpreting
that extension's memory source.
Mode selection:
- INIT phase: existing artifacts are missing/empty (especially `memory_summary.md`