From a813bf8ab6906eafb1fd9d9c3408ec8b5bb2bfc8 Mon Sep 17 00:00:00 2001
From: Owen Lin <owen@openai.com>
Date: Mon, 4 May 2026 09:35:05 -0700
Subject: [PATCH] fix flaky test
 load_history_uses_live_writer_rollout_path_for_archived_source

---
 codex-rs/rollout/src/state_db.rs              | 22 +++++++++++++++++++
 codex-rs/thread-store/src/local/mod.rs        | 12 ++++++++++
 .../thread-store/src/local/read_thread.rs     | 20 +++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs
index 41b59c9760..68cec9e0b7 100644
--- a/codex-rs/rollout/src/state_db.rs
+++ b/codex-rs/rollout/src/state_db.rs
@@ -1,3 +1,4 @@
+use crate::ARCHIVED_SESSIONS_SUBDIR;
 use crate::config::RolloutConfig;
 use crate::config::RolloutConfigView;
 use crate::list::Cursor;
@@ -14,6 +15,7 @@ pub use codex_state::LogEntry;
 use codex_state::ThreadMetadataBuilder;
 use codex_utils_path::normalize_for_path_comparison;
 use serde_json::Value;
+use std::ffi::OsStr;
 use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -534,6 +536,10 @@ pub async fn apply_rollout_items(
     }
     builder.rollout_path = rollout_path.to_path_buf();
     builder.cwd = normalize_cwd_for_state_db(&builder.cwd);
+    let updated_at = updated_at_override.unwrap_or_else(Utc::now);
+    if rollout_path_is_archived(rollout_path) && builder.archived_at.is_none() {
+        builder.archived_at = Some(updated_at);
+    }
     if let Err(err) = ctx
         .apply_rollout_items(&builder, items, new_thread_memory_mode, updated_at_override)
         .await
@@ -542,7 +548,23 @@ pub async fn apply_rollout_items(
             "state db apply_rollout_items failed during {stage} for {}: {err}",
             rollout_path.display()
         );
+        return;
     }
+    if rollout_path_is_archived(rollout_path)
+        && let Err(err) = ctx
+            .mark_archived(builder.id, rollout_path, updated_at)
+            .await
+    {
+        warn!(
+            "state db apply_rollout_items failed to preserve archived status during {stage} for {}: {err}",
+            rollout_path.display()
+        );
+    }
+}
+
+fn rollout_path_is_archived(path: &Path) -> bool {
+    path.components()
+        .any(|component| component.as_os_str() == OsStr::new(ARCHIVED_SESSIONS_SUBDIR))
 }
 
 pub async fn touch_thread_updated_at(
diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs
index 04dd8b2490..5f18c92000 100644
--- a/codex-rs/thread-store/src/local/mod.rs
+++ b/codex-rs/thread-store/src/local/mod.rs
@@ -664,6 +664,18 @@ mod tests {
             .flush_thread(thread_id)
             .await
             .expect("flush live thread");
+        assert!(
+            store
+                .state_db()
+                .await
+                .expect("state db")
+                .get_thread(thread_id)
+                .await
+                .expect("read state db metadata")
+                .expect("state db metadata")
+                .archived_at
+                .is_some()
+        );
 
         let err = store
             .read_thread(ReadThreadParams {
diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs
index 8b3d3160db..daa20fae40 100644
--- a/codex-rs/thread-store/src/local/read_thread.rs
+++ b/codex-rs/thread-store/src/local/read_thread.rs
@@ -31,6 +31,19 @@ pub(super) async fn read_thread(
     params: ReadThreadParams,
 ) -> ThreadStoreResult<StoredThread> {
     let thread_id = params.thread_id;
+    if let Ok(rollout_path) = live_writer::rollout_path(store, thread_id).await {
+        if !params.include_archived
+            && rollout_path_is_archived(store.config.codex_home.as_path(), rollout_path.as_path())
+        {
+            return Err(ThreadStoreError::InvalidRequest {
+                message: format!("thread {thread_id} is archived"),
+            });
+        }
+        let mut thread = read_thread_from_rollout_path(store, rollout_path).await?;
+        attach_history_if_requested(&mut thread, params.include_history).await?;
+        return Ok(thread);
+    }
+
     if let Some(metadata) = read_sqlite_metadata(store, thread_id).await
         && (params.include_archived
             || (metadata.archived_at.is_none()
@@ -69,6 +82,13 @@ pub(super) async fn read_thread(
         .ok_or_else(|| ThreadStoreError::InvalidRequest {
             message: format!("no rollout found for thread id {thread_id}"),
         })?;
+    if !params.include_archived
+        && rollout_path_is_archived(store.config.codex_home.as_path(), path.as_path())
+    {
+        return Err(ThreadStoreError::InvalidRequest {
+            message: format!("thread {thread_id} is archived"),
+        });
+    }
 
     let mut thread = read_thread_from_rollout_path(store, path).await?;
     attach_history_if_requested(&mut thread, params.include_history).await?;