diff --git a/codex-rs/memories/write/src/phase2.rs b/codex-rs/memories/write/src/phase2.rs index 0c588a607072..1c7f35b5ae3f 100644 --- a/codex-rs/memories/write/src/phase2.rs +++ b/codex-rs/memories/write/src/phase2.rs @@ -12,6 +12,7 @@ use crate::sync_rollout_summaries_from_memories; use crate::workspace::memory_workspace_diff; use crate::workspace::prepare_memory_workspace; use crate::workspace::reset_memory_workspace_baseline; +use crate::workspace::validate_consolidation_artifacts; use crate::workspace::write_workspace_diff; use codex_config::Constrained; use codex_core::config::Config; @@ -140,7 +141,7 @@ pub async fn run(context: Arc, config: Arc) { return; } }; - if !workspace_diff.has_changes() { + if !workspace_diff.has_changes() && validate_consolidation_artifacts(&root).await.is_ok() { tracing::error!("Phase 2 no changes"); // We check only after sync of the file system. job::succeed( @@ -378,14 +379,30 @@ mod agent { let final_status = loop_agent(db.clone(), claim.token.clone(), thread_id, &thread).await; - if matches!(final_status, AgentStatus::Completed(_)) { - if let Some(token_usage) = thread + let agent_completed = matches!(final_status, AgentStatus::Completed(_)); + if agent_completed + && let Some(token_usage) = thread .token_usage_info() .await .map(|info| info.total_token_usage) - { - emit_token_usage_metrics(context.as_ref(), &token_usage); + { + emit_token_usage_metrics(context.as_ref(), &token_usage); + } + let artifacts_valid = if agent_completed { + match validate_consolidation_artifacts(&memory_root).await { + Ok(()) => true, + Err(err) => { + tracing::error!("memory consolidation artifacts are invalid: {err}"); + job::failed(context.as_ref(), &db, &claim, "failed_invalid_artifacts") + .await; + false + } } + } else { + false + }; + + if agent_completed && artifacts_valid { // Do not reset the workspace baseline if we lost the lock. let still_owns_lock = match db .memories() @@ -431,7 +448,7 @@ mod agent { ); } } - } else { + } else if !agent_completed { job::failed(context.as_ref(), &db, &claim, "failed_agent").await; } diff --git a/codex-rs/memories/write/src/startup_tests.rs b/codex-rs/memories/write/src/startup_tests.rs index fad1ce14a425..4154b852bf82 100644 --- a/codex-rs/memories/write/src/startup_tests.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -4,6 +4,8 @@ use crate::phase1; use crate::phase2; use crate::runtime::MemoryStartupContext; use crate::start_memories_startup_task; +use crate::storage::rebuild_raw_memories_file_from_memories; +use crate::storage::sync_rollout_summaries_from_memories; use codex_config::types::MemoriesConfig; use codex_features::Feature; use codex_git_utils::diff_since_latest_init; @@ -27,6 +29,7 @@ use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionSource; +use codex_state::Phase2JobClaimOutcome; use core_test_support::responses::ResponseMock; use core_test_support::responses::ResponsesRequest; use core_test_support::responses::ev_assistant_message; @@ -91,6 +94,7 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> anyhow:: "git_branch: branch-rollout-a\n\nrollout summary A\n", ) .await?; + seed_required_memory_artifacts(&memory_root).await?; reset_git_repository(&memory_root).await?; let _thread_b = seed_stage1_output( @@ -149,6 +153,59 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> anyhow:: Ok(()) } +#[tokio::test] +async fn phase2_retries_when_clean_workspace_is_missing_artifacts() -> anyhow::Result<()> { + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let db = init_state_db(&home).await?; + let memory_root = home.path().join("memories"); + seed_stage1_output( + db.as_ref(), + home.path(), + chrono::Utc::now(), + "raw memory", + "rollout summary", + "missing-artifacts", + ) + .await?; + let raw_memories = db + .memories() + .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 1) + .await?; + sync_rollout_summaries_from_memories(&memory_root, &raw_memories, raw_memories.len()).await?; + rebuild_raw_memories_file_from_memories(&memory_root, &raw_memories, raw_memories.len()) + .await?; + seed_extension_instructions(&memory_root).await?; + reset_git_repository(&memory_root).await?; + let phase2 = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-phase2-missing-artifacts"), + ev_assistant_message("msg-phase2-missing-artifacts", "phase2 complete"), + ev_completed("resp-phase2-missing-artifacts"), + ]), + ) + .await; + let test = build_test_codex(&server, home.clone()).await?; + + trigger_memories_startup(&test).await; + wait_for_single_request(&phase2).await; + + assert_eq!( + wait_for_phase2_job_to_finish(db.as_ref()).await?, + Phase2JobClaimOutcome::SkippedRetryUnavailable + ); + assert!(!memory_root.join("MEMORY.md").exists()); + assert!(!memory_root.join("memory_summary.md").exists()); + assert_eq!( + tokio::fs::read_to_string(memory_root.join("phase2_workspace_diff.md")).await?, + "# Memory Workspace Diff\n\nGenerated by Codex before Phase 2 memory consolidation. Read this file first and do not edit it.\n\n## Status\n- none\n" + ); + + shutdown_test_codex(&test).await?; + Ok(()) +} + #[tokio::test] async fn memories_startup_phase2_prunes_old_extension_resources() -> anyhow::Result<()> { let server = start_mock_server().await; @@ -183,6 +240,7 @@ async fn memories_startup_phase2_prunes_old_extension_resources() -> anyhow::Res (now - chrono::Duration::days(6)).format("%Y-%m-%dT%H-%M-%S") )); tokio::fs::write(&recent_file, "recent resource").await?; + seed_required_memory_artifacts(&home.path().join("memories")).await?; let phase2 = mount_sse_once( &server, @@ -243,6 +301,7 @@ async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_i (now - chrono::Duration::days(8)).format("%Y-%m-%dT%H-%M-%S") )); tokio::fs::write(&old_file, "old resource").await?; + seed_required_memory_artifacts(&home.path().join("memories")).await?; let phase2 = mount_sse_once( &server, @@ -494,6 +553,7 @@ async fn run_memory_phase_two_model_request_test( let root = memory_root(&config.codex_home); tokio::fs::create_dir_all(&root).await?; seed_extension_instructions(&root).await?; + seed_required_memory_artifacts(&root).await?; phase2::run(context, config).await; let request = wait_for_single_request(&response).await; wait_for_phase2_workspace_reset(&home.path().join("memories")).await?; @@ -803,6 +863,33 @@ async fn wait_for_phase2_workspace_reset(memory_root: &Path) -> anyhow::Result<( } } +async fn wait_for_phase2_job_to_finish( + db: &codex_state::StateRuntime, +) -> anyhow::Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let outcome = db + .memories() + .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) + .await?; + if outcome != Phase2JobClaimOutcome::SkippedRunning { + return Ok(outcome); + } + anyhow::ensure!( + Instant::now() < deadline, + "timed out waiting for phase-2 job to finish" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +async fn seed_required_memory_artifacts(root: &Path) -> anyhow::Result<()> { + tokio::fs::create_dir_all(root).await?; + tokio::fs::write(root.join("MEMORY.md"), "memory\n").await?; + tokio::fs::write(root.join("memory_summary.md"), "v1\n\nsummary\n").await?; + Ok(()) +} + async fn seed_stage1_output_for_existing_thread( db: &codex_state::StateRuntime, thread_id: ThreadId, diff --git a/codex-rs/memories/write/src/workspace.rs b/codex-rs/memories/write/src/workspace.rs index 17b8e5a95287..9f80693e4768 100644 --- a/codex-rs/memories/write/src/workspace.rs +++ b/codex-rs/memories/write/src/workspace.rs @@ -45,6 +45,34 @@ pub async fn reset_memory_workspace_baseline(root: &Path) -> anyhow::Result<()> reset_git_repository(root).await } +/// Verifies that a completed consolidation run left the required memory artifacts in place. +pub async fn validate_consolidation_artifacts(root: &Path) -> anyhow::Result<()> { + let memory_path = root.join("MEMORY.md"); + let memory_metadata = tokio::fs::metadata(&memory_path).await.with_context(|| { + format!( + "read consolidated memory artifact {}", + memory_path.display() + ) + })?; + anyhow::ensure!( + memory_metadata.is_file(), + "consolidated memory artifact is not a file: {}", + memory_path.display() + ); + + let summary_path = root.join("memory_summary.md"); + let summary = tokio::fs::read_to_string(&summary_path) + .await + .with_context(|| format!("read memory summary artifact {}", summary_path.display()))?; + anyhow::ensure!( + summary.lines().next() == Some("v1"), + "memory summary artifact does not start with v1: {}", + summary_path.display() + ); + + Ok(()) +} + /// 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 diff --git a/codex-rs/memories/write/src/workspace_tests.rs b/codex-rs/memories/write/src/workspace_tests.rs index e46576c0d687..8c157273a60a 100644 --- a/codex-rs/memories/write/src/workspace_tests.rs +++ b/codex-rs/memories/write/src/workspace_tests.rs @@ -76,3 +76,18 @@ fn previous_char_boundary_handles_multibyte_text() { let text = "aƩ"; assert_eq!(previous_char_boundary(text, /*max_bytes*/ 2), 1); } + +#[tokio::test] +async fn validate_consolidation_artifacts_rejects_invalid_summary() { + let home = TempDir::new().expect("tempdir"); + let root = home.path().join("memories"); + fs::create_dir_all(&root).expect("create memory root"); + fs::write(root.join("MEMORY.md"), "memory").expect("write memory"); + fs::write(root.join("memory_summary.md"), "outdated\n").expect("write summary"); + + let err = validate_consolidation_artifacts(&root) + .await + .expect_err("invalid summary should fail validation"); + + assert!(err.to_string().contains("does not start with v1")); +}