mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
chore: split memories part 1 (#19818)
Extract memories into 2 different crates
This commit is contained in:
jif-oai
@@ -0,0 +1,44 @@
|
||||
use std::path::Path;
|
||||
|
||||
pub async fn clear_memory_roots_contents(codex_home: &Path) -> std::io::Result<()> {
|
||||
for memory_root in [
|
||||
codex_home.join("memories"),
|
||||
codex_home.join("memories_extensions"),
|
||||
] {
|
||||
clear_memory_root_contents(memory_root.as_path()).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn clear_memory_root_contents(memory_root: &Path) -> std::io::Result<()> {
|
||||
match tokio::fs::symlink_metadata(memory_root).await {
|
||||
Ok(metadata) if metadata.file_type().is_symlink() => {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"refusing to clear symlinked memory root {}",
|
||||
memory_root.display()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
|
||||
Err(err) => return Err(err),
|
||||
}
|
||||
|
||||
tokio::fs::create_dir_all(memory_root).await?;
|
||||
|
||||
let mut entries = tokio::fs::read_dir(memory_root).await?;
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
let file_type = entry.file_type().await?;
|
||||
if file_type.is_dir() {
|
||||
tokio::fs::remove_dir_all(path).await?;
|
||||
} else {
|
||||
tokio::fs::remove_file(path).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
use crate::memory_extensions_root;
|
||||
use chrono::DateTime;
|
||||
use chrono::Duration;
|
||||
use chrono::NaiveDateTime;
|
||||
use chrono::Utc;
|
||||
use std::path::Path;
|
||||
use tracing::warn;
|
||||
|
||||
const FILENAME_TS_FORMAT: &str = "%Y-%m-%dT%H-%M-%S";
|
||||
const EXTENSION_RESOURCE_RETENTION_DAYS: i64 = 7;
|
||||
|
||||
pub async fn prune_old_extension_resources(memory_root: &Path) {
|
||||
prune_old_extension_resources_with_now(memory_root, Utc::now()).await
|
||||
}
|
||||
|
||||
async fn prune_old_extension_resources_with_now(memory_root: &Path, now: DateTime<Utc>) {
|
||||
let cutoff = now - Duration::days(EXTENSION_RESOURCE_RETENTION_DAYS);
|
||||
let extensions_root = memory_extensions_root(memory_root);
|
||||
let mut extensions = match tokio::fs::read_dir(&extensions_root).await {
|
||||
Ok(extensions) => extensions,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"failed reading memory extensions root {}: {err}",
|
||||
extensions_root.display()
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
while let Ok(Some(extension_entry)) = extensions.next_entry().await {
|
||||
let extension_path = extension_entry.path();
|
||||
let Ok(file_type) = extension_entry.file_type().await else {
|
||||
continue;
|
||||
};
|
||||
if !file_type.is_dir()
|
||||
|| !tokio::fs::try_exists(extension_path.join("instructions.md"))
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let resources_path = extension_path.join("resources");
|
||||
let mut resources = match tokio::fs::read_dir(&resources_path).await {
|
||||
Ok(resources) => resources,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => continue,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"failed reading memory extension resources {}: {err}",
|
||||
resources_path.display()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
while let Ok(Some(resource_entry)) = resources.next_entry().await {
|
||||
let resource_file_path = resource_entry.path();
|
||||
let Ok(file_type) = resource_entry.file_type().await else {
|
||||
continue;
|
||||
};
|
||||
if !file_type.is_file() {
|
||||
continue;
|
||||
}
|
||||
let Some(file_name) = resource_file_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if !file_name.ends_with(".md") {
|
||||
continue;
|
||||
}
|
||||
let Some(resource_timestamp) = resource_timestamp(file_name) else {
|
||||
continue;
|
||||
};
|
||||
if resource_timestamp > cutoff {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Err(err) = tokio::fs::remove_file(&resource_file_path).await
|
||||
&& err.kind() != std::io::ErrorKind::NotFound
|
||||
{
|
||||
warn!(
|
||||
"failed pruning old memory extension resource {}: {err}",
|
||||
resource_file_path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn resource_timestamp(file_name: &str) -> Option<DateTime<Utc>> {
|
||||
let timestamp = file_name.get(..19)?;
|
||||
let naive = NaiveDateTime::parse_from_str(timestamp, FILENAME_TS_FORMAT).ok()?;
|
||||
Some(DateTime::from_naive_utc_and_offset(naive, Utc))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "extensions_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,81 @@
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[tokio::test]
|
||||
async fn prunes_only_old_resources_from_extensions_with_instructions() {
|
||||
let codex_home = TempDir::new().expect("create temp codex home");
|
||||
let memory_root = codex_home.path().join("memories");
|
||||
let extensions_root = memory_extensions_root(&memory_root);
|
||||
let chronicle_resources = extensions_root.join("chronicle/resources");
|
||||
tokio::fs::create_dir_all(&chronicle_resources)
|
||||
.await
|
||||
.expect("create chronicle resources");
|
||||
tokio::fs::write(
|
||||
extensions_root.join("chronicle/instructions.md"),
|
||||
"instructions",
|
||||
)
|
||||
.await
|
||||
.expect("write chronicle instructions");
|
||||
|
||||
let now = DateTime::from_naive_utc_and_offset(
|
||||
NaiveDateTime::parse_from_str("2026-04-14T12-00-00", FILENAME_TS_FORMAT)
|
||||
.expect("parse now"),
|
||||
Utc,
|
||||
);
|
||||
let old_file = chronicle_resources.join("2026-04-06T11-59-59-abcd-10min-old.md");
|
||||
let exact_cutoff_file = chronicle_resources.join("2026-04-07T12-00-00-abcd-10min-cutoff.md");
|
||||
let recent_file = chronicle_resources.join("2026-04-08T12-00-00-abcd-10min-recent.md");
|
||||
let invalid_file = chronicle_resources.join("not-a-timestamp.md");
|
||||
for file in [&old_file, &exact_cutoff_file, &recent_file, &invalid_file] {
|
||||
tokio::fs::write(file, "resource")
|
||||
.await
|
||||
.expect("write chronicle resource");
|
||||
}
|
||||
|
||||
let ignored_resources = extensions_root.join("ignored/resources");
|
||||
tokio::fs::create_dir_all(&ignored_resources)
|
||||
.await
|
||||
.expect("create ignored resources");
|
||||
let ignored_old_file = ignored_resources.join("2026-04-06T11-59-59-abcd-10min-old.md");
|
||||
tokio::fs::write(&ignored_old_file, "ignored")
|
||||
.await
|
||||
.expect("write ignored resource");
|
||||
|
||||
prune_old_extension_resources_with_now(&memory_root, now).await;
|
||||
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&old_file)
|
||||
.await
|
||||
.expect("check old file")
|
||||
);
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&exact_cutoff_file)
|
||||
.await
|
||||
.expect("check cutoff file")
|
||||
);
|
||||
assert!(
|
||||
tokio::fs::try_exists(&recent_file)
|
||||
.await
|
||||
.expect("check recent file")
|
||||
);
|
||||
assert!(
|
||||
tokio::fs::try_exists(&invalid_file)
|
||||
.await
|
||||
.expect("check invalid file")
|
||||
);
|
||||
assert!(
|
||||
tokio::fs::try_exists(&ignored_old_file)
|
||||
.await
|
||||
.expect("check ignored file")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_timestamp_prefix_from_resource_file_name() {
|
||||
let parsed = resource_timestamp("2026-04-06T11-59-59-abcd-10min-old.md")
|
||||
.expect("timestamp should parse");
|
||||
|
||||
assert_eq!(parsed.timestamp(), 1_775_476_799);
|
||||
assert!(resource_timestamp("not-a-timestamp.md").is_none());
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
//! Write-path helpers for Codex memories.
|
||||
//!
|
||||
//! This crate owns the file-backed memory artifact helpers, Phase 1 and Phase
|
||||
//! 2 prompt rendering, extension pruning, and workspace diffing. Runtime
|
||||
//! orchestration for Phase 1 and Phase 2 remains in `codex-core`.
|
||||
|
||||
mod control;
|
||||
mod extensions;
|
||||
mod prompts;
|
||||
mod storage;
|
||||
pub mod workspace;
|
||||
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
pub use control::clear_memory_roots_contents;
|
||||
pub use extensions::prune_old_extension_resources;
|
||||
pub use prompts::build_consolidation_prompt;
|
||||
pub use prompts::build_stage_one_input_message;
|
||||
pub use storage::rebuild_raw_memories_file_from_memories;
|
||||
pub use storage::rollout_summary_file_stem;
|
||||
pub use storage::sync_rollout_summaries_from_memories;
|
||||
|
||||
/// Prompt used for phase 1 extraction.
|
||||
pub const STAGE_ONE_PROMPT: &str = include_str!("../templates/memories/stage_one_system.md");
|
||||
|
||||
/// Fallback stage-1 rollout truncation limit (tokens) when model metadata
|
||||
/// does not include a valid context window.
|
||||
pub const DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT: usize = 150_000;
|
||||
|
||||
/// Portion of the model effective input window reserved for the stage-1
|
||||
/// rollout input.
|
||||
///
|
||||
/// Keeping this below 100% leaves room for system instructions, prompt framing,
|
||||
/// and model output.
|
||||
pub const STAGE_ONE_CONTEXT_WINDOW_PERCENT: i64 = 70;
|
||||
|
||||
mod artifacts {
|
||||
pub(super) const EXTENSIONS_SUBDIR: &str = "extensions";
|
||||
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
|
||||
pub(super) const RAW_MEMORIES_FILENAME: &str = "raw_memories.md";
|
||||
}
|
||||
|
||||
pub fn memory_root(codex_home: &AbsolutePathBuf) -> AbsolutePathBuf {
|
||||
codex_home.join("memories")
|
||||
}
|
||||
|
||||
pub fn rollout_summaries_dir(root: &Path) -> PathBuf {
|
||||
root.join(artifacts::ROLLOUT_SUMMARIES_SUBDIR)
|
||||
}
|
||||
|
||||
pub fn memory_extensions_root(root: &Path) -> PathBuf {
|
||||
root.join(artifacts::EXTENSIONS_SUBDIR)
|
||||
}
|
||||
|
||||
pub fn raw_memories_file(root: &Path) -> PathBuf {
|
||||
root.join(artifacts::RAW_MEMORIES_FILENAME)
|
||||
}
|
||||
|
||||
pub async fn ensure_layout(root: &Path) -> std::io::Result<()> {
|
||||
tokio::fs::create_dir_all(rollout_summaries_dir(root)).await
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
use crate::DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT;
|
||||
use crate::STAGE_ONE_CONTEXT_WINDOW_PERCENT;
|
||||
use crate::memory_extensions_root;
|
||||
use crate::workspace::WORKSPACE_DIFF_FILENAME;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
use codex_utils_output_truncation::TruncationPolicy;
|
||||
use codex_utils_output_truncation::truncate_text;
|
||||
use codex_utils_template::Template;
|
||||
use std::path::Path;
|
||||
use std::sync::LazyLock;
|
||||
use tracing::warn;
|
||||
|
||||
static CONSOLIDATION_PROMPT_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
|
||||
parse_embedded_template(
|
||||
include_str!("../templates/memories/consolidation.md"),
|
||||
"memories/consolidation.md",
|
||||
)
|
||||
});
|
||||
static STAGE_ONE_INPUT_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
|
||||
parse_embedded_template(
|
||||
include_str!("../templates/memories/stage_one_input.md"),
|
||||
"memories/stage_one_input.md",
|
||||
)
|
||||
});
|
||||
static MEMORY_EXTENSIONS_FOLDER_STRUCTURE_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
|
||||
parse_embedded_template(
|
||||
MEMORY_EXTENSIONS_FOLDER_STRUCTURE,
|
||||
"memories/extensions_folder_structure.md",
|
||||
)
|
||||
});
|
||||
static MEMORY_EXTENSIONS_PRIMARY_INPUTS_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
|
||||
parse_embedded_template(
|
||||
MEMORY_EXTENSIONS_PRIMARY_INPUTS,
|
||||
"memories/extensions_primary_inputs.md",
|
||||
)
|
||||
});
|
||||
|
||||
fn parse_embedded_template(source: &'static str, template_name: &str) -> Template {
|
||||
match Template::parse(source) {
|
||||
Ok(template) => template,
|
||||
Err(err) => panic!("embedded template {template_name} is invalid: {err}"),
|
||||
}
|
||||
}
|
||||
|
||||
const MEMORY_EXTENSIONS_FOLDER_STRUCTURE: &str = r#"
|
||||
Memory extensions (under {{ memory_extensions_root }}/):
|
||||
|
||||
- <extension_name>/instructions.md
|
||||
- Source-specific guidance for interpreting additional memory signals. If an
|
||||
extension folder exists, you must read its instructions.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 the workspace diff shows deleted extension resource files,
|
||||
remove stale memories derived only from those resources. If it has no extension folders, continue
|
||||
with the standard memory inputs only.
|
||||
"#;
|
||||
|
||||
const MEMORY_EXTENSIONS_PRIMARY_INPUTS: &str = r#"
|
||||
Optional source-specific inputs:
|
||||
Under `{{ memory_extensions_root }}/`:
|
||||
|
||||
- `<extension_name>/instructions.md`
|
||||
- If extension folders exist, read each instructions.md first and follow it when interpreting
|
||||
that extension's memory source.
|
||||
|
||||
If the workspace diff shows deleted memory extension resources, use that extension-specific deletion
|
||||
signal to remove stale memories derived only from those resources.
|
||||
"#;
|
||||
|
||||
/// Builds the consolidation subagent prompt for a specific memory root.
|
||||
pub fn build_consolidation_prompt(memory_root: &Path) -> String {
|
||||
let memory_extensions_root = memory_extensions_root(memory_root);
|
||||
let memory_extensions_exist = memory_extensions_root.is_dir();
|
||||
let memory_root = memory_root.display().to_string();
|
||||
let memory_extensions_root = memory_extensions_root.display().to_string();
|
||||
let phase2_workspace_diff_file = WORKSPACE_DIFF_FILENAME.to_string();
|
||||
let memory_extensions_folder_structure = if memory_extensions_exist {
|
||||
render_memory_extensions_block(
|
||||
&MEMORY_EXTENSIONS_FOLDER_STRUCTURE_TEMPLATE,
|
||||
&memory_extensions_root,
|
||||
)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
let memory_extensions_primary_inputs = if memory_extensions_exist {
|
||||
render_memory_extensions_block(
|
||||
&MEMORY_EXTENSIONS_PRIMARY_INPUTS_TEMPLATE,
|
||||
&memory_extensions_root,
|
||||
)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
CONSOLIDATION_PROMPT_TEMPLATE
|
||||
.render([
|
||||
("memory_root", memory_root.as_str()),
|
||||
(
|
||||
"memory_extensions_folder_structure",
|
||||
memory_extensions_folder_structure.as_str(),
|
||||
),
|
||||
(
|
||||
"memory_extensions_primary_inputs",
|
||||
memory_extensions_primary_inputs.as_str(),
|
||||
),
|
||||
(
|
||||
"phase2_workspace_diff_file",
|
||||
phase2_workspace_diff_file.as_str(),
|
||||
),
|
||||
])
|
||||
.unwrap_or_else(|err| {
|
||||
warn!("failed to render memories consolidation prompt template: {err}");
|
||||
format!(
|
||||
"## Memory Phase 2 (Consolidation)\nConsolidate Codex memories in: {memory_root}\n\nRead {phase2_workspace_diff_file} first."
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn render_memory_extensions_block(template: &Template, memory_extensions_root: &str) -> String {
|
||||
template
|
||||
.render([("memory_extensions_root", memory_extensions_root)])
|
||||
.unwrap_or_else(|err| {
|
||||
warn!("failed to render memories extension prompt block: {err}");
|
||||
String::new()
|
||||
})
|
||||
}
|
||||
|
||||
/// Builds the stage-1 user message containing rollout metadata and content.
|
||||
///
|
||||
/// Large rollout payloads are truncated to 70% of the active model's effective
|
||||
/// input window token budget while keeping both head and tail context.
|
||||
pub fn build_stage_one_input_message(
|
||||
model_info: &ModelInfo,
|
||||
rollout_path: &Path,
|
||||
rollout_cwd: &Path,
|
||||
rollout_contents: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
let rollout_token_limit = model_info
|
||||
.resolved_context_window()
|
||||
.and_then(|limit| (limit > 0).then_some(limit))
|
||||
.map(|limit| limit.saturating_mul(model_info.effective_context_window_percent) / 100)
|
||||
.map(|limit| (limit.saturating_mul(STAGE_ONE_CONTEXT_WINDOW_PERCENT) / 100).max(1))
|
||||
.and_then(|limit| usize::try_from(limit).ok())
|
||||
.unwrap_or(DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT);
|
||||
let truncated_rollout_contents = truncate_text(
|
||||
rollout_contents,
|
||||
TruncationPolicy::Tokens(rollout_token_limit),
|
||||
);
|
||||
|
||||
let rollout_path = rollout_path.display().to_string();
|
||||
let rollout_cwd = rollout_cwd.display().to_string();
|
||||
Ok(STAGE_ONE_INPUT_TEMPLATE.render([
|
||||
("rollout_path", rollout_path.as_str()),
|
||||
("rollout_cwd", rollout_cwd.as_str()),
|
||||
("rollout_contents", truncated_rollout_contents.as_str()),
|
||||
])?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "prompts_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,71 @@
|
||||
use super::*;
|
||||
use codex_models_manager::model_info::model_info_from_slug;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn build_stage_one_input_message_truncates_rollout_using_model_context_window() {
|
||||
let input = format!("{}{}{}", "a".repeat(700_000), "middle", "z".repeat(700_000));
|
||||
let mut model_info = model_info_from_slug("gpt-5.3-codex");
|
||||
model_info.context_window = Some(123_000);
|
||||
let expected_rollout_token_limit = usize::try_from(
|
||||
((123_000_i64 * model_info.effective_context_window_percent) / 100)
|
||||
* STAGE_ONE_CONTEXT_WINDOW_PERCENT
|
||||
/ 100,
|
||||
)
|
||||
.unwrap();
|
||||
let expected_truncated = truncate_text(
|
||||
&input,
|
||||
TruncationPolicy::Tokens(expected_rollout_token_limit),
|
||||
);
|
||||
let message = build_stage_one_input_message(
|
||||
&model_info,
|
||||
Path::new("/tmp/rollout.jsonl"),
|
||||
Path::new("/tmp"),
|
||||
&input,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(expected_truncated.contains("tokens truncated"));
|
||||
assert!(expected_truncated.starts_with('a'));
|
||||
assert!(expected_truncated.ends_with('z'));
|
||||
assert!(message.contains(&expected_truncated));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_stage_one_input_message_uses_default_limit_when_model_context_window_missing() {
|
||||
let input = format!("{}{}{}", "a".repeat(700_000), "middle", "z".repeat(700_000));
|
||||
let mut model_info = model_info_from_slug("gpt-5.3-codex");
|
||||
model_info.context_window = None;
|
||||
model_info.max_context_window = None;
|
||||
let expected_truncated = truncate_text(
|
||||
&input,
|
||||
TruncationPolicy::Tokens(DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT),
|
||||
);
|
||||
let message = build_stage_one_input_message(
|
||||
&model_info,
|
||||
Path::new("/tmp/rollout.jsonl"),
|
||||
Path::new("/tmp"),
|
||||
&input,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(message.contains(&expected_truncated));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_consolidation_prompt_points_to_workspace_diff_and_extension_tree() {
|
||||
let temp = tempdir().unwrap();
|
||||
let memory_root = temp.path().join("memories");
|
||||
let memory_extensions_root = memory_root.join("extensions");
|
||||
std::fs::create_dir_all(&memory_extensions_root).unwrap();
|
||||
|
||||
let prompt = build_consolidation_prompt(&memory_root);
|
||||
|
||||
assert!(prompt.contains("Memory workspace diff:"));
|
||||
assert!(prompt.contains("phase2_workspace_diff.md"));
|
||||
assert!(prompt.contains(&format!(
|
||||
"Memory extensions (under {}/):",
|
||||
memory_extensions_root.display()
|
||||
)));
|
||||
assert!(prompt.contains("workspace diff shows deleted extension resource files"));
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
use codex_state::Stage1Output;
|
||||
use std::collections::HashSet;
|
||||
use std::fmt::Write as _;
|
||||
use std::path::Path;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ensure_layout;
|
||||
use crate::raw_memories_file;
|
||||
use crate::rollout_summaries_dir;
|
||||
|
||||
/// Rebuild `raw_memories.md` from DB-backed stage-1 outputs.
|
||||
pub async fn rebuild_raw_memories_file_from_memories(
|
||||
root: &Path,
|
||||
memories: &[Stage1Output],
|
||||
max_raw_memories_for_consolidation: usize,
|
||||
) -> std::io::Result<()> {
|
||||
ensure_layout(root).await?;
|
||||
rebuild_raw_memories_file(root, memories, max_raw_memories_for_consolidation).await
|
||||
}
|
||||
|
||||
/// Syncs canonical rollout summary files from DB-backed stage-1 output rows.
|
||||
pub async fn sync_rollout_summaries_from_memories(
|
||||
root: &Path,
|
||||
memories: &[Stage1Output],
|
||||
max_raw_memories_for_consolidation: usize,
|
||||
) -> std::io::Result<()> {
|
||||
ensure_layout(root).await?;
|
||||
|
||||
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
|
||||
let keep = retained
|
||||
.iter()
|
||||
.map(rollout_summary_file_stem)
|
||||
.collect::<HashSet<_>>();
|
||||
prune_rollout_summaries(root, &keep).await?;
|
||||
|
||||
for memory in retained {
|
||||
write_rollout_summary_for_thread(root, memory).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn rebuild_raw_memories_file(
|
||||
root: &Path,
|
||||
memories: &[Stage1Output],
|
||||
max_raw_memories_for_consolidation: usize,
|
||||
) -> std::io::Result<()> {
|
||||
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
|
||||
let mut body = String::from("# Raw Memories\n\n");
|
||||
|
||||
if retained.is_empty() {
|
||||
body.push_str("No raw memories yet.\n");
|
||||
return tokio::fs::write(raw_memories_file(root), body).await;
|
||||
}
|
||||
|
||||
body.push_str("Merged stage-1 raw memories (latest first):\n\n");
|
||||
for memory in retained {
|
||||
writeln!(body, "## Thread `{}`", memory.thread_id).map_err(raw_memories_format_error)?;
|
||||
writeln!(
|
||||
body,
|
||||
"updated_at: {}",
|
||||
memory.source_updated_at.to_rfc3339()
|
||||
)
|
||||
.map_err(raw_memories_format_error)?;
|
||||
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(raw_memories_format_error)?;
|
||||
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
|
||||
.map_err(raw_memories_format_error)?;
|
||||
let rollout_summary_file = format!("{}.md", rollout_summary_file_stem(memory));
|
||||
writeln!(body, "rollout_summary_file: {rollout_summary_file}")
|
||||
.map_err(raw_memories_format_error)?;
|
||||
writeln!(body).map_err(raw_memories_format_error)?;
|
||||
body.push_str(memory.raw_memory.trim());
|
||||
body.push_str("\n\n");
|
||||
}
|
||||
|
||||
tokio::fs::write(raw_memories_file(root), body).await
|
||||
}
|
||||
|
||||
async fn prune_rollout_summaries(root: &Path, keep: &HashSet<String>) -> std::io::Result<()> {
|
||||
let dir_path = rollout_summaries_dir(root);
|
||||
let mut dir = match tokio::fs::read_dir(&dir_path).await {
|
||||
Ok(dir) => dir,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
|
||||
while let Some(entry) = dir.next_entry().await? {
|
||||
let path = entry.path();
|
||||
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
let Some(stem) = file_name.strip_suffix(".md") else {
|
||||
continue;
|
||||
};
|
||||
if !keep.contains(stem)
|
||||
&& let Err(err) = tokio::fs::remove_file(&path).await
|
||||
&& err.kind() != std::io::ErrorKind::NotFound
|
||||
{
|
||||
warn!(
|
||||
"failed pruning outdated rollout summary {}: {err}",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_rollout_summary_for_thread(
|
||||
root: &Path,
|
||||
memory: &Stage1Output,
|
||||
) -> std::io::Result<()> {
|
||||
let file_stem = rollout_summary_file_stem(memory);
|
||||
let path = rollout_summaries_dir(root).join(format!("{file_stem}.md"));
|
||||
|
||||
let mut body = String::new();
|
||||
writeln!(body, "thread_id: {}", memory.thread_id).map_err(rollout_summary_format_error)?;
|
||||
writeln!(
|
||||
body,
|
||||
"updated_at: {}",
|
||||
memory.source_updated_at.to_rfc3339()
|
||||
)
|
||||
.map_err(rollout_summary_format_error)?;
|
||||
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
|
||||
.map_err(rollout_summary_format_error)?;
|
||||
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(rollout_summary_format_error)?;
|
||||
if let Some(git_branch) = memory.git_branch.as_deref() {
|
||||
writeln!(body, "git_branch: {git_branch}").map_err(rollout_summary_format_error)?;
|
||||
}
|
||||
writeln!(body).map_err(rollout_summary_format_error)?;
|
||||
body.push_str(&memory.rollout_summary);
|
||||
body.push('\n');
|
||||
|
||||
tokio::fs::write(path, body).await
|
||||
}
|
||||
|
||||
fn retained_memories(
|
||||
memories: &[Stage1Output],
|
||||
max_raw_memories_for_consolidation: usize,
|
||||
) -> &[Stage1Output] {
|
||||
&memories[..memories.len().min(max_raw_memories_for_consolidation)]
|
||||
}
|
||||
|
||||
fn raw_memories_format_error(err: std::fmt::Error) -> std::io::Error {
|
||||
std::io::Error::other(format!("format raw memories: {err}"))
|
||||
}
|
||||
|
||||
fn rollout_summary_format_error(err: std::fmt::Error) -> std::io::Error {
|
||||
std::io::Error::other(format!("format rollout summary: {err}"))
|
||||
}
|
||||
|
||||
pub fn rollout_summary_file_stem(memory: &Stage1Output) -> String {
|
||||
rollout_summary_file_stem_from_parts(
|
||||
memory.thread_id,
|
||||
memory.source_updated_at,
|
||||
memory.rollout_slug.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn rollout_summary_file_stem_from_parts(
|
||||
thread_id: codex_protocol::ThreadId,
|
||||
source_updated_at: chrono::DateTime<chrono::Utc>,
|
||||
rollout_slug: Option<&str>,
|
||||
) -> String {
|
||||
const ROLLOUT_SLUG_MAX_LEN: usize = 60;
|
||||
const SHORT_HASH_ALPHABET: &[u8; 62] =
|
||||
b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
|
||||
const SHORT_HASH_SPACE: u32 = 14_776_336;
|
||||
|
||||
let thread_id = thread_id.to_string();
|
||||
let (timestamp_fragment, short_hash_seed) = match Uuid::parse_str(&thread_id) {
|
||||
Ok(thread_uuid) => {
|
||||
let timestamp = thread_uuid
|
||||
.get_timestamp()
|
||||
.and_then(|uuid_timestamp| {
|
||||
let (seconds, nanos) = uuid_timestamp.to_unix();
|
||||
i64::try_from(seconds).ok().and_then(|secs| {
|
||||
chrono::DateTime::<chrono::Utc>::from_timestamp(secs, nanos)
|
||||
})
|
||||
})
|
||||
.unwrap_or(source_updated_at);
|
||||
let short_hash_seed = (thread_uuid.as_u128() & 0xFFFF_FFFF) as u32;
|
||||
(
|
||||
timestamp.format("%Y-%m-%dT%H-%M-%S").to_string(),
|
||||
short_hash_seed,
|
||||
)
|
||||
}
|
||||
Err(_) => {
|
||||
let mut short_hash_seed = 0u32;
|
||||
for byte in thread_id.bytes() {
|
||||
short_hash_seed = short_hash_seed
|
||||
.wrapping_mul(31)
|
||||
.wrapping_add(u32::from(byte));
|
||||
}
|
||||
(
|
||||
source_updated_at.format("%Y-%m-%dT%H-%M-%S").to_string(),
|
||||
short_hash_seed,
|
||||
)
|
||||
}
|
||||
};
|
||||
let mut short_hash_value = short_hash_seed % SHORT_HASH_SPACE;
|
||||
let mut short_hash_chars = ['0'; 4];
|
||||
for idx in (0..short_hash_chars.len()).rev() {
|
||||
let alphabet_idx = (short_hash_value % SHORT_HASH_ALPHABET.len() as u32) as usize;
|
||||
short_hash_chars[idx] = SHORT_HASH_ALPHABET[alphabet_idx] as char;
|
||||
short_hash_value /= SHORT_HASH_ALPHABET.len() as u32;
|
||||
}
|
||||
let short_hash: String = short_hash_chars.iter().collect();
|
||||
let file_prefix = format!("{timestamp_fragment}-{short_hash}");
|
||||
|
||||
let Some(raw_slug) = rollout_slug else {
|
||||
return file_prefix;
|
||||
};
|
||||
|
||||
let mut slug = String::with_capacity(ROLLOUT_SLUG_MAX_LEN);
|
||||
for ch in raw_slug.chars() {
|
||||
if slug.len() >= ROLLOUT_SLUG_MAX_LEN {
|
||||
break;
|
||||
}
|
||||
|
||||
if ch.is_ascii_alphanumeric() {
|
||||
slug.push(ch.to_ascii_lowercase());
|
||||
} else {
|
||||
slug.push('_');
|
||||
}
|
||||
}
|
||||
|
||||
while slug.ends_with('_') {
|
||||
slug.pop();
|
||||
}
|
||||
|
||||
if slug.is_empty() {
|
||||
file_prefix
|
||||
} else {
|
||||
format!("{file_prefix}-{slug}")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "storage_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,61 @@
|
||||
use super::rollout_summary_file_stem;
|
||||
use chrono::TimeZone;
|
||||
use chrono::Utc;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_state::Stage1Output;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::PathBuf;
|
||||
const FIXED_PREFIX: &str = "2025-02-11T15-35-19-jqmb";
|
||||
|
||||
fn stage1_output_with_slug(thread_id: ThreadId, rollout_slug: Option<&str>) -> Stage1Output {
|
||||
Stage1Output {
|
||||
thread_id,
|
||||
source_updated_at: Utc.timestamp_opt(123, 0).single().expect("timestamp"),
|
||||
raw_memory: "raw memory".to_string(),
|
||||
rollout_summary: "summary".to_string(),
|
||||
rollout_slug: rollout_slug.map(ToString::to_string),
|
||||
rollout_path: PathBuf::from("/tmp/rollout.jsonl"),
|
||||
cwd: PathBuf::from("/tmp/workspace"),
|
||||
git_branch: None,
|
||||
generated_at: Utc.timestamp_opt(124, 0).single().expect("timestamp"),
|
||||
}
|
||||
}
|
||||
|
||||
fn fixed_thread_id() -> ThreadId {
|
||||
ThreadId::try_from("0194f5a6-89ab-7cde-8123-456789abcdef").expect("valid thread id")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_missing() {
|
||||
let thread_id = fixed_thread_id();
|
||||
let memory = stage1_output_with_slug(thread_id, /*rollout_slug*/ None);
|
||||
|
||||
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rollout_summary_file_stem_sanitizes_and_truncates_slug() {
|
||||
let thread_id = fixed_thread_id();
|
||||
let memory = stage1_output_with_slug(
|
||||
thread_id,
|
||||
Some("Unsafe Slug/With Spaces & Symbols + EXTRA_LONG_12345_67890_ABCDE_fghij_klmno"),
|
||||
);
|
||||
|
||||
let stem = rollout_summary_file_stem(&memory);
|
||||
let slug = stem
|
||||
.strip_prefix(&format!("{FIXED_PREFIX}-"))
|
||||
.expect("slug suffix should be present");
|
||||
assert_eq!(slug.len(), 60);
|
||||
assert_eq!(
|
||||
slug,
|
||||
"unsafe_slug_with_spaces___symbols___extra_long_12345_67890_a"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_is_empty() {
|
||||
let thread_id = fixed_thread_id();
|
||||
let memory = stage1_output_with_slug(thread_id, Some(""));
|
||||
|
||||
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
use anyhow::Context;
|
||||
use codex_git_utils::GitBaselineDiff;
|
||||
use codex_git_utils::diff_since_latest_init;
|
||||
use codex_git_utils::ensure_git_baseline_repository;
|
||||
use codex_git_utils::reset_git_repository;
|
||||
use std::path::Path;
|
||||
|
||||
/// Generated diff file the Phase 2 consolidation agent reads before editing memories.
|
||||
pub const WORKSPACE_DIFF_FILENAME: &str = "phase2_workspace_diff.md";
|
||||
|
||||
const WORKSPACE_DIFF_MAX_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
/// Prepares the memory directory for git-baseline diffing.
|
||||
///
|
||||
/// This keeps an existing usable `.git/` baseline intact. It initializes a new git baseline when the
|
||||
/// metadata is missing or unusable, and removes any stale generated `phase2_workspace_diff.md` file
|
||||
/// so that the next diff does not include a previous prompt artifact.
|
||||
pub async fn prepare_memory_workspace(root: &Path) -> anyhow::Result<()> {
|
||||
tokio::fs::create_dir_all(root)
|
||||
.await
|
||||
.with_context(|| format!("create memory workspace {}", root.display()))?;
|
||||
remove_workspace_diff(root).await?;
|
||||
ensure_git_baseline_repository(root).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the current workspace diff after removing any stale generated diff artifact.
|
||||
///
|
||||
/// The removed file is only `phase2_workspace_diff.md`; memory artifacts and `.git/` metadata are
|
||||
/// left intact.
|
||||
pub async fn memory_workspace_diff(root: &Path) -> anyhow::Result<GitBaselineDiff> {
|
||||
remove_workspace_diff(root).await?;
|
||||
diff_since_latest_init(root).await
|
||||
}
|
||||
|
||||
/// Writes `phase2_workspace_diff.md` with a bounded git-style diff from the current baseline.
|
||||
pub async fn write_workspace_diff(root: &Path, diff: &GitBaselineDiff) -> anyhow::Result<()> {
|
||||
let path = root.join(WORKSPACE_DIFF_FILENAME);
|
||||
tokio::fs::write(&path, render_workspace_diff_file(diff))
|
||||
.await
|
||||
.with_context(|| format!("write memory workspace diff file {}", path.display()))
|
||||
}
|
||||
|
||||
/// Marks the current memory root as the new baseline.
|
||||
///
|
||||
/// The generated diff file is removed before resetting the baseline so deleted memory content is
|
||||
/// not retained in the prompt artifact or in unreachable git objects.
|
||||
pub async fn reset_memory_workspace_baseline(root: &Path) -> anyhow::Result<()> {
|
||||
remove_workspace_diff(root).await?;
|
||||
reset_git_repository(root).await
|
||||
}
|
||||
|
||||
/// Removes the generated `phase2_workspace_diff.md` prompt artifact.
|
||||
///
|
||||
/// This does not remove `.git/`, reset the baseline, or delete memory content. It is used before
|
||||
/// diffing and before baseline reset so the generated diff file itself is not treated as memory
|
||||
/// workspace input.
|
||||
pub(super) async fn remove_workspace_diff(root: &Path) -> anyhow::Result<()> {
|
||||
let path = root.join(WORKSPACE_DIFF_FILENAME);
|
||||
match tokio::fs::remove_file(&path).await {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(err) => Err(err)
|
||||
.with_context(|| format!("remove memory workspace diff file {}", path.display())),
|
||||
}
|
||||
}
|
||||
|
||||
fn render_workspace_diff_file(diff: &GitBaselineDiff) -> String {
|
||||
let mut rendered = String::from(
|
||||
"# Memory Workspace Diff\n\n\
|
||||
Generated by Codex before Phase 2 memory consolidation. Read this file first and do not edit it.\n\n\
|
||||
## Status\n",
|
||||
);
|
||||
|
||||
if !diff.has_changes() {
|
||||
rendered.push_str("- none\n");
|
||||
return rendered;
|
||||
}
|
||||
|
||||
for change in &diff.changes {
|
||||
rendered.push_str(&format!("- {} {}\n", change.status.label(), change.path));
|
||||
}
|
||||
rendered.push_str("\n## Diff\n\n```diff\n");
|
||||
append_bounded_diff(&mut rendered, &diff.unified_diff);
|
||||
rendered.push_str("```\n");
|
||||
rendered
|
||||
}
|
||||
|
||||
fn append_bounded_diff(rendered: &mut String, diff: &str) {
|
||||
if diff.len() <= WORKSPACE_DIFF_MAX_BYTES {
|
||||
rendered.push_str(diff);
|
||||
if !diff.ends_with('\n') {
|
||||
rendered.push('\n');
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let boundary = previous_char_boundary(diff, WORKSPACE_DIFF_MAX_BYTES);
|
||||
rendered.push_str(&diff[..boundary]);
|
||||
if !rendered.ends_with('\n') {
|
||||
rendered.push('\n');
|
||||
}
|
||||
rendered.push_str(&format!(
|
||||
"\n[workspace diff truncated at {WORKSPACE_DIFF_MAX_BYTES} bytes]\n"
|
||||
));
|
||||
}
|
||||
|
||||
fn previous_char_boundary(value: &str, max_bytes: usize) -> usize {
|
||||
if max_bytes >= value.len() {
|
||||
return value.len();
|
||||
}
|
||||
let mut index = max_bytes;
|
||||
while !value.is_char_boundary(index) {
|
||||
index -= 1;
|
||||
}
|
||||
index
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "workspace_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,78 @@
|
||||
use super::*;
|
||||
use codex_git_utils::GitBaselineChange;
|
||||
use codex_git_utils::GitBaselineChangeStatus;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn render_workspace_diff_file_bounds_large_diff() {
|
||||
let diff = GitBaselineDiff {
|
||||
changes: vec![GitBaselineChange {
|
||||
status: GitBaselineChangeStatus::Modified,
|
||||
path: "MEMORY.md".to_string(),
|
||||
}],
|
||||
unified_diff: "a".repeat(WORKSPACE_DIFF_MAX_BYTES + 128),
|
||||
};
|
||||
|
||||
let rendered = render_workspace_diff_file(&diff);
|
||||
|
||||
assert!(rendered.contains("- M MEMORY.md"));
|
||||
assert!(rendered.contains("[workspace diff truncated at 4194304 bytes]"));
|
||||
assert!(rendered.ends_with("```\n"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reset_memory_workspace_baseline_removes_generated_diff() {
|
||||
let home = TempDir::new().expect("tempdir");
|
||||
let root = home.path().join("memories");
|
||||
prepare_memory_workspace(&root)
|
||||
.await
|
||||
.expect("prepare memory workspace");
|
||||
fs::write(root.join("MEMORY.md"), "memory").expect("write memory");
|
||||
write_workspace_diff(
|
||||
&root,
|
||||
&GitBaselineDiff {
|
||||
changes: vec![GitBaselineChange {
|
||||
status: GitBaselineChangeStatus::Added,
|
||||
path: "MEMORY.md".to_string(),
|
||||
}],
|
||||
unified_diff: "+memory\n".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("write workspace diff");
|
||||
|
||||
reset_memory_workspace_baseline(&root)
|
||||
.await
|
||||
.expect("reset baseline");
|
||||
|
||||
assert!(!root.join(WORKSPACE_DIFF_FILENAME).exists());
|
||||
let diff = memory_workspace_diff(&root)
|
||||
.await
|
||||
.expect("load workspace diff");
|
||||
assert_eq!(diff.changes, Vec::new());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_memory_workspace_recovers_unusable_git_dir() {
|
||||
let home = TempDir::new().expect("tempdir");
|
||||
let root = home.path().join("memories");
|
||||
fs::create_dir_all(root.join(".git")).expect("create unusable git dir");
|
||||
fs::write(root.join("MEMORY.md"), "memory").expect("write memory");
|
||||
|
||||
prepare_memory_workspace(&root)
|
||||
.await
|
||||
.expect("prepare memory workspace");
|
||||
|
||||
let diff = memory_workspace_diff(&root)
|
||||
.await
|
||||
.expect("load workspace diff");
|
||||
assert_eq!(diff.changes, Vec::new());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn previous_char_boundary_handles_multibyte_text() {
|
||||
let text = "aé";
|
||||
assert_eq!(previous_char_boundary(text, /*max_bytes*/ 2), 1);
|
||||
}
|
||||
Reference in New Issue
Block a user