use super::LocalThreadStore; 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 super::helpers::validated_rollout_file_name; use crate::ArchiveThreadsParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; use chrono::Utc; use codex_rollout::RolloutReferenceIndex; use tracing::warn; use super::thread_rollout_resolver; pub(super) async fn archive_threads( store: &LocalThreadStore, params: ArchiveThreadsParams, ) -> ThreadStoreResult> { let thread_ids = params.thread_ids; if thread_ids.is_empty() { return Ok(Vec::new()); } let mut lock_thread_ids = params.writer_lock_thread_ids; lock_thread_ids.extend(thread_ids.iter().copied()); 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); } let mut _live_writer_guards = Vec::with_capacity(lock_thread_ids.len()); for thread_id in &lock_thread_ids { _live_writer_guards.push(store.live_writer_locks.lock(*thread_id).await); if store.live_recorders.lock().await.contains_key(thread_id) { return Err(ThreadStoreError::Conflict { message: format!("thread {thread_id} already has an active writer"), }); } } 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 { 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!( "failed to archive spawned descendant thread {thread_id} while archiving {parent_thread_id}: {err}" ), } } Ok(archived_thread_ids) } async fn archive_thread_with_paths( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, mut rollout_paths: Vec, ) -> ThreadStoreResult<()> { let state_db_ctx = store.state_db().await; 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 .codex_home .join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR); std::fs::create_dir_all(&archive_folder).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()); } if !rollout_moves .iter() .any(|(source, _)| source == &canonical_rollout_path) { 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}"), })?; 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 { 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(()) } #[cfg(test)] mod tests { use std::time::Duration; use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; use codex_utils_absolute_path::test_support::PathExt; use pretty_assertions::assert_eq; use tempfile::TempDir; use uuid::Uuid; use super::*; use crate::ArchiveThreadParams; use crate::ListThreadsParams; use crate::ThreadSortKey; use crate::ThreadStore; use crate::local::LocalThreadStore; use crate::local::test_support::test_config; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with_history_mode; #[tokio::test] async fn archive_waits_for_fork_reservation_without_holding_writer_lock() { let home = TempDir::new().expect("temp dir"); let store = 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 active_path = write_session_file_with_history_mode( home.path(), "2025-01-03T12-00-00", uuid, ThreadHistoryMode::Paginated, ) .expect("session file"); let reservation = store.live_writer_locks.reserve_lifecycle(thread_id).await; let mut archive = Box::pin(store.archive_thread(ArchiveThreadParams { thread_id })); tokio::select! { biased; result = &mut archive => panic!("archive completed while the source was reserved: {result:?}"), _ = tokio::task::yield_now() => {} } let writer_guard = tokio::time::timeout( Duration::from_secs(1), store.live_writer_locks.lock(thread_id), ) .await .expect("pending archive should not hold the writer lock"); drop(writer_guard); drop(reservation); archive.await.expect("archive reserved thread"); assert!(!active_path.exists()); } #[tokio::test] async fn archive_threads_rejects_owned_descendants_before_archiving_anything() { 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); for (parent_uuid, child_uuid, history_mode) in [ ( Uuid::from_u128(203), Uuid::from_u128(204), ThreadHistoryMode::Legacy, ), ( Uuid::from_u128(206), Uuid::from_u128(207), ThreadHistoryMode::Paginated, ), ] { let parent_thread_id = ThreadId::from_string(&parent_uuid.to_string()).expect("valid parent thread id"); let parent_path = write_session_file_with_history_mode( home.path(), "2025-01-03T12-00-00", parent_uuid, history_mode, ) .expect("parent session file"); let child_thread_id = ThreadId::from_string(&child_uuid.to_string()).expect("valid child thread id"); let child_path = write_session_file_with_history_mode( home.path(), "2025-01-03T12-00-01", child_uuid, history_mode, ) .expect("child session file"); let _owner_guard = owner .writer_lock_coordinator .acquire(child_thread_id) .expect("acquire child writer lock"); let error = store .archive_threads(ArchiveThreadsParams { thread_ids: vec![parent_thread_id, child_thread_id], writer_lock_thread_ids: Vec::new(), }) .await .expect_err("owned descendant should block archive"); assert!(matches!(error, ThreadStoreError::Conflict { .. })); assert!(parent_path.exists()); assert!(child_path.exists()); } } #[tokio::test] async fn archive_thread_moves_rollout_to_archived_collection() { let home = TempDir::new().expect("temp dir"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let uuid = Uuid::from_u128(201); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); store .archive_thread(ArchiveThreadParams { thread_id }) .await .expect("archive thread"); assert!(!active_path.exists()); let archived_path = home .path() .join(ARCHIVED_SESSIONS_SUBDIR) .join(active_path.file_name().expect("file name")); assert!(archived_path.exists()); let archived = store .list_threads(ListThreadsParams { page_size: 10, cursor: None, sort_key: ThreadSortKey::CreatedAt, sort_direction: crate::SortDirection::Desc, allowed_sources: Vec::new(), model_providers: None, cwd_filters: None, section: None, project_id: None, archived: true, search_term: None, relation_filter: None, use_state_db_only: false, }) .await .expect("archived listing"); assert_eq!(archived.items.len(), 1); assert_eq!(archived.items[0].thread_id, thread_id); assert_eq!(archived.items[0].rollout_path, Some(archived_path)); assert_eq!( archived.items[0].archived_at, Some(archived.items[0].updated_at) ); } #[tokio::test] async fn archive_thread_deduplicates_rollout_paths_and_updates_sqlite_metadata() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let uuid = Uuid::from_u128(202); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); let alternate_directory = active_path .parent() .expect("session directory") .join("alternate"); std::fs::create_dir(&alternate_directory).expect("alternate session directory"); let selected_rollout_path = alternate_directory .join("..") .join(active_path.file_name().expect("file name")); let runtime = codex_state::StateRuntime::init( codex_state::SqliteConfig::new_for_testing(home.path().abs()), config.default_model_provider_id.clone(), ) .await .expect("state db should initialize"); let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); runtime .mark_backfill_complete(/*last_watermark*/ None) .await .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, selected_rollout_path, Utc::now(), SessionSource::Cli, ); builder.model_provider = Some(config.default_model_provider_id.clone()); builder.cwd = home.path().to_path_buf(); builder.cli_version = Some("test_version".to_string()); let metadata = builder.build(config.default_model_provider_id.as_str()); runtime .upsert_thread(&metadata) .await .expect("state db upsert should succeed"); store .archive_thread(ArchiveThreadParams { thread_id }) .await .expect("archive thread"); let archived_path = home .path() .join(ARCHIVED_SESSIONS_SUBDIR) .join(active_path.file_name().expect("file name")); let updated = runtime .get_thread(thread_id) .await .expect("state db read should succeed") .expect("thread metadata should exist"); assert_eq!(updated.rollout_path, archived_path); assert!(updated.archived_at.is_some()); assert_eq!(updated.recency_at, metadata.recency_at); } }