diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 67e9995294..751eae93f7 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2406,6 +2406,7 @@ dependencies = [ "codex-state", "codex-stdio-to-uds", "codex-terminal-detection", + "codex-thread-store", "codex-tui", "codex-utils-absolute-path", "codex-utils-cargo-bin", diff --git a/codex-rs/app-server/tests/suite/v2/rollout_migration.rs b/codex-rs/app-server/tests/suite/v2/rollout_migration.rs index 59d182ffde..431534a175 100644 --- a/codex-rs/app-server/tests/suite/v2/rollout_migration.rs +++ b/codex-rs/app-server/tests/suite/v2/rollout_migration.rs @@ -83,7 +83,7 @@ async fn migrated_legacy_thread_cold_resume_preserves_model_context() -> Result< let report = store .migrate_rollouts(RolloutMigrationOptions { mode: RolloutMigrationMode::Apply, - max_mib_per_second: 1024, + max_mib_per_second: Some(1024), ..RolloutMigrationOptions::default() }) .await?; diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 5adc7a6cfe..0fe910d272 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -67,6 +67,7 @@ codex-skills-extension = { workspace = true } codex-state = { workspace = true } codex-stdio-to-uds = { workspace = true } codex-terminal-detection = { workspace = true } +codex-thread-store = { workspace = true } codex-tui = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-path = { workspace = true } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 4c88718a20..b22ff80b1e 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -51,6 +51,7 @@ mod doctor; mod exec_server_telemetry; mod marketplace_cmd; mod mcp_cmd; +mod migrate_rollouts; mod plugin_cmd; mod remote_control_cmd; #[cfg(target_os = "windows")] @@ -190,6 +191,9 @@ enum Subcommand { /// Permanently delete a saved session by id or session name. Delete(DeleteCommand), + /// Inspect or migrate legacy local sessions to paginated thread history. + MigrateRollouts(migrate_rollouts::MigrateRolloutsCommand), + /// Unarchive a saved session by id or session name. Unarchive(SessionArchiveCommand), @@ -1339,6 +1343,14 @@ async fn cli_main( .await?; println!("{output}"); } + Some(Subcommand::MigrateRollouts(command)) => { + reject_remote_mode_for_subcommand( + root_remote.as_deref(), + root_remote_auth_token_env.as_deref(), + "migrate-rollouts", + )?; + migrate_rollouts::run(command, root_config_overrides).await?; + } Some(Subcommand::Unarchive(cmd)) => { let output = run_session_archive_cli_command( codex_tui::SessionArchiveAction::Unarchive, @@ -2245,6 +2257,7 @@ fn unsupported_subcommand_name_for_strict_config( Some(Subcommand::RemoteControl(remote_control)) => Some(remote_control.subcommand_name()), Some(Subcommand::Mcp(_)) => Some("mcp"), Some(Subcommand::Plugin(_)) => Some("plugin"), + Some(Subcommand::MigrateRollouts(_)) => Some("migrate-rollouts"), #[cfg(any(target_os = "macos", target_os = "windows"))] Some(Subcommand::App(_)) => Some("app"), Some(Subcommand::Login(_)) => Some("login"), diff --git a/codex-rs/cli/src/migrate_rollouts.rs b/codex-rs/cli/src/migrate_rollouts.rs new file mode 100644 index 0000000000..fec035e55c --- /dev/null +++ b/codex-rs/cli/src/migrate_rollouts.rs @@ -0,0 +1,369 @@ +use std::io; +use std::io::IsTerminal; +use std::io::Write; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Context; +use clap::Parser; +use codex_core::config::ConfigBuilder; +use codex_protocol::ThreadId; +use codex_thread_store::LocalThreadStore; +use codex_thread_store::LocalThreadStoreConfig; +use codex_thread_store::RolloutMigrationMode; +use codex_thread_store::RolloutMigrationOptions; +use codex_thread_store::RolloutMigrationProgress; +use codex_thread_store::RolloutMigrationReport; +use codex_thread_store::RolloutMigrationStatus; +use codex_utils_cli::CliConfigOverrides; + +#[derive(Debug, Parser)] +pub(crate) struct MigrateRolloutsCommand { + /// Publish the migration. Without this flag the command only reports eligible sessions. + #[arg(long)] + apply: bool, + + /// Restrict inspection or migration to one or more thread IDs. + #[arg(long, value_name = "THREAD_ID", value_parser = ThreadId::from_string)] + thread: Vec, + + /// Limit aggregate rollout read and write throughput, in MiB per second. + #[arg( + long, + value_name = "MIB", + value_parser = clap::value_parser!(u64).range(1..) + )] + max_mib_per_second: Option, + + /// Emit the complete per-thread report as JSON. + #[arg(long)] + json: bool, + + /// Print one line for every inspected rollout. + #[arg(long)] + verbose: bool, +} + +pub(crate) async fn run( + command: MigrateRolloutsCommand, + config_overrides: CliConfigOverrides, +) -> anyhow::Result<()> { + let overrides = config_overrides + .parse_overrides() + .map_err(anyhow::Error::msg)?; + let config = ConfigBuilder::default() + .cli_overrides(overrides) + .build() + .await?; + let otel = codex_core::otel_init::build_provider( + &config, + env!("CARGO_PKG_VERSION"), + /*service_name_override*/ None, + /*default_analytics_enabled*/ true, + ) + .unwrap_or_else(|error| { + eprintln!("Could not create otel exporter: {error}"); + None + }); + codex_core::otel_init::record_process_start(otel.as_ref(), "codex_migrate_rollouts"); + let mode = if command.apply { + RolloutMigrationMode::Apply + } else { + RolloutMigrationMode::DryRun + }; + let json = command.json; + let verbose = command.verbose; + let state_db = if mode == RolloutMigrationMode::Apply { + Some( + codex_rollout::state_db::try_init(&config) + .await + .context("failed to initialize local thread metadata")?, + ) + } else { + None + }; + let store = LocalThreadStore::new(LocalThreadStoreConfig::from_config(&config), state_db); + let mut progress = MigrationProgress::new(mode, json); + progress.begin(); + let result = store + .migrate_rollouts_with_progress( + RolloutMigrationOptions { + mode, + thread_ids: command.thread, + max_mib_per_second: command.max_mib_per_second, + }, + |update| progress.update(update), + ) + .await; + progress.finish(); + let report = result?; + + if json { + println!("{}", serde_json::to_string_pretty(&report)?); + } else { + print_human_report(&report, mode, verbose, progress.elapsed()); + } + + if report + .outcomes + .iter() + .any(|outcome| outcome.status == RolloutMigrationStatus::Failed) + { + anyhow::bail!("one or more rollout migrations failed"); + } + Ok(()) +} + +const TTY_PROGRESS_INTERVAL: Duration = Duration::from_millis(250); +const NON_TTY_PROGRESS_INTERVAL: usize = 1_000; +const MAX_EXCEPTION_DETAILS: usize = 20; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ProgressOutput { + Quiet, + Tty, + Plain, +} + +struct MigrationProgress { + mode: RolloutMigrationMode, + output: ProgressOutput, + started_at: Instant, + last_rendered_at: Instant, + last_plain_processed: usize, + counts: MigrationCounts, + wrote_tty_line: bool, +} + +impl MigrationProgress { + fn new(mode: RolloutMigrationMode, json: bool) -> Self { + let now = Instant::now(); + let output = if json { + ProgressOutput::Quiet + } else if io::stderr().is_terminal() + && std::env::var("TERM").ok().as_deref() != Some("dumb") + { + ProgressOutput::Tty + } else { + ProgressOutput::Plain + }; + Self { + mode, + output, + started_at: now, + last_rendered_at: now, + last_plain_processed: 0, + counts: MigrationCounts::default(), + wrote_tty_line: false, + } + } + + fn begin(&self) { + if self.output != ProgressOutput::Quiet { + eprintln!("Scanning local rollouts..."); + } + } + + fn update(&mut self, update: RolloutMigrationProgress) { + if let Some(status) = update.outcome_status { + self.counts.observe(status); + } + match self.output { + ProgressOutput::Quiet => {} + ProgressOutput::Tty + if update.processed_paths == update.total_paths + || self.last_rendered_at.elapsed() >= TTY_PROGRESS_INTERVAL => + { + let line = self.line(update); + let mut stderr = io::stderr().lock(); + let _ = write!(stderr, "\r\x1b[2K{line}"); + let _ = stderr.flush(); + self.last_rendered_at = Instant::now(); + self.wrote_tty_line = true; + } + ProgressOutput::Plain + if update.processed_paths == update.total_paths + || update + .processed_paths + .saturating_sub(self.last_plain_processed) + >= NON_TTY_PROGRESS_INTERVAL => + { + eprintln!("{}", self.line(update)); + self.last_plain_processed = update.processed_paths; + } + ProgressOutput::Tty | ProgressOutput::Plain => {} + } + } + + fn finish(&mut self) { + if self.output != ProgressOutput::Tty || !self.wrote_tty_line { + return; + } + let mut stderr = io::stderr().lock(); + let _ = write!(stderr, "\r\x1b[2K"); + let _ = stderr.flush(); + self.wrote_tty_line = false; + } + + fn elapsed(&self) -> Duration { + self.started_at.elapsed() + } + + fn line(&self, update: RolloutMigrationProgress) -> String { + let percent = update + .processed_paths + .saturating_mul(100) + .checked_div(update.total_paths) + .unwrap_or(100); + let action = match self.mode { + RolloutMigrationMode::DryRun => "Checking", + RolloutMigrationMode::Apply => "Migrating", + }; + let status_counts = match self.mode { + RolloutMigrationMode::DryRun => format!( + "{} eligible • {} already paginated", + self.counts.eligible, self.counts.already_paginated + ), + RolloutMigrationMode::Apply => format!( + "{} migrated • {} already paginated", + self.counts.migrated, self.counts.already_paginated + ), + }; + format!( + "{action} rollouts {}/{} ({percent}%) • {status_counts} • {} skipped • {} failed • {}", + update.processed_paths, + update.total_paths, + self.counts.skipped(), + self.counts.failed, + format_elapsed(self.elapsed()), + ) + } +} + +#[derive(Default)] +struct MigrationCounts { + eligible: usize, + migrated: usize, + already_paginated: usize, + skipped_empty: usize, + skipped_busy: usize, + failed: usize, +} + +impl MigrationCounts { + fn observe(&mut self, status: RolloutMigrationStatus) { + match status { + RolloutMigrationStatus::Eligible => self.eligible += 1, + RolloutMigrationStatus::Migrated => self.migrated += 1, + RolloutMigrationStatus::AlreadyPaginated => self.already_paginated += 1, + RolloutMigrationStatus::SkippedEmpty => self.skipped_empty += 1, + RolloutMigrationStatus::SkippedBusy => self.skipped_busy += 1, + RolloutMigrationStatus::Failed => self.failed += 1, + } + } + + fn skipped(&self) -> usize { + self.skipped_empty + self.skipped_busy + } +} + +fn print_human_report( + report: &RolloutMigrationReport, + mode: RolloutMigrationMode, + verbose: bool, + elapsed: Duration, +) { + let mut counts = MigrationCounts::default(); + for outcome in &report.outcomes { + counts.observe(outcome.status); + } + let completion = match mode { + RolloutMigrationMode::DryRun => "Scan complete", + RolloutMigrationMode::Apply => "Migration complete", + }; + println!("{completion} in {}.", format_elapsed(elapsed)); + match mode { + RolloutMigrationMode::DryRun => println!( + "Scanned {} rollout(s): {} eligible, {} already paginated, {} skipped ({} empty, {} busy), {} failed.", + report.outcomes.len(), + counts.eligible, + counts.already_paginated, + counts.skipped(), + counts.skipped_empty, + counts.skipped_busy, + counts.failed, + ), + RolloutMigrationMode::Apply => println!( + "Scanned {} rollout(s): {} migrated, {} already paginated, {} skipped ({} empty, {} busy), {} failed.", + report.outcomes.len(), + counts.migrated, + counts.already_paginated, + counts.skipped(), + counts.skipped_empty, + counts.skipped_busy, + counts.failed, + ), + } + if mode == RolloutMigrationMode::DryRun && counts.eligible > 0 { + println!("Run `codex migrate-rollouts --apply` to migrate eligible sessions."); + } + + if verbose { + for outcome in &report.outcomes { + print_outcome(outcome); + } + return; + } + + let exceptions = report.outcomes.iter().filter(|outcome| { + matches!( + outcome.status, + RolloutMigrationStatus::SkippedBusy | RolloutMigrationStatus::Failed + ) + }); + let exception_count = exceptions.clone().count(); + if exception_count == 0 { + return; + } + println!(); + for outcome in exceptions.take(MAX_EXCEPTION_DETAILS) { + print_outcome(outcome); + } + if exception_count > MAX_EXCEPTION_DETAILS { + println!( + "... and {} more; rerun with --json for the complete report.", + exception_count - MAX_EXCEPTION_DETAILS + ); + } +} + +fn print_outcome(outcome: &codex_thread_store::RolloutMigrationOutcome) { + let status = match outcome.status { + RolloutMigrationStatus::Eligible => "eligible", + RolloutMigrationStatus::Migrated => "migrated", + RolloutMigrationStatus::AlreadyPaginated => "already paginated", + RolloutMigrationStatus::SkippedEmpty => "skipped empty", + RolloutMigrationStatus::SkippedBusy => "skipped busy", + RolloutMigrationStatus::Failed => "failed", + }; + let thread_id = outcome + .thread_id + .map_or_else(|| "unknown".to_string(), |thread_id| thread_id.to_string()); + match &outcome.message { + Some(message) => println!("{status}\t{thread_id}\t{message}"), + None => println!("{status}\t{thread_id}"), + } +} + +fn format_elapsed(elapsed: Duration) -> String { + let seconds = elapsed.as_secs(); + let minutes = seconds / 60; + if minutes == 0 { + return format!("{seconds}s"); + } + let hours = minutes / 60; + if hours == 0 { + return format!("{minutes}m{:02}s", seconds % 60); + } + format!("{hours}h{:02}m{:02}s", minutes % 60, seconds % 60) +} diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 8c0a10169a..8177d0f677 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -456,6 +456,9 @@ "auth_elicitation": { "type": "boolean" }, + "background_paginated_rollout_migration": { + "type": "boolean" + }, "browser_use": { "type": "boolean" }, @@ -5074,6 +5077,9 @@ "auth_elicitation": { "type": "boolean" }, + "background_paginated_rollout_migration": { + "type": "boolean" + }, "browser_use": { "type": "boolean" }, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 95795c28d6..8053df8a34 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -349,16 +349,32 @@ pub fn thread_store_from_config( ) -> Arc { match &config.experimental_thread_store { ThreadStoreConfig::Local => { - if config + let compression_enabled = config .features - .enabled(Feature::LocalThreadStoreCompression) - { - codex_rollout::spawn_rollout_compression_worker(config.codex_home.to_path_buf()); - } - Arc::new(LocalThreadStore::new( + .enabled(Feature::LocalThreadStoreCompression); + let background_migration_enabled = config + .features + .enabled(Feature::BackgroundPaginatedRolloutMigration); + let has_state_db = state_db.is_some(); + let store = Arc::new(LocalThreadStore::new( LocalThreadStoreConfig::from_config(config), state_db, - )) + )); + if has_state_db && background_migration_enabled { + let startup_store = Arc::clone(&store); + let codex_home = config.codex_home.to_path_buf(); + tokio::spawn(async move { + if let Err(err) = startup_store.migrate_rollouts_on_startup().await { + warn!("failed to migrate legacy rollouts on startup: {err}"); + } + if compression_enabled { + codex_rollout::spawn_rollout_compression_worker(codex_home); + } + }); + } else if compression_enabled { + codex_rollout::spawn_rollout_compression_worker(config.codex_home.to_path_buf()); + } + store } ThreadStoreConfig::InMemory { id } => InMemoryThreadStore::for_id(id), } diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 51f77749cd..06ce37e1e6 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -147,6 +147,8 @@ pub enum Feature { ExternalAgentMemoryImport, /// Compress cold local thread-store rollout files. LocalThreadStoreCompression, + /// Migrate legacy local rollout files to paginated history in the background. + BackgroundPaginatedRolloutMigration, /// Enable the Chronicle sidecar for passive screen-context memories. Chronicle, /// Compress request bodies (zstd) when sending streaming requests to codex-backend. @@ -1009,6 +1011,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::BackgroundPaginatedRolloutMigration, + key: "background_paginated_rollout_migration", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::Chronicle, key: "chronicle", diff --git a/codex-rs/state/migrations/0047_rollout_migration_state.sql b/codex-rs/state/migrations/0047_rollout_migration_state.sql new file mode 100644 index 0000000000..db92256c0e --- /dev/null +++ b/codex-rs/state/migrations/0047_rollout_migration_state.sql @@ -0,0 +1,16 @@ +CREATE TABLE rollout_migration_state ( + migration_id TEXT PRIMARY KEY, + last_checked_thread_created_at INTEGER, + last_checked_thread_id TEXT, + updated_at INTEGER NOT NULL +); + +CREATE TABLE rollout_migration_skipped_rollouts ( + migration_id TEXT NOT NULL, + rollout_path TEXT NOT NULL, + rollout_size_bytes INTEGER NOT NULL, + rollout_modified_at_ns INTEGER NOT NULL, + skip_reason TEXT NOT NULL, + skipped_at INTEGER NOT NULL, + PRIMARY KEY (migration_id, rollout_path) +); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index c852fa2355..43731d39e4 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -24,6 +24,9 @@ pub use model::LogQuery; pub use model::LogRow; pub use model::Phase2JobClaimOutcome; pub use model::QueuedUserSubmissionRecord; +pub use model::RolloutMigrationCursor; +pub use model::RolloutMigrationSkippedRollout; +pub use model::RolloutMigrationState; /// Preferred entrypoint: owns configuration and metrics. pub use runtime::StateRuntime; pub use sqlite::SqliteConfig; diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 76ad9db254..0f4c4cf6d7 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -3,6 +3,7 @@ mod graph; mod log; mod memories; mod queued_item; +mod rollout_migration_state; mod thread_goal; mod thread_metadata; @@ -18,6 +19,9 @@ pub use memories::Stage1JobClaimOutcome; pub use memories::Stage1Output; pub use memories::Stage1StartupClaimParams; pub use queued_item::QueuedUserSubmissionRecord; +pub use rollout_migration_state::RolloutMigrationCursor; +pub use rollout_migration_state::RolloutMigrationSkippedRollout; +pub use rollout_migration_state::RolloutMigrationState; pub use thread_goal::ThreadGoal; pub use thread_goal::ThreadGoalStatus; pub use thread_metadata::Anchor; diff --git a/codex-rs/state/src/model/rollout_migration_state.rs b/codex-rs/state/src/model/rollout_migration_state.rs new file mode 100644 index 0000000000..37400a22b9 --- /dev/null +++ b/codex-rs/state/src/model/rollout_migration_state.rs @@ -0,0 +1,67 @@ +//! SQLite-backed bookkeeping types shared by rollout migrations. +//! +//! A migration cursor says “everything through this creation-ordered thread was checked.” A +//! skipped rollout stores enough file fingerprint information to tell whether a previously +//! unmigratable rollout is still unchanged. +//! +//! These types are intentionally generic so future rollout migrations can reuse the same state +//! shape without depending on legacy -> paginated policy. + +use anyhow::Result; +use sqlx::Row; +use sqlx::sqlite::SqliteRow; + +/// Creation-ordered thread frontier checked by one rollout migration. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct RolloutMigrationCursor { + pub thread_created_at: i64, + pub thread_id: String, +} + +/// Persisted progress for one rollout migration. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RolloutMigrationState { + pub last_checked_thread: Option, +} + +impl RolloutMigrationState { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + let thread_created_at = row.try_get("last_checked_thread_created_at")?; + let thread_id = row.try_get("last_checked_thread_id")?; + let last_checked_thread = match (thread_created_at, thread_id) { + (Some(thread_created_at), Some(thread_id)) => Some(RolloutMigrationCursor { + thread_created_at, + thread_id, + }), + (None, None) => None, + _ => { + return Err(anyhow::anyhow!( + "rollout migration state has incomplete last checked thread" + )); + } + }; + Ok(Self { + last_checked_thread, + }) + } +} + +/// An unchanged rollout that one migration can safely skip. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RolloutMigrationSkippedRollout { + pub rollout_path: String, + pub rollout_size_bytes: i64, + pub rollout_modified_at_ns: i64, + pub skip_reason: String, +} + +impl RolloutMigrationSkippedRollout { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + Ok(Self { + rollout_path: row.try_get("rollout_path")?, + rollout_size_bytes: row.try_get("rollout_size_bytes")?, + rollout_modified_at_ns: row.try_get("rollout_modified_at_ns")?, + skip_reason: row.try_get("skip_reason")?, + }) + } +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 29528f3ce9..dd9c84a822 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -47,6 +47,7 @@ mod memories; mod queued_items; mod recovery; mod remote_control; +mod rollout_migration; #[cfg(test)] pub(crate) mod test_support; mod thread_section_order; diff --git a/codex-rs/state/src/runtime/rollout_migration.rs b/codex-rs/state/src/runtime/rollout_migration.rs new file mode 100644 index 0000000000..bc1c55720e --- /dev/null +++ b/codex-rs/state/src/runtime/rollout_migration.rs @@ -0,0 +1,149 @@ +//! Persists generic rollout-migration progress in the state database. +//! +//! It stores a monotonic creation-ordered cursor for each migration and records fingerprinted +//! rollouts that the migration could not process. Thread-store uses this state to avoid rescanning +//! old rollouts on every startup while still retrying skipped files that later change. + +use super::*; + +impl StateRuntime { + pub async fn get_rollout_migration_state( + &self, + migration_id: &str, + ) -> anyhow::Result> { + let row = sqlx::query( + r#" +SELECT last_checked_thread_created_at, last_checked_thread_id +FROM rollout_migration_state +WHERE migration_id = ? + "#, + ) + .bind(migration_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.as_ref() + .map(crate::RolloutMigrationState::try_from_row) + .transpose() + } + + /// Advance one migration's checked frontier without letting concurrent startup checks move + /// it backward. + pub async fn advance_rollout_migration_state( + &self, + migration_id: &str, + last_checked_thread: Option<&crate::RolloutMigrationCursor>, + ) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + let (thread_created_at, thread_id) = last_checked_thread.map_or((None, None), |cursor| { + ( + Some(cursor.thread_created_at), + Some(cursor.thread_id.as_str()), + ) + }); + sqlx::query( + r#" +INSERT INTO rollout_migration_state ( + migration_id, + last_checked_thread_created_at, + last_checked_thread_id, + updated_at +) +VALUES (?, ?, ?, ?) +ON CONFLICT(migration_id) DO UPDATE SET + last_checked_thread_created_at = excluded.last_checked_thread_created_at, + last_checked_thread_id = excluded.last_checked_thread_id, + updated_at = excluded.updated_at +WHERE excluded.last_checked_thread_created_at IS NOT NULL + AND ( + rollout_migration_state.last_checked_thread_created_at IS NULL + OR excluded.last_checked_thread_created_at + > rollout_migration_state.last_checked_thread_created_at + OR ( + excluded.last_checked_thread_created_at + = rollout_migration_state.last_checked_thread_created_at + AND excluded.last_checked_thread_id > rollout_migration_state.last_checked_thread_id + ) + ) + "#, + ) + .bind(migration_id) + .bind(thread_created_at) + .bind(thread_id) + .bind(now) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn list_rollout_migration_skipped_rollouts( + &self, + migration_id: &str, + ) -> anyhow::Result> { + let rows = sqlx::query( + r#" +SELECT rollout_path, rollout_size_bytes, rollout_modified_at_ns, skip_reason +FROM rollout_migration_skipped_rollouts +WHERE migration_id = ? + "#, + ) + .bind(migration_id) + .fetch_all(self.pool.as_ref()) + .await?; + rows.iter() + .map(crate::RolloutMigrationSkippedRollout::try_from_row) + .collect() + } + + pub async fn record_rollout_migration_skip( + &self, + migration_id: &str, + skipped_rollout: &crate::RolloutMigrationSkippedRollout, + ) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + sqlx::query( + r#" +INSERT INTO rollout_migration_skipped_rollouts ( + migration_id, + rollout_path, + rollout_size_bytes, + rollout_modified_at_ns, + skip_reason, + skipped_at +) +VALUES (?, ?, ?, ?, ?, ?) +ON CONFLICT(migration_id, rollout_path) DO UPDATE SET + rollout_size_bytes = excluded.rollout_size_bytes, + rollout_modified_at_ns = excluded.rollout_modified_at_ns, + skip_reason = excluded.skip_reason, + skipped_at = excluded.skipped_at + "#, + ) + .bind(migration_id) + .bind(skipped_rollout.rollout_path.as_str()) + .bind(skipped_rollout.rollout_size_bytes) + .bind(skipped_rollout.rollout_modified_at_ns) + .bind(skipped_rollout.skip_reason.as_str()) + .bind(now) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn remove_rollout_migration_skip( + &self, + migration_id: &str, + rollout_path: &str, + ) -> anyhow::Result<()> { + sqlx::query( + r#" +DELETE FROM rollout_migration_skipped_rollouts +WHERE migration_id = ? AND rollout_path = ? + "#, + ) + .bind(migration_id) + .bind(rollout_path) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } +} diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 3d5eecb62f..812a1b1bdf 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -27,6 +27,7 @@ pub use local::LocalThreadStoreConfig; pub use local::RolloutMigrationMode; pub use local::RolloutMigrationOptions; pub use local::RolloutMigrationOutcome; +pub use local::RolloutMigrationProgress; pub use local::RolloutMigrationReport; pub use local::RolloutMigrationStatus; pub use queue_store::LocalQueueStore; diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index cb8e8d9699..9213acd5be 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -82,6 +82,7 @@ use crate::local::writer_lock::WriterLockGuard; pub use rollout_migration::RolloutMigrationMode; pub use rollout_migration::RolloutMigrationOptions; pub use rollout_migration::RolloutMigrationOutcome; +pub use rollout_migration::RolloutMigrationProgress; pub use rollout_migration::RolloutMigrationReport; pub use rollout_migration::RolloutMigrationStatus; diff --git a/codex-rs/thread-store/src/local/rollout_migration.rs b/codex-rs/thread-store/src/local/rollout_migration.rs index facfa952d5..201c184607 100644 --- a/codex-rs/thread-store/src/local/rollout_migration.rs +++ b/codex-rs/thread-store/src/local/rollout_migration.rs @@ -45,7 +45,9 @@ mod publish; mod rollback; mod rollback_plan; mod rollback_replay; +mod startup; mod subagent; +mod telemetry; use canonicalizer::LegacyRolloutCanonicalizer; use publish::compress_rollout_to_path; @@ -62,6 +64,8 @@ use publish::sync_parent_directory; use publish::write_migration_journal; use rollback_plan::RollbackPlan; use rollback_plan::RollbackPlanner; +use telemetry::RolloutMigrationTelemetry; +use telemetry::RolloutMigrationTrigger; const PROJECTION_BATCH_BYTES: u64 = 256 * 1024; const MAX_ROLLOUT_LINE_BYTES: usize = 16 * 1024 * 1024; @@ -98,7 +102,9 @@ pub enum RolloutMigrationMode { pub struct RolloutMigrationOptions { pub mode: RolloutMigrationMode, pub thread_ids: Vec, - pub max_mib_per_second: u64, + /// Optional aggregate rollout I/O limit. Without one, migration runs as fast as local I/O + /// allows while still yielding between projection-sized batches. + pub max_mib_per_second: Option, } impl Default for RolloutMigrationOptions { @@ -106,7 +112,7 @@ impl Default for RolloutMigrationOptions { Self { mode: RolloutMigrationMode::DryRun, thread_ids: Vec::new(), - max_mib_per_second: 8, + max_mib_per_second: None, } } } @@ -139,10 +145,19 @@ pub struct RolloutMigrationReport { pub outcomes: Vec, } +/// Incremental progress emitted while rollout migration scans discovered paths. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RolloutMigrationProgress { + pub processed_paths: usize, + pub total_paths: usize, + /// None means the scanned path was excluded by the requested thread filter. + pub outcome_status: Option, +} + struct RolloutMigrationRateLimiter { started_at: Instant, bytes_processed: u64, - bytes_per_second: u64, + bytes_per_second: Option, bytes_since_yield: u64, } @@ -158,13 +173,17 @@ enum RolloutMigrationKind { } impl RolloutMigrationRateLimiter { - fn new(max_mib_per_second: u64) -> ThreadStoreResult { + fn new(max_mib_per_second: Option) -> ThreadStoreResult { let bytes_per_second = max_mib_per_second - .checked_mul(1024 * 1024) - .filter(|rate| *rate > 0) - .ok_or_else(|| ThreadStoreError::InvalidRequest { - message: "--max-mib-per-second must be a positive supported integer".to_string(), - })?; + .map(|rate| { + rate.checked_mul(1024 * 1024) + .filter(|rate| *rate > 0) + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: "--max-mib-per-second must be a positive supported integer" + .to_string(), + }) + }) + .transpose()?; Ok(Self { started_at: Instant::now(), bytes_processed: 0, @@ -180,8 +199,12 @@ impl RolloutMigrationRateLimiter { return; } self.bytes_since_yield = 0; + let Some(bytes_per_second) = self.bytes_per_second else { + tokio::task::yield_now().await; + return; + }; let expected = - Duration::from_secs_f64(self.bytes_processed as f64 / self.bytes_per_second as f64); + Duration::from_secs_f64(self.bytes_processed as f64 / bytes_per_second as f64); if let Some(remaining) = expected.checked_sub(self.started_at.elapsed()) { tokio::time::sleep(remaining).await; } else { @@ -191,10 +214,57 @@ impl RolloutMigrationRateLimiter { } impl LocalThreadStore { + /// Check whether startup needs to migrate legacy rollouts, then migrate in the background + /// when it does. + pub async fn migrate_rollouts_on_startup(&self) -> ThreadStoreResult<()> { + startup::migrate_rollouts_on_startup(self).await + } + /// Inspect or migrate eligible legacy rollout files beneath active and archived sessions. pub async fn migrate_rollouts( &self, options: RolloutMigrationOptions, + ) -> ThreadStoreResult { + self.migrate_rollouts_with_progress_for_trigger( + options, + |_| {}, + RolloutMigrationTrigger::Manual, + ) + .await + } + + /// Inspect or migrate rollouts while reporting each discovered path after it is processed. + pub async fn migrate_rollouts_with_progress( + &self, + options: RolloutMigrationOptions, + on_progress: impl FnMut(RolloutMigrationProgress), + ) -> ThreadStoreResult { + self.migrate_rollouts_with_progress_for_trigger( + options, + on_progress, + RolloutMigrationTrigger::Manual, + ) + .await + } + + async fn migrate_rollouts_with_progress_for_trigger( + &self, + options: RolloutMigrationOptions, + mut on_progress: impl FnMut(RolloutMigrationProgress), + trigger: RolloutMigrationTrigger, + ) -> ThreadStoreResult { + let telemetry = RolloutMigrationTelemetry::new(trigger, &options); + let result = self + .migrate_rollouts_with_progress_inner(options, &mut on_progress) + .await; + telemetry.finish(&result); + result + } + + async fn migrate_rollouts_with_progress_inner( + &self, + options: RolloutMigrationOptions, + on_progress: &mut impl FnMut(RolloutMigrationProgress), ) -> ThreadStoreResult { let mut limiter = RolloutMigrationRateLimiter::new(options.max_mib_per_second)?; let _maintenance_guard = match options.mode { @@ -227,15 +297,22 @@ impl LocalThreadStore { .is_some_and(|thread_id| pending_thread_ids.contains(&thread_id)) }); } + let total_paths = paths.len(); let mut report = RolloutMigrationReport::default(); - for path in paths { - if let Some(outcome) = self + for (index, path) in paths.into_iter().enumerate() { + let outcome = self .migrate_rollout_path(path, &options, &mut limiter) - .await? - { + .await?; + let outcome_status = outcome.as_ref().map(|outcome| outcome.status); + if let Some(outcome) = outcome { report.outcomes.push(outcome); } + on_progress(RolloutMigrationProgress { + processed_paths: index + 1, + total_paths, + outcome_status, + }); } Ok(report) @@ -301,6 +378,7 @@ impl LocalThreadStore { if metadata.meta.history_mode == ThreadHistoryMode::Paginated { let bytes_before = limiter.bytes_processed; let result = if pending_published_migration { + let _live_writer_guard = self.live_writer_locks.lock(thread_id).await; match self .recover_published_migration(thread_id, &path, &journal_path, limiter) .await @@ -341,6 +419,7 @@ impl LocalThreadStore { ))); } + let _live_writer_guard = self.live_writer_locks.lock(thread_id).await; let _writer_guard = match self.writer_lock_coordinator.acquire(thread_id) { Ok(guard) => guard, Err(ThreadStoreError::Conflict { message }) => { diff --git a/codex-rs/thread-store/src/local/rollout_migration/startup.rs b/codex-rs/thread-store/src/local/rollout_migration/startup.rs new file mode 100644 index 0000000000..18c467526c --- /dev/null +++ b/codex-rs/thread-store/src/local/rollout_migration/startup.rs @@ -0,0 +1,332 @@ +//! Decides whether startup needs to invoke legacy -> paginated rollout migration. +//! +//! It keeps startup cheap by storing a creation-ordered cursor in SQLite and checking only newer +//! rollout files on later launches. When it finds legacy history or a pending recovery marker, it +//! invokes the existing full migration path. +//! +//! Empty or malformed rollouts are fingerprinted so they do not block the cursor forever. If one +//! changes later, startup retries it instead of trusting the old skip. + +use std::collections::HashSet; +use std::io::ErrorKind; +use std::path::Path; +use std::path::PathBuf; +use std::time::SystemTime; + +use chrono::NaiveDateTime; +use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_rollout::StateDbHandle; +use codex_state::RolloutMigrationCursor; +use codex_state::RolloutMigrationSkippedRollout; + +use super::LocalThreadStore; +use super::RolloutMigrationMode; +use super::RolloutMigrationOptions; +use super::RolloutMigrationStatus; +use super::find_rollout_paths; +use super::migration_error; +use super::publish::pending_migration_thread_ids; +use super::telemetry::RolloutMigrationTrigger; +use crate::ThreadStoreResult; + +const LEGACY_TO_PAGINATED_MIGRATION_ID: &str = "legacy_to_paginated_v1"; +const EMPTY_SKIP_REASON: &str = "empty"; +const MALFORMED_SESSION_META_SKIP_REASON: &str = "malformed_session_meta"; +const CURSOR_LOOKBACK_SECONDS: i64 = 48 * 60 * 60; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct RolloutFingerprint { + size_bytes: i64, + modified_at_ns: i64, +} + +enum StartupInspection { + Paginated, + Legacy, + Skipped, + Unresolved, +} + +pub(super) async fn migrate_rollouts_on_startup(store: &LocalThreadStore) -> ThreadStoreResult<()> { + let Some(state_db) = store.state_db.as_ref() else { + return Ok(()); + }; + let paths = find_all_rollout_paths(store).await?; + let skipped_rollouts = state_db + .list_rollout_migration_skipped_rollouts(LEGACY_TO_PAGINATED_MIGRATION_ID) + .await + .map_err(migration_error)?; + if !pending_migration_thread_ids(&store.config.codex_home) + .await? + .is_empty() + { + return migrate_all_rollouts(store, paths, skipped_rollouts.as_slice()).await; + } + let (unchanged_skips, invalidated_skip) = + revalidate_skipped_rollouts(store, skipped_rollouts.as_slice()).await?; + let state = state_db + .get_rollout_migration_state(LEGACY_TO_PAGINATED_MIGRATION_ID) + .await + .map_err(migration_error)?; + + if state.is_none() || invalidated_skip { + return migrate_all_rollouts(store, paths, skipped_rollouts.as_slice()).await; + } + + let last_checked_thread = state.and_then(|state| state.last_checked_thread); + let lookback_created_at = last_checked_thread.as_ref().map(|cursor| { + cursor + .thread_created_at + .saturating_sub(CURSOR_LOOKBACK_SECONDS) + }); + let candidates = paths + .iter() + .filter(|path| { + let relative_path = relative_rollout_path(store, path); + !unchanged_skips.contains(relative_path.as_str()) + && thread_creation_cursor(path).is_none_or(|cursor| { + lookback_created_at.is_none_or(|lookback_created_at| { + cursor.thread_created_at >= lookback_created_at + }) + }) + }) + .collect::>(); + if candidates.is_empty() { + return Ok(()); + } + + let mut unresolved = false; + for path in candidates { + match inspect_rollout_path(store, path).await? { + StartupInspection::Paginated | StartupInspection::Skipped => {} + StartupInspection::Legacy => { + return migrate_all_rollouts(store, paths, skipped_rollouts.as_slice()).await; + } + StartupInspection::Unresolved => unresolved = true, + } + } + if unresolved { + return Ok(()); + } + + advance_last_checked_thread(store, paths.as_slice()).await +} + +async fn migrate_all_rollouts( + store: &LocalThreadStore, + paths_before_migration: Vec, + existing_skips: &[RolloutMigrationSkippedRollout], +) -> ThreadStoreResult<()> { + let report = store + .migrate_rollouts_with_progress_for_trigger( + RolloutMigrationOptions { + mode: RolloutMigrationMode::Apply, + thread_ids: Vec::new(), + max_mib_per_second: None, + }, + |_| {}, + RolloutMigrationTrigger::Startup, + ) + .await?; + let existing_skip_paths = existing_skips + .iter() + .map(|skipped_rollout| skipped_rollout.rollout_path.as_str()) + .collect::>(); + let mut terminal = true; + let mut reported_paths = HashSet::new(); + for outcome in &report.outcomes { + let relative_path = relative_rollout_path(store, &outcome.rollout_path); + reported_paths.insert(relative_path.clone()); + match outcome.status { + RolloutMigrationStatus::Migrated | RolloutMigrationStatus::AlreadyPaginated => { + if existing_skip_paths.contains(relative_path.as_str()) { + remove_skip(store, relative_path.as_str()).await?; + } + } + RolloutMigrationStatus::SkippedEmpty | RolloutMigrationStatus::Failed => { + if !matches!( + inspect_rollout_path(store, &outcome.rollout_path).await?, + StartupInspection::Skipped + ) { + terminal = false; + } + } + RolloutMigrationStatus::Eligible | RolloutMigrationStatus::SkippedBusy => { + terminal = false + } + } + } + if !terminal { + return Ok(()); + } + for skipped_rollout in existing_skips { + if !reported_paths.contains(skipped_rollout.rollout_path.as_str()) { + remove_skip(store, skipped_rollout.rollout_path.as_str()).await?; + } + } + // Only mark the pre-migration snapshot; newer rollouts wait for the next startup check. + advance_last_checked_thread(store, paths_before_migration.as_slice()).await +} + +async fn inspect_rollout_path( + store: &LocalThreadStore, + path: &Path, +) -> ThreadStoreResult { + let before = rollout_fingerprint(path).await?; + match codex_rollout::read_session_meta_line(path).await { + Ok(metadata) if metadata.meta.history_mode == ThreadHistoryMode::Legacy => { + Ok(StartupInspection::Legacy) + } + Ok(_) => Ok(StartupInspection::Paginated), + Err(error) => { + let after = rollout_fingerprint(path).await?; + if before != after || !matches!(error.kind(), ErrorKind::Other | ErrorKind::InvalidData) + { + return Ok(StartupInspection::Unresolved); + } + record_skip(store, path, before).await?; + Ok(StartupInspection::Skipped) + } + } +} + +async fn record_skip( + store: &LocalThreadStore, + path: &Path, + fingerprint: RolloutFingerprint, +) -> ThreadStoreResult<()> { + let state_db = startup_state_db(store)?; + let skipped_rollout = RolloutMigrationSkippedRollout { + rollout_path: relative_rollout_path(store, path), + rollout_size_bytes: fingerprint.size_bytes, + rollout_modified_at_ns: fingerprint.modified_at_ns, + skip_reason: if fingerprint.size_bytes == 0 { + EMPTY_SKIP_REASON.to_string() + } else { + MALFORMED_SESSION_META_SKIP_REASON.to_string() + }, + }; + state_db + .record_rollout_migration_skip(LEGACY_TO_PAGINATED_MIGRATION_ID, &skipped_rollout) + .await + .map_err(migration_error) +} + +async fn remove_skip(store: &LocalThreadStore, rollout_path: &str) -> ThreadStoreResult<()> { + startup_state_db(store)? + .remove_rollout_migration_skip(LEGACY_TO_PAGINATED_MIGRATION_ID, rollout_path) + .await + .map_err(migration_error) +} + +async fn revalidate_skipped_rollouts( + store: &LocalThreadStore, + skipped_rollouts: &[RolloutMigrationSkippedRollout], +) -> ThreadStoreResult<(HashSet, bool)> { + let mut unchanged_skips = HashSet::new(); + let mut invalidated_skip = false; + for skipped_rollout in skipped_rollouts { + let path = store.config.codex_home.join(&skipped_rollout.rollout_path); + let fingerprint = match rollout_fingerprint(&path).await { + Ok(fingerprint) => fingerprint, + Err(_) => { + invalidated_skip = true; + continue; + } + }; + if fingerprint.size_bytes == skipped_rollout.rollout_size_bytes + && fingerprint.modified_at_ns == skipped_rollout.rollout_modified_at_ns + { + unchanged_skips.insert(skipped_rollout.rollout_path.clone()); + } else { + invalidated_skip = true; + } + } + Ok((unchanged_skips, invalidated_skip)) +} + +async fn advance_last_checked_thread( + store: &LocalThreadStore, + paths: &[PathBuf], +) -> ThreadStoreResult<()> { + let last_checked_thread = paths + .iter() + .filter_map(|path| thread_creation_cursor(path)) + .max(); + startup_state_db(store)? + .advance_rollout_migration_state( + LEGACY_TO_PAGINATED_MIGRATION_ID, + last_checked_thread.as_ref(), + ) + .await + .map_err(migration_error) +} + +fn startup_state_db(store: &LocalThreadStore) -> ThreadStoreResult<&StateDbHandle> { + store + .state_db + .as_ref() + .ok_or_else(|| migration_error("startup migration requires state db")) +} + +async fn find_all_rollout_paths(store: &LocalThreadStore) -> ThreadStoreResult> { + let mut paths = + find_rollout_paths(&store.config.codex_home.join(codex_rollout::SESSIONS_SUBDIR)).await?; + paths.extend( + find_rollout_paths( + &store + .config + .codex_home + .join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR), + ) + .await?, + ); + Ok(paths) +} + +fn thread_creation_cursor(path: &Path) -> Option { + let name = path.file_name()?.to_str()?; + let stem = name + .strip_suffix(".jsonl.zst") + .or_else(|| name.strip_suffix(".jsonl"))? + .strip_prefix("rollout-")?; + let separator = stem.len().checked_sub(37)?; + let thread_id = stem.get(separator + 1..)?; + ThreadId::from_string(thread_id).ok()?; + let timestamp = NaiveDateTime::parse_from_str(stem.get(..separator)?, "%Y-%m-%dT%H-%M-%S") + .ok()? + .and_utc() + .timestamp(); + Some(RolloutMigrationCursor { + thread_created_at: timestamp, + thread_id: thread_id.to_string(), + }) +} + +fn relative_rollout_path(store: &LocalThreadStore, path: &Path) -> String { + path.strip_prefix(&store.config.codex_home) + .unwrap_or(path) + .to_string_lossy() + .replace('\\', "/") +} + +async fn rollout_fingerprint(path: &Path) -> ThreadStoreResult { + let metadata = tokio::fs::metadata(path).await.map_err(migration_error)?; + let size_bytes = i64::try_from(metadata.len()).map_err(migration_error)?; + let modified_at_ns = metadata + .modified() + .map_err(migration_error)? + .duration_since(SystemTime::UNIX_EPOCH) + .map_err(migration_error)? + .as_nanos(); + let modified_at_ns = i64::try_from(modified_at_ns).map_err(migration_error)?; + Ok(RolloutFingerprint { + size_bytes, + modified_at_ns, + }) +} + +#[cfg(test)] +#[path = "startup_tests.rs"] +mod tests; diff --git a/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs new file mode 100644 index 0000000000..e33c2f4133 --- /dev/null +++ b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs @@ -0,0 +1,333 @@ +use std::fs; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; +use std::time::Duration; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutConfig; +use codex_rollout::RolloutRecorder; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +use super::super::publish::migration_journal_path; +use super::super::publish::write_migration_journal; +use super::super::thread_history; +use super::LocalThreadStore; +use crate::local::test_support::test_config; + +const TIMESTAMP: &str = "2025-01-03T12:00:00Z"; + +fn write_rollout(home: &Path, thread_id: ThreadId, history_mode: ThreadHistoryMode) -> PathBuf { + let directory = home.join("sessions/2025/01/03"); + fs::create_dir_all(&directory).expect("create rollout directory"); + let path = directory.join(format!("rollout-2025-01-03T12-00-00-{thread_id}.jsonl")); + let mut file = fs::File::create(&path).expect("create legacy rollout"); + let paginated = history_mode == ThreadHistoryMode::Paginated; + let metadata = SessionMeta { + session_id: thread_id.into(), + id: thread_id, + timestamp: TIMESTAMP.to_string(), + cwd: home.to_path_buf(), + originator: "test-originator".to_string(), + cli_version: "0.0.0".to_string(), + source: SessionSource::Cli, + model_provider: Some("test-provider".to_string()), + history_mode, + ..SessionMeta::default() + }; + for (ordinal, item) in [ + RolloutItem::SessionMeta(SessionMetaLine { + meta: metadata, + git: None, + }), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "question".to_string(), + ..UserMessageEvent::default() + })), + ] + .into_iter() + .enumerate() + { + let line = RolloutLine { + timestamp: TIMESTAMP.to_string(), + ordinal: paginated.then_some(ordinal as u64), + item, + }; + writeln!( + file, + "{}", + serde_json::to_string(&line).expect("serialize legacy record") + ) + .expect("write legacy record"); + } + path +} + +fn move_to_timestamp( + home: &Path, + path: PathBuf, + session_day: &str, + filename_timestamp: &str, +) -> PathBuf { + let directory = home.join(format!("sessions/{session_day}")); + fs::create_dir_all(&directory).expect("create rollout directory"); + let stem = path + .file_name() + .and_then(|name| name.to_str()) + .and_then(|name| name.strip_suffix(".jsonl")) + .expect("rollout filename"); + let thread_id = stem + .get(stem.len().checked_sub(36).expect("thread id offset")..) + .expect("thread id in rollout filename"); + let moved_path = directory.join(format!("rollout-{filename_timestamp}-{thread_id}.jsonl")); + fs::rename(path, &moved_path).expect("move rollout timestamp"); + moved_path +} + +async fn indexed_store(home: &Path) -> LocalThreadStore { + let config = test_config(home); + let rollout_config = RolloutConfig { + codex_home: config.codex_home.clone(), + sqlite: config.sqlite.clone(), + cwd: home.to_path_buf(), + model_provider_id: config.default_model_provider_id.clone(), + generate_memories: false, + }; + let state_db = codex_rollout::state_db::try_init(&rollout_config) + .await + .expect("backfill legacy thread metadata"); + LocalThreadStore::new(config, Some(state_db)) +} + +#[tokio::test] +async fn records_and_advances_checked_thread() { + let home = TempDir::new().expect("create Codex home"); + let legacy_thread_id = ThreadId::new(); + let legacy_path = write_rollout(home.path(), legacy_thread_id, ThreadHistoryMode::Legacy); + let store = indexed_store(home.path()).await; + + store + .migrate_rollouts_on_startup() + .await + .expect("migrate startup rollouts"); + assert_eq!( + codex_rollout::read_session_meta_line(&legacy_path) + .await + .expect("read migrated metadata") + .meta + .history_mode, + ThreadHistoryMode::Paginated + ); + + let newer_thread_id = ThreadId::new(); + move_to_timestamp( + home.path(), + write_rollout(home.path(), newer_thread_id, ThreadHistoryMode::Paginated), + "2025/01/04", + "2025-01-04T12-00-00", + ); + store + .migrate_rollouts_on_startup() + .await + .expect("check newer paginated rollout"); + + let state = store + .state_db() + .await + .expect("state db") + .get_rollout_migration_state(super::LEGACY_TO_PAGINATED_MIGRATION_ID) + .await + .expect("read migration state") + .expect("migration state"); + assert_eq!( + state + .last_checked_thread + .expect("last checked thread") + .thread_id, + newer_thread_id.to_string() + ); +} + +#[tokio::test] +async fn checks_rollouts_within_the_cursor_lookback() { + let home = TempDir::new().expect("create Codex home"); + let older_thread_id = ThreadId::new(); + let older_path = move_to_timestamp( + home.path(), + write_rollout(home.path(), older_thread_id, ThreadHistoryMode::Legacy), + "2025/01/02", + "2025-01-02T12-00-00", + ); + let newer_path = move_to_timestamp( + home.path(), + write_rollout(home.path(), ThreadId::new(), ThreadHistoryMode::Paginated), + "2025/01/03", + "2025-01-03T12-00-00", + ); + let store = indexed_store(home.path()).await; + let cursor = super::thread_creation_cursor(&newer_path).expect("newer rollout cursor"); + store + .state_db() + .await + .expect("state db") + .advance_rollout_migration_state(super::LEGACY_TO_PAGINATED_MIGRATION_ID, Some(&cursor)) + .await + .expect("seed migration cursor"); + + store + .migrate_rollouts_on_startup() + .await + .expect("check rollout behind cursor"); + + assert_eq!( + codex_rollout::read_session_meta_line(&older_path) + .await + .expect("read migrated metadata") + .meta + .history_mode, + ThreadHistoryMode::Paginated + ); +} + +#[tokio::test] +async fn recovers_pending_migrations_behind_the_checked_thread() { + let home = TempDir::new().expect("create Codex home"); + let thread_id = ThreadId::new(); + let path = write_rollout(home.path(), thread_id, ThreadHistoryMode::Legacy); + let store = indexed_store(home.path()).await; + + store + .migrate_rollouts_on_startup() + .await + .expect("migrate and advance startup cursor"); + thread_history::delete_thread(&store, thread_id) + .await + .expect("simulate missing projection"); + let journal_path = migration_journal_path(home.path(), thread_id); + write_migration_journal(&journal_path) + .await + .expect("simulate pending migration marker"); + + store + .migrate_rollouts_on_startup() + .await + .expect("recover pending migration behind cursor"); + + assert!(!journal_path.exists()); + assert!( + thread_history::projection_state(&store, thread_id) + .await + .expect("read repaired projection") + .is_some() + ); + assert_eq!( + codex_rollout::read_session_meta_line(&path) + .await + .expect("read migrated metadata") + .meta + .history_mode, + ThreadHistoryMode::Paginated + ); +} + +#[tokio::test] +async fn waits_for_a_live_writer_before_migrating() { + let home = TempDir::new().expect("create Codex home"); + let thread_id = ThreadId::new(); + let path = write_rollout(home.path(), thread_id, ThreadHistoryMode::Legacy); + let store = indexed_store(home.path()).await; + let live_writer_guard = store.live_writer_locks.lock(thread_id).await; + let migration_store = store.clone(); + let mut migration = tokio::spawn(async move { + migration_store + .migrate_rollouts_on_startup() + .await + .expect("migrate startup rollouts"); + }); + + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut migration) + .await + .is_err(), + "migration should wait for the live writer" + ); + drop(live_writer_guard); + migration.await.expect("join startup migration"); + + assert_eq!( + codex_rollout::read_session_meta_line(&path) + .await + .expect("read migrated metadata") + .meta + .history_mode, + ThreadHistoryMode::Paginated + ); +} + +#[tokio::test] +async fn rechecks_changed_empty_rollouts() { + let home = TempDir::new().expect("create Codex home"); + write_rollout(home.path(), ThreadId::new(), ThreadHistoryMode::Legacy); + let empty_thread_id = ThreadId::new(); + let empty_path = move_to_timestamp( + home.path(), + write_rollout(home.path(), empty_thread_id, ThreadHistoryMode::Legacy), + "2025/01/04", + "2025-01-04T12-00-00", + ); + let restored_contents = fs::read(&empty_path).expect("read rollout before emptying"); + fs::write(&empty_path, []).expect("empty rollout"); + let store = indexed_store(home.path()).await; + + store + .migrate_rollouts_on_startup() + .await + .expect("record empty rollout"); + fs::write(&empty_path, restored_contents).expect("restore rollout"); + let (items, _, _) = RolloutRecorder::load_rollout_items(&empty_path) + .await + .expect("load restored rollout"); + let metadata = codex_rollout::builder_from_items(items.as_slice(), &empty_path) + .expect("build restored metadata") + .build("test-provider"); + store + .state_db() + .await + .expect("state db") + .upsert_thread(&metadata) + .await + .expect("seed restored metadata"); + + store + .migrate_rollouts_on_startup() + .await + .expect("migrate changed rollout"); + + assert_eq!( + codex_rollout::read_session_meta_line(&empty_path) + .await + .expect("read migrated metadata") + .meta + .history_mode, + ThreadHistoryMode::Paginated + ); + assert!( + store + .state_db() + .await + .expect("state db") + .list_rollout_migration_skipped_rollouts(super::LEGACY_TO_PAGINATED_MIGRATION_ID) + .await + .expect("read skipped rollouts") + .is_empty() + ); +} diff --git a/codex-rs/thread-store/src/local/rollout_migration/telemetry.rs b/codex-rs/thread-store/src/local/rollout_migration/telemetry.rs new file mode 100644 index 0000000000..8a8d672b30 --- /dev/null +++ b/codex-rs/thread-store/src/local/rollout_migration/telemetry.rs @@ -0,0 +1,133 @@ +//! Records low-cardinality migration metrics shared by manual and startup runs. +//! +//! Migration owns the outcome stream, so it is the one place that can count successful, +//! skipped, and failed rollouts consistently. Callers only tell it why the run started. + +use std::time::Instant; + +use super::RolloutMigrationMode; +use super::RolloutMigrationOptions; +use super::RolloutMigrationReport; +use super::RolloutMigrationStatus; +use crate::ThreadStoreResult; + +const RUN_METRIC: &str = "codex.rollout_migration.run"; +const RUN_DURATION_METRIC: &str = "codex.rollout_migration.run.duration_ms"; +const RUN_IO_BYTES_METRIC: &str = "codex.rollout_migration.run.io_bytes"; +const THREAD_METRIC: &str = "codex.rollout_migration.thread"; + +#[derive(Clone, Copy)] +pub(super) enum RolloutMigrationTrigger { + Manual, + Startup, +} + +#[derive(Clone, Copy)] +enum RolloutMigrationScope { + All, + Selected, +} + +pub(super) struct RolloutMigrationTelemetry { + trigger: RolloutMigrationTrigger, + mode: RolloutMigrationMode, + scope: RolloutMigrationScope, + started_at: Instant, +} + +impl RolloutMigrationTelemetry { + pub(super) fn new(trigger: RolloutMigrationTrigger, options: &RolloutMigrationOptions) -> Self { + Self { + trigger, + mode: options.mode, + scope: if options.thread_ids.is_empty() { + RolloutMigrationScope::All + } else { + RolloutMigrationScope::Selected + }, + started_at: Instant::now(), + } + } + + pub(super) fn finish(&self, result: &ThreadStoreResult) { + let Some(metrics) = codex_otel::global() else { + return; + }; + let mut io_bytes = 0_u64; + if let Ok(report) = result { + for outcome in &report.outcomes { + let tags = [ + ("trigger", self.trigger.tag()), + ("mode", self.mode.tag()), + ("scope", self.scope.tag()), + ("status", outcome.status.tag()), + ]; + let _ = metrics.counter(THREAD_METRIC, /*inc*/ 1, &tags); + io_bytes = io_bytes.saturating_add(outcome.bytes_processed); + } + } + let result = match result { + Err(_) => "error", + Ok(report) + if report + .outcomes + .iter() + .any(|outcome| outcome.status == RolloutMigrationStatus::Failed) => + { + "partial_failure" + } + Ok(_) => "success", + }; + let tags = [ + ("trigger", self.trigger.tag()), + ("mode", self.mode.tag()), + ("scope", self.scope.tag()), + ("result", result), + ]; + let duration_ms = i64::try_from(self.started_at.elapsed().as_millis()).unwrap_or(i64::MAX); + let io_bytes = i64::try_from(io_bytes).unwrap_or(i64::MAX); + let _ = metrics.counter(RUN_METRIC, /*inc*/ 1, &tags); + let _ = metrics.histogram(RUN_DURATION_METRIC, duration_ms, &tags); + let _ = metrics.histogram(RUN_IO_BYTES_METRIC, io_bytes, &tags); + } +} + +impl RolloutMigrationTrigger { + fn tag(self) -> &'static str { + match self { + Self::Manual => "manual", + Self::Startup => "startup", + } + } +} + +impl RolloutMigrationMode { + fn tag(self) -> &'static str { + match self { + Self::DryRun => "dry_run", + Self::Apply => "apply", + } + } +} + +impl RolloutMigrationScope { + fn tag(self) -> &'static str { + match self { + Self::All => "all", + Self::Selected => "selected", + } + } +} + +impl RolloutMigrationStatus { + fn tag(self) -> &'static str { + match self { + Self::Eligible => "eligible", + Self::Migrated => "migrated", + Self::AlreadyPaginated => "already_paginated", + Self::SkippedEmpty => "skipped_empty", + Self::SkippedBusy => "skipped_busy", + Self::Failed => "failed", + } + } +} diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 7761ca382b..bdad98164f 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -41,6 +41,7 @@ use tempfile::TempDir; use super::LocalThreadStore; use super::RolloutMigrationMode; use super::RolloutMigrationOptions; +use super::RolloutMigrationProgress; use super::RolloutMigrationStatus; #[cfg(unix)] use super::decompress_rollout_to_path; @@ -241,7 +242,7 @@ fn read_rollout(path: &Path) -> Vec { fn apply_options() -> RolloutMigrationOptions { RolloutMigrationOptions { mode: RolloutMigrationMode::Apply, - max_mib_per_second: 1024, + max_mib_per_second: Some(1024), ..RolloutMigrationOptions::default() } } @@ -1516,8 +1517,11 @@ async fn dry_run_reports_migration_order() { let original = fs::read(&root).expect("read original root rollout"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let mut progress = Vec::new(); let report = store - .migrate_rollouts(RolloutMigrationOptions::default()) + .migrate_rollouts_with_progress(RolloutMigrationOptions::default(), |update| { + progress.push(update); + }) .await .expect("inspect legacy rollouts"); @@ -1540,6 +1544,14 @@ async fn dry_run_reports_migration_order() { .collect::>(), expected, ); + assert_eq!( + progress.last(), + Some(&RolloutMigrationProgress { + processed_paths: 5, + total_paths: 5, + outcome_status: Some(RolloutMigrationStatus::Eligible), + }) + ); assert_eq!(fs::read(&root).expect("read inspected rollout"), original); let selected = store