codex/codex-rs/rollout/src/metadata.rs

use crate::ARCHIVED_SESSIONS_SUBDIR;
use crate::SESSIONS_SUBDIR;
use crate::compression;
use crate::list::parse_timestamp_uuid_from_filename;
use crate::recorder::RolloutRecorder;
use crate::state_db::normalize_cwd_for_state_db;
use chrono::DateTime;
use chrono::NaiveDateTime;
use chrono::Timelike;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
use codex_state::BackfillState;
use codex_state::BackfillStats;
use codex_state::BackfillStatus;
use codex_state::DB_ERROR_METRIC;
use codex_state::DB_METRIC_BACKFILL;
use codex_state::DB_METRIC_BACKFILL_DURATION_MS;
use codex_state::ExtractionOutcome;
use codex_state::ThreadMetadataBuilder;
use codex_state::apply_rollout_item;
use std::path::Path;
use std::path::PathBuf;
use tracing::info;
use tracing::warn;

const BACKFILL_BATCH_SIZE: usize = 200;
#[cfg(not(test))]
const BACKFILL_LEASE_SECONDS: i64 = 900;
#[cfg(test)]
const BACKFILL_LEASE_SECONDS: i64 = 1;

pub(crate) fn builder_from_session_meta(
    session_meta: &SessionMetaLine,
    rollout_path: &Path,
) -> Option<ThreadMetadataBuilder> {
    let created_at = parse_timestamp_to_utc(session_meta.meta.timestamp.as_str())?;
    let mut builder = ThreadMetadataBuilder::new(
        session_meta.meta.id,
        rollout_path.to_path_buf(),
        created_at,
        session_meta.meta.source.clone(),
    );
    builder.history_mode = session_meta.meta.history_mode;
    builder.model_provider = session_meta.meta.model_provider.clone();
    builder.agent_nickname = session_meta.meta.agent_nickname.clone();
    builder.agent_role = session_meta.meta.agent_role.clone();
    builder.agent_path = session_meta.meta.agent_path.clone();
    builder.cwd = session_meta.meta.cwd.clone();
    builder.cli_version = Some(session_meta.meta.cli_version.clone());
    builder.sandbox_policy = SandboxPolicy::new_read_only_policy();
    builder.approval_mode = AskForApproval::OnRequest;
    if let Some(git) = session_meta.git.as_ref() {
        builder.git_sha = git.commit_hash.as_ref().map(|sha| sha.0.clone());
        builder.git_branch = git.branch.clone();
        builder.git_origin_url = git.repository_url.clone();
    }
    Some(builder)
}

pub fn builder_from_items(
    items: &[RolloutItem],
    rollout_path: &Path,
) -> Option<ThreadMetadataBuilder> {
    if let Some(session_meta) = items.iter().find_map(|item| match item {
        RolloutItem::SessionMeta(meta_line) => Some(meta_line),
        RolloutItem::ResponseItem(_)
        | RolloutItem::InterAgentCommunication(_)
        | RolloutItem::InterAgentCommunicationMetadata { .. }
        | RolloutItem::Compacted(_)
        | RolloutItem::TurnContext(_)
        | RolloutItem::WorldState(_)
        | RolloutItem::EventMsg(_) => None,
    }) && let Some(builder) = builder_from_session_meta(session_meta, rollout_path)
    {
        return Some(builder);
    }

    let file_name = rollout_path.file_name()?.to_str()?;
    let file_name = compression::parse_rollout_file_name(file_name)?;
    let (created_ts, uuid) = parse_timestamp_uuid_from_filename(file_name)?;
    let created_at =
        DateTime::<Utc>::from_timestamp(created_ts.unix_timestamp(), 0)?.with_nanosecond(0)?;
    let id = ThreadId::from_string(&uuid.to_string()).ok()?;
    Some(ThreadMetadataBuilder::new(
        id,
        rollout_path.to_path_buf(),
        created_at,
        SessionSource::default(),
    ))
}

pub async fn extract_metadata_from_rollout(
    rollout_path: &Path,
    default_provider: &str,
) -> anyhow::Result<ExtractionOutcome> {
    let (items, _thread_id, parse_errors) =
        RolloutRecorder::load_rollout_items(rollout_path).await?;
    if items.is_empty() {
        return Err(anyhow::anyhow!(
            "empty session file: {}",
            rollout_path.display()
        ));
    }
    let builder = builder_from_items(items.as_slice(), rollout_path).ok_or_else(|| {
        anyhow::anyhow!(
            "rollout missing metadata builder: {}",
            rollout_path.display()
        )
    })?;
    let mut metadata = builder.build(default_provider);
    for item in &items {
        apply_rollout_item(&mut metadata, item, default_provider);
    }
    if let Some(updated_at) = file_modified_time_utc(rollout_path).await {
        metadata.updated_at = updated_at;
        metadata.recency_at = updated_at;
    }
    Ok(ExtractionOutcome {
        metadata,
        memory_mode: items.iter().rev().find_map(|item| match item {
            RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(),
            RolloutItem::ResponseItem(_)
            | RolloutItem::InterAgentCommunication(_)
            | RolloutItem::InterAgentCommunicationMetadata { .. }
            | RolloutItem::Compacted(_)
            | RolloutItem::TurnContext(_)
            | RolloutItem::WorldState(_)
            | RolloutItem::EventMsg(_) => None,
        }),
        parse_errors,
    })
}

