use crate::RolloutItem; use crate::config::RolloutConfig; use crate::config::RolloutConfigView; use crate::list::Cursor; use crate::list::SortDirection; use crate::list::ThreadSortKey; use crate::metadata; use crate::sqlite_metrics; use anyhow::Context; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; pub use codex_state::LogEntry; use codex_state::SqliteConfig; use codex_state::ThreadMetadataBuilder; use codex_utils_path::normalize_for_path_comparison; use serde_json::Value; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use std::time::Instant; use tracing::info; use tracing::warn; /// Core-facing handle to the SQLite-backed state runtime. pub type StateDbHandle = Arc; #[cfg(not(test))] const STARTUP_BACKFILL_POLL_INTERVAL: Duration = Duration::from_secs(1); #[cfg(test)] const STARTUP_BACKFILL_POLL_INTERVAL: Duration = Duration::from_millis(10); #[cfg(not(test))] const STARTUP_BACKFILL_WAIT_TIMEOUT: Duration = Duration::from_secs(30); #[cfg(test)] const STARTUP_BACKFILL_WAIT_TIMEOUT: Duration = Duration::from_secs(2); /// Initialize the state runtime for thread state persistence. /// /// This is the process entry point for local state: it opens the SQLite-backed /// runtime, applies rollout metadata backfills as needed, and returns the /// initialized handle. pub async fn init(config: &impl RolloutConfigView) -> Option { let config = RolloutConfig::from_view(config); match try_init_with_roots(config.codex_home, config.sqlite, config.model_provider_id).await { Ok(runtime) => Some(runtime), Err(err) => { emit_startup_warning(&format!("failed to initialize state runtime: {err:#}")); None } } } /// Initialize the state runtime and return any initialization error to the caller. /// /// Prefer [`init`] unless the caller needs to surface the exact failure after /// tracing or UI setup has completed. pub async fn try_init(config: &impl RolloutConfigView) -> anyhow::Result { let config = RolloutConfig::from_view(config); try_init_with_roots(config.codex_home, config.sqlite, config.model_provider_id).await } async fn try_init_with_roots( codex_home: PathBuf, sqlite: SqliteConfig, default_model_provider_id: String, ) -> anyhow::Result { try_init_with_roots_inner( codex_home, sqlite, default_model_provider_id, /*backfill_lease_seconds*/ None, ) .await } #[cfg(test)] async fn try_init_with_roots_and_backfill_lease( codex_home: PathBuf, sqlite: SqliteConfig, default_model_provider_id: String, backfill_lease_seconds: i64, ) -> anyhow::Result { try_init_with_roots_inner( codex_home, sqlite, default_model_provider_id, Some(backfill_lease_seconds), ) .await } async fn try_init_with_roots_inner( codex_home: PathBuf, sqlite: SqliteConfig, default_model_provider_id: String, backfill_lease_seconds: Option, ) -> anyhow::Result { let runtime = codex_state::StateRuntime::init(sqlite.clone(), default_model_provider_id.clone()) .await .with_context(|| { format!( "failed to initialize state runtime at {}", sqlite.home().display() ) })?; let backfill_gate_started = Instant::now(); let backfill_gate_result = wait_for_backfill_gate( runtime.as_ref(), codex_home.as_path(), default_model_provider_id.as_str(), backfill_lease_seconds, ) .await; codex_state::record_backfill_gate( /*telemetry*/ None, backfill_gate_started.elapsed(), &backfill_gate_result, ); if let Err(err) = backfill_gate_result { runtime.close().await; return Err(err); } Ok(runtime) } async fn wait_for_backfill_gate( runtime: &codex_state::StateRuntime, codex_home: &Path, default_model_provider_id: &str, backfill_lease_seconds: Option, ) -> anyhow::Result<()> { let wait_started = Instant::now(); let mut reported_wait = false; loop { let backfill_state = runtime.get_backfill_state().await.map_err(|err| { anyhow::anyhow!( "failed to read backfill state at {}: {err}", codex_home.display() ) })?; if backfill_state.status == codex_state::BackfillStatus::Complete { return Ok(()); } if let Some(backfill_lease_seconds) = backfill_lease_seconds { metadata::backfill_sessions_with_lease( runtime, codex_home, default_model_provider_id, backfill_lease_seconds, ) .await; } else { metadata::backfill_sessions(runtime, codex_home, default_model_provider_id).await; } let backfill_state = runtime.get_backfill_state().await.map_err(|err| { anyhow::anyhow!( "failed to read backfill state at {} after startup backfill: {err}", codex_home.display() ) })?; if backfill_state.status == codex_state::BackfillStatus::Complete { return Ok(()); } if wait_started.elapsed() >= STARTUP_BACKFILL_WAIT_TIMEOUT { return Err(anyhow::anyhow!( "timed out waiting for state db backfill at {} after {:?} (status: {})", codex_home.display(), STARTUP_BACKFILL_WAIT_TIMEOUT, backfill_state.status.as_str() )); } let message = format!( "state db backfill is {} at {}; waiting up to {:?} before retrying startup initialization", backfill_state.status.as_str(), codex_home.display(), STARTUP_BACKFILL_WAIT_TIMEOUT, ); if reported_wait { info!("{message}"); } else { emit_startup_warning(&message); reported_wait = true; } tokio::time::sleep(STARTUP_BACKFILL_POLL_INTERVAL).await; } } fn emit_startup_warning(message: &str) { warn!("{message}"); if !tracing::dispatcher::has_been_set() { #[allow(clippy::print_stderr)] { eprintln!("{message}"); } } } /// Open the DB if it exists and its startup backfill has already completed. /// /// Unlike [`init`], this helper does not run rollout backfill. It is for /// optional local reads from non-owning contexts such as remote app-server mode. pub async fn get_state_db(config: &impl RolloutConfigView) -> Option { let state_path = config.sqlite_config().state_db_path(); if !tokio::fs::try_exists(&state_path).await.unwrap_or(false) { codex_state::record_fallback( "get_state_db", "db_unavailable", /*telemetry_override*/ None, ); return None; } let runtime = match codex_state::StateRuntime::init( config.sqlite_config().clone(), config.model_provider_id().to_string(), ) .await { Ok(runtime) => runtime, Err(_) => { codex_state::record_fallback( "get_state_db", "db_error", /*telemetry_override*/ None, ); return None; } }; require_backfill_complete(runtime, config.sqlite_config().home()).await } /// Build a SQLite telemetry recorder backed by an OTEL metrics client. pub fn sqlite_telemetry_recorder( metrics: codex_otel::MetricsClient, originator: &str, ) -> codex_state::DbTelemetryHandle { sqlite_metrics::recorder(metrics, originator) } async fn require_backfill_complete( runtime: StateDbHandle, codex_home: &Path, ) -> Option { match runtime.get_backfill_state().await { Ok(state) if state.status == codex_state::BackfillStatus::Complete => Some(runtime), Ok(state) => { warn!( "state db backfill not complete at {} (status: {})", codex_home.display(), state.status.as_str() ); codex_state::record_fallback( "get_state_db", "backfill_incomplete", /*telemetry_override*/ None, ); None } Err(err) => { warn!( "failed to read backfill state at {}: {err}", codex_home.display() ); codex_state::record_fallback( "get_state_db", "db_error", /*telemetry_override*/ None, ); None } } } fn cursor_to_anchor(cursor: Option<&Cursor>) -> Option { let cursor = cursor?; let millis = cursor.timestamp().unix_timestamp_nanos() / 1_000_000; let millis = i64::try_from(millis).ok()?; let ts = chrono::DateTime::::from_timestamp_millis(millis)?; Some(codex_state::Anchor { ts, id: cursor.thread_id(), }) } pub fn normalize_cwd_for_state_db(cwd: &Path) -> PathBuf { normalize_for_path_comparison(cwd).unwrap_or_else(|_| cwd.to_path_buf()) } /// List thread ids from SQLite for parity checks without rollout scanning. #[allow(clippy::too_many_arguments)] pub async fn list_thread_ids_db( context: Option<&codex_state::StateRuntime>, sqlite: &codex_state::SqliteConfig, page_size: usize, cursor: Option<&Cursor>, sort_key: ThreadSortKey, allowed_sources: &[SessionSource], model_providers: Option<&[String]>, archived_only: bool, stage: &str, ) -> Option> { let ctx = context?; if ctx.sqlite() != sqlite { warn!( "state db SQLite home mismatch: expected {}, got {}", sqlite.home().display(), ctx.sqlite().home().display() ); return None; } let anchor = cursor_to_anchor(cursor); let allowed_sources: Vec = allowed_sources .iter() .map(|value| match serde_json::to_value(value) { Ok(Value::String(s)) => s, Ok(other) => other.to_string(), Err(_) => String::new(), }) .collect(); let model_providers = model_providers.map(<[String]>::to_vec); match ctx .list_thread_ids( page_size, anchor.as_ref(), match sort_key { ThreadSortKey::CreatedAt => codex_state::SortKey::CreatedAt, ThreadSortKey::UpdatedAt => codex_state::SortKey::UpdatedAt, ThreadSortKey::RecencyAt => codex_state::SortKey::RecencyAt, }, allowed_sources.as_slice(), model_providers.as_deref(), archived_only, ) .await { Ok(ids) => Some(ids), Err(err) => { warn!("state db list_thread_ids failed during {stage}: {err}"); None } } } /// List thread metadata from SQLite without rollout directory traversal. #[allow(clippy::too_many_arguments)] pub async fn list_threads_db( context: Option<&codex_state::StateRuntime>, sqlite: &codex_state::SqliteConfig, page_size: usize, cursor: Option<&Cursor>, sort_key: ThreadSortKey, sort_direction: SortDirection, allowed_sources: &[SessionSource], model_providers: Option<&[String]>, cwd_filters: Option<&[PathBuf]>, relation_filter: Option, archived: bool, section: Option>, project_id: Option>, search_term: Option<&str>, ) -> Option { let ctx = context?; if ctx.sqlite() != sqlite { warn!( "state db SQLite home mismatch: expected {}, got {}", sqlite.home().display(), ctx.sqlite().home().display() ); return None; } let mut anchor = cursor_to_anchor(cursor); let allowed_sources: Vec = allowed_sources .iter() .map(|value| match serde_json::to_value(value) { Ok(Value::String(s)) => s, Ok(other) => other.to_string(), Err(_) => String::new(), }) .collect(); let model_providers = model_providers.map(<[String]>::to_vec); let normalized_cwd_filters = cwd_filters.map(|filters| { filters .iter() .map(|cwd| normalize_cwd_for_state_db(cwd)) .collect::>() }); let state_sort_key = match sort_key { ThreadSortKey::CreatedAt => codex_state::SortKey::CreatedAt, ThreadSortKey::UpdatedAt => codex_state::SortKey::UpdatedAt, ThreadSortKey::RecencyAt => codex_state::SortKey::RecencyAt, }; let state_sort_direction = match sort_direction { SortDirection::Asc => codex_state::SortDirection::Asc, SortDirection::Desc => codex_state::SortDirection::Desc, }; if let Some(relation_filter) = relation_filter { let filters = codex_state::ThreadFilterOptions { archived_only: archived, section, allowed_sources: allowed_sources.as_slice(), model_providers: model_providers.as_deref(), cwd_filters: normalized_cwd_filters.as_deref(), project_id, anchor: anchor.as_ref(), sort_key: state_sort_key, sort_direction: state_sort_direction, search_term, }; return match ctx .list_threads_by_relation(page_size, relation_filter, filters) .await { Ok(page) => Some(page), Err(err) => { warn!("state db list_threads failed: {err}"); None } }; } let mut items = Vec::with_capacity(page_size); let mut next_anchor = None; let mut num_scanned_rows = 0usize; while items.len() < page_size { let filters = codex_state::ThreadFilterOptions { archived_only: archived, section, allowed_sources: allowed_sources.as_slice(), model_providers: model_providers.as_deref(), cwd_filters: normalized_cwd_filters.as_deref(), project_id, anchor: anchor.as_ref(), sort_key: state_sort_key, sort_direction: state_sort_direction, search_term, }; let page = match ctx.list_threads(page_size - items.len(), filters).await { Ok(page) => page, Err(err) => { warn!("state db list_threads failed: {err}"); return None; } }; num_scanned_rows = num_scanned_rows.saturating_add(page.num_scanned_rows); next_anchor = page.next_anchor; for item in page.items { if let Some(existing_path) = crate::compression::existing_rollout_path(item.rollout_path.as_path()).await { let mut item = item; item.rollout_path = existing_path; items.push(item); } else { warn!( "state db list_threads returned stale rollout path for thread {}: {}", item.id, item.rollout_path.display() ); warn!("state db discrepancy during list_threads_db: stale_db_path_retained"); } } if items.len() == page_size || next_anchor.is_none() { break; } anchor = next_anchor.clone(); } Some(codex_state::ThreadsPage { items, parent_thread_ids: Default::default(), next_anchor, num_scanned_rows, }) } /// Look up the rollout path for a thread id using SQLite. pub async fn find_rollout_path_by_id( context: Option<&codex_state::StateRuntime>, thread_id: ThreadId, archived_only: Option, stage: &str, ) -> Option { let ctx = context?; ctx.find_rollout_path_by_id(thread_id, archived_only) .await .unwrap_or_else(|err| { warn!("state db find_rollout_path_by_id failed during {stage}: {err}"); None }) } pub async fn mark_thread_memory_mode_polluted( context: Option<&codex_state::StateRuntime>, thread_id: ThreadId, stage: &str, ) { let Some(ctx) = context else { return; }; if let Err(err) = ctx .memories() .mark_thread_memory_mode_polluted(thread_id) .await { warn!("memories db mark_thread_memory_mode_polluted failed during {stage}: {err}"); } } /// Reconcile rollout items into SQLite, falling back to scanning the rollout file. pub async fn reconcile_rollout( context: Option<&codex_state::StateRuntime>, rollout_path: &Path, default_provider: &str, builder: Option<&ThreadMetadataBuilder>, items: &[RolloutItem], archived_only: Option, new_thread_memory_mode: Option<&str>, ) { let Some(ctx) = context else { return; }; if builder.is_some() || !items.is_empty() { apply_rollout_items( Some(ctx), rollout_path, default_provider, builder, items, "reconcile_rollout", new_thread_memory_mode, /*updated_at_override*/ None, ) .await; return; } let outcome = match metadata::extract_metadata_from_rollout(rollout_path, default_provider).await { Ok(outcome) => outcome, Err(err) => { warn!( "state db reconcile_rollout extraction failed {}: {err}", rollout_path.display() ); return; } }; let mut metadata = outcome.metadata; let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string()); metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd); let existing_metadata = ctx.get_thread(metadata.id).await.ok().flatten(); // Filesystem repair may seed a missing row, but it must not change an existing row's // selected rollout path. After `thread/revert`, a scan can find multiple immutable // rollouts for one thread and cannot know which one SQLite selected. if existing_metadata .as_ref() .is_some_and(|existing| existing.rollout_path.as_path() != rollout_path) { return; } // Paginated metadata updates are SQLite-only. Use the rollout mode to seed a // missing row, then keep the value from SQLite. let restore_memory_mode_from_rollout = existing_metadata.is_none() || matches!(metadata.history_mode, ThreadHistoryMode::Legacy); if let Some(existing_metadata) = existing_metadata.as_ref() { metadata.prefer_existing_git_info(existing_metadata); metadata.prefer_existing_explicit_title(existing_metadata); } match archived_only { Some(true) if metadata.archived_at.is_none() => { metadata.archived_at = Some(metadata.updated_at); } Some(false) => { metadata.archived_at = None; } Some(true) | None => {} } if let Err(err) = ctx.upsert_thread(&metadata).await { warn!( "state db reconcile_rollout upsert failed {}: {err}", rollout_path.display() ); return; } if restore_memory_mode_from_rollout && let Err(err) = ctx .set_thread_memory_mode(metadata.id, memory_mode.as_str()) .await { warn!( "state db reconcile_rollout memory_mode update failed {}: {err}", rollout_path.display() ); } } /// Repair a thread's rollout path after filesystem fallback succeeds. pub async fn read_repair_rollout_path( context: Option<&codex_state::StateRuntime>, thread_id: Option, archived_only: Option, rollout_path: &Path, ) { let Some(ctx) = context else { return; }; // Fast path: update an existing metadata row in place, but avoid writes when // read-repair computes no effective change. let mut saw_existing_metadata = false; if let Some(thread_id) = thread_id && let Ok(Some(metadata)) = ctx.get_thread(thread_id).await { // Filesystem repair may seed a missing row, but it must not change an existing row's // selected rollout path. After `thread/revert`, a scan can find multiple immutable // rollouts for one thread and cannot know which one SQLite selected. if metadata.rollout_path.as_path() != rollout_path { return; } saw_existing_metadata = true; let mut repaired = metadata.clone(); repaired.rollout_path = rollout_path.to_path_buf(); repaired.cwd = normalize_cwd_for_state_db(&repaired.cwd); match archived_only { Some(true) if repaired.archived_at.is_none() => { repaired.archived_at = Some(repaired.updated_at); } Some(false) => { repaired.archived_at = None; } Some(true) | None => {} } if repaired == metadata { return; } warn!("state db discrepancy during read_repair_rollout_path: upsert_needed (fast path)"); if let Err(err) = ctx.upsert_thread(&repaired).await { warn!( "state db read-repair upsert failed for {}: {err}", rollout_path.display() ); } else { return; } } // Slow path: when the row is missing/unreadable (or direct upsert failed), // rebuild metadata from rollout contents and reconcile it into SQLite. if !saw_existing_metadata { warn!("state db discrepancy during read_repair_rollout_path: upsert_needed (slow path)"); } let default_provider = crate::list::read_session_meta_line(rollout_path) .await .ok() .and_then(|meta| meta.meta.model_provider) .unwrap_or_default(); reconcile_rollout( Some(ctx), rollout_path, default_provider.as_str(), /*builder*/ None, &[], archived_only, /*new_thread_memory_mode*/ None, ) .await; } /// Apply rollout items incrementally to SQLite. #[allow(clippy::too_many_arguments)] pub async fn apply_rollout_items( context: Option<&codex_state::StateRuntime>, rollout_path: &Path, default_provider: &str, builder: Option<&ThreadMetadataBuilder>, items: &[RolloutItem], stage: &str, new_thread_memory_mode: Option<&str>, updated_at_override: Option>, ) { let Some(ctx) = context else { return; }; let mut builder = match builder { Some(builder) => builder.clone(), None => match metadata::builder_from_items(items, rollout_path) { Some(builder) => builder, None => { warn!( "state db apply_rollout_items missing builder during {stage}: {}", rollout_path.display() ); warn!("state db discrepancy during apply_rollout_items: {stage}, missing_builder"); return; } }, }; if builder.model_provider.is_none() { builder.model_provider = Some(default_provider.to_string()); } builder.rollout_path = rollout_path.to_path_buf(); builder.cwd = normalize_cwd_for_state_db(&builder.cwd); if let Err(err) = ctx .apply_rollout_items(&builder, items, new_thread_memory_mode, updated_at_override) .await { warn!( "state db apply_rollout_items failed during {stage} for {}: {err}", rollout_path.display() ); } } pub async fn touch_thread_updated_at( context: Option<&codex_state::StateRuntime>, thread_id: Option, updated_at: DateTime, stage: &str, ) -> bool { let Some(ctx) = context else { return false; }; let Some(thread_id) = thread_id else { return false; }; ctx.touch_thread_updated_at(thread_id, updated_at) .await .unwrap_or_else(|err| { warn!("state db touch_thread_updated_at failed during {stage} for {thread_id}: {err}"); false }) } #[cfg(test)] #[path = "state_db_tests.rs"] mod tests;