diff --git a/codex-rs/state/BUILD.bazel b/codex-rs/state/BUILD.bazel index 8b7c454fdb..ce4900a135 100644 --- a/codex-rs/state/BUILD.bazel +++ b/codex-rs/state/BUILD.bazel @@ -7,6 +7,7 @@ codex_rust_crate( "logs_migrations/**", "memory_migrations/**", "migrations/**", + "thread_history_migrations/**", ]), crate_name = "codex_state", ) diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index cdfcd29c50..cdac70a4eb 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -85,6 +85,8 @@ pub use runtime::sqlite_error_detail_is_lock; pub use runtime::sqlite_integrity_check; pub use runtime::state_db_filename; pub use runtime::state_db_path; +pub use runtime::thread_history_db_filename; +pub use runtime::thread_history_db_path; pub use telemetry::DbTelemetry; pub use telemetry::DbTelemetryHandle; pub use telemetry::install_process_db_telemetry; @@ -98,6 +100,7 @@ pub const LOGS_DB_FILENAME: &str = "logs_2.sqlite"; pub const GOALS_DB_FILENAME: &str = "goals_1.sqlite"; pub const MEMORIES_DB_FILENAME: &str = "memories_1.sqlite"; pub const STATE_DB_FILENAME: &str = "state_5.sqlite"; +pub const THREAD_HISTORY_DB_FILENAME: &str = "thread_history_1.sqlite"; /// Errors encountered during DB operations. Tags: [stage] pub const DB_ERROR_METRIC: &str = "codex.db.error"; diff --git a/codex-rs/state/src/migrations.rs b/codex-rs/state/src/migrations.rs index 187e1fb33d..03371f7a70 100644 --- a/codex-rs/state/src/migrations.rs +++ b/codex-rs/state/src/migrations.rs @@ -7,6 +7,7 @@ pub(crate) static STATE_MIGRATOR: Migrator = sqlx::migrate!("./migrations"); pub(crate) static LOGS_MIGRATOR: Migrator = sqlx::migrate!("./logs_migrations"); pub(crate) static GOALS_MIGRATOR: Migrator = sqlx::migrate!("./goals_migrations"); pub(crate) static MEMORIES_MIGRATOR: Migrator = sqlx::migrate!("./memory_migrations"); +pub(crate) static THREAD_HISTORY_MIGRATOR: Migrator = sqlx::migrate!("./thread_history_migrations"); /// Allow an older Codex binary to open a database that has already been /// migrated by a newer binary running in parallel. @@ -41,6 +42,10 @@ pub(crate) fn runtime_memories_migrator() -> Migrator { runtime_migrator(&MEMORIES_MIGRATOR) } +pub(crate) fn runtime_thread_history_migrator() -> Migrator { + runtime_migrator(&THREAD_HISTORY_MIGRATOR) +} + pub(crate) async fn repair_legacy_recency_migration_version( pool: &SqlitePool, migrator: &Migrator, diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 4b1d264d0b..7ed4e80ba8 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -13,6 +13,7 @@ use crate::LogRow; use crate::MEMORIES_DB_FILENAME; use crate::STATE_DB_FILENAME; use crate::SortKey; +use crate::THREAD_HISTORY_DB_FILENAME; use crate::ThreadMetadata; use crate::ThreadMetadataBuilder; use crate::ThreadsPage; @@ -22,6 +23,7 @@ use crate::migrations::runtime_goals_migrator; use crate::migrations::runtime_logs_migrator; use crate::migrations::runtime_memories_migrator; use crate::migrations::runtime_state_migrator; +use crate::migrations::runtime_thread_history_migrator; use crate::model::AgentJobRow; use crate::model::ThreadRow; use crate::model::anchor_from_item; @@ -144,7 +146,16 @@ const MEMORIES_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_memories", }; -const RUNTIME_DBS: [RuntimeDbSpec; 4] = [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB]; +const THREAD_HISTORY_DB: RuntimeDbSpec = RuntimeDbSpec { + label: "thread history DB", + filename: THREAD_HISTORY_DB_FILENAME, + kind: DbKind::ThreadHistory, + open_phase: "open_thread_history", + migrate_phase: "migrate_thread_history", +}; + +const RUNTIME_DBS: [RuntimeDbSpec; 5] = + [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB, THREAD_HISTORY_DB]; #[derive(Clone, Debug, Eq, PartialEq)] pub struct RuntimeDbPath { @@ -158,6 +169,7 @@ pub struct StateRuntime { default_provider: String, pool: Arc, logs_pool: Arc, + thread_history_pool: Arc, thread_goals: GoalStore, memories: MemoryStore, thread_updated_at_millis: Arc, @@ -167,9 +179,9 @@ pub struct StateRuntime { impl StateRuntime { /// Initialize the state runtime using the provided Codex home and default provider. /// - /// This opens (and migrates) the SQLite databases under `codex_home`, - /// keeping logs in a dedicated file to reduce lock contention with the - /// rest of the state store. + /// This opens (and migrates) the SQLite databases under `codex_home`. + /// Logs and paginated thread history live in dedicated files to reduce + /// lock contention with the rest of the state store. pub async fn init(codex_home: PathBuf, default_provider: String) -> anyhow::Result> { Self::init_inner( codex_home, @@ -198,10 +210,12 @@ impl StateRuntime { let logs_migrator = runtime_logs_migrator(); let goals_migrator = runtime_goals_migrator(); let memories_migrator = runtime_memories_migrator(); + let thread_history_migrator = runtime_thread_history_migrator(); let state_path = STATE_DB.path(codex_home.as_path()); let logs_path = LOGS_DB.path(codex_home.as_path()); let goals_path = GOALS_DB.path(codex_home.as_path()); let memories_path = MEMORIES_DB.path(codex_home.as_path()); + let thread_history_path = THREAD_HISTORY_DB.path(codex_home.as_path()); let pool = match open_state_sqlite(&state_path, &state_migrator, telemetry_override).await { Ok(db) => Arc::new(db), Err(err) => { @@ -244,6 +258,29 @@ impl StateRuntime { return Err(err); } }; + let thread_history_pool = match open_thread_history_sqlite( + &thread_history_path, + &thread_history_migrator, + telemetry_override, + ) + .await + { + Ok(db) => Arc::new(db), + Err(err) => { + warn!( + "failed to open thread history db at {}: {err}", + thread_history_path.display() + ); + close_sqlite_pools(&[ + pool.as_ref(), + logs_pool.as_ref(), + goals_pool.as_ref(), + memories_pool.as_ref(), + ]) + .await; + return Err(err); + } + }; let started = Instant::now(); let backfill_state_result = ensure_backfill_state_row_in_pool(pool.as_ref()).await; crate::telemetry::record_init_result( @@ -259,6 +296,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + thread_history_pool.as_ref(), ]) .await; return Err(err); @@ -287,6 +325,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + thread_history_pool.as_ref(), ]) .await; return Err(err); @@ -299,6 +338,7 @@ impl StateRuntime { memories: MemoryStore::new(Arc::clone(&memories_pool), Arc::clone(&pool)), pool, logs_pool, + thread_history_pool, codex_home, default_provider, thread_updated_at_millis: Arc::new(AtomicI64::new(thread_updated_at_millis)), @@ -330,6 +370,7 @@ impl StateRuntime { pub async fn close(&self) { self.memories.close().await; self.thread_goals.close().await; + self.thread_history_pool.close().await; self.logs_pool.close().await; self.pool.close().await; } @@ -404,6 +445,14 @@ async fn open_memories_sqlite( open_sqlite(path, migrator, MEMORIES_DB, telemetry_override).await } +async fn open_thread_history_sqlite( + path: &Path, + migrator: &Migrator, + telemetry_override: Option<&dyn DbTelemetry>, +) -> anyhow::Result { + open_sqlite(path, migrator, THREAD_HISTORY_DB, telemetry_override).await +} + async fn open_sqlite( path: &Path, migrator: &Migrator, @@ -508,6 +557,14 @@ pub fn memories_db_path(codex_home: &Path) -> PathBuf { MEMORIES_DB.path(codex_home) } +pub fn thread_history_db_filename() -> String { + THREAD_HISTORY_DB.filename.to_string() +} + +pub fn thread_history_db_path(codex_home: &Path) -> PathBuf { + THREAD_HISTORY_DB.path(codex_home) +} + pub fn runtime_db_paths(codex_home: &Path) -> Vec { RUNTIME_DBS .iter() @@ -726,6 +783,8 @@ mod tests { "migrate_goals", "open_memories", "migrate_memories", + "open_thread_history", + "migrate_thread_history", "ensure_backfill_state", "post_init_query", ] @@ -734,8 +793,7 @@ mod tests { .collect::>(); assert_eq!(phases, expected); - runtime.pool.close().await; - runtime.logs_pool.close().await; + runtime.close().await; let _ = tokio::fs::remove_dir_all(codex_home).await; } } diff --git a/codex-rs/state/src/telemetry.rs b/codex-rs/state/src/telemetry.rs index 1a9e4f9951..57a6c71314 100644 --- a/codex-rs/state/src/telemetry.rs +++ b/codex-rs/state/src/telemetry.rs @@ -41,6 +41,7 @@ pub(crate) enum DbKind { Logs, Goals, Memories, + ThreadHistory, } impl DbKind { @@ -50,6 +51,7 @@ impl DbKind { Self::Logs => "logs", Self::Goals => "goals", Self::Memories => "memories", + Self::ThreadHistory => "thread_history", } } } diff --git a/codex-rs/state/thread_history_migrations/0001_thread_history.sql b/codex-rs/state/thread_history_migrations/0001_thread_history.sql new file mode 100644 index 0000000000..cdf1279119 --- /dev/null +++ b/codex-rs/state/thread_history_migrations/0001_thread_history.sql @@ -0,0 +1,34 @@ +CREATE TABLE thread_turns ( + turn_id TEXT PRIMARY KEY, + thread_id TEXT NOT NULL, + turn_created_at_ms INTEGER NOT NULL, + metadata_json TEXT NOT NULL, + summary_first_user_item_key TEXT, + summary_final_agent_item_key TEXT +); + +CREATE INDEX idx_thread_turns_page + ON thread_turns(thread_id, turn_created_at_ms, turn_id); + +CREATE TABLE thread_items ( + thread_id TEXT NOT NULL, + turn_id TEXT NOT NULL, + item_key TEXT NOT NULL, + item_ordinal INTEGER NOT NULL, + item_created_at_ms INTEGER NOT NULL, + item_type TEXT NOT NULL, + materialized_thread_item_json TEXT NOT NULL, + PRIMARY KEY (thread_id, turn_id, item_key), + UNIQUE (thread_id, item_ordinal) +); + +CREATE INDEX idx_thread_items_page + ON thread_items(thread_id, item_ordinal); + +CREATE INDEX idx_thread_items_by_turn_page + ON thread_items(thread_id, turn_id, item_ordinal); + +CREATE TABLE thread_history_projection_state ( + thread_id TEXT PRIMARY KEY, + next_rollout_byte_offset INTEGER NOT NULL +);