diff --git a/codex-rs/rollout/src/rollout_reference_index.rs b/codex-rs/rollout/src/rollout_reference_index.rs index bd23b086b6..1e94bd03df 100644 --- a/codex-rs/rollout/src/rollout_reference_index.rs +++ b/codex-rs/rollout/src/rollout_reference_index.rs @@ -1,13 +1,15 @@ //! Indexes direct fork references found in local rollout files. use std::collections::HashMap; -use std::collections::HashSet; +use std::collections::hash_map::Entry; use std::io; use std::path::Path; +use std::path::PathBuf; use std::time::Duration; use std::time::Instant; use codex_protocol::RolloutId; +use codex_protocol::ThreadId; use codex_protocol::protocol::HistoryPosition; use crate::ARCHIVED_SESSIONS_SUBDIR; @@ -20,10 +22,17 @@ use crate::compression::RolloutFile; /// cheap inverse-reference questions without each reimplementing rollout discovery. #[derive(Debug, Default)] pub struct RolloutReferenceIndex { - history_base_by_rollout: HashMap, + rollouts_by_id: HashMap, reference_counts_by_rollout: HashMap, } +#[derive(Debug)] +struct IndexedRollout { + thread_id: ThreadId, + path: PathBuf, + history_base: Option, +} + impl RolloutReferenceIndex { /// Scans active and archived local rollout metadata without a deadline. pub async fn scan(codex_home: &Path) -> io::Result { @@ -64,15 +73,27 @@ impl RolloutReferenceIndex { /// Returns the direct history-base edge for `rollout_id`, if one was discovered. pub fn history_base(&self, rollout_id: RolloutId) -> Option<&HistoryPosition> { - self.history_base_by_rollout.get(&rollout_id) + self.rollouts_by_id + .get(&rollout_id) + .and_then(|rollout| rollout.history_base.as_ref()) + } + + /// Returns rollout IDs and paths whose session metadata belongs to `thread_id`. + pub fn rollouts_for_thread( + &self, + thread_id: ThreadId, + ) -> impl Iterator { + self.rollouts_by_id + .iter() + .filter(move |(_, rollout)| rollout.thread_id == thread_id) + .map(|(rollout_id, rollout)| (*rollout_id, rollout.path.as_path())) } async fn scan_with_deadline( codex_home: &Path, deadline: ScanDeadline, ) -> io::Result> { - let mut history_base_by_rollout = HashMap::new(); - let mut seen_rollout_ids = HashSet::new(); + let mut rollouts_by_id = HashMap::new(); let mut stack = vec![ codex_home.join(ARCHIVED_SESSIONS_SUBDIR), codex_home.join(SESSIONS_SUBDIR), @@ -111,17 +132,21 @@ impl RolloutReferenceIndex { let Ok(meta) = crate::read_session_meta_line(rollout_file.path()).await else { continue; }; - if !seen_rollout_ids.insert(rollout_id) { - continue; - } - if let Some(history_base) = meta.meta.history_base { - history_base_by_rollout.insert(rollout_id, history_base); + if let Entry::Vacant(entry) = rollouts_by_id.entry(rollout_id) { + entry.insert(IndexedRollout { + thread_id: meta.meta.id, + path: rollout_file.into_path(), + history_base: meta.meta.history_base, + }); } } } let mut reference_counts_by_rollout = HashMap::new(); - for (rollout_id, history_base) in &history_base_by_rollout { + for (rollout_id, rollout) in &rollouts_by_id { + let Some(history_base) = rollout.history_base else { + continue; + }; if history_base.thread_id == *rollout_id { continue; } @@ -130,7 +155,7 @@ impl RolloutReferenceIndex { .or_default() += 1; } Ok(Some(Self { - history_base_by_rollout, + rollouts_by_id, reference_counts_by_rollout, })) } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 20710bc2d4..dc9a1c154a 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -373,6 +373,26 @@ ON CONFLICT(child_thread_id) DO NOTHING .map(PathBuf::from)) } + /// Swap one thread's rollout path only when it still matches the expected path. + /// + /// This intentionally updates only the physical path. The logical thread metadata remains + /// attached to the stable thread id. + pub async fn replace_rollout_path_if_current( + &self, + id: ThreadId, + expected: &Path, + replacement: &Path, + ) -> anyhow::Result { + let result = + sqlx::query("UPDATE threads SET rollout_path = ? WHERE id = ? AND rollout_path = ?") + .bind(replacement.display().to_string()) + .bind(id.to_string()) + .bind(expected.display().to_string()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() == 1) + } + /// Find the newest thread whose user-facing title exactly matches `title`. #[allow(clippy::too_many_arguments)] pub async fn find_thread_by_exact_title( diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index f7bf9baefb..8c1e2fbb13 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -63,6 +63,7 @@ pub use types::PreparedFork; pub use types::ReadThreadByRolloutPathParams; pub use types::ReadThreadParams; pub use types::ResumeThreadParams; +pub use types::RevertThreadParams; pub use types::SearchTextRange; pub use types::SearchThreadOccurrencesParams; pub use types::SearchThreadsParams; diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 91af9dbdd3..66f42315cf 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -1,14 +1,17 @@ use super::LocalThreadStore; -use super::helpers::matching_rollout_file_name; +use super::helpers::owned_rollout_paths_from_index; +use super::helpers::restore_rollout_moves; +use super::helpers::rollout_path_is_archived; use super::helpers::scoped_rollout_path; -use crate::ArchiveThreadParams; +use super::helpers::validated_rollout_file_name; use crate::ArchiveThreadsParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; use chrono::Utc; -use codex_rollout::find_thread_path_by_id_str; +use codex_rollout::RolloutReferenceIndex; use tracing::warn; +use super::thread_rollout_resolver; pub(super) async fn archive_threads( store: &LocalThreadStore, params: ArchiveThreadsParams, @@ -36,11 +39,17 @@ pub(super) async fn archive_threads( } } let _writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?; + let reference_index = RolloutReferenceIndex::scan(store.config.codex_home.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to scan thread rollout files: {err}"), + })?; let parent_thread_id = thread_ids[0]; let mut archived_thread_ids = Vec::new(); for thread_id in thread_ids { - match archive_thread(store, ArchiveThreadParams { thread_id }).await { + let rollout_paths = owned_rollout_paths_from_index(&reference_index, thread_id); + match archive_thread_with_paths(store, thread_id, rollout_paths).await { Ok(()) => archived_thread_ids.push(thread_id), Err(err) if archived_thread_ids.is_empty() => return Err(err), Err(err) => warn!( @@ -51,35 +60,18 @@ pub(super) async fn archive_threads( Ok(archived_thread_ids) } -pub(super) async fn archive_thread( +async fn archive_thread_with_paths( store: &LocalThreadStore, - params: ArchiveThreadParams, + thread_id: codex_protocol::ThreadId, + mut rollout_paths: Vec, ) -> ThreadStoreResult<()> { - let thread_id = params.thread_id; let state_db_ctx = store.state_db().await; - let rollout_path = find_thread_path_by_id_str( - store.config.codex_home.as_path(), - &thread_id.to_string(), - state_db_ctx.as_deref(), - ) - .await - .map_err(|err| ThreadStoreError::InvalidRequest { - message: format!("failed to locate thread id {thread_id}: {err}"), - })? - .ok_or_else(|| ThreadStoreError::InvalidRequest { - message: format!("no rollout found for thread id {thread_id}"), - })?; - - let canonical_rollout_path = scoped_rollout_path( - store.config.codex_home.join(codex_rollout::SESSIONS_SUBDIR), - rollout_path.as_path(), - "sessions", - )?; - let file_name = matching_rollout_file_name( - canonical_rollout_path.as_path(), - thread_id, - rollout_path.as_path(), - )?; + let selected_rollout_path = thread_rollout_resolver::resolve_current(store, thread_id) + .await? + .map(|resolved| resolved.path) + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("no rollout found for thread id {thread_id}"), + })?; let archive_folder = store .config @@ -88,17 +80,62 @@ pub(super) async fn archive_thread( std::fs::create_dir_all(&archive_folder).map_err(|err| ThreadStoreError::Internal { message: format!("failed to archive thread: {err}"), })?; - let archived_path = archive_folder.join(&file_name); - std::fs::rename(&canonical_rollout_path, &archived_path).map_err(|err| { - ThreadStoreError::Internal { - message: format!("failed to archive thread: {err}"), + if !rollout_paths.contains(&selected_rollout_path) { + rollout_paths.push(selected_rollout_path.clone()); + } + let mut archived_path = None; + let mut rollout_moves = Vec::new(); + for rollout_path in rollout_paths { + if rollout_path_is_archived(store.config.codex_home.as_path(), rollout_path.as_path()) { + continue; } + let canonical_rollout_path = scoped_rollout_path( + store.config.codex_home.join(codex_rollout::SESSIONS_SUBDIR), + rollout_path.as_path(), + "sessions", + )?; + let file_name = + validated_rollout_file_name(canonical_rollout_path.as_path(), rollout_path.as_path())?; + let destination = archive_folder.join(&file_name); + if rollout_path == selected_rollout_path { + archived_path = Some(destination.clone()); + } + rollout_moves.push((canonical_rollout_path, destination)); + } + let archived_path = archived_path.ok_or_else(|| ThreadStoreError::Internal { + message: format!("failed to archive selected rollout for thread {thread_id}"), })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx + for (index, (source, destination)) in rollout_moves.iter().enumerate() { + if let Err(err) = std::fs::rename(source, destination) { + if let Err(restore_err) = restore_rollout_moves(&rollout_moves[..index]) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to archive thread: {err}; failed to restore moved rollouts: {restore_err}" + ), + }); + } + return Err(ThreadStoreError::Internal { + message: format!("failed to archive thread: {err}"), + }); + } + } + + if let Some(ctx) = state_db_ctx + && let Err(err) = ctx .mark_archived(thread_id, archived_path.as_path(), Utc::now()) - .await; + .await + { + if let Err(restore_err) = restore_rollout_moves(&rollout_moves) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to update archived thread metadata: {err}; failed to restore moved rollouts: {restore_err}" + ), + }); + } + return Err(ThreadStoreError::Internal { + message: format!("failed to update archived thread metadata: {err}"), + }); } Ok(()) } @@ -118,6 +155,7 @@ mod tests { use uuid::Uuid; use super::*; + use crate::ArchiveThreadParams; use crate::ListThreadsParams; use crate::ThreadSortKey; use crate::ThreadStore; diff --git a/codex-rs/thread-store/src/local/delete_thread.rs b/codex-rs/thread-store/src/local/delete_thread.rs index 69c0ee9802..74fafd5447 100644 --- a/codex-rs/thread-store/src/local/delete_thread.rs +++ b/codex-rs/thread-store/src/local/delete_thread.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::io::ErrorKind; use std::path::Path; +use std::path::PathBuf; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; use codex_rollout::RolloutReferenceIndex; @@ -17,13 +18,49 @@ use codex_rollout::find_thread_path_by_id_str; use codex_rollout::remove_thread_name_entries; use super::LocalThreadStore; -use super::helpers::matching_rollout_file_name; use super::helpers::scoped_rollout_path; +use super::helpers::validated_rollout_file_name; use crate::DeleteThreadParams; use crate::DeleteThreadsParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; +struct ThreadRollouts { + thread_id: codex_protocol::ThreadId, + rollout_ids: HashSet, + paths: Vec, +} + +impl ThreadRollouts { + fn from_index( + reference_index: &RolloutReferenceIndex, + thread_id: codex_protocol::ThreadId, + ) -> Self { + let mut rollout_ids = HashSet::new(); + let paths = reference_index + .rollouts_for_thread(thread_id) + .map(|(rollout_id, path)| { + rollout_ids.insert(rollout_id); + path.to_path_buf() + }) + .collect(); + Self { + thread_id, + rollout_ids, + paths, + } + } + + fn add_path(&mut self, path: PathBuf) { + if let Some(rollout_id) = codex_rollout::rollout_id_from_path(path.as_path()) { + self.rollout_ids.insert(rollout_id); + } + if !self.paths.contains(&path) { + self.paths.push(path); + } + } +} + pub(super) async fn delete_thread( store: &LocalThreadStore, params: DeleteThreadParams, @@ -32,11 +69,10 @@ pub(super) async fn delete_thread( let _lifecycle_guard = store.live_writer_locks.lock_lifecycle(thread_id).await; let _live_writer_guard = store.live_writer_locks.lock(thread_id).await; let reference_index = scan_reference_index(store).await?; - if reference_index.reference_count(thread_id) > 0 { - return Err(referenced_thread_error(thread_id)); - } + let thread_rollouts = ThreadRollouts::from_index(&reference_index, thread_id); + ensure_no_external_references(&reference_index, std::slice::from_ref(&thread_rollouts))?; let mut writer_guards = store.acquire_writer_locks(&[thread_id]).await?; - delete_thread_after_reference_check(store, thread_id, &mut writer_guards).await + delete_thread_after_reference_check(store, thread_rollouts, &mut writer_guards).await } pub(super) async fn delete_threads( @@ -48,9 +84,9 @@ pub(super) async fn delete_threads( return Ok(()); } - let deletion_set: HashSet<_> = thread_ids.iter().copied().collect(); - let mut lock_thread_ids: Vec<_> = deletion_set.iter().copied().collect(); + let mut lock_thread_ids = thread_ids.clone(); lock_thread_ids.sort_unstable_by_key(ToString::to_string); + lock_thread_ids.dedup(); let mut _lifecycle_guards = Vec::with_capacity(lock_thread_ids.len()); for thread_id in &lock_thread_ids { _lifecycle_guards.push(store.live_writer_locks.lock_lifecycle(*thread_id).await); @@ -61,34 +97,51 @@ pub(super) async fn delete_threads( } let reference_index = scan_reference_index(store).await?; - // References from children in this delete set are removed by the same request, so only - // references from children outside the set should block it. + let thread_rollouts = thread_ids + .iter() + .map(|thread_id| ThreadRollouts::from_index(&reference_index, *thread_id)) + .collect::>(); + ensure_no_external_references(&reference_index, thread_rollouts.as_slice())?; + + let mut writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?; + for thread_rollouts in thread_rollouts { + match delete_thread_after_reference_check(store, thread_rollouts, &mut writer_guards).await + { + Ok(()) | Err(ThreadStoreError::ThreadNotFound { .. }) => {} + Err(err) => return Err(err), + } + } + Ok(()) +} + +fn ensure_no_external_references( + reference_index: &RolloutReferenceIndex, + thread_rollouts: &[ThreadRollouts], +) -> ThreadStoreResult<()> { + let deletion_rollout_ids = thread_rollouts + .iter() + .flat_map(|thread_rollouts| thread_rollouts.rollout_ids.iter().copied()) + .collect::>(); let mut internal_reference_counts = HashMap::new(); - for child_thread_id in &deletion_set { - if let Some(history_base) = reference_index.history_base(*child_thread_id) - && history_base.thread_id != *child_thread_id - && deletion_set.contains(&history_base.thread_id) + for source_rollout_id in &deletion_rollout_ids { + if let Some(history_base) = reference_index.history_base(*source_rollout_id) + && history_base.thread_id != *source_rollout_id + && deletion_rollout_ids.contains(&history_base.thread_id) { *internal_reference_counts .entry(history_base.thread_id) .or_default() += 1; } } - for thread_id in &thread_ids { - let internal_reference_count = internal_reference_counts - .get(thread_id) - .copied() - .unwrap_or_default(); - if reference_index.reference_count(*thread_id) > internal_reference_count { - return Err(referenced_thread_error(*thread_id)); - } - } - - let mut writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?; - for thread_id in thread_ids { - match delete_thread_after_reference_check(store, thread_id, &mut writer_guards).await { - Ok(()) | Err(ThreadStoreError::ThreadNotFound { .. }) => {} - Err(err) => return Err(err), + for thread_rollouts in thread_rollouts { + if thread_rollouts.rollout_ids.iter().any(|rollout_id| { + let internal_reference_count = internal_reference_counts + .get(rollout_id) + .copied() + .unwrap_or_default(); + reference_index.reference_count(*rollout_id) > internal_reference_count + }) { + return Err(referenced_thread_error(thread_rollouts.thread_id)); } } Ok(()) @@ -112,12 +165,12 @@ fn referenced_thread_error(thread_id: codex_protocol::ThreadId) -> ThreadStoreEr async fn delete_thread_after_reference_check( store: &LocalThreadStore, - thread_id: codex_protocol::ThreadId, + mut thread_rollouts: ThreadRollouts, writer_guards: &mut Vec, ) -> ThreadStoreResult<()> { + let thread_id = thread_rollouts.thread_id; let thread_id_str = thread_id.to_string(); let state_db_ctx = store.state_db().await; - let mut rollout_paths = Vec::new(); match find_thread_path_by_id_str( store.config.codex_home.as_path(), thread_id_str.as_str(), @@ -125,7 +178,7 @@ async fn delete_thread_after_reference_check( ) .await { - Ok(Some(path)) => rollout_paths.push(path), + Ok(Some(path)) => thread_rollouts.add_path(path), Ok(None) => {} Err(err) => { return Err(ThreadStoreError::InvalidRequest { @@ -140,11 +193,7 @@ async fn delete_thread_after_reference_check( ) .await { - Ok(Some(path)) => { - if !rollout_paths.contains(&path) { - rollout_paths.push(path); - } - } + Ok(Some(path)) => thread_rollouts.add_path(path), Ok(None) => {} Err(err) => { return Err(ThreadStoreError::InvalidRequest { @@ -152,15 +201,18 @@ async fn delete_thread_after_reference_check( }); } } - super::thread_history::delete_thread(store, thread_id).await?; + thread_rollouts.rollout_ids.insert(thread_id); + for rollout_id in thread_rollouts.rollout_ids { + super::thread_history::delete_thread(store, rollout_id).await?; + } // Drop the recorder before removing files, but retain its writer lock until cleanup finishes. if let Some(entry) = store.live_recorders.lock().await.remove(&thread_id) { writer_guards.push(entry.writer_lock); } - let found_rollout_path = !rollout_paths.is_empty(); - for rollout_path in rollout_paths { - delete_rollout_file(store, rollout_path.as_path(), thread_id)?; + let found_rollout_path = !thread_rollouts.paths.is_empty(); + for rollout_path in thread_rollouts.paths { + delete_rollout_file(store, rollout_path.as_path())?; } remove_thread_name_entries(store.config.codex_home.as_path(), thread_id) .await @@ -175,23 +227,15 @@ async fn delete_thread_after_reference_check( Ok(()) } -fn delete_rollout_file( - store: &LocalThreadStore, - rollout_path: &Path, - thread_id: codex_protocol::ThreadId, -) -> ThreadStoreResult { +fn delete_rollout_file(store: &LocalThreadStore, rollout_path: &Path) -> ThreadStoreResult { let plain_path = codex_rollout::plain_rollout_path(rollout_path); let compressed_path = plain_path.with_extension("jsonl.zst"); - let deleted_plain = delete_rollout_path(store, plain_path.as_path(), thread_id)?; - let deleted_compressed = delete_rollout_path(store, compressed_path.as_path(), thread_id)?; + let deleted_plain = delete_rollout_path(store, plain_path.as_path())?; + let deleted_compressed = delete_rollout_path(store, compressed_path.as_path())?; Ok(deleted_plain || deleted_compressed) } -fn delete_rollout_path( - store: &LocalThreadStore, - rollout_path: &Path, - thread_id: codex_protocol::ThreadId, -) -> ThreadStoreResult { +fn delete_rollout_path(store: &LocalThreadStore, rollout_path: &Path) -> ThreadStoreResult { let canonical_rollout_path = scoped_rollout_path( store.config.codex_home.join(SESSIONS_SUBDIR), rollout_path, @@ -208,7 +252,7 @@ fn delete_rollout_path( Ok(false) => Ok(rollout_path.to_path_buf()), Ok(true) | Err(_) => Err(err), })?; - matching_rollout_file_name(&canonical_rollout_path, thread_id, rollout_path)?; + validated_rollout_file_name(&canonical_rollout_path, rollout_path)?; match std::fs::remove_file(&canonical_rollout_path) { Ok(()) => Ok(true), Err(err) if err.kind() == ErrorKind::NotFound => Ok(false), @@ -328,6 +372,73 @@ mod tests { assert!(source_path.exists()); } + #[tokio::test] + async fn delete_thread_rejects_reference_to_reverted_rollout() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let thread_uuid = Uuid::from_u128(320); + let thread_id = ThreadId::from_string(&thread_uuid.to_string()).expect("thread id"); + let source_path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T12-00-00", + thread_uuid, + ThreadHistoryMode::Paginated, + ) + .expect("source session file"); + let replacement_uuid = Uuid::from_u128(321); + let replacement_rollout_id = + ThreadId::from_string(&replacement_uuid.to_string()).expect("replacement rollout id"); + let replacement_path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T12-00-01", + replacement_uuid, + ThreadHistoryMode::Paginated, + ) + .expect("replacement session file"); + set_thread_id(replacement_path.as_path(), thread_id); + set_history_base( + replacement_path.as_path(), + HistoryPosition { + thread_id, + end_ordinal_exclusive: 1, + end_byte_offset: std::fs::metadata(source_path.as_path()) + .expect("source rollout metadata") + .len(), + }, + ); + let child_path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T12-00-02", + Uuid::from_u128(322), + ThreadHistoryMode::Paginated, + ) + .expect("child session file"); + set_history_base( + child_path.as_path(), + HistoryPosition { + thread_id: replacement_rollout_id, + end_ordinal_exclusive: 1, + end_byte_offset: std::fs::metadata(replacement_path.as_path()) + .expect("replacement rollout metadata") + .len(), + }, + ); + + let err = store + .delete_thread(DeleteThreadParams { thread_id }) + .await + .expect_err("referenced replacement should not be deleted"); + + assert_eq!( + err.to_string(), + format!( + "invalid thread-store request: cannot delete thread {thread_id}: forked history still references it" + ) + ); + assert!(source_path.exists()); + assert!(replacement_path.exists()); + } + #[tokio::test] async fn delete_thread_ignores_unreadable_reference_metadata() { let home = TempDir::new().expect("temp dir"); @@ -498,12 +609,11 @@ mod tests { let home = TempDir::new().expect("temp dir"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let uuid = Uuid::from_u128(305); - let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); std::fs::remove_file(&path).expect("remove session file"); - assert!(!delete_rollout_file(&store, path.as_path(), thread_id).expect("delete rollout")); + assert!(!delete_rollout_file(&store, path.as_path()).expect("delete rollout")); } #[tokio::test] @@ -696,4 +806,17 @@ SELECT serde_json::to_value(history_base).expect("serialize history base"); std::fs::write(path, format!("{session_meta}\n")).expect("write session file"); } + + fn set_thread_id(path: &Path, thread_id: ThreadId) { + let mut session_meta: serde_json::Value = serde_json::from_str( + std::fs::read_to_string(path) + .expect("read session file") + .lines() + .next() + .expect("session metadata"), + ) + .expect("parse session metadata"); + session_meta["payload"]["id"] = serde_json::to_value(thread_id).expect("thread id"); + std::fs::write(path, format!("{session_meta}\n")).expect("write session file"); + } } diff --git a/codex-rs/thread-store/src/local/helpers.rs b/codex-rs/thread-store/src/local/helpers.rs index a583ea8864..c5a362f73c 100644 --- a/codex-rs/thread-store/src/local/helpers.rs +++ b/codex-rs/thread-store/src/local/helpers.rs @@ -19,6 +19,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::RolloutReferenceIndex; use codex_rollout::ThreadItem; use codex_rollout::find_thread_names_by_ids; use codex_state::ThreadMetadata; @@ -66,9 +67,31 @@ pub(super) fn rollout_path_is_archived(codex_home: &Path, path: &Path) -> bool { .any(|component| component.as_os_str() == OsStr::new(ARCHIVED_SESSIONS_SUBDIR)) } -pub(super) fn matching_rollout_file_name( - rollout_path: &Path, +/// Returns rollout files whose session metadata belongs to `thread_id`. +pub(super) async fn owned_rollout_paths( + store: &LocalThreadStore, thread_id: ThreadId, +) -> ThreadStoreResult> { + RolloutReferenceIndex::scan(store.config.codex_home.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to scan thread rollout files: {err}"), + }) + .map(|index| owned_rollout_paths_from_index(&index, thread_id)) +} + +pub(super) fn owned_rollout_paths_from_index( + index: &RolloutReferenceIndex, + thread_id: ThreadId, +) -> Vec { + index + .rollouts_for_thread(thread_id) + .map(|(_, path)| path.to_path_buf()) + .collect() +} + +pub(super) fn validated_rollout_file_name( + rollout_path: &Path, display_path: &Path, ) -> ThreadStoreResult { let Some(file_name) = rollout_path.file_name().map(OsStr::to_owned) else { @@ -79,17 +102,12 @@ pub(super) fn matching_rollout_file_name( ), }); }; - let required_plain_suffix = format!("{thread_id}.jsonl"); - let required_compressed_suffix = format!("{required_plain_suffix}.zst"); - let file_name_str = file_name.to_string_lossy(); - if file_name_str.ends_with(required_plain_suffix.as_str()) - || file_name_str.ends_with(required_compressed_suffix.as_str()) - { + if codex_rollout::rollout_id_from_path(rollout_path).is_some() { Ok(file_name) } else { Err(ThreadStoreError::InvalidRequest { message: format!( - "rollout path `{}` does not match thread id {thread_id}", + "rollout path `{}` has an invalid filename", display_path.display() ), }) @@ -101,6 +119,13 @@ pub(super) fn touch_modified_time(path: &Path) -> std::io::Result<()> { OpenOptions::new().append(true).open(path)?.set_times(times) } +pub(super) fn restore_rollout_moves(moves: &[(PathBuf, PathBuf)]) -> std::io::Result<()> { + for (source, destination) in moves.iter().rev() { + std::fs::rename(destination, source)?; + } + Ok(()) +} + pub(super) fn stored_thread_from_rollout_item( item: ThreadItem, archived: bool, diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index bb3b3967ef..0512e9cfc3 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -8,6 +8,7 @@ mod model_context; mod move_thread_to_section; mod paginated_fork; mod read_thread; +mod revert_thread; mod rollout_migration; // This lands before the reader PRs that consume the shared lineage resolver. #[allow(dead_code)] @@ -62,6 +63,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::RenameThreadSectionParams; use crate::ResumeThreadParams; +use crate::RevertThreadParams; use crate::SearchThreadOccurrencesParams; use crate::SearchThreadsParams; use crate::StoredModelContext; @@ -445,6 +447,10 @@ impl ThreadStore for LocalThreadStore { Box::pin(async move { paginated_fork::prepare(self, params).await }) } + fn revert_thread(&self, params: RevertThreadParams) -> ThreadStoreFuture<'_, ()> { + Box::pin(async move { revert_thread::revert(self, params).await }) + } + fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread> { Box::pin(async move { read_thread::read_thread(self, params).await }) } diff --git a/codex-rs/thread-store/src/local/revert_thread.rs b/codex-rs/thread-store/src/local/revert_thread.rs new file mode 100644 index 0000000000..1bb40be400 --- /dev/null +++ b/codex-rs/thread-store/src/local/revert_thread.rs @@ -0,0 +1,176 @@ +use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_rollout::RolloutConfig; +use codex_rollout::RolloutRecorder; +use codex_rollout::RolloutRecorderParams; + +use super::LocalThreadStore; +use super::paginated_fork; +use super::thread_rollout_resolver; +use crate::ForkBoundary; +use crate::RevertThreadParams; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +/// Revert an unloaded paginated thread by creating a new immutable rollout file. +/// +/// Old rollouts stay intact. The new file references the retained prefix, and the only mutable +/// cutover is the existing SQLite rollout-path pointer for the thread. +pub(super) async fn revert( + store: &LocalThreadStore, + params: RevertThreadParams, +) -> ThreadStoreResult<()> { + let RevertThreadParams { + thread_id, + before_turn_id, + } = params; + let state_db = store + .state_db() + .await + .ok_or(ThreadStoreError::Unsupported { + operation: "revert_thread", + })?; + let _lifecycle_guard = store.live_writer_locks.lock_lifecycle(thread_id).await; + let _live_writer_guard = store.live_writer_locks.lock(thread_id).await; + store.ensure_live_recorder_absent(thread_id).await?; + let _writer_lock = store.writer_lock_coordinator.acquire(thread_id)?; + + // Resolution may return a compressed sibling. Keep SQLite's exact stored path for the CAS. + let expected_sqlite_path = state_db + .get_thread(thread_id) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to read thread metadata for {thread_id}: {err}"), + })? + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })? + .rollout_path; + let current_rollout = thread_rollout_resolver::resolve_current(store, thread_id) + .await? + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })?; + let source_path = current_rollout.path; + let source_meta = codex_rollout::read_session_meta_line(source_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to read current paginated rollout {}: {err}", + source_path.display() + ), + })? + .meta; + if source_meta.id != current_rollout.thread_id { + return Err(ThreadStoreError::InvalidRequest { + message: format!("current rollout for {thread_id} belongs to another thread"), + }); + } + if source_meta.history_mode != ThreadHistoryMode::Paginated { + return Err(ThreadStoreError::InvalidRequest { + message: format!("thread {thread_id} does not use paginated history"), + }); + } + + // HistoryPosition stores byte offsets in plain JSONL, so materialize any compressed + // immutable lineage before deriving the retained boundary. + let mut lineage = store.resolve_rollout_lineage(thread_id).await?; + for segment in &mut lineage.segments { + segment.rollout_path = + codex_rollout::materialize_rollout_for_reference(segment.rollout_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to materialize rollout {} for revert: {err}", + segment.rollout_path.display() + ), + })?; + super::thread_history_materialization::materialize_to_sqlite( + store, + segment.rollout_id(), + segment.rollout_path.as_path(), + ) + .await?; + } + let history_base = paginated_fork::history_base_at_boundary( + store, + thread_id, + ForkBoundary::BeforeTurn(before_turn_id), + &lineage, + ) + .await?; + + let rollout_id = ThreadId::new(); + let recorder = + create_replacement_recorder(store, source_meta, rollout_id, history_base).await?; + let replacement_path = recorder.rollout_path().to_path_buf(); + recorder.persist().await.map_err(thread_store_io_error)?; + recorder.shutdown().await.map_err(thread_store_io_error)?; + + let replaced = state_db + .replace_rollout_path_if_current( + thread_id, + expected_sqlite_path.as_path(), + replacement_path.as_path(), + ) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to switch thread {thread_id} to reverted rollout: {err}"), + })?; + if !replaced { + let _ = tokio::fs::remove_file(replacement_path.as_path()).await; + return Err(ThreadStoreError::Conflict { + message: format!("thread {thread_id} changed while it was being reverted"), + }); + } + Ok(()) +} + +async fn create_replacement_recorder( + store: &LocalThreadStore, + source_meta: codex_rollout::SessionMeta, + rollout_id: ThreadId, + history_base: Option, +) -> ThreadStoreResult { + let config = RolloutConfig { + codex_home: store.config.codex_home.clone(), + sqlite: store.config.sqlite.clone(), + cwd: source_meta.cwd.clone(), + model_provider_id: source_meta + .model_provider + .clone() + .unwrap_or_else(|| store.config.default_model_provider_id.clone()), + generate_memories: source_meta.memory_mode.as_deref() != Some("disabled"), + }; + let mut params = RolloutRecorderParams::new( + source_meta.id, + source_meta.forked_from_id, + source_meta.parent_thread_id, + source_meta.source, + source_meta.thread_source, + source_meta.originator, + source_meta.base_instructions.unwrap_or_default(), + source_meta.dynamic_tools.unwrap_or_default(), + ) + .with_session_id(source_meta.session_id) + .with_rollout_id(rollout_id) + .with_selected_capability_roots(source_meta.selected_capability_roots) + .with_multi_agent_version(source_meta.multi_agent_version) + .with_history_mode(ThreadHistoryMode::Paginated) + .with_history_base(history_base) + .with_subagent_history_start_ordinal(source_meta.subagent_history_start_ordinal); + if let Some(context_window) = source_meta.context_window { + params = params.with_initial_window_id(context_window.window_id); + } + RolloutRecorder::new(&config, params) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to create reverted rollout: {err}"), + }) +} + +fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError { + ThreadStoreError::Internal { + message: err.to_string(), + } +} + +#[cfg(test)] +#[path = "revert_thread_tests.rs"] +mod tests; diff --git a/codex-rs/thread-store/src/local/revert_thread_tests.rs b/codex-rs/thread-store/src/local/revert_thread_tests.rs new file mode 100644 index 0000000000..1bf99b62e0 --- /dev/null +++ b/codex-rs/thread-store/src/local/revert_thread_tests.rs @@ -0,0 +1,223 @@ +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_rollout::RolloutItem; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +use super::super::LocalThreadStore; +use super::super::test_support::test_config; +use crate::AppendThreadItemsParams; +use crate::ArchiveThreadParams; +use crate::CreateThreadParams; +use crate::DeleteThreadParams; +use crate::ListTurnsParams; +use crate::RevertThreadParams; +use crate::SortDirection; +use crate::StoredTurnItemsView; +use crate::ThreadPersistenceMetadata; +use crate::ThreadStore; + +#[tokio::test] +async fn revert_keeps_thread_id_and_hides_suffix_across_repeated_reverts() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let state_db = codex_state::StateRuntime::init( + config.sqlite.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("initialize state database"); + let store = LocalThreadStore::new(config, Some(state_db.clone())); + let thread_id = ThreadId::new(); + create_paginated_thread(&store, thread_id).await; + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![ + turn_started("turn-1"), + turn_completed("turn-1"), + turn_started("turn-2"), + turn_completed("turn-2"), + ], + }) + .await + .expect("append turns"); + let original_path = store + .live_rollout_path(thread_id) + .await + .expect("source rollout path"); + store + .shutdown_thread(thread_id) + .await + .expect("close source writer"); + codex_rollout::state_db::reconcile_rollout( + Some(state_db.as_ref()), + original_path.as_path(), + "test-provider", + /*builder*/ None, + &[], + /*archived_only*/ Some(false), + /*new_thread_memory_mode*/ None, + ) + .await; + compress_rollout(original_path.as_path()); + + store + .revert_thread(RevertThreadParams { + thread_id, + before_turn_id: "turn-2".to_string(), + }) + .await + .expect("revert before second turn"); + let first_replacement_path = state_db + .get_thread(thread_id) + .await + .expect("read metadata") + .expect("thread metadata") + .rollout_path; + assert_ne!(first_replacement_path, original_path); + assert_ne!( + codex_rollout::rollout_id_from_path(first_replacement_path.as_path()), + Some(thread_id) + ); + let replacement_meta = codex_rollout::read_session_meta_line(first_replacement_path.as_path()) + .await + .expect("read replacement metadata") + .meta; + assert_eq!(replacement_meta.id, thread_id); + assert_eq!(replacement_meta.memory_mode, None); + assert_eq!(turn_ids(&store, thread_id).await, vec!["turn-1"]); + + store + .revert_thread(RevertThreadParams { + thread_id, + before_turn_id: "turn-1".to_string(), + }) + .await + .expect("revert before first turn"); + assert_eq!(turn_ids(&store, thread_id).await, Vec::::new()); + + store + .archive_thread(ArchiveThreadParams { thread_id }) + .await + .expect("archive reverted thread"); + assert!( + rollout_paths_for_thread(home.path(), thread_id) + .await + .iter() + .all(|path| path.starts_with(home.path().join("archived_sessions"))) + ); + store + .unarchive_thread(ArchiveThreadParams { thread_id }) + .await + .expect("unarchive reverted thread"); + let owned_rollout_paths = rollout_paths_for_thread(home.path(), thread_id).await; + assert_eq!(owned_rollout_paths.len(), 3); + assert!( + owned_rollout_paths + .iter() + .all(|path| path.starts_with(home.path().join("sessions"))) + ); + + store + .delete_thread(DeleteThreadParams { thread_id }) + .await + .expect("delete reverted thread"); + for rollout_path in owned_rollout_paths { + assert!(!rollout_path.exists()); + } +} + +async fn rollout_paths_for_thread( + home: &std::path::Path, + thread_id: ThreadId, +) -> Vec { + codex_rollout::RolloutReferenceIndex::scan(home) + .await + .expect("scan rollout references") + .rollouts_for_thread(thread_id) + .map(|(_, path)| path.to_path_buf()) + .collect() +} + +async fn create_paginated_thread(store: &LocalThreadStore, thread_id: ThreadId) { + store + .create_thread(CreateThreadParams { + session_id: thread_id.into(), + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), + multi_agent_version: None, + history_mode: ThreadHistoryMode::Paginated, + history_base: None, + subagent_history_start_ordinal: None, + initial_window_id: "window-1".to_string(), + metadata: ThreadPersistenceMetadata { + cwd: Some(std::env::current_dir().expect("cwd")), + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + }) + .await + .expect("create paginated thread"); +} + +async fn turn_ids(store: &LocalThreadStore, thread_id: ThreadId) -> Vec { + store + .list_turns(ListTurnsParams { + thread_id, + include_archived: false, + cursor: None, + page_size: 10, + sort_direction: SortDirection::Asc, + items_view: StoredTurnItemsView::NotLoaded, + }) + .await + .expect("list turns") + .turns + .into_iter() + .map(|turn| turn.turn_id) + .collect() +} + +fn turn_started(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: Some(10), + model_context_window: None, + collaboration_mode_kind: Default::default(), + })) +} + +fn turn_completed(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + error: None, + started_at: Some(10), + completed_at: Some(20), + duration_ms: Some(10_000), + time_to_first_token_ms: None, + })) +} + +fn compress_rollout(path: &std::path::Path) { + let contents = std::fs::read(path).expect("read rollout"); + let compressed = zstd::stream::encode_all(contents.as_slice(), 3).expect("compress rollout"); + std::fs::write(path.with_extension("jsonl.zst"), compressed).expect("write compressed rollout"); + std::fs::remove_file(path).expect("remove plain rollout"); +} diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index c743c00d3d..50c932f136 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -1,81 +1,129 @@ use super::LocalThreadStore; -use super::helpers::matching_rollout_file_name; +use super::helpers::owned_rollout_paths; +use super::helpers::restore_rollout_moves; +use super::helpers::rollout_path_is_archived; use super::helpers::scoped_rollout_path; use super::helpers::touch_modified_time; +use super::helpers::validated_rollout_file_name; use crate::ArchiveThreadParams; use crate::ReadThreadParams; use crate::StoredThread; use crate::ThreadStoreError; use crate::ThreadStoreResult; -use codex_rollout::find_archived_thread_path_by_id_str; use codex_rollout::rollout_date_parts; +use super::thread_rollout_resolver; +use super::thread_rollout_resolver::RolloutLocation; + pub(super) async fn unarchive_thread( store: &LocalThreadStore, params: ArchiveThreadParams, ) -> ThreadStoreResult { let thread_id = params.thread_id; let _lifecycle_guard = store.live_writer_locks.lock_lifecycle(thread_id).await; + // Archive, delete, and revert use the same cross-process lock while moving or selecting + // rollout files. Unarchive must participate before it moves those files back. + let _writer_lock = store.writer_lock_coordinator.acquire(thread_id)?; let state_db_ctx = store.state_db().await; - let archived_path = find_archived_thread_path_by_id_str( - store.config.codex_home.as_path(), - &thread_id.to_string(), - state_db_ctx.as_deref(), - ) - .await - .map_err(|err| ThreadStoreError::InvalidRequest { - message: format!("failed to locate archived thread id {thread_id}: {err}"), - })? - .ok_or_else(|| ThreadStoreError::InvalidRequest { - message: format!("no archived rollout found for thread id {thread_id}"), - })?; + let selected_archived_path = + thread_rollout_resolver::resolve_current_including_archived(store, thread_id) + .await? + .filter(|resolved| resolved.location == RolloutLocation::Archived) + .map(|resolved| resolved.path) + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("no archived rollout found for thread id {thread_id}"), + })?; - let canonical_archived_path = scoped_rollout_path( - store + let mut rollout_paths = owned_rollout_paths(store, thread_id).await?; + if !rollout_paths.contains(&selected_archived_path) { + rollout_paths.push(selected_archived_path.clone()); + } + let mut restored_path = None; + let mut rollout_moves = Vec::new(); + for rollout_path in rollout_paths { + if !rollout_path_is_archived(store.config.codex_home.as_path(), rollout_path.as_path()) { + continue; + } + let canonical_archived_path = scoped_rollout_path( + store + .config + .codex_home + .join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR), + rollout_path.as_path(), + "archived", + )?; + let file_name = + validated_rollout_file_name(canonical_archived_path.as_path(), rollout_path.as_path())?; + let Some((year, month, day)) = rollout_date_parts(&file_name) else { + return Err(ThreadStoreError::InvalidRequest { + message: format!( + "rollout path `{}` missing filename timestamp", + rollout_path.display() + ), + }); + }; + let dest_dir = store .config .codex_home - .join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR), - archived_path.as_path(), - "archived", - )?; - let file_name = matching_rollout_file_name( - canonical_archived_path.as_path(), - thread_id, - archived_path.as_path(), - )?; - let Some((year, month, day)) = rollout_date_parts(&file_name) else { - return Err(ThreadStoreError::InvalidRequest { - message: format!( - "rollout path `{}` missing filename timestamp", - archived_path.display() - ), - }); - }; - - let dest_dir = store - .config - .codex_home - .join(codex_rollout::SESSIONS_SUBDIR) - .join(year) - .join(month) - .join(day); - std::fs::create_dir_all(&dest_dir).map_err(|err| ThreadStoreError::Internal { - message: format!("failed to unarchive thread: {err}"), - })?; - let restored_path = dest_dir.join(&file_name); - std::fs::rename(&canonical_archived_path, &restored_path).map_err(|err| { - ThreadStoreError::Internal { + .join(codex_rollout::SESSIONS_SUBDIR) + .join(year) + .join(month) + .join(day); + std::fs::create_dir_all(&dest_dir).map_err(|err| ThreadStoreError::Internal { message: format!("failed to unarchive thread: {err}"), + })?; + let destination = dest_dir.join(&file_name); + if rollout_path == selected_archived_path { + restored_path = Some(destination.clone()); } - })?; - touch_modified_time(restored_path.as_path()).map_err(|err| ThreadStoreError::Internal { - message: format!("failed to update unarchived thread timestamp: {err}"), + rollout_moves.push((canonical_archived_path, destination)); + } + let restored_path = restored_path.ok_or_else(|| ThreadStoreError::Internal { + message: format!("failed to unarchive selected rollout for thread {thread_id}"), })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx + for (index, (source, destination)) in rollout_moves.iter().enumerate() { + if let Err(err) = std::fs::rename(source, destination) { + if let Err(restore_err) = restore_rollout_moves(&rollout_moves[..index]) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to unarchive thread: {err}; failed to restore moved rollouts: {restore_err}" + ), + }); + } + return Err(ThreadStoreError::Internal { + message: format!("failed to unarchive thread: {err}"), + }); + } + } + if let Err(err) = touch_modified_time(restored_path.as_path()) { + if let Err(restore_err) = restore_rollout_moves(&rollout_moves) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to update unarchived thread timestamp: {err}; failed to restore moved rollouts: {restore_err}" + ), + }); + } + return Err(ThreadStoreError::Internal { + message: format!("failed to update unarchived thread timestamp: {err}"), + }); + } + + if let Some(ctx) = state_db_ctx + && let Err(err) = ctx .mark_unarchived(thread_id, restored_path.as_path()) - .await; + .await + { + if let Err(restore_err) = restore_rollout_moves(&rollout_moves) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to update unarchived thread metadata: {err}; failed to restore moved rollouts: {restore_err}" + ), + }); + } + return Err(ThreadStoreError::Internal { + message: format!("failed to update unarchived thread metadata: {err}"), + }); } super::read_thread::read_thread( @@ -135,6 +183,34 @@ mod tests { ); } + #[tokio::test] + async fn unarchive_thread_rejects_a_cross_process_writer() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let owner = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let uuid = Uuid::from_u128(205); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) + .expect("archived session file"); + let writer_guard = owner + .writer_lock_coordinator + .acquire(thread_id) + .expect("acquire writer lock"); + + let error = store + .unarchive_thread(ArchiveThreadParams { thread_id }) + .await + .expect_err("active writer should block unarchive"); + + assert!(matches!(error, ThreadStoreError::Conflict { .. })); + assert!(archived_path.exists()); + drop(writer_guard); + store + .unarchive_thread(ArchiveThreadParams { thread_id }) + .await + .expect("unarchive after writer exits"); + } + #[tokio::test] async fn unarchive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 9d88724489..d0927b5132 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -25,6 +25,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::RenameThreadSectionParams; use crate::ResumeThreadParams; +use crate::RevertThreadParams; use crate::SearchThreadOccurrencesParams; use crate::SearchThreadsParams; use crate::StoredModelContext; @@ -132,6 +133,21 @@ pub trait ThreadStore: Any + Send + Sync { }) } + /// Reverts a paginated thread's durable history so it ends immediately before + /// `before_turn_id`. + /// + /// Callers must close the thread's live writer first. The logical thread id and semantic + /// metadata stay unchanged. + /// + /// Stores without paginated revert support can retain this default implementation. + fn revert_thread(&self, _params: RevertThreadParams) -> ThreadStoreFuture<'_, ()> { + Box::pin(async { + Err(ThreadStoreError::Unsupported { + operation: "revert_thread", + }) + }) + } + /// Reads a thread summary and optionally its persisted history. fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread>; diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 39a8bc3075..d129d7e65e 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -198,6 +198,15 @@ pub struct PrepareForkParams { pub boundary: ForkBoundary, } +/// Parameters for reverting a paginated thread's durable history. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RevertThreadParams { + /// Stable logical thread to revert. + pub thread_id: ThreadId, + /// First turn excluded from the retained history. + pub before_turn_id: String, +} + /// Frozen source history and model context for a reference-backed fork. #[derive(Debug)] pub struct PreparedFork {