feat: memories forgetting (#12900)

Add diff based memory forgetting

codex-rs/state/src/lib.rs

@@ -14,6 +14,7 @@ mod runtime;
 pub use model::LogEntry;
 pub use model::LogQuery;
 pub use model::LogRow;
+pub use model::Phase2InputSelection;
 pub use model::Phase2JobClaimOutcome;
 /// Preferred entrypoint: owns configuration and metrics.
 pub use runtime::StateRuntime;
@@ -38,6 +39,7 @@ pub use model::SortKey;
 pub use model::Stage1JobClaim;
 pub use model::Stage1JobClaimOutcome;
 pub use model::Stage1Output;
+pub use model::Stage1OutputRef;
 pub use model::Stage1StartupClaimParams;
 pub use model::ThreadMetadata;
 pub use model::ThreadMetadataBuilder;

codex-rs/state/src/model/memories.rs

@@ -21,6 +21,21 @@ pub struct Stage1Output {
     pub generated_at: DateTime<Utc>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Stage1OutputRef {
+    pub thread_id: ThreadId,
+    pub source_updated_at: DateTime<Utc>,
+    pub rollout_slug: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct Phase2InputSelection {
+    pub selected: Vec<Stage1Output>,
+    pub previous_selected: Vec<Stage1Output>,
+    pub retained_thread_ids: Vec<ThreadId>,
+    pub removed: Vec<Stage1OutputRef>,
+}
+
 #[derive(Debug)]
 pub(crate) struct Stage1OutputRow {
     thread_id: String,
@@ -70,6 +85,18 @@ fn epoch_seconds_to_datetime(secs: i64) -> Result<DateTime<Utc>> {
         .ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}"))
 }
 
+pub(crate) fn stage1_output_ref_from_parts(
+    thread_id: String,
+    source_updated_at: i64,
+    rollout_slug: Option<String>,
+) -> Result<Stage1OutputRef> {
+    Ok(Stage1OutputRef {
+        thread_id: ThreadId::try_from(thread_id)?,
+        source_updated_at: epoch_seconds_to_datetime(source_updated_at)?,
+        rollout_slug,
+    })
+}
+
 /// Result of trying to claim a stage-1 memory extraction job.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Stage1JobClaimOutcome {

codex-rs/state/src/model/mod.rs

@@ -16,10 +16,12 @@ pub use backfill_state::BackfillStatus;
 pub use log::LogEntry;
 pub use log::LogQuery;
 pub use log::LogRow;
+pub use memories::Phase2InputSelection;
 pub use memories::Phase2JobClaimOutcome;
 pub use memories::Stage1JobClaim;
 pub use memories::Stage1JobClaimOutcome;
 pub use memories::Stage1Output;
+pub use memories::Stage1OutputRef;
 pub use memories::Stage1StartupClaimParams;
 pub use thread_metadata::Anchor;
 pub use thread_metadata::BackfillStats;
@@ -32,6 +34,7 @@ pub use thread_metadata::ThreadsPage;
 pub(crate) use agent_job::AgentJobItemRow;
 pub(crate) use agent_job::AgentJobRow;
 pub(crate) use memories::Stage1OutputRow;
+pub(crate) use memories::stage1_output_ref_from_parts;
 pub(crate) use thread_metadata::ThreadRow;
 pub(crate) use thread_metadata::anchor_from_item;
 pub(crate) use thread_metadata::datetime_to_epoch_seconds;

codex-rs/state/src/runtime.rs

@@ -2773,7 +2773,11 @@ WHERE kind = 'memory_stage1'
         assert_eq!(phase2_input_watermark, 100);
         assert!(
             runtime
-                .mark_global_phase2_job_succeeded(phase2_token.as_str(), phase2_input_watermark)
+                .mark_global_phase2_job_succeeded(
+                    phase2_token.as_str(),
+                    phase2_input_watermark,
+                    &[],
+                )
                 .await
                 .expect("mark initial phase2 succeeded"),
             "initial phase2 success should clear global dirty state"
@@ -2819,7 +2823,11 @@ WHERE kind = 'memory_stage1'
         assert_eq!(phase2_input_watermark, 101);
         assert!(
             runtime
-                .mark_global_phase2_job_succeeded(phase2_token.as_str(), phase2_input_watermark)
+                .mark_global_phase2_job_succeeded(
+                    phase2_token.as_str(),
+                    phase2_input_watermark,
+                    &[],
+                )
                 .await
                 .expect("mark phase2 succeeded after no-output delete")
         );
@@ -2936,7 +2944,7 @@ WHERE kind = 'memory_stage1'
         };
         assert!(
             runtime
-                .mark_global_phase2_job_succeeded(ownership_token.as_str(), input_watermark)
+                .mark_global_phase2_job_succeeded(ownership_token.as_str(), input_watermark, &[],)
                 .await
                 .expect("mark phase2 succeeded"),
             "phase2 success should finalize for current token"
@@ -3124,6 +3132,646 @@ VALUES (?, ?, ?, ?, ?)
         let _ = tokio::fs::remove_dir_all(codex_home).await;
     }
 
+    #[tokio::test]
+    async fn get_phase2_input_selection_reports_added_retained_and_removed_rows() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("initialize runtime");
+
+        let thread_id_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let thread_id_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let thread_id_c = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
+
+        for (thread_id, workspace) in [
+            (thread_id_a, "workspace-a"),
+            (thread_id_b, "workspace-b"),
+            (thread_id_c, "workspace-c"),
+        ] {
+            runtime
+                .upsert_thread(&test_thread_metadata(
+                    &codex_home,
+                    thread_id,
+                    codex_home.join(workspace),
+                ))
+                .await
+                .expect("upsert thread");
+        }
+
+        for (thread_id, updated_at, slug) in [
+            (thread_id_a, 100, Some("rollout-a")),
+            (thread_id_b, 101, Some("rollout-b")),
+            (thread_id_c, 102, Some("rollout-c")),
+        ] {
+            let claim = runtime
+                .try_claim_stage1_job(thread_id, owner, updated_at, 3600, 64)
+                .await
+                .expect("claim stage1");
+            let ownership_token = match claim {
+                Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+                other => panic!("unexpected stage1 claim outcome: {other:?}"),
+            };
+            assert!(
+                runtime
+                    .mark_stage1_job_succeeded(
+                        thread_id,
+                        ownership_token.as_str(),
+                        updated_at,
+                        &format!("raw-{updated_at}"),
+                        &format!("summary-{updated_at}"),
+                        slug,
+                    )
+                    .await
+                    .expect("mark stage1 succeeded"),
+                "stage1 success should persist output"
+            );
+        }
+
+        let claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim phase2");
+        let (ownership_token, input_watermark) = match claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected phase2 claim outcome: {other:?}"),
+        };
+        assert_eq!(input_watermark, 102);
+        let selected_outputs = runtime
+            .list_stage1_outputs_for_global(10)
+            .await
+            .expect("list stage1 outputs for global")
+            .into_iter()
+            .filter(|output| output.thread_id == thread_id_c || output.thread_id == thread_id_a)
+            .collect::<Vec<_>>();
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    ownership_token.as_str(),
+                    input_watermark,
+                    &selected_outputs,
+                )
+                .await
+                .expect("mark phase2 success with selection"),
+            "phase2 success should persist selected rows"
+        );
+
+        let selection = runtime
+            .get_phase2_input_selection(2)
+            .await
+            .expect("load phase2 input selection");
+
+        assert_eq!(selection.selected.len(), 2);
+        assert_eq!(selection.previous_selected.len(), 2);
+        assert_eq!(selection.selected[0].thread_id, thread_id_c);
+        assert_eq!(
+            selection.selected[0].rollout_path,
+            codex_home.join(format!("rollout-{thread_id_c}.jsonl"))
+        );
+        assert_eq!(selection.selected[1].thread_id, thread_id_b);
+        assert_eq!(selection.retained_thread_ids, vec![thread_id_c]);
+
+        assert_eq!(selection.removed.len(), 1);
+        assert_eq!(selection.removed[0].thread_id, thread_id_a);
+        assert_eq!(
+            selection.removed[0].rollout_slug.as_deref(),
+            Some("rollout-a")
+        );
+
+        let _ = tokio::fs::remove_dir_all(codex_home).await;
+    }
+
+    #[tokio::test]
+    async fn get_phase2_input_selection_treats_regenerated_selected_rows_as_added() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("initialize runtime");
+
+        let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
+        runtime
+            .upsert_thread(&test_thread_metadata(
+                &codex_home,
+                thread_id,
+                codex_home.join("workspace"),
+            ))
+            .await
+            .expect("upsert thread");
+
+        let first_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 100, 3600, 64)
+            .await
+            .expect("claim initial stage1");
+        let first_token = match first_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    first_token.as_str(),
+                    100,
+                    "raw-100",
+                    "summary-100",
+                    Some("rollout-100"),
+                )
+                .await
+                .expect("mark initial stage1 success"),
+            "initial stage1 success should persist output"
+        );
+
+        let phase2_claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim phase2");
+        let (phase2_token, input_watermark) = match phase2_claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected phase2 claim outcome: {other:?}"),
+        };
+        let selected_outputs = runtime
+            .list_stage1_outputs_for_global(1)
+            .await
+            .expect("list selected outputs");
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    phase2_token.as_str(),
+                    input_watermark,
+                    &selected_outputs,
+                )
+                .await
+                .expect("mark phase2 success"),
+            "phase2 success should persist selected rows"
+        );
+
+        let refreshed_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 101, 3600, 64)
+            .await
+            .expect("claim refreshed stage1");
+        let refreshed_token = match refreshed_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    refreshed_token.as_str(),
+                    101,
+                    "raw-101",
+                    "summary-101",
+                    Some("rollout-101"),
+                )
+                .await
+                .expect("mark refreshed stage1 success"),
+            "refreshed stage1 success should persist output"
+        );
+
+        let selection = runtime
+            .get_phase2_input_selection(1)
+            .await
+            .expect("load phase2 input selection");
+        assert_eq!(selection.selected.len(), 1);
+        assert_eq!(selection.previous_selected.len(), 1);
+        assert_eq!(selection.selected[0].thread_id, thread_id);
+        assert_eq!(selection.selected[0].source_updated_at.timestamp(), 101);
+        assert!(selection.retained_thread_ids.is_empty());
+        assert!(selection.removed.is_empty());
+
+        let (selected_for_phase2, selected_for_phase2_source_updated_at) =
+            sqlx::query_as::<_, (i64, Option<i64>)>(
+                "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?",
+            )
+        .bind(thread_id.to_string())
+        .fetch_one(runtime.pool.as_ref())
+        .await
+        .expect("load selected_for_phase2");
+        assert_eq!(selected_for_phase2, 1);
+        assert_eq!(selected_for_phase2_source_updated_at, Some(100));
+
+        let _ = tokio::fs::remove_dir_all(codex_home).await;
+    }
+
+    #[tokio::test]
+    async fn get_phase2_input_selection_reports_regenerated_previous_selection_as_removed() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("initialize runtime");
+
+        let thread_id_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread a");
+        let thread_id_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread b");
+        let thread_id_c = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread c");
+        let thread_id_d = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread d");
+        let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
+
+        for (thread_id, workspace) in [
+            (thread_id_a, "workspace-a"),
+            (thread_id_b, "workspace-b"),
+            (thread_id_c, "workspace-c"),
+            (thread_id_d, "workspace-d"),
+        ] {
+            runtime
+                .upsert_thread(&test_thread_metadata(
+                    &codex_home,
+                    thread_id,
+                    codex_home.join(workspace),
+                ))
+                .await
+                .expect("upsert thread");
+        }
+
+        for (thread_id, updated_at, slug) in [
+            (thread_id_a, 100, Some("rollout-a-100")),
+            (thread_id_b, 101, Some("rollout-b-101")),
+            (thread_id_c, 99, Some("rollout-c-99")),
+            (thread_id_d, 98, Some("rollout-d-98")),
+        ] {
+            let claim = runtime
+                .try_claim_stage1_job(thread_id, owner, updated_at, 3600, 64)
+                .await
+                .expect("claim initial stage1");
+            let ownership_token = match claim {
+                Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+                other => panic!("unexpected stage1 claim outcome: {other:?}"),
+            };
+            assert!(
+                runtime
+                    .mark_stage1_job_succeeded(
+                        thread_id,
+                        ownership_token.as_str(),
+                        updated_at,
+                        &format!("raw-{updated_at}"),
+                        &format!("summary-{updated_at}"),
+                        slug,
+                    )
+                    .await
+                    .expect("mark stage1 succeeded"),
+                "stage1 success should persist output"
+            );
+        }
+
+        let phase2_claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim phase2");
+        let (phase2_token, input_watermark) = match phase2_claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected phase2 claim outcome: {other:?}"),
+        };
+        let selected_outputs = runtime
+            .list_stage1_outputs_for_global(2)
+            .await
+            .expect("list selected outputs");
+        assert_eq!(
+            selected_outputs
+                .iter()
+                .map(|output| output.thread_id)
+                .collect::<Vec<_>>(),
+            vec![thread_id_b, thread_id_a]
+        );
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    phase2_token.as_str(),
+                    input_watermark,
+                    &selected_outputs,
+                )
+                .await
+                .expect("mark phase2 success"),
+            "phase2 success should persist selected rows"
+        );
+
+        for (thread_id, updated_at, slug) in [
+            (thread_id_a, 102, Some("rollout-a-102")),
+            (thread_id_c, 103, Some("rollout-c-103")),
+            (thread_id_d, 104, Some("rollout-d-104")),
+        ] {
+            let claim = runtime
+                .try_claim_stage1_job(thread_id, owner, updated_at, 3600, 64)
+                .await
+                .expect("claim refreshed stage1");
+            let ownership_token = match claim {
+                Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+                other => panic!("unexpected stage1 claim outcome: {other:?}"),
+            };
+            assert!(
+                runtime
+                    .mark_stage1_job_succeeded(
+                        thread_id,
+                        ownership_token.as_str(),
+                        updated_at,
+                        &format!("raw-{updated_at}"),
+                        &format!("summary-{updated_at}"),
+                        slug,
+                    )
+                    .await
+                    .expect("mark refreshed stage1 success"),
+                "refreshed stage1 success should persist output"
+            );
+        }
+
+        let selection = runtime
+            .get_phase2_input_selection(2)
+            .await
+            .expect("load phase2 input selection");
+        assert_eq!(
+            selection
+                .selected
+                .iter()
+                .map(|output| output.thread_id)
+                .collect::<Vec<_>>(),
+            vec![thread_id_d, thread_id_c]
+        );
+        assert_eq!(
+            selection
+                .previous_selected
+                .iter()
+                .map(|output| output.thread_id)
+                .collect::<Vec<_>>(),
+            vec![thread_id_a, thread_id_b]
+        );
+        assert!(selection.retained_thread_ids.is_empty());
+        assert_eq!(
+            selection
+                .removed
+                .iter()
+                .map(|output| (output.thread_id, output.source_updated_at.timestamp()))
+                .collect::<Vec<_>>(),
+            vec![(thread_id_a, 102), (thread_id_b, 101)]
+        );
+
+        let _ = tokio::fs::remove_dir_all(codex_home).await;
+    }
+
+    #[tokio::test]
+    async fn mark_global_phase2_job_succeeded_updates_selected_snapshot_timestamp() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("initialize runtime");
+
+        let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
+        runtime
+            .upsert_thread(&test_thread_metadata(
+                &codex_home,
+                thread_id,
+                codex_home.join("workspace"),
+            ))
+            .await
+            .expect("upsert thread");
+
+        let initial_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 100, 3600, 64)
+            .await
+            .expect("claim initial stage1");
+        let initial_token = match initial_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    initial_token.as_str(),
+                    100,
+                    "raw-100",
+                    "summary-100",
+                    Some("rollout-100"),
+                )
+                .await
+                .expect("mark initial stage1 success"),
+            "initial stage1 success should persist output"
+        );
+
+        let first_phase2_claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim first phase2");
+        let (first_phase2_token, first_input_watermark) = match first_phase2_claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected first phase2 claim outcome: {other:?}"),
+        };
+        let first_selected_outputs = runtime
+            .list_stage1_outputs_for_global(1)
+            .await
+            .expect("list first selected outputs");
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    first_phase2_token.as_str(),
+                    first_input_watermark,
+                    &first_selected_outputs,
+                )
+                .await
+                .expect("mark first phase2 success"),
+            "first phase2 success should persist selected rows"
+        );
+
+        let refreshed_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 101, 3600, 64)
+            .await
+            .expect("claim refreshed stage1");
+        let refreshed_token = match refreshed_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected refreshed stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    refreshed_token.as_str(),
+                    101,
+                    "raw-101",
+                    "summary-101",
+                    Some("rollout-101"),
+                )
+                .await
+                .expect("mark refreshed stage1 success"),
+            "refreshed stage1 success should persist output"
+        );
+
+        let second_phase2_claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim second phase2");
+        let (second_phase2_token, second_input_watermark) = match second_phase2_claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected second phase2 claim outcome: {other:?}"),
+        };
+        let second_selected_outputs = runtime
+            .list_stage1_outputs_for_global(1)
+            .await
+            .expect("list second selected outputs");
+        assert_eq!(
+            second_selected_outputs[0].source_updated_at.timestamp(),
+            101
+        );
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    second_phase2_token.as_str(),
+                    second_input_watermark,
+                    &second_selected_outputs,
+                )
+                .await
+                .expect("mark second phase2 success"),
+            "second phase2 success should persist selected rows"
+        );
+
+        let selection = runtime
+            .get_phase2_input_selection(1)
+            .await
+            .expect("load phase2 input selection after refresh");
+        assert_eq!(selection.retained_thread_ids, vec![thread_id]);
+
+        let (selected_for_phase2, selected_for_phase2_source_updated_at) =
+            sqlx::query_as::<_, (i64, Option<i64>)>(
+                "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?",
+            )
+            .bind(thread_id.to_string())
+            .fetch_one(runtime.pool.as_ref())
+            .await
+            .expect("load selected snapshot after phase2");
+        assert_eq!(selected_for_phase2, 1);
+        assert_eq!(selected_for_phase2_source_updated_at, Some(101));
+
+        let _ = tokio::fs::remove_dir_all(codex_home).await;
+    }
+
+    #[tokio::test]
+    async fn mark_global_phase2_job_succeeded_only_marks_exact_selected_snapshots() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("initialize runtime");
+
+        let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
+        let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
+        runtime
+            .upsert_thread(&test_thread_metadata(
+                &codex_home,
+                thread_id,
+                codex_home.join("workspace"),
+            ))
+            .await
+            .expect("upsert thread");
+
+        let initial_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 100, 3600, 64)
+            .await
+            .expect("claim initial stage1");
+        let initial_token = match initial_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    initial_token.as_str(),
+                    100,
+                    "raw-100",
+                    "summary-100",
+                    Some("rollout-100"),
+                )
+                .await
+                .expect("mark initial stage1 success"),
+            "initial stage1 success should persist output"
+        );
+
+        let phase2_claim = runtime
+            .try_claim_global_phase2_job(owner, 3600)
+            .await
+            .expect("claim phase2");
+        let (phase2_token, input_watermark) = match phase2_claim {
+            Phase2JobClaimOutcome::Claimed {
+                ownership_token,
+                input_watermark,
+            } => (ownership_token, input_watermark),
+            other => panic!("unexpected phase2 claim outcome: {other:?}"),
+        };
+        let selected_outputs = runtime
+            .list_stage1_outputs_for_global(1)
+            .await
+            .expect("list selected outputs");
+        assert_eq!(selected_outputs[0].source_updated_at.timestamp(), 100);
+
+        let refreshed_claim = runtime
+            .try_claim_stage1_job(thread_id, owner, 101, 3600, 64)
+            .await
+            .expect("claim refreshed stage1");
+        let refreshed_token = match refreshed_claim {
+            Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
+            other => panic!("unexpected stage1 claim outcome: {other:?}"),
+        };
+        assert!(
+            runtime
+                .mark_stage1_job_succeeded(
+                    thread_id,
+                    refreshed_token.as_str(),
+                    101,
+                    "raw-101",
+                    "summary-101",
+                    Some("rollout-101"),
+                )
+                .await
+                .expect("mark refreshed stage1 success"),
+            "refreshed stage1 success should persist output"
+        );
+
+        assert!(
+            runtime
+                .mark_global_phase2_job_succeeded(
+                    phase2_token.as_str(),
+                    input_watermark,
+                    &selected_outputs,
+                )
+                .await
+                .expect("mark phase2 success"),
+            "phase2 success should still complete"
+        );
+
+        let (selected_for_phase2, selected_for_phase2_source_updated_at) =
+            sqlx::query_as::<_, (i64, Option<i64>)>(
+                "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?",
+            )
+            .bind(thread_id.to_string())
+            .fetch_one(runtime.pool.as_ref())
+            .await
+            .expect("load selected_for_phase2");
+        assert_eq!(selected_for_phase2, 0);
+        assert_eq!(selected_for_phase2_source_updated_at, None);
+
+        let selection = runtime
+            .get_phase2_input_selection(1)
+            .await
+            .expect("load phase2 input selection");
+        assert_eq!(selection.selected.len(), 1);
+        assert_eq!(selection.selected[0].source_updated_at.timestamp(), 101);
+        assert!(selection.retained_thread_ids.is_empty());
+
+        let _ = tokio::fs::remove_dir_all(codex_home).await;
+    }
+
     #[tokio::test]
     async fn record_stage1_output_usage_updates_usage_metadata() {
         let codex_home = unique_temp_dir();
@@ -3395,7 +4043,7 @@ VALUES (?, ?, ?, ?, ?)
 
         assert_eq!(
             runtime
-                .mark_global_phase2_job_succeeded(token_a.as_str(), 300)
+                .mark_global_phase2_job_succeeded(token_a.as_str(), 300, &[])
                 .await
                 .expect("mark stale owner success result"),
             false,
@@ -3403,7 +4051,7 @@ VALUES (?, ?, ?, ?, ?)
         );
         assert!(
             runtime
-                .mark_global_phase2_job_succeeded(token_b.as_str(), 300)
+                .mark_global_phase2_job_succeeded(token_b.as_str(), 300, &[])
                 .await
                 .expect("mark takeover owner success"),
             "takeover owner should finalize consolidation"
@@ -3440,7 +4088,7 @@ VALUES (?, ?, ?, ?, ?)
         };
         assert!(
             runtime
-                .mark_global_phase2_job_succeeded(token_a.as_str(), 500)
+                .mark_global_phase2_job_succeeded(token_a.as_str(), 500, &[])
                 .await
                 .expect("mark initial phase2 success"),
             "initial phase2 success should finalize"

codex-rs/state/src/runtime/memories.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::model::Phase2InputSelection;
 use crate::model::Phase2JobClaimOutcome;
 use crate::model::Stage1JobClaim;
 use crate::model::Stage1JobClaimOutcome;
@@ -6,10 +7,12 @@ use crate::model::Stage1Output;
 use crate::model::Stage1OutputRow;
 use crate::model::Stage1StartupClaimParams;
 use crate::model::ThreadRow;
+use crate::model::stage1_output_ref_from_parts;
 use chrono::Duration;
 use sqlx::Executor;
 use sqlx::QueryBuilder;
 use sqlx::Sqlite;
+use std::collections::HashSet;
 
 const JOB_KIND_MEMORY_STAGE1: &str = "memory_stage1";
 const JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL: &str = "memory_consolidate_global";
@@ -257,6 +260,117 @@ LIMIT ?
             .collect::<Result<Vec<_>, _>>()
     }
 
+    /// Returns the current phase-2 input set along with its diff against the
+    /// last successful phase-2 selection.
+    ///
+    /// Query behavior:
+    /// - current selection is the latest `n` non-empty stage-1 outputs ordered
+    ///   by `source_updated_at DESC, thread_id DESC`
+    /// - previously selected rows are identified by `selected_for_phase2 = 1`
+    /// - `previous_selected` contains the current persisted rows that belonged
+    ///   to the last successful phase-2 baseline
+    /// - `retained_thread_ids` records which current rows still match the exact
+    ///   snapshot selected in the last successful phase-2 run
+    /// - removed rows are previously selected rows that are still present in
+    ///   `stage1_outputs` but fall outside the current top-`n` selection
+    pub async fn get_phase2_input_selection(
+        &self,
+        n: usize,
+    ) -> anyhow::Result<Phase2InputSelection> {
+        if n == 0 {
+            return Ok(Phase2InputSelection::default());
+        }
+
+        let current_rows = sqlx::query(
+            r#"
+SELECT
+    so.thread_id,
+    COALESCE(t.rollout_path, '') AS rollout_path,
+    so.source_updated_at,
+    so.raw_memory,
+    so.rollout_summary,
+    so.rollout_slug,
+    so.generated_at,
+    so.selected_for_phase2,
+    so.selected_for_phase2_source_updated_at,
+    COALESCE(t.cwd, '') AS cwd
+FROM stage1_outputs AS so
+LEFT JOIN threads AS t
+    ON t.id = so.thread_id
+WHERE length(trim(so.raw_memory)) > 0 OR length(trim(so.rollout_summary)) > 0
+ORDER BY so.source_updated_at DESC, so.thread_id DESC
+LIMIT ?
+            "#,
+        )
+        .bind(n as i64)
+        .fetch_all(self.pool.as_ref())
+        .await?;
+
+        let mut current_thread_ids = HashSet::with_capacity(current_rows.len());
+        let mut selected = Vec::with_capacity(current_rows.len());
+        let mut retained_thread_ids = Vec::new();
+        for row in current_rows {
+            let thread_id = row.try_get::<String, _>("thread_id")?;
+            current_thread_ids.insert(thread_id.clone());
+            let source_updated_at = row.try_get::<i64, _>("source_updated_at")?;
+            if row.try_get::<i64, _>("selected_for_phase2")? != 0
+                && row.try_get::<Option<i64>, _>("selected_for_phase2_source_updated_at")?
+                    == Some(source_updated_at)
+            {
+                retained_thread_ids.push(ThreadId::try_from(thread_id.clone())?);
+            }
+            selected.push(Stage1Output::try_from(Stage1OutputRow::try_from_row(
+                &row,
+            )?)?);
+        }
+
+        let previous_rows = sqlx::query(
+            r#"
+SELECT
+    so.thread_id,
+    COALESCE(t.rollout_path, '') AS rollout_path,
+    so.source_updated_at,
+    so.raw_memory,
+    so.rollout_summary,
+    so.rollout_slug
+  , so.generated_at
+  , COALESCE(t.cwd, '') AS cwd
+FROM stage1_outputs AS so
+LEFT JOIN threads AS t
+    ON t.id = so.thread_id
+WHERE so.selected_for_phase2 = 1
+ORDER BY so.source_updated_at DESC, so.thread_id DESC
+            "#,
+        )
+        .fetch_all(self.pool.as_ref())
+        .await?;
+
+        let previous_selected = previous_rows
+            .iter()
+            .map(Stage1OutputRow::try_from_row)
+            .map(|row| row.and_then(Stage1Output::try_from))
+            .collect::<Result<Vec<_>, _>>()?;
+        let mut removed = Vec::new();
+        for row in previous_rows {
+            let thread_id = row.try_get::<String, _>("thread_id")?;
+            if current_thread_ids.contains(thread_id.as_str()) {
+                continue;
+            }
+            removed.push(stage1_output_ref_from_parts(
+                thread_id,
+                row.try_get("source_updated_at")?,
+                row.try_get("rollout_slug")?,
+            )?);
+        }
+
+        Ok(Phase2InputSelection {
+            selected,
+            previous_selected,
+            retained_thread_ids,
+            removed,
+        })
+    }
+
     /// Attempts to claim a stage-1 job for a thread at `source_updated_at`.
     ///
     /// Claim semantics:
@@ -454,6 +568,9 @@ WHERE kind = ? AND job_key = ?
     /// - sets `status='done'` and `last_success_watermark = input_watermark`
     /// - upserts `stage1_outputs` for the thread, replacing existing output only
     ///   when `source_updated_at` is newer or equal
+    /// - preserves any existing `selected_for_phase2` baseline until the next
+    ///   successful phase-2 run rewrites the baseline selection, including the
+    ///   snapshot timestamp chosen during that run
     /// - persists optional `rollout_slug` for rollout summary artifact naming
     /// - enqueues/advances the global phase-2 job watermark using
     ///   `source_updated_at`
@@ -806,12 +923,18 @@ WHERE kind = ? AND job_key = ?
     /// - sets `status='done'`, clears lease/errors
     /// - advances `last_success_watermark` to
     ///   `max(existing_last_success_watermark, completed_watermark)`
+    /// - rewrites `selected_for_phase2` so only the exact selected stage-1
+    ///   snapshots remain marked as part of the latest successful phase-2
+    ///   selection, and persists each selected snapshot's
+    ///   `source_updated_at` for future retained-vs-added diffing
     pub async fn mark_global_phase2_job_succeeded(
         &self,
         ownership_token: &str,
         completed_watermark: i64,
+        selected_outputs: &[Stage1Output],
     ) -> anyhow::Result<bool> {
         let now = Utc::now().timestamp();
+        let mut tx = self.pool.begin().await?;
         let rows_affected = sqlx::query(
             r#"
 UPDATE jobs
@@ -830,11 +953,46 @@ WHERE kind = ? AND job_key = ?
         .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL)
         .bind(MEMORY_CONSOLIDATION_JOB_KEY)
         .bind(ownership_token)
-        .execute(self.pool.as_ref())
+        .execute(&mut *tx)
         .await?
         .rows_affected();
 
-        Ok(rows_affected > 0)
+        if rows_affected == 0 {
+            tx.commit().await?;
+            return Ok(false);
+        }
+
+        sqlx::query(
+            r#"
+UPDATE stage1_outputs
+SET
+    selected_for_phase2 = 0,
+    selected_for_phase2_source_updated_at = NULL
+WHERE selected_for_phase2 != 0 OR selected_for_phase2_source_updated_at IS NOT NULL
+            "#,
+        )
+        .execute(&mut *tx)
+        .await?;
+
+        for output in selected_outputs {
+            sqlx::query(
+                r#"
+UPDATE stage1_outputs
+SET
+    selected_for_phase2 = 1,
+    selected_for_phase2_source_updated_at = ?
+WHERE thread_id = ? AND source_updated_at = ?
+                "#,
+            )
+            .bind(output.source_updated_at.timestamp())
+            .bind(output.thread_id.to_string())
+            .bind(output.source_updated_at.timestamp())
+            .execute(&mut *tx)
+            .await?;
+        }
+
+        tx.commit().await?;
+        Ok(true)
     }
 
     /// Marks the owned running global phase-2 job as failed and schedules retry.