codex/codex-rs/thread-store/src/local/archive_thread.rs

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<Vec<codex_protocol::ThreadId>> {
    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?;
    // Only inspect active files whose names belong to the threads being archived.
    let reference_index = RolloutReferenceIndex::scan_unarchived_threads(
        store.config.codex_home.as_path(),
        &thread_ids,
    )
    .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<std::path::PathBuf>,
) -> 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
                .acquire_writer_lock(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);
    }
}