feat: align memory phase 1 and make it stronger (#11300)

## Align with the new phase-1 design Basically we know run phase 1 in parallel by considering: * Max 64 rollouts * Max 1 month old * Consider the most recent first This PR also adds stronger parallelization capabilities by detecting stale jobs, retry policies, ownership of computation to prevent double computations etc etc
2026-07-01 00:31:56 +08:00 · 2026-02-10 13:42:09 +00:00
parent 223fadc760
commit 1d5eba0090
10 changed files with 1553 additions and 259 deletions
@@ -17,13 +17,20 @@ use std::path::PathBuf;
 /// Subagent source label used to identify consolidation tasks.
 pub(crate) const MEMORY_CONSOLIDATION_SUBAGENT_LABEL: &str = "memory_consolidation";
 /// Maximum number of rollout candidates processed per startup pass.
-pub(crate) const MAX_ROLLOUTS_PER_STARTUP: usize = 8;
+pub(crate) const MAX_ROLLOUTS_PER_STARTUP: usize = 64;
 /// Concurrency cap for startup memory extraction and consolidation scheduling.
 pub(crate) const PHASE_ONE_CONCURRENCY_LIMIT: usize = MAX_ROLLOUTS_PER_STARTUP;
-/// Maximum number of recent raw memories retained per working directory.
-pub(crate) const MAX_RAW_MEMORIES_PER_CWD: usize = 10;
+/// Maximum number of recent raw memories retained per scope.
+pub(crate) const MAX_RAW_MEMORIES_PER_SCOPE: usize = 64;
+/// Maximum rollout age considered for phase-1 extraction.
+pub(crate) const PHASE_ONE_MAX_ROLLOUT_AGE_DAYS: i64 = 30;
+/// Lease duration (seconds) for phase-1 job ownership.
+pub(crate) const PHASE_ONE_JOB_LEASE_SECONDS: i64 = 3_600;
 /// Lease duration (seconds) for per-cwd consolidation locks.
 pub(crate) const CONSOLIDATION_LOCK_LEASE_SECONDS: i64 = 600;
+pub(crate) const MEMORY_SCOPE_KIND_CWD: &str = "cwd";
+pub(crate) const MEMORY_SCOPE_KIND_USER: &str = "user";
+pub(crate) const MEMORY_SCOPE_KEY_USER: &str = "user";

 const MEMORY_SUBDIR: &str = "memory";
 const RAW_MEMORIES_SUBDIR: &str = "raw_memories";
@@ -31,6 +38,7 @@ const MEMORY_SUMMARY_FILENAME: &str = "memory_summary.md";
 const MEMORY_REGISTRY_FILENAME: &str = "MEMORY.md";
 const LEGACY_CONSOLIDATED_FILENAME: &str = "consolidated.md";
 const SKILLS_SUBDIR: &str = "skills";
+const CWD_MEMORY_BUCKET_HEX_LEN: usize = 16;

 pub(crate) use phase_one::RAW_MEMORY_PROMPT;
 pub(crate) use phase_one::parse_stage_one_output;
@@ -43,8 +51,9 @@ pub(crate) use rollout::StageOneRolloutFilter;
 pub(crate) use rollout::serialize_filtered_rollout_response_items;
 pub(crate) use selection::select_rollout_candidates_from_db;
 pub(crate) use storage::prune_to_recent_memories_and_rebuild_summary;
+pub(crate) use storage::rebuild_memory_summary_from_memories;
+pub(crate) use storage::sync_raw_memories_from_memories;
 pub(crate) use storage::wipe_consolidation_outputs;
-pub(crate) use storage::write_raw_memory;
 pub(crate) use types::RolloutCandidate;

 /// Returns the on-disk memory root directory for a given working directory.
@@ -56,6 +65,21 @@ pub(crate) fn memory_root_for_cwd(codex_home: &Path, cwd: &Path) -> PathBuf {
    codex_home.join("memories").join(bucket).join(MEMORY_SUBDIR)
 }

+/// Returns the DB scope key for a cwd-scoped memory entry.
+///
+/// This uses the same normalization/fallback behavior as cwd bucket derivation.
+pub(crate) fn memory_scope_key_for_cwd(cwd: &Path) -> String {
+    normalize_cwd_for_memory(cwd).display().to_string()
+}
+
+/// Returns the on-disk user-shared memory root directory.
+pub(crate) fn memory_root_for_user(codex_home: &Path) -> PathBuf {
+    codex_home
+        .join("memories")
+        .join(MEMORY_SCOPE_KEY_USER)
+        .join(MEMORY_SUBDIR)
+}
+
 fn raw_memories_dir(root: &Path) -> PathBuf {
    root.join(RAW_MEMORIES_SUBDIR)
 }
@@ -70,9 +94,14 @@ pub(crate) async fn ensure_layout(root: &Path) -> std::io::Result<()> {
 }

 fn memory_bucket_for_cwd(cwd: &Path) -> String {
-    let normalized = normalize_for_path_comparison(cwd).unwrap_or_else(|_| cwd.to_path_buf());
+    let normalized = normalize_cwd_for_memory(cwd);
    let normalized = normalized.to_string_lossy();
    let mut hasher = Sha256::new();
    hasher.update(normalized.as_bytes());
-    format!("{:x}", hasher.finalize())
+    let full_hash = format!("{:x}", hasher.finalize());
+    full_hash[..CWD_MEMORY_BUCKET_HEX_LEN].to_string()
+}
+
+fn normalize_cwd_for_memory(cwd: &Path) -> PathBuf {
+    normalize_for_path_comparison(cwd).unwrap_or_else(|_| cwd.to_path_buf())
 }
@@ -1,28 +1,25 @@
+use chrono::Duration;
+use chrono::Utc;
 use codex_protocol::ThreadId;
-use codex_state::ThreadMemory;
 use codex_state::ThreadMetadata;
-use std::collections::BTreeMap;

 use super::types::RolloutCandidate;

 /// Selects rollout candidates that need stage-1 memory extraction.
 ///
-/// A rollout is selected when it is not the active thread and has no memory yet
-/// (or the stored memory is older than the thread metadata timestamp).
+/// A rollout is selected when it is not the active thread and was updated
+/// within the configured max age window.
 pub(crate) fn select_rollout_candidates_from_db(
    items: &[ThreadMetadata],
    current_thread_id: ThreadId,
-    existing_memories: &[ThreadMemory],
    max_items: usize,
+    max_age_days: i64,
 ) -> Vec<RolloutCandidate> {
    if max_items == 0 {
        return Vec::new();
    }

-    let memory_updated_by_thread = existing_memories
-        .iter()
-        .map(|memory| (memory.thread_id.to_string(), memory.updated_at))
-        .collect::<BTreeMap<_, _>>();
+    let cutoff = Utc::now() - Duration::days(max_age_days.max(0));

    let mut candidates = Vec::new();

@@ -30,10 +27,7 @@ pub(crate) fn select_rollout_candidates_from_db(
        if item.id == current_thread_id {
            continue;
        }
-
-        let memory_updated_at = memory_updated_by_thread.get(&item.id.to_string());
-        if memory_updated_at.is_some_and(|memory_updated_at| *memory_updated_at >= item.updated_at)
-        {
+        if item.updated_at < cutoff {
            continue;
        }

@@ -41,7 +35,6 @@ pub(crate) fn select_rollout_candidates_from_db(
            thread_id: item.id,
            rollout_path: item.rollout_path.clone(),
            cwd: item.cwd.clone(),
-            title: item.title.clone(),
            updated_at: Some(item.updated_at.to_rfc3339()),
        });

@@ -6,47 +6,12 @@ use std::path::PathBuf;
 use tracing::warn;

 use super::LEGACY_CONSOLIDATED_FILENAME;
-use super::MAX_RAW_MEMORIES_PER_CWD;
+use super::MAX_RAW_MEMORIES_PER_SCOPE;
 use super::MEMORY_REGISTRY_FILENAME;
 use super::SKILLS_SUBDIR;
 use super::ensure_layout;
 use super::memory_summary_file;
 use super::raw_memories_dir;
-use super::types::RolloutCandidate;
-
-/// Writes (or replaces) the per-thread markdown raw memory on disk.
-///
-/// This also removes older files for the same thread id to keep one canonical
-/// raw memory file per thread.
-pub(crate) async fn write_raw_memory(
-    root: &Path,
-    candidate: &RolloutCandidate,
-    raw_memory: &str,
-) -> std::io::Result<PathBuf> {
-    let slug = build_memory_slug(&candidate.title);
-    let filename = format!("{}_{}.md", candidate.thread_id, slug);
-    let path = raw_memories_dir(root).join(filename);
-
-    remove_outdated_thread_raw_memories(root, &candidate.thread_id.to_string(), &path).await?;
-
-    let mut body = String::new();
-    writeln!(body, "thread_id: {}", candidate.thread_id)
-        .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
-    writeln!(body, "cwd: {}", candidate.cwd.display())
-        .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
-    writeln!(body, "rollout_path: {}", candidate.rollout_path.display())
-        .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
-    if let Some(updated_at) = candidate.updated_at.as_deref() {
-        writeln!(body, "updated_at: {updated_at}")
-            .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
-    }
-    writeln!(body).map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
-    body.push_str(raw_memory.trim());
-    body.push('\n');
-
-    tokio::fs::write(&path, body).await?;
-    Ok(path)
-}

 /// Prunes stale raw memory files and rebuilds the routing summary for recent memories.
 pub(crate) async fn prune_to_recent_memories_and_rebuild_summary(
@@ -57,7 +22,7 @@ pub(crate) async fn prune_to_recent_memories_and_rebuild_summary(

    let keep = memories
        .iter()
-        .take(MAX_RAW_MEMORIES_PER_CWD)
+        .take(MAX_RAW_MEMORIES_PER_SCOPE)
        .map(|memory| memory.thread_id.to_string())
        .collect::<BTreeSet<_>>();

@@ -65,6 +30,38 @@ pub(crate) async fn prune_to_recent_memories_and_rebuild_summary(
    rebuild_memory_summary(root, memories).await
 }

+/// Rebuild `memory_summary.md` for a scope without pruning raw memory files.
+pub(crate) async fn rebuild_memory_summary_from_memories(
+    root: &Path,
+    memories: &[ThreadMemory],
+) -> std::io::Result<()> {
+    ensure_layout(root).await?;
+    rebuild_memory_summary(root, memories).await
+}
+
+/// Syncs canonical raw memory files from DB-backed memory rows.
+pub(crate) async fn sync_raw_memories_from_memories(
+    root: &Path,
+    memories: &[ThreadMemory],
+) -> std::io::Result<()> {
+    ensure_layout(root).await?;
+
+    let retained = memories
+        .iter()
+        .take(MAX_RAW_MEMORIES_PER_SCOPE)
+        .collect::<Vec<_>>();
+    let keep = retained
+        .iter()
+        .map(|memory| memory.thread_id.to_string())
+        .collect::<BTreeSet<_>>();
+    prune_raw_memories(root, &keep).await?;
+
+    for memory in retained {
+        write_raw_memory_for_thread(root, memory).await?;
+    }
+    Ok(())
+}
+
 /// Clears consolidation outputs so a fresh consolidation run can regenerate them.
 ///
 /// Phase-1 artifacts (`raw_memories/` and `memory_summary.md`) are preserved.
@@ -103,7 +100,7 @@ async fn rebuild_memory_summary(root: &Path, memories: &[ThreadMemory]) -> std::
    }

    body.push_str("Map of concise summaries to thread IDs (latest first):\n\n");
-    for memory in memories.iter().take(MAX_RAW_MEMORIES_PER_CWD) {
+    for memory in memories.iter().take(MAX_RAW_MEMORIES_PER_SCOPE) {
        let summary = compact_summary_for_index(&memory.memory_summary);
        writeln!(body, "- {summary} (thread: `{}`)", memory.thread_id)
            .map_err(|err| std::io::Error::other(format!("format memory summary: {err}")))?;
@@ -179,27 +176,25 @@ async fn remove_outdated_thread_raw_memories(
    Ok(())
 }

-fn build_memory_slug(value: &str) -> String {
-    let mut slug = String::new();
-    let mut last_was_sep = false;
+async fn write_raw_memory_for_thread(
+    root: &Path,
+    memory: &ThreadMemory,
+) -> std::io::Result<PathBuf> {
+    let path = raw_memories_dir(root).join(format!("{}.md", memory.thread_id));

-    for ch in value.chars() {
-        let normalized = ch.to_ascii_lowercase();
-        if normalized.is_ascii_alphanumeric() {
-            slug.push(normalized);
-            last_was_sep = false;
-        } else if !last_was_sep {
-            slug.push('_');
-            last_was_sep = true;
-        }
-    }
+    remove_outdated_thread_raw_memories(root, &memory.thread_id.to_string(), &path).await?;

-    let slug = slug.trim_matches('_').to_string();
-    if slug.is_empty() {
-        "memory".to_string()
-    } else {
-        slug.chars().take(64).collect()
-    }
+    let mut body = String::new();
+    writeln!(body, "thread_id: {}", memory.thread_id)
+        .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
+    writeln!(body, "updated_at: {}", memory.updated_at.to_rfc3339())
+        .map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
+    writeln!(body).map_err(|err| std::io::Error::other(format!("format raw memory: {err}")))?;
+    body.push_str(memory.raw_memory.trim());
+    body.push('\n');
+
+    tokio::fs::write(&path, body).await?;
+    Ok(path)
 }

 fn compact_summary_for_index(summary: &str) -> String {
@@ -208,10 +203,15 @@ fn compact_summary_for_index(summary: &str) -> String {

 fn extract_thread_id_from_summary_filename(file_name: &str) -> Option<&str> {
    let stem = file_name.strip_suffix(".md")?;
-    let (thread_id, _) = stem.split_once('_')?;
-    if thread_id.is_empty() {
+    if stem.is_empty() {
        None
+    } else if let Some((thread_id, _legacy_slug)) = stem.split_once('_') {
+        if thread_id.is_empty() {
+            None
+        } else {
+            Some(thread_id)
+        }
    } else {
-        Some(thread_id)
+        Some(stem)
    }
 }
@@ -1,7 +1,10 @@
+use super::MEMORY_SCOPE_KIND_CWD;
+use super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS;
 use super::StageOneResponseItemKinds;
 use super::StageOneRolloutFilter;
 use super::ensure_layout;
 use super::memory_root_for_cwd;
+use super::memory_scope_key_for_cwd;
 use super::memory_summary_file;
 use super::parse_stage_one_output;
 use super::prune_to_recent_memories_and_rebuild_summary;
@@ -77,7 +80,7 @@ fn memory_root_varies_by_cwd() {
        .and_then(std::path::Path::file_name)
        .and_then(std::ffi::OsStr::to_str)
        .expect("cwd bucket");
-    assert_eq!(bucket_a.len(), 64);
+    assert_eq!(bucket_a.len(), 16);
    assert!(bucket_a.chars().all(|ch| ch.is_ascii_hexdigit()));
 }

@@ -97,6 +100,22 @@ fn memory_root_encoding_avoids_component_collisions() {
    assert!(!root_hash.display().to_string().contains("workspace"));
 }

+#[test]
+fn memory_scope_key_uses_normalized_cwd() {
+    let dir = tempdir().expect("tempdir");
+    let workspace = dir.path().join("workspace");
+    std::fs::create_dir_all(&workspace).expect("mkdir workspace");
+    std::fs::create_dir_all(workspace.join("nested")).expect("mkdir nested");
+
+    let alias = workspace.join("nested").join("..");
+    let normalized = workspace
+        .canonicalize()
+        .expect("canonical workspace path should resolve");
+    let alias_key = memory_scope_key_for_cwd(&alias);
+    let normalized_key = memory_scope_key_for_cwd(&normalized);
+    assert_eq!(alias_key, normalized_key);
+}
+
 #[test]
 fn parse_stage_one_output_accepts_fenced_json() {
    let raw = "```json\n{\"rawMemory\":\"abc\",\"summary\":\"short\"}\n```";
@@ -206,64 +225,58 @@ fn serialize_filtered_rollout_response_items_filters_by_response_item_kind() {
 }

 #[test]
-fn select_rollout_candidates_uses_db_memory_recency() {
+fn select_rollout_candidates_filters_by_age_window() {
    let dir = tempdir().expect("tempdir");
    let cwd_a = dir.path().join("workspace-a");
    let cwd_b = dir.path().join("workspace-b");
    std::fs::create_dir_all(&cwd_a).expect("mkdir cwd a");
    std::fs::create_dir_all(&cwd_b).expect("mkdir cwd b");

+    let now = Utc::now().timestamp();
    let current_thread_id = ThreadId::default();
-    let stale_thread_id = ThreadId::default();
-    let fresh_thread_id = ThreadId::default();
-    let missing_thread_id = ThreadId::default();
+    let recent_thread_id = ThreadId::default();
+    let old_thread_id = ThreadId::default();
+    let recent_two_thread_id = ThreadId::default();

    let current = thread_metadata(
        current_thread_id,
        dir.path().join("current.jsonl"),
        cwd_a.clone(),
        "current",
-        500,
+        now,
    );
-    let fresh = thread_metadata(
-        fresh_thread_id,
-        dir.path().join("fresh.jsonl"),
+    let recent = thread_metadata(
+        recent_thread_id,
+        dir.path().join("recent.jsonl"),
        cwd_a,
-        "fresh",
-        400,
+        "recent",
+        now - 10,
    );
-    let stale = thread_metadata(
-        stale_thread_id,
-        dir.path().join("stale.jsonl"),
+    let old = thread_metadata(
+        old_thread_id,
+        dir.path().join("old.jsonl"),
        cwd_b.clone(),
-        "stale",
-        300,
+        "old",
+        now - (PHASE_ONE_MAX_ROLLOUT_AGE_DAYS + 1) * 24 * 60 * 60,
    );
-    let missing = thread_metadata(
-        missing_thread_id,
-        dir.path().join("missing.jsonl"),
+    let recent_two = thread_metadata(
+        recent_two_thread_id,
+        dir.path().join("recent-two.jsonl"),
        cwd_b,
-        "missing",
-        200,
+        "recent-two",
+        now - 20,
    );

-    let memories = vec![ThreadMemory {
-        thread_id: fresh_thread_id,
-        raw_memory: "raw memory".to_string(),
-        memory_summary: "memory".to_string(),
-        updated_at: Utc.timestamp_opt(450, 0).single().expect("timestamp"),
-    }];
-
    let candidates = select_rollout_candidates_from_db(
-        &[current, fresh, stale, missing],
+        &[current, recent, old, recent_two],
        current_thread_id,
-        &memories,
        5,
+        PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
    );

    assert_eq!(candidates.len(), 2);
-    assert_eq!(candidates[0].thread_id, stale_thread_id);
-    assert_eq!(candidates[1].thread_id, missing_thread_id);
+    assert_eq!(candidates[0].thread_id, recent_thread_id);
+    assert_eq!(candidates[1].thread_id, recent_two_thread_id);
 }

 #[tokio::test]
@@ -274,8 +287,8 @@ async fn prune_and_rebuild_summary_keeps_latest_memories_only() {

    let keep_id = ThreadId::default().to_string();
    let drop_id = ThreadId::default().to_string();
-    let keep_path = raw_memories_dir(&root).join(format!("{keep_id}_keep.md"));
-    let drop_path = raw_memories_dir(&root).join(format!("{drop_id}_drop.md"));
+    let keep_path = raw_memories_dir(&root).join(format!("{keep_id}.md"));
+    let drop_path = raw_memories_dir(&root).join(format!("{drop_id}.md"));
    tokio::fs::write(&keep_path, "keep")
        .await
        .expect("write keep");
@@ -285,9 +298,15 @@ async fn prune_and_rebuild_summary_keeps_latest_memories_only() {

    let memories = vec![ThreadMemory {
        thread_id: ThreadId::try_from(keep_id.clone()).expect("thread id"),
+        scope_kind: MEMORY_SCOPE_KIND_CWD.to_string(),
+        scope_key: "scope".to_string(),
        raw_memory: "raw memory".to_string(),
        memory_summary: "short summary".to_string(),
        updated_at: Utc.timestamp_opt(100, 0).single().expect("timestamp"),
+        last_used_at: None,
+        used_count: 0,
+        invalidated_at: None,
+        invalid_reason: None,
    }];

    prune_to_recent_memories_and_rebuild_summary(&root, &memories)
@@ -11,8 +11,6 @@ pub(crate) struct RolloutCandidate {
    pub(crate) rollout_path: PathBuf,
    /// Thread working directory used for per-project memory bucketing.
    pub(crate) cwd: PathBuf,
-    /// Best-effort thread title used to build readable memory filenames.
-    pub(crate) title: String,
    /// Last observed thread update timestamp (RFC3339), if available.
    pub(crate) updated_at: Option<String>,
 }