Wait for memory consolidation shutdown before finishing jobs (#36140)

## Why

A phase-two memory job could release its lease while its consolidation agent was
still shutting down, allowing another worker to race the existing agent.

## What changed

- Await consolidation-agent shutdown before validating artifacts and completing
  the job.
- Keep the lease until it expires if shutdown fails, and remove the agent from
  the thread manager only after shutdown succeeds.

## Testing

Update startup tests to wait for the phase-two job lifecycle to finish and verify
that the consolidation thread has been removed.

GitOrigin-RevId: 7ac3a7369529822ac95fd2ee9dcb27d39c191d20

codex-rs/memories/write/src/phase2.rs

@@ -406,6 +406,17 @@ mod agent {
             {
                 emit_token_usage_metrics(context.as_ref(), &token_usage);
             }
+
+            if let Err(err) = context
+                .shutdown_consolidation_agent(SpawnedConsolidationAgent { thread_id, thread })
+                .await
+            {
+                warn!("failed to auto-close global memory consolidation agent {thread_id}: {err}");
+                // Keep the existing lease until it expires so another worker cannot race a
+                // consolidation agent whose shutdown has not completed.
+                return;
+            }
+
             let artifacts_valid = if agent_completed {
                 match validate_consolidation_artifacts(&memory_root).await {
                     Ok(()) => true,
@@ -469,18 +480,6 @@ mod agent {
             } else if !agent_completed {
                 job::failed(context.as_ref(), &db, &claim, "failed_agent").await;
             }
-
-            let cleanup_context = Arc::clone(&context);
-            tokio::spawn(async move {
-                if let Err(err) = cleanup_context
-                    .shutdown_consolidation_agent(SpawnedConsolidationAgent { thread_id, thread })
-                    .await
-                {
-                    warn!(
-                        "failed to auto-close global memory consolidation agent {thread_id}: {err}"
-                    );
-                }
-            });
         });
     }
 

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

@@ -361,18 +361,14 @@ impl MemoryStartupContext {
         agent: SpawnedConsolidationAgent,
     ) -> anyhow::Result<()> {
         let SpawnedConsolidationAgent { thread_id, thread } = agent;
-        let thread = self
-            .thread_manager
-            .remove_thread(&thread_id)
-            .await
-            .unwrap_or(thread);
-
         tokio::time::timeout(Duration::from_secs(10), thread.shutdown_and_wait())
             .await
             .map_err(|_| {
                 anyhow::anyhow!("memory consolidation agent {thread_id} shutdown timed out")
             })??;
 
+        self.thread_manager.remove_thread(&thread_id).await;
+
         Ok(())
     }
 }

codex-rs/memories/write/src/startup_tests.rs

@@ -128,7 +128,7 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> anyhow::
         "expected workspace diff file in prompt: {prompt}"
     );
 
-    wait_for_phase2_workspace_reset(&memory_root).await?;
+    wait_for_phase2_workspace_reset(db.as_ref(), &memory_root).await?;
     let raw_memories = tokio::fs::read_to_string(memory_root.join("raw_memories.md")).await?;
     assert!(raw_memories.contains("raw memory B"));
     assert!(!raw_memories.contains("raw memory A"));
@@ -263,7 +263,7 @@ async fn memories_startup_phase2_prunes_old_extension_resources() -> anyhow::Res
         "expected workspace diff file in prompt: {prompt}"
     );
 
-    wait_for_phase2_workspace_reset(&home.path().join("memories")).await?;
+    wait_for_phase2_workspace_reset(db.as_ref(), &home.path().join("memories")).await?;
     wait_for_file_removed(&old_file).await?;
     assert!(
         !tokio::fs::try_exists(&old_file).await?,
@@ -324,8 +324,8 @@ async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_i
         "expected workspace diff file in prompt: {prompt}"
     );
 
+    wait_for_phase2_workspace_reset(db.as_ref(), &home.path().join("memories")).await?;
     wait_for_file_removed(&old_file).await?;
-    wait_for_phase2_workspace_reset(&home.path().join("memories")).await?;
 
     shutdown_test_codex(&test).await?;
     Ok(())
@@ -558,8 +558,20 @@ async fn run_memory_phase_two_model_request_test(
     let parent_permission_profile = config.permissions.effective_permission_profile();
     phase2::run(context, config, parent_permission_profile).await;
     let request = wait_for_single_request(&response).await;
-    wait_for_phase2_workspace_reset(&home.path().join("memories")).await?;
+    let consolidation_thread_id = ThreadId::from_string(
-    shutdown_test_codex(&test).await?;
+        request.body_json()["client_metadata"]["thread_id"]
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("phase-2 request should include the child thread id"))?,
+    )?;
+    wait_for_phase2_workspace_reset(db.as_ref(), &home.path().join("memories")).await?;
+    assert!(
+        test.thread_manager
+            .remove_thread(&consolidation_thread_id)
+            .await
+            .is_none(),
+        "phase-2 consolidation agent should be removed after shutdown"
+    );
+    test.codex.shutdown_and_wait().await?;
     Ok(request)
 }
 
@@ -854,21 +866,17 @@ fn phase2_prompt_text(request: &ResponsesRequest) -> String {
         .expect("phase2 prompt text")
 }
 
-async fn wait_for_phase2_workspace_reset(memory_root: &Path) -> anyhow::Result<()> {
+async fn wait_for_phase2_workspace_reset(
-    wait_for_file_removed(&memory_root.join("phase2_workspace_diff.md")).await?;
+    db: &codex_state::StateRuntime,
-    let deadline = Instant::now() + Duration::from_secs(10);
+    memory_root: &Path,
-    loop {
+) -> anyhow::Result<()> {
-        if let Ok(diff) = diff_since_latest_init(memory_root).await
+    assert_eq!(
-            && !diff.has_changes()
+        wait_for_phase2_job_to_finish(db).await?,
-        {
+        Phase2JobClaimOutcome::SkippedCooldown
-            return Ok(());
+    );
-        }
+    assert!(!tokio::fs::try_exists(memory_root.join("phase2_workspace_diff.md")).await?);
-        assert!(
+    assert!(!diff_since_latest_init(memory_root).await?.has_changes());
-            Instant::now() < deadline,
+    Ok(())
-            "timed out waiting for clean memory workspace baseline"
-        );
-        tokio::time::sleep(Duration::from_millis(50)).await;
-    }
 }
 
 async fn wait_for_phase2_job_to_finish(