pub(crate) async fn backfill_sessions(
    runtime: &codex_state::StateRuntime,
    codex_home: &Path,
    default_provider: &str,
) {
    backfill_sessions_with_lease(
        runtime,
        codex_home,
        default_provider,
        BACKFILL_LEASE_SECONDS,
    )
    .await;
}

pub(crate) async fn backfill_sessions_with_lease(
    runtime: &codex_state::StateRuntime,
    codex_home: &Path,
    default_provider: &str,
    backfill_lease_seconds: i64,
) {
    let metric_client = codex_otel::global();
    let timer = metric_client
        .as_ref()
        .and_then(|otel| otel.start_timer(DB_METRIC_BACKFILL_DURATION_MS, &[]).ok());
    let backfill_state = match runtime.get_backfill_state().await {
        Ok(state) => state,
        Err(err) => {
            warn!(
                "failed to read backfill state at {}: {err}",
                codex_home.display()
            );
            BackfillState::default()
        }
    };
    if backfill_state.status == BackfillStatus::Complete {
        return;
    }
    let claimed = match runtime.try_claim_backfill(backfill_lease_seconds).await {
        Ok(claimed) => claimed,
        Err(err) => {
            warn!(
                "failed to claim backfill worker at {}: {err}",
                codex_home.display()
            );
            return;
        }
    };
    if !claimed {
        info!(
            "state db backfill already running at {}; skipping duplicate worker",
            codex_home.display()
        );
        return;
    }
    let mut backfill_state = match runtime.get_backfill_state().await {
        Ok(state) => state,
        Err(err) => {
            warn!(
                "failed to read claimed backfill state at {}: {err}",
                codex_home.display()
            );
            BackfillState {
                status: BackfillStatus::Running,
                ..Default::default()
            }
        }
    };
    if backfill_state.status != BackfillStatus::Running {
        if let Err(err) = runtime.mark_backfill_running().await {
            warn!(
                "failed to mark backfill running at {}: {err}",
                codex_home.display()
            );
        } else {
            backfill_state.status = BackfillStatus::Running;
        }
    }

    let sessions_root = codex_home.join(SESSIONS_SUBDIR);
    let archived_root = codex_home.join(ARCHIVED_SESSIONS_SUBDIR);
    let mut rollout_paths: Vec<BackfillRolloutPath> = Vec::new();
    for (root, archived) in [(sessions_root, false), (archived_root, true)] {
        if !tokio::fs::try_exists(&root).await.unwrap_or(false) {
            continue;
        }
        match collect_rollout_paths(&root).await {
            Ok(paths) => {
                rollout_paths.extend(paths.into_iter().map(|path| BackfillRolloutPath {
                    watermark: backfill_watermark_for_path(codex_home, &path),
                    path,
                    archived,
                }));
            }
            Err(err) => {
                warn!(
                    "failed to collect rollout paths under {}: {err}",
                    root.display()
                );
            }
        }
    }
    rollout_paths.sort_by(|a, b| a.watermark.cmp(&b.watermark));
    if let Some(last_watermark) = backfill_state.last_watermark.as_deref() {
        rollout_paths.retain(|entry| entry.watermark.as_str() > last_watermark);
    }

    let mut stats = BackfillStats {
        scanned: 0,
        upserted: 0,
        failed: 0,
    };
    let mut last_watermark = backfill_state.last_watermark.clone();
    for batch in rollout_paths.chunks(BACKFILL_BATCH_SIZE) {
        for rollout in batch {
            stats.scanned = stats.scanned.saturating_add(1);
            match extract_metadata_from_rollout(&rollout.path, default_provider).await {
                Ok(outcome) => {
                    if outcome.parse_errors > 0
                        && let Some(ref metric_client) = metric_client
                    {
                        let _ = metric_client.counter(
                            DB_ERROR_METRIC,
                            outcome.parse_errors as i64,
                            &[("stage", "backfill_sessions")],
                        );
                    }
                    let mut metadata = outcome.metadata;
                    metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd);
                    let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string());
                    if let Ok(Some(existing_metadata)) = runtime.get_thread(metadata.id).await {
                        metadata.prefer_existing_git_info(&existing_metadata);
                        metadata.prefer_existing_explicit_title(&existing_metadata);
                    }
                    if rollout.archived && metadata.archived_at.is_none() {
                        let fallback_archived_at = metadata.updated_at;
                        metadata.archived_at = file_modified_time_utc(&rollout.path)
                            .await
                            .or(Some(fallback_archived_at));
                    }
                    if let Err(err) = runtime.upsert_thread(&metadata).await {
                        stats.failed = stats.failed.saturating_add(1);
                        warn!("failed to upsert rollout {}: {err}", rollout.path.display());
                    } else {
                        if let Err(err) = runtime
                            .set_thread_memory_mode(metadata.id, memory_mode.as_str())
                            .await
                        {
                            stats.failed = stats.failed.saturating_add(1);
                            warn!(
                                "failed to restore memory mode for {}: {err}",
                                rollout.path.display()
                            );
                            continue;
                        }
                        stats.upserted = stats.upserted.saturating_add(1);
                    }
                }
                Err(err) => {
                    stats.failed = stats.failed.saturating_add(1);
                    warn!(
                        "failed to extract rollout {}: {err}",
                        rollout.path.display()
                    );
                }
            }
        }

        if let Some(last_entry) = batch.last() {
            if let Err(err) = runtime
                .checkpoint_backfill(last_entry.watermark.as_str())
                .await
            {
                warn!(
                    "failed to checkpoint backfill at {}: {err}",
                    codex_home.display()
                );
            } else {
                last_watermark = Some(last_entry.watermark.clone());
            }
        }
    }
    if let Err(err) = runtime
        .mark_backfill_complete(last_watermark.as_deref())
        .await
    {
        warn!(
            "failed to mark backfill complete at {}: {err}",
            codex_home.display()
        );
    }

    info!(
        "state db backfill scanned={}, upserted={}, failed={}",
        stats.scanned, stats.upserted, stats.failed
    );
    if let Some(metric_client) = metric_client {
        let _ = metric_client.counter(
            DB_METRIC_BACKFILL,
            stats.upserted as i64,
            &[("status", "upserted")],
        );
        let _ = metric_client.counter(
            DB_METRIC_BACKFILL,
            stats.failed as i64,
            &[("status", "failed")],
        );
    }
    if let Some(timer) = timer.as_ref() {
        let status = if stats.failed == 0 {
            "success"
        } else if stats.upserted == 0 {
            "failed"
        } else {
            "partial_failure"
        };
        let _ = timer.record(&[("status", status)]);
    }
}

