feat: async backfill (#10089)

This commit is contained in:
jif-oai
2026-01-29 10:57:50 +01:00
committed by GitHub
parent 780482da84
commit 714dc8d8bd
4 changed files with 58 additions and 27 deletions

View File

@@ -11,7 +11,6 @@ use codex_otel::OtelManager;
use codex_protocol::ThreadId;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SessionSource;
use codex_state::DB_METRIC_BACKFILL;
pub use codex_state::LogEntry;
use codex_state::STATE_DB_FILENAME;
use codex_state::ThreadMetadataBuilder;
@@ -19,15 +18,18 @@ use serde_json::Value;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::info;
use tracing::warn;
use uuid::Uuid;
/// Core-facing handle to the optional SQLite-backed state runtime.
pub type StateDbHandle = Arc<codex_state::StateRuntime>;
/// Initialize the state runtime when the `sqlite` feature flag is enabled.
pub async fn init_if_enabled(config: &Config, otel: Option<&OtelManager>) -> Option<StateDbHandle> {
/// Initialize the state runtime when the `sqlite` feature flag is enabled. To only be used
/// inside `core`. The initialization should not be done anywhere else.
pub(crate) async fn init_if_enabled(
config: &Config,
otel: Option<&OtelManager>,
) -> Option<StateDbHandle> {
let state_path = config.codex_home.join(STATE_DB_FILENAME);
if !config.features.enabled(Feature::Sqlite) {
// We delete the file on best effort basis to maintain retro-compatibility in the future.
@@ -59,27 +61,38 @@ pub async fn init_if_enabled(config: &Config, otel: Option<&OtelManager>) -> Opt
}
};
if !existed {
let stats = metadata::backfill_sessions(runtime.as_ref(), config, otel).await;
info!(
"state db backfill scanned={}, upserted={}, failed={}",
stats.scanned, stats.upserted, stats.failed
);
if let Some(otel) = otel {
otel.counter(
DB_METRIC_BACKFILL,
stats.upserted as i64,
&[("status", "upserted")],
);
otel.counter(
DB_METRIC_BACKFILL,
stats.failed as i64,
&[("status", "failed")],
);
}
let runtime_for_backfill = Arc::clone(&runtime);
let config_for_backfill = config.clone();
let otel_for_backfill = otel.cloned();
tokio::task::spawn(async move {
metadata::backfill_sessions(
runtime_for_backfill.as_ref(),
&config_for_backfill,
otel_for_backfill.as_ref(),
)
.await;
});
}
Some(runtime)
}
/// Get the DB if the feature is enabled and the DB exists.
pub async fn get_state_db(config: &Config, otel: Option<&OtelManager>) -> Option<StateDbHandle> {
let state_path = config.codex_home.join(STATE_DB_FILENAME);
if !config.features.enabled(Feature::Sqlite)
|| !tokio::fs::try_exists(&state_path).await.unwrap_or(false)
{
return None;
}
codex_state::StateRuntime::init(
config.codex_home.clone(),
config.model_provider_id.clone(),
otel.cloned(),
)
.await
.ok()
}
/// Open the state runtime when the SQLite file exists, without feature gating.
///
/// This is used for parity checks during the SQLite migration phase.