From 2793c826e8f5a04e0619a18e9bbdae91c0435cd5 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Mon, 20 Jul 2026 17:35:17 +0000 Subject: [PATCH] Enable memories for paginated threads (#34386) ## Why Paginated threads keep metadata updates in SQLite, while their rollout retains the initial `memory_mode`. Reconciliation could therefore overwrite the current setting with a stale value, and memory processing excluded these threads entirely. ## What changed - Include paginated threads in stage 1 memory job selection and global memory output lookup. - Preserve the SQLite `memory_mode` when reconciling or backfilling an existing paginated thread, while continuing to seed missing rows and restore legacy threads from rollouts. - Omit `memory_mode` alongside Git metadata when flushing resumed paginated history. ## Testing Added coverage for preserving disabled memory mode during reconciliation and backfill, selecting eligible paginated threads, and omitting initial metadata on paginated resume. GitOrigin-RevId: 2a6e16068e69680728757fbec27aeefae45b8110 --- codex-rs/rollout/src/metadata.rs | 19 +++++-- codex-rs/rollout/src/metadata_tests.rs | 51 ++++++++++++++++- codex-rs/rollout/src/state_db.rs | 19 +++++-- codex-rs/rollout/src/state_db_tests.rs | 55 ++++++++++++++++++- codex-rs/state/src/runtime/memories.rs | 32 +++++++---- .../thread-store/src/thread_metadata_sync.rs | 7 ++- 6 files changed, 154 insertions(+), 29 deletions(-) diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index 127a8af223..b62e1e463d 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -14,6 +14,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_state::BackfillState; use codex_state::BackfillStats; use codex_state::BackfillStatus; @@ -265,9 +266,14 @@ pub(crate) async fn backfill_sessions_with_lease( let mut metadata = outcome.metadata; metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd); let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string()); - if let Ok(Some(existing_metadata)) = runtime.get_thread(metadata.id).await { - metadata.prefer_existing_git_info(&existing_metadata); - metadata.prefer_existing_explicit_title(&existing_metadata); + let existing_metadata = runtime.get_thread(metadata.id).await.ok().flatten(); + // Paginated metadata updates are SQLite-only. Use the rollout mode to seed a + // missing row, then keep the value from SQLite. + let restore_memory_mode_from_rollout = existing_metadata.is_none() + || matches!(metadata.history_mode, ThreadHistoryMode::Legacy); + if let Some(existing_metadata) = existing_metadata.as_ref() { + metadata.prefer_existing_git_info(existing_metadata); + metadata.prefer_existing_explicit_title(existing_metadata); } if rollout.archived && metadata.archived_at.is_none() { let fallback_archived_at = metadata.updated_at; @@ -279,9 +285,10 @@ pub(crate) async fn backfill_sessions_with_lease( stats.failed = stats.failed.saturating_add(1); warn!("failed to upsert rollout {}: {err}", rollout.path.display()); } else { - if let Err(err) = runtime - .set_thread_memory_mode(metadata.id, memory_mode.as_str()) - .await + if restore_memory_mode_from_rollout + && let Err(err) = runtime + .set_thread_memory_mode(metadata.id, memory_mode.as_str()) + .await { stats.failed = stats.failed.saturating_add(1); warn!( diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index bea534bf0c..03677a8b56 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -345,6 +345,52 @@ async fn backfill_sessions_preserves_existing_git_branch_and_fills_missing_git_f ); } +#[tokio::test] +async fn backfill_sessions_preserves_existing_paginated_memory_mode() { + let dir = tempdir().expect("tempdir"); + let codex_home = dir.path().to_path_buf(); + let thread_uuid = Uuid::new_v4(); + let rollout_path = write_rollout_in_sessions_with_cwd( + codex_home.as_path(), + "2026-01-27T12-34-56", + "2026-01-27T12:34:56Z", + thread_uuid, + codex_home.clone(), + /*git*/ None, + ThreadHistoryMode::Paginated, + ); + + let runtime = codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + let thread_id = ThreadId::from_string(&thread_uuid.to_string()).expect("thread id"); + let existing = extract_metadata_from_rollout(&rollout_path, "test-provider") + .await + .expect("extract") + .metadata; + runtime + .upsert_thread(&existing) + .await + .expect("existing metadata upsert"); + assert!( + runtime + .set_thread_memory_mode(thread_id, "disabled") + .await + .expect("disable memory mode") + ); + + backfill_sessions(runtime.as_ref(), codex_home.as_path(), "test-provider").await; + + assert_eq!( + runtime + .get_thread_memory_mode(thread_id) + .await + .expect("get memory mode") + .as_deref(), + Some("disabled") + ); +} + #[tokio::test] async fn backfill_sessions_normalizes_cwd_before_upsert() { let dir = tempdir().expect("tempdir"); @@ -358,6 +404,7 @@ async fn backfill_sessions_normalizes_cwd_before_upsert() { thread_uuid, session_cwd.clone(), /*git*/ None, + ThreadHistoryMode::Legacy, ); let runtime = codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) @@ -391,6 +438,7 @@ fn write_rollout_in_sessions( thread_uuid, codex_home.to_path_buf(), git, + ThreadHistoryMode::Legacy, ) } @@ -401,6 +449,7 @@ fn write_rollout_in_sessions_with_cwd( thread_uuid: Uuid, cwd: PathBuf, git: Option, + history_mode: ThreadHistoryMode, ) -> PathBuf { let id = ThreadId::from_string(&thread_uuid.to_string()).expect("thread id"); let sessions_dir = codex_home.join("sessions"); @@ -425,7 +474,7 @@ fn write_rollout_in_sessions_with_cwd( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, - history_mode: Default::default(), + history_mode, history_base: None, subagent_history_start_ordinal: None, multi_agent_version: None, diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs index f22019af84..213fb65d0f 100644 --- a/codex-rs/rollout/src/state_db.rs +++ b/codex-rs/rollout/src/state_db.rs @@ -11,6 +11,7 @@ use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; pub use codex_state::LogEntry; use codex_state::ThreadMetadataBuilder; use codex_utils_path::normalize_for_path_comparison; @@ -528,9 +529,14 @@ pub async fn reconcile_rollout( let mut metadata = outcome.metadata; let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string()); metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd); - if let Ok(Some(existing_metadata)) = ctx.get_thread(metadata.id).await { - metadata.prefer_existing_git_info(&existing_metadata); - metadata.prefer_existing_explicit_title(&existing_metadata); + let existing_metadata = ctx.get_thread(metadata.id).await.ok().flatten(); + // Paginated metadata updates are SQLite-only. Use the rollout mode to seed a + // missing row, then keep the value from SQLite. + let restore_memory_mode_from_rollout = + existing_metadata.is_none() || matches!(metadata.history_mode, ThreadHistoryMode::Legacy); + if let Some(existing_metadata) = existing_metadata.as_ref() { + metadata.prefer_existing_git_info(existing_metadata); + metadata.prefer_existing_explicit_title(existing_metadata); } match archived_only { Some(true) if metadata.archived_at.is_none() => { @@ -548,9 +554,10 @@ pub async fn reconcile_rollout( ); return; } - if let Err(err) = ctx - .set_thread_memory_mode(metadata.id, memory_mode.as_str()) - .await + if restore_memory_mode_from_rollout + && let Err(err) = ctx + .set_thread_memory_mode(metadata.id, memory_mode.as_str()) + .await { warn!( "state db reconcile_rollout memory_mode update failed {}: {err}", diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 33eae943ba..5cc22fa654 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -10,6 +10,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::UserMessageEvent; use pretty_assertions::assert_eq; use std::path::Path; @@ -111,7 +112,8 @@ async fn try_init_times_out_waiting_for_stuck_startup_backfill() -> anyhow::Resu async fn reconcile_rollout_preserves_existing_explicit_title() -> anyhow::Result<()> { let home = TempDir::new().expect("temp dir"); let thread_id = ThreadId::new(); - let rollout_path = write_rollout_with_user_message(home.path(), thread_id, "Hey")?; + let rollout_path = + write_rollout_with_user_message(home.path(), thread_id, "Hey", ThreadHistoryMode::Legacy)?; let runtime = codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".to_string()) .await?; @@ -145,10 +147,59 @@ async fn reconcile_rollout_preserves_existing_explicit_title() -> anyhow::Result Ok(()) } +#[tokio::test] +async fn reconcile_rollout_preserves_existing_paginated_memory_mode() -> anyhow::Result<()> { + let home = TempDir::new().expect("temp dir"); + let thread_id = ThreadId::new(); + let rollout_path = write_rollout_with_user_message( + home.path(), + thread_id, + "Hey", + ThreadHistoryMode::Paginated, + )?; + let runtime = + codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".to_string()) + .await?; + + reconcile_rollout( + Some(runtime.as_ref()), + rollout_path.as_path(), + "test-provider", + /*builder*/ None, + &[], + /*archived_only*/ None, + /*new_thread_memory_mode*/ None, + ) + .await; + assert!( + runtime + .set_thread_memory_mode(thread_id, "disabled") + .await? + ); + + reconcile_rollout( + Some(runtime.as_ref()), + rollout_path.as_path(), + "test-provider", + /*builder*/ None, + &[], + /*archived_only*/ None, + /*new_thread_memory_mode*/ None, + ) + .await; + + assert_eq!( + runtime.get_thread_memory_mode(thread_id).await?.as_deref(), + Some("disabled") + ); + Ok(()) +} + fn write_rollout_with_user_message( home: &Path, thread_id: ThreadId, message: &str, + history_mode: ThreadHistoryMode, ) -> anyhow::Result { let dir = home.join("sessions/2026/06/01"); std::fs::create_dir_all(dir.as_path())?; @@ -177,7 +228,7 @@ fn write_rollout_with_user_message( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, - history_mode: Default::default(), + history_mode, history_base: None, subagent_history_start_ordinal: None, multi_agent_version: None, diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index fb82a82390..33103584c6 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -136,7 +136,6 @@ WHERE kind = ? AND job_key = ? /// - starts from `threads` filtered to active threads and allowed sources /// (`push_thread_filters`) /// - excludes threads with `memory_mode != 'enabled'` - /// - excludes paginated threads because stage 1 still full-loads rollout JSONL /// - excludes the current thread id /// - keeps only threads whose millisecond `updated_at` is in the age window /// - checks memory staleness against the memories DB @@ -216,7 +215,7 @@ FROM threads }, /*include_thread_id_tiebreaker*/ false, ); - builder.push(" AND threads.memory_mode = 'enabled' AND threads.history_mode = 'legacy'"); + builder.push(" AND threads.memory_mode = 'enabled'"); builder .push(" AND threads.id != ") .push_bind(current_thread_id.as_str()); @@ -574,7 +573,7 @@ SELECT threads.git_branch, threads.git_origin_url FROM threads -WHERE threads.id = ? AND threads.memory_mode = 'enabled' AND threads.history_mode = 'legacy' +WHERE threads.id = ? AND threads.memory_mode = 'enabled' "#, ) .bind(thread_id.to_string()) @@ -2167,7 +2166,7 @@ mod tests { } #[tokio::test] - async fn claim_stage1_jobs_skips_threads_without_legacy_enabled_memory() { + async fn claim_stage1_jobs_skips_threads_without_enabled_memory() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await @@ -2198,6 +2197,7 @@ mod tests { test_thread_metadata(&codex_home, disabled_thread_id, codex_home.join("disabled")); disabled.created_at = eligible_at; disabled.updated_at = eligible_at; + disabled.history_mode = ThreadHistoryMode::Paginated; runtime .upsert_thread(&disabled) .await @@ -2246,8 +2246,17 @@ mod tests { .await .expect("claim stage1 startup jobs"); - assert_eq!(claims.len(), 1); - assert_eq!(claims[0].thread.id, enabled_thread_id); + let mut claimed_ids = claims + .iter() + .map(|claim| claim.thread.id.to_string()) + .collect::>(); + claimed_ids.sort(); + let mut expected_ids = vec![ + enabled_thread_id.to_string(), + paginated_thread_id.to_string(), + ]; + expected_ids.sort(); + assert_eq!(claimed_ids, expected_ids); let _ = tokio::fs::remove_dir_all(codex_home).await; } @@ -3301,7 +3310,7 @@ VALUES (?, ?, ?, ?, ?) } #[tokio::test] - async fn list_stage1_outputs_for_global_skips_polluted_threads() { + async fn list_stage1_outputs_for_global_includes_paginated_and_skips_polluted_threads() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await @@ -3317,12 +3326,11 @@ VALUES (?, ?, ?, ?, ?) (thread_id_enabled, "workspace-enabled"), (thread_id_polluted, "workspace-polluted"), ] { + let mut metadata = + test_thread_metadata(&codex_home, thread_id, codex_home.join(workspace)); + metadata.history_mode = ThreadHistoryMode::Paginated; runtime - .upsert_thread(&test_thread_metadata( - &codex_home, - thread_id, - codex_home.join(workspace), - )) + .upsert_thread(&metadata) .await .expect("upsert thread"); diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index da12f7b304..44e502eacc 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -201,8 +201,9 @@ impl ThreadMetadataSync { ThreadHistoryMode::Paginated ) { // Paginated rollouts never append metadata-only SessionMeta updates. Do not reapply - // the initial Git tuple when resume history is flushed after the first append. + // initial metadata when resume history is flushed after the first append. update.git_info = None; + update.memory_mode = None; } Some(update) } @@ -675,10 +676,11 @@ mod tests { } #[test] - fn paginated_resume_history_does_not_reapply_initial_git_info() { + fn paginated_resume_history_does_not_reapply_initial_metadata() { let thread_id = ThreadId::new(); let mut meta = session_meta(thread_id); meta.meta.history_mode = ThreadHistoryMode::Paginated; + meta.meta.memory_mode = Some("disabled".to_string()); meta.git = Some(GitInfo { commit_hash: None, branch: Some("stale-branch".to_string()), @@ -694,6 +696,7 @@ mod tests { let update = sync.take_pending_update().expect("pending metadata update"); assert_eq!(update.patch.git_info, None); + assert_eq!(update.patch.memory_mode, None); assert_eq!(update.patch.preview.as_deref(), Some("hello metadata")); }