#[derive(Debug, Clone)]
struct BackfillRolloutPath {
    watermark: String,
    path: PathBuf,
    archived: bool,
}

fn backfill_watermark_for_path(codex_home: &Path, path: &Path) -> String {
    path.strip_prefix(codex_home)
        .unwrap_or(path)
        .to_string_lossy()
        .replace('\\', "/")
}

async fn file_modified_time_utc(path: &Path) -> Option<DateTime<Utc>> {
    let modified = compression::file_modified_time(path).await.ok()??;
    DateTime::<Utc>::from_timestamp(modified.unix_timestamp(), modified.nanosecond())
}

fn parse_timestamp_to_utc(ts: &str) -> Option<DateTime<Utc>> {
    const FILENAME_TS_FORMAT: &str = "%Y-%m-%dT%H-%M-%S";
    if let Ok(naive) = NaiveDateTime::parse_from_str(ts, FILENAME_TS_FORMAT) {
        let dt = DateTime::<Utc>::from_naive_utc_and_offset(naive, Utc);
        return dt.with_nanosecond(0);
    }
    if let Ok(dt) = DateTime::parse_from_rfc3339(ts) {
        return Some(dt.with_timezone(&Utc));
    }
    None
}

async fn collect_rollout_paths(root: &Path) -> std::io::Result<Vec<PathBuf>> {
    let mut stack = vec![root.to_path_buf()];
    let mut paths = Vec::new();
    while let Some(dir) = stack.pop() {
        let mut read_dir = match tokio::fs::read_dir(&dir).await {
            Ok(read_dir) => read_dir,
            Err(err) => {
                warn!("failed to read directory {}: {err}", dir.display());
                continue;
            }
        };
        loop {
            let next_entry = match read_dir.next_entry().await {
                Ok(next_entry) => next_entry,
                Err(err) => {
                    warn!(
                        "failed to read directory entry under {}: {err}",
                        dir.display()
                    );
                    continue;
                }
            };
            let Some(entry) = next_entry else {
                break;
            };
            let path = entry.path();
            let file_type = match entry.file_type().await {
                Ok(file_type) => file_type,
                Err(err) => {
                    warn!("failed to read file type for {}: {err}", path.display());
                    continue;
                }
            };
            if file_type.is_dir() {
                stack.push(path);
                continue;
            }
            if !file_type.is_file() {
                continue;
            }
            if let Some(rollout_file) = compression::RolloutFile::from_path(path) {
                paths.push(rollout_file.into_path());
            }
        }
    }
    Ok(paths)
}

#[cfg(test)]
#[path = "metadata_tests.rs"]
mod tests;