Move memory state to a dedicated SQLite DB (#24591)

## Summary

Generated memory rows and their stage-one/stage-two job state currently
live in `state_5.sqlite` alongside thread metadata. That makes memory
cleanup and regeneration share the main state schema even though those
rows are memory-pipeline data and can be rebuilt independently from the
durable thread records.

This PR moves the memory-owned tables into a dedicated
`memories_1.sqlite` runtime database while keeping thread metadata in
`state_5.sqlite`.

## Changes

- Adds a separate memories DB runtime, migrator, path helpers, telemetry
kind, and Bazel compile data for `state/memory_migrations`.
- Introduces `MemoryStore` behind `StateRuntime::memories()` and moves
memory table/job operations onto that store.
- Drops the old memory tables from the state DB and recreates their
schema in `state/memory_migrations/0001_memories.sql`.
- Updates memory startup, citation usage tracking, rollout pollution
handling, `debug clear-memories`, and app-server `memory/reset` to
operate through the memories DB.
- Preserves cross-DB behavior by hydrating thread metadata from the
state DB when selecting visible memory outputs and checking stage-one
staleness.

## Verification

- Added/updated `codex-state` tests for deleted-thread memory visibility
and already-polluted phase-two enqueue behavior.
- Updated `debug clear-memories`, app-server `memory/reset`, and
memories startup tests to seed and assert memory rows through
`memories_1.sqlite`.
This commit is contained in:
jif-oai authored and GitHub committed 2026-05-26 20:07:25 +02:00
1 parent 823381e867
commit aad59a0916
22 files changed
+1014 -280

No files matched your search

@@ -1513,9 +1513,13 @@ impl ThreadRequestProcessor {
.clone() .clone()
.ok_or_else(|| internal_error("sqlite state db unavailable for memory reset"))?; .ok_or_else(|| internal_error("sqlite state db unavailable for memory reset"))?;
state_db.clear_memory_data().await.map_err(|err| { state_db
internal_error(format!("failed to clear memory rows in state db: {err}")) .memories()
})?; .clear_memory_data()
.await
.map_err(|err| {
internal_error(format!("failed to clear memory rows in memories db: {err}"))
})?;
clear_memory_roots_contents(&self.config.codex_home) clear_memory_roots_contents(&self.config.codex_home)
.await .await
@@ -49,7 +49,10 @@ async fn memory_reset_clears_memory_files_and_rows_preserves_threads() -> Result
.await??; .await??;
let _: MemoryResetResponse = to_response::<MemoryResetResponse>(response)?; let _: MemoryResetResponse = to_response::<MemoryResetResponse>(response)?;
let stage1_outputs = state_db.list_stage1_outputs_for_global(/*n*/ 10).await?; let stage1_outputs = state_db
.memories()
.list_stage1_outputs_for_global(/*n*/ 10)
.await?;
assert_eq!(stage1_outputs, Vec::new()); assert_eq!(stage1_outputs, Vec::new());
assert_eq!( assert_eq!(
state_db.get_thread_memory_mode(thread_id).await?.as_deref(), state_db.get_thread_memory_mode(thread_id).await?.as_deref(),
@@ -81,6 +84,7 @@ async fn seed_stage1_output(state_db: &Arc<StateRuntime>, codex_home: &Path) ->
state_db.upsert_thread(&metadata).await?; state_db.upsert_thread(&metadata).await?;
let claim = state_db let claim = state_db
.memories()
.try_claim_stage1_job( .try_claim_stage1_job(
thread_id, thread_id,
worker_id, worker_id,
@@ -94,6 +98,7 @@ async fn seed_stage1_output(state_db: &Arc<StateRuntime>, codex_home: &Path) ->
}; };
assert!( assert!(
state_db state_db
.memories()
.mark_stage1_job_succeeded( .mark_stage1_job_succeeded(
thread_id, thread_id,
ownership_token.as_str(), ownership_token.as_str(),
@@ -106,6 +111,7 @@ async fn seed_stage1_output(state_db: &Arc<StateRuntime>, codex_home: &Path) ->
"stage1 success should be recorded" "stage1 success should be recorded"
); );
state_db state_db
.memories()
.enqueue_global_consolidation(now.timestamp()) .enqueue_global_consolidation(now.timestamp())
.await?; .await?;
+32
View File
@@ -702,6 +702,7 @@ fn state_summary(check: &DoctorCheck) -> String {
"state DB integrity", "state DB integrity",
"log DB integrity", "log DB integrity",
"goals DB integrity", "goals DB integrity",
"memories DB integrity",
] ]
.into_iter() .into_iter()
.all(|label| detail::detail_value(check, label).is_some_and(|value| value == "ok")); .all(|label| detail::detail_value(check, label).is_some_and(|value| value == "ok"));
@@ -1363,6 +1364,37 @@ Run codex doctor without --summary for detailed diagnostics.
); );
} }
#[test]
fn render_human_report_includes_memories_db_in_state_health_summary() {
let report = DoctorReport {
schema_version: 1,
generated_at: "0s since unix epoch".to_string(),
overall_status: CheckStatus::Ok,
codex_version: "0.0.0".to_string(),
checks: vec![
DoctorCheck::new(
"state.paths",
"state",
CheckStatus::Ok,
"state paths inspectable",
)
.detail("state DB: /tmp/state.sqlite")
.detail("state DB integrity: ok")
.detail("log DB: /tmp/logs.sqlite")
.detail("log DB integrity: ok")
.detail("goals DB: /tmp/goals.sqlite")
.detail("goals DB integrity: ok")
.detail("memories DB: /tmp/memories.sqlite")
.detail("memories DB integrity: ok"),
],
};
let rendered = render_human_report(&report, detailed_no_color_unicode_options());
assert!(rendered.contains("✓ state databases healthy"));
assert!(rendered.contains("memories DB /tmp/memories.sqlite · integrity ok"));
}
#[test] #[test]
fn render_human_report_supports_ascii_output() { fn render_human_report_supports_ascii_output() {
let rendered = render_human_report( let rendered = render_human_report(
+3
View File
@@ -413,6 +413,7 @@ fn state_details(parsed: &[ParsedDetail]) -> Vec<HumanDetail> {
push_database_row(&mut out, parsed, "state DB"); push_database_row(&mut out, parsed, "state DB");
push_database_row(&mut out, parsed, "log DB"); push_database_row(&mut out, parsed, "log DB");
push_database_row(&mut out, parsed, "goals DB"); push_database_row(&mut out, parsed, "goals DB");
push_database_row(&mut out, parsed, "memories DB");
for (source, label) in [ for (source, label) in [
("active rollout files", "active rollouts"), ("active rollout files", "active rollouts"),
@@ -440,6 +441,8 @@ fn state_details(parsed: &[ParsedDetail]) -> Vec<HumanDetail> {
"state DB integrity", "state DB integrity",
"log DB integrity", "log DB integrity",
"goals DB integrity", "goals DB integrity",
"memories DB",
"memories DB integrity",
"active rollout files", "active rollout files",
"archived rollout files", "archived rollout files",
], ],
+7 -13
View File
@@ -27,7 +27,7 @@ use codex_responses_api_proxy::Args as ResponsesApiProxyArgs;
use codex_rollout_trace::REDUCED_STATE_FILE_NAME; use codex_rollout_trace::REDUCED_STATE_FILE_NAME;
use codex_rollout_trace::replay_bundle; use codex_rollout_trace::replay_bundle;
use codex_state::StateRuntime; use codex_state::StateRuntime;
use codex_state::state_db_path; use codex_state::memories_db_path;
use codex_tui::AppExitInfo; use codex_tui::AppExitInfo;
use codex_tui::Cli as TuiCli; use codex_tui::Cli as TuiCli;
use codex_tui::ExitReason; use codex_tui::ExitReason;
@@ -1751,22 +1751,16 @@ async fn run_debug_clear_memories_command(
.build() .build()
.await?; .await?;
let state_path = state_db_path(config.sqlite_home.as_path()); let memories_path = memories_db_path(config.sqlite_home.as_path());
let mut cleared_state_db = false; let cleared_memories_db =
if tokio::fs::try_exists(&state_path).await? { StateRuntime::clear_memory_data_in_sqlite_home(config.sqlite_home.as_path()).await?;
let state_db =
StateRuntime::init(config.sqlite_home.clone(), config.model_provider_id.clone())
.await?;
state_db.clear_memory_data().await?;
cleared_state_db = true;
}
clear_memory_roots_contents(&config.codex_home).await?; clear_memory_roots_contents(&config.codex_home).await?;
let mut message = if cleared_state_db { let mut message = if cleared_memories_db {
format!("Cleared memory state from {}.", state_path.display()) format!("Cleared memory state from {}.", memories_path.display())
} else { } else {
format!("No state db found at {}.", state_path.display()) format!("No memories db found at {}.", memories_path.display())
}; };
message.push_str(&format!( message.push_str(&format!(
" Cleared memory directories under {}.", " Cleared memory directories under {}.",
+59 -3
View File
@@ -2,6 +2,7 @@ use std::path::Path;
use anyhow::Result; use anyhow::Result;
use codex_state::StateRuntime; use codex_state::StateRuntime;
use codex_state::memories_db_path;
use codex_state::state_db_path; use codex_state::state_db_path;
use predicates::str::contains; use predicates::str::contains;
use sqlx::SqlitePool; use sqlx::SqlitePool;
@@ -23,6 +24,9 @@ async fn debug_clear_memories_resets_state_and_removes_memory_dir() -> Result<()
let thread_id = "00000000-0000-0000-0000-000000000123"; let thread_id = "00000000-0000-0000-0000-000000000123";
let db_path = state_db_path(codex_home.path()); let db_path = state_db_path(codex_home.path());
let pool = SqlitePool::connect(&format!("sqlite://{}", db_path.display())).await?; let pool = SqlitePool::connect(&format!("sqlite://{}", db_path.display())).await?;
let memories_db_path = memories_db_path(codex_home.path());
let memories_pool =
SqlitePool::connect(&format!("sqlite://{}", memories_db_path.display())).await?;
sqlx::query( sqlx::query(
r#" r#"
@@ -74,7 +78,7 @@ INSERT INTO stage1_outputs (
"#, "#,
) )
.bind(thread_id) .bind(thread_id)
.execute(&pool) .execute(&memories_pool)
.await?; .await?;
sqlx::query( sqlx::query(
@@ -99,13 +103,14 @@ INSERT INTO jobs (
"#, "#,
) )
.bind(thread_id) .bind(thread_id)
.execute(&pool) .execute(&memories_pool)
.await?; .await?;
let memory_root = codex_home.path().join("memories"); let memory_root = codex_home.path().join("memories");
std::fs::create_dir_all(&memory_root)?; std::fs::create_dir_all(&memory_root)?;
std::fs::write(memory_root.join("memory_summary.md"), "stale memory")?; std::fs::write(memory_root.join("memory_summary.md"), "stale memory")?;
pool.close().await; pool.close().await;
memories_pool.close().await;
let mut cmd = codex_command(codex_home.path())?; let mut cmd = codex_command(codex_home.path())?;
cmd.args(["debug", "clear-memories"]) cmd.args(["debug", "clear-memories"])
@@ -113,7 +118,7 @@ INSERT INTO jobs (
.success() .success()
.stdout(contains("Cleared memory state")); .stdout(contains("Cleared memory state"));
let pool = SqlitePool::connect(&format!("sqlite://{}", db_path.display())).await?; let pool = SqlitePool::connect(&format!("sqlite://{}", memories_db_path.display())).await?;
let stage1_outputs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs") let stage1_outputs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs")
.fetch_one(&pool) .fetch_one(&pool)
.await?; .await?;
@@ -131,3 +136,54 @@ INSERT INTO jobs (
Ok(()) Ok(())
} }
#[tokio::test]
async fn debug_clear_memories_resets_memories_db_without_state_db() -> Result<()> {
let codex_home = TempDir::new()?;
let runtime =
StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()).await?;
drop(runtime);
let db_path = state_db_path(codex_home.path());
let memories_db_path = memories_db_path(codex_home.path());
let memories_pool =
SqlitePool::connect(&format!("sqlite://{}", memories_db_path.display())).await?;
sqlx::query(
r#"
INSERT INTO stage1_outputs (
thread_id,
source_updated_at,
raw_memory,
rollout_summary,
generated_at,
rollout_slug,
usage_count,
last_usage,
selected_for_phase2,
selected_for_phase2_source_updated_at
) VALUES ('00000000-0000-0000-0000-000000000123', 1, 'raw', 'summary', 1, NULL, 0, NULL, 0, NULL)
"#,
)
.execute(&memories_pool)
.await?;
memories_pool.close().await;
std::fs::remove_file(&db_path)?;
let mut cmd = codex_command(codex_home.path())?;
cmd.args(["debug", "clear-memories"])
.assert()
.success()
.stdout(contains("Cleared memory state"));
let pool = SqlitePool::connect(&format!("sqlite://{}", memories_db_path.display())).await?;
let stage1_outputs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs")
.fetch_one(&pool)
.await?;
assert_eq!(stage1_outputs_count, 0);
pool.close().await;
assert!(!db_path.exists());
Ok(())
}
+1 -1
View File
@@ -233,7 +233,7 @@ async fn record_stage1_output_usage_for_memory_citation(
} }
if let Some(db) = state_db_ctx { if let Some(db) = state_db_ctx {
let _ = db.record_stage1_output_usage(&thread_ids).await; let _ = db.memories().record_stage1_output_usage(&thread_ids).await;
} }
true true
} }
+9 -2
View File
@@ -112,6 +112,7 @@ pub async fn prune(context: &MemoryStartupContext, config: &Config) {
if let Some(db) = context.state_db() { if let Some(db) = context.state_db() {
let max_unused_days = config.memories.max_unused_days; let max_unused_days = config.memories.max_unused_days;
match db match db
.memories()
.prune_stage1_outputs_for_retention(max_unused_days, crate::stage_one::PRUNE_BATCH_SIZE) .prune_stage1_outputs_for_retention(max_unused_days, crate::stage_one::PRUNE_BATCH_SIZE)
.await .await
{ {
@@ -124,7 +125,7 @@ pub async fn prune(context: &MemoryStartupContext, config: &Config) {
} }
Err(err) => { Err(err) => {
warn!( warn!(
"state db prune_stage1_outputs_for_retention failed during memories startup: {err}" "memories db prune_stage1_outputs_for_retention failed during memories startup: {err}"
); );
} }
} }
@@ -161,6 +162,7 @@ async fn claim_startup_jobs(
.collect::<Vec<_>>(); .collect::<Vec<_>>();
match state_db match state_db
.memories()
.claim_stage1_jobs_for_startup( .claim_stage1_jobs_for_startup(
context.thread_id(), context.thread_id(),
codex_state::Stage1StartupClaimParams { codex_state::Stage1StartupClaimParams {
@@ -176,7 +178,9 @@ async fn claim_startup_jobs(
{ {
Ok(claims) => Some(claims), Ok(claims) => Some(claims),
Err(err) => { Err(err) => {
warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}"); warn!(
"memories db claim_stage1_jobs_for_startup failed during memories startup: {err}"
);
None None
} }
} }
@@ -329,6 +333,7 @@ mod job {
tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}"); tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}");
if let Some(state_db) = context.state_db() { if let Some(state_db) = context.state_db() {
let _ = state_db let _ = state_db
.memories()
.mark_stage1_job_failed( .mark_stage1_job_failed(
thread_id, thread_id,
ownership_token, ownership_token,
@@ -349,6 +354,7 @@ mod job {
}; };
if state_db if state_db
.memories()
.mark_stage1_job_succeeded_no_output(thread_id, ownership_token) .mark_stage1_job_succeeded_no_output(thread_id, ownership_token)
.await .await
.unwrap_or(false) .unwrap_or(false)
@@ -373,6 +379,7 @@ mod job {
}; };
if state_db if state_db
.memories()
.mark_stage1_job_succeeded( .mark_stage1_job_succeeded(
thread_id, thread_id,
ownership_token, ownership_token,
+14 -7
View File
@@ -91,6 +91,7 @@ pub async fn run(context: Arc<MemoryStartupContext>, config: Arc<Config>) {
// 4. Load current DB-backed Phase 2 inputs. // 4. Load current DB-backed Phase 2 inputs.
let raw_memories = match db let raw_memories = match db
.memories()
.get_phase2_input_selection(max_raw_memories, max_unused_days) .get_phase2_input_selection(max_raw_memories, max_unused_days)
.await .await
{ {
@@ -217,6 +218,7 @@ mod job {
db: &StateRuntime, db: &StateRuntime,
) -> Result<Claim, &'static str> { ) -> Result<Claim, &'static str> {
let claim = db let claim = db
.memories()
.try_claim_global_phase2_job(context.thread_id(), crate::stage_two::JOB_LEASE_SECONDS) .try_claim_global_phase2_job(context.thread_id(), crate::stage_two::JOB_LEASE_SECONDS)
.await .await
.map_err(|e| { .map_err(|e| {
@@ -255,15 +257,17 @@ mod job {
) { ) {
context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]);
if matches!( if matches!(
db.mark_global_phase2_job_failed( db.memories()
&claim.token, .mark_global_phase2_job_failed(
reason, &claim.token,
crate::stage_two::JOB_RETRY_DELAY_SECONDS, reason,
) crate::stage_two::JOB_RETRY_DELAY_SECONDS,
.await, )
.await,
Ok(false) Ok(false)
) { ) {
let _ = db let _ = db
.memories()
.mark_global_phase2_job_failed_if_unowned( .mark_global_phase2_job_failed_if_unowned(
&claim.token, &claim.token,
reason, reason,
@@ -282,7 +286,8 @@ mod job {
reason: &'static str, reason: &'static str,
) -> bool { ) -> bool {
context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]);
db.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs) db.memories()
.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs)
.await .await
.unwrap_or(false) .unwrap_or(false)
} }
@@ -382,6 +387,7 @@ mod agent {
} }
// Do not reset the workspace baseline if we lost the lock. // Do not reset the workspace baseline if we lost the lock.
let still_owns_lock = match db let still_owns_lock = match db
.memories()
.heartbeat_global_phase2_job( .heartbeat_global_phase2_job(
&claim.token, &claim.token,
crate::stage_two::JOB_LEASE_SECONDS, crate::stage_two::JOB_LEASE_SECONDS,
@@ -479,6 +485,7 @@ mod agent {
} }
_ = heartbeat_interval.tick() => { _ = heartbeat_interval.tick() => {
match db match db
.memories()
.heartbeat_global_phase2_job( .heartbeat_global_phase2_job(
&token, &token,
crate::stage_two::JOB_LEASE_SECONDS, crate::stage_two::JOB_LEASE_SECONDS,
+13 -10
View File
@@ -190,7 +190,8 @@ async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_i
let server = start_mock_server().await; let server = start_mock_server().await;
let home = Arc::new(TempDir::new()?); let home = Arc::new(TempDir::new()?);
let db = init_state_db(&home).await?; let db = init_state_db(&home).await?;
db.enqueue_global_consolidation(/*input_watermark*/ 1) db.memories()
.enqueue_global_consolidation(/*input_watermark*/ 1)
.await?; .await?;
let now = chrono::Utc::now(); let now = chrono::Utc::now();
@@ -445,6 +446,7 @@ async fn seed_stage1_output_for_existing_thread(
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
let owner = ThreadId::new(); let owner = ThreadId::new();
let claim = db let claim = db
.memories()
.try_claim_stage1_job( .try_claim_stage1_job(
thread_id, owner, updated_at, /*lease_seconds*/ 3_600, thread_id, owner, updated_at, /*lease_seconds*/ 3_600,
/*max_running_jobs*/ 64, /*max_running_jobs*/ 64,
@@ -456,15 +458,16 @@ async fn seed_stage1_output_for_existing_thread(
}; };
assert!( assert!(
db.mark_stage1_job_succeeded( db.memories()
thread_id, .mark_stage1_job_succeeded(
&ownership_token, thread_id,
updated_at, &ownership_token,
raw_memory, updated_at,
rollout_summary, raw_memory,
rollout_slug, rollout_summary,
) rollout_slug,
.await?, )
.await?,
"stage-1 success should enqueue global consolidation" "stage-1 success should enqueue global consolidation"
); );
+6 -2
View File
@@ -493,8 +493,12 @@ pub async fn mark_thread_memory_mode_polluted(
let Some(ctx) = context else { let Some(ctx) = context else {
return; return;
}; };
if let Err(err) = ctx.mark_thread_memory_mode_polluted(thread_id).await { if let Err(err) = ctx
warn!("state db mark_thread_memory_mode_polluted failed during {stage}: {err}"); .memories()
.mark_thread_memory_mode_polluted(thread_id)
.await
{
warn!("memories db mark_thread_memory_mode_polluted failed during {stage}: {err}");
} }
} }
+6 -1
View File
@@ -3,5 +3,10 @@ load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate( codex_rust_crate(
name = "state", name = "state",
crate_name = "codex_state", crate_name = "codex_state",
compile_data = glob(["goals_migrations/**", "logs_migrations/**", "migrations/**"]), compile_data = glob([
"goals_migrations/**",
"logs_migrations/**",
"memory_migrations/**",
"migrations/**",
]),
) )
@@ -0,0 +1,35 @@
CREATE TABLE stage1_outputs (
thread_id TEXT PRIMARY KEY,
source_updated_at INTEGER NOT NULL,
raw_memory TEXT NOT NULL,
rollout_summary TEXT NOT NULL,
rollout_slug TEXT,
generated_at INTEGER NOT NULL,
usage_count INTEGER,
last_usage INTEGER,
selected_for_phase2 INTEGER NOT NULL DEFAULT 0,
selected_for_phase2_source_updated_at INTEGER
);
CREATE INDEX idx_stage1_outputs_source_updated_at
ON stage1_outputs(source_updated_at DESC, thread_id DESC);
CREATE TABLE jobs (
kind TEXT NOT NULL,
job_key TEXT NOT NULL,
status TEXT NOT NULL,
worker_id TEXT,
ownership_token TEXT,
started_at INTEGER,
finished_at INTEGER,
lease_until INTEGER,
retry_at INTEGER,
retry_remaining INTEGER NOT NULL,
last_error TEXT,
input_watermark INTEGER,
last_success_watermark INTEGER,
PRIMARY KEY (kind, job_key)
);
CREATE INDEX idx_jobs_kind_status_retry_lease
ON jobs(kind, status, retry_at, lease_until);
@@ -0,0 +1,2 @@
DROP TABLE IF EXISTS jobs;
DROP TABLE IF EXISTS stage1_outputs;
+4
View File
@@ -55,6 +55,7 @@ pub use runtime::GoalAccountingMode;
pub use runtime::GoalAccountingOutcome; pub use runtime::GoalAccountingOutcome;
pub use runtime::GoalStore; pub use runtime::GoalStore;
pub use runtime::GoalUpdate; pub use runtime::GoalUpdate;
pub use runtime::MemoryStore;
pub use runtime::RemoteControlEnrollmentRecord; pub use runtime::RemoteControlEnrollmentRecord;
pub use runtime::RuntimeDbPath; pub use runtime::RuntimeDbPath;
pub use runtime::ThreadFilterOptions; pub use runtime::ThreadFilterOptions;
@@ -62,6 +63,8 @@ pub use runtime::goals_db_filename;
pub use runtime::goals_db_path; pub use runtime::goals_db_path;
pub use runtime::logs_db_filename; pub use runtime::logs_db_filename;
pub use runtime::logs_db_path; pub use runtime::logs_db_path;
pub use runtime::memories_db_filename;
pub use runtime::memories_db_path;
pub use runtime::runtime_db_paths; pub use runtime::runtime_db_paths;
pub use runtime::sqlite_integrity_check; pub use runtime::sqlite_integrity_check;
pub use runtime::state_db_filename; pub use runtime::state_db_filename;
@@ -77,6 +80,7 @@ pub const SQLITE_HOME_ENV: &str = "CODEX_SQLITE_HOME";
pub const LOGS_DB_FILENAME: &str = "logs_2.sqlite"; pub const LOGS_DB_FILENAME: &str = "logs_2.sqlite";
pub const GOALS_DB_FILENAME: &str = "goals_1.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 STATE_DB_FILENAME: &str = "state_5.sqlite";
/// Errors encountered during DB operations. Tags: [stage] /// Errors encountered during DB operations. Tags: [stage]
+5
View File
@@ -5,6 +5,7 @@ use sqlx::migrate::Migrator;
pub(crate) static STATE_MIGRATOR: Migrator = sqlx::migrate!("./migrations"); pub(crate) static STATE_MIGRATOR: Migrator = sqlx::migrate!("./migrations");
pub(crate) static LOGS_MIGRATOR: Migrator = sqlx::migrate!("./logs_migrations"); pub(crate) static LOGS_MIGRATOR: Migrator = sqlx::migrate!("./logs_migrations");
pub(crate) static GOALS_MIGRATOR: Migrator = sqlx::migrate!("./goals_migrations"); pub(crate) static GOALS_MIGRATOR: Migrator = sqlx::migrate!("./goals_migrations");
pub(crate) static MEMORIES_MIGRATOR: Migrator = sqlx::migrate!("./memory_migrations");
/// Allow an older Codex binary to open a database that has already been /// Allow an older Codex binary to open a database that has already been
/// migrated by a newer binary running in parallel. /// migrated by a newer binary running in parallel.
@@ -32,3 +33,7 @@ pub(crate) fn runtime_logs_migrator() -> Migrator {
pub(crate) fn runtime_goals_migrator() -> Migrator { pub(crate) fn runtime_goals_migrator() -> Migrator {
runtime_migrator(&GOALS_MIGRATOR) runtime_migrator(&GOALS_MIGRATOR)
} }
pub(crate) fn runtime_memories_migrator() -> Migrator {
runtime_migrator(&MEMORIES_MIGRATOR)
}
-55
View File
@@ -1,9 +1,6 @@
use anyhow::Result;
use chrono::DateTime; use chrono::DateTime;
use chrono::Utc; use chrono::Utc;
use codex_protocol::ThreadId; use codex_protocol::ThreadId;
use sqlx::Row;
use sqlx::sqlite::SqliteRow;
use std::path::PathBuf; use std::path::PathBuf;
use super::ThreadMetadata; use super::ThreadMetadata;
@@ -22,58 +19,6 @@ pub struct Stage1Output {
pub generated_at: DateTime<Utc>, pub generated_at: DateTime<Utc>,
} }
#[derive(Debug)]
pub(crate) struct Stage1OutputRow {
thread_id: String,
rollout_path: String,
source_updated_at: i64,
raw_memory: String,
rollout_summary: String,
rollout_slug: Option<String>,
cwd: String,
git_branch: Option<String>,
generated_at: i64,
}
impl Stage1OutputRow {
pub(crate) fn try_from_row(row: &SqliteRow) -> Result<Self> {
Ok(Self {
thread_id: row.try_get("thread_id")?,
rollout_path: row.try_get("rollout_path")?,
source_updated_at: row.try_get("source_updated_at")?,
raw_memory: row.try_get("raw_memory")?,
rollout_summary: row.try_get("rollout_summary")?,
rollout_slug: row.try_get("rollout_slug")?,
cwd: row.try_get("cwd")?,
git_branch: row.try_get("git_branch")?,
generated_at: row.try_get("generated_at")?,
})
}
}
impl TryFrom<Stage1OutputRow> for Stage1Output {
type Error = anyhow::Error;
fn try_from(row: Stage1OutputRow) -> std::result::Result<Self, Self::Error> {
Ok(Self {
thread_id: ThreadId::try_from(row.thread_id)?,
rollout_path: PathBuf::from(row.rollout_path),
source_updated_at: epoch_seconds_to_datetime(row.source_updated_at)?,
raw_memory: row.raw_memory,
rollout_summary: row.rollout_summary,
rollout_slug: row.rollout_slug,
cwd: PathBuf::from(row.cwd),
git_branch: row.git_branch,
generated_at: epoch_seconds_to_datetime(row.generated_at)?,
})
}
}
fn epoch_seconds_to_datetime(secs: i64) -> Result<DateTime<Utc>> {
DateTime::<Utc>::from_timestamp(secs, 0)
.ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}"))
}
/// Result of trying to claim a stage-1 memory extraction job. /// Result of trying to claim a stage-1 memory extraction job.
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
pub enum Stage1JobClaimOutcome { pub enum Stage1JobClaimOutcome {
-1
View File
@@ -37,7 +37,6 @@ pub use thread_metadata::ThreadsPage;
pub(crate) use agent_job::AgentJobItemRow; pub(crate) use agent_job::AgentJobItemRow;
pub(crate) use agent_job::AgentJobRow; pub(crate) use agent_job::AgentJobRow;
pub(crate) use memories::Stage1OutputRow;
pub(crate) use thread_goal::ThreadGoalRow; pub(crate) use thread_goal::ThreadGoalRow;
pub(crate) use thread_metadata::ThreadRow; pub(crate) use thread_metadata::ThreadRow;
pub(crate) use thread_metadata::anchor_from_item; pub(crate) use thread_metadata::anchor_from_item;
+72 -1
View File
@@ -10,6 +10,7 @@ use crate::LOGS_DB_FILENAME;
use crate::LogEntry; use crate::LogEntry;
use crate::LogQuery; use crate::LogQuery;
use crate::LogRow; use crate::LogRow;
use crate::MEMORIES_DB_FILENAME;
use crate::STATE_DB_FILENAME; use crate::STATE_DB_FILENAME;
use crate::SortKey; use crate::SortKey;
use crate::ThreadMetadata; use crate::ThreadMetadata;
@@ -18,6 +19,7 @@ use crate::ThreadsPage;
use crate::apply_rollout_item; use crate::apply_rollout_item;
use crate::migrations::runtime_goals_migrator; use crate::migrations::runtime_goals_migrator;
use crate::migrations::runtime_logs_migrator; use crate::migrations::runtime_logs_migrator;
use crate::migrations::runtime_memories_migrator;
use crate::migrations::runtime_state_migrator; use crate::migrations::runtime_state_migrator;
use crate::model::AgentJobRow; use crate::model::AgentJobRow;
use crate::model::ThreadRow; use crate::model::ThreadRow;
@@ -70,6 +72,7 @@ pub use goals::GoalAccountingMode;
pub use goals::GoalAccountingOutcome; pub use goals::GoalAccountingOutcome;
pub use goals::GoalStore; pub use goals::GoalStore;
pub use goals::GoalUpdate; pub use goals::GoalUpdate;
pub use memories::MemoryStore;
pub use remote_control::RemoteControlEnrollmentRecord; pub use remote_control::RemoteControlEnrollmentRecord;
pub use threads::ThreadFilterOptions; pub use threads::ThreadFilterOptions;
@@ -121,7 +124,15 @@ const GOALS_DB: RuntimeDbSpec = RuntimeDbSpec {
migrate_phase: "migrate_goals", migrate_phase: "migrate_goals",
}; };
const RUNTIME_DBS: [RuntimeDbSpec; 3] = [STATE_DB, LOGS_DB, GOALS_DB]; const MEMORIES_DB: RuntimeDbSpec = RuntimeDbSpec {
label: "memories DB",
filename: MEMORIES_DB_FILENAME,
kind: DbKind::Memories,
open_phase: "open_memories",
migrate_phase: "migrate_memories",
};
const RUNTIME_DBS: [RuntimeDbSpec; 4] = [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB];
#[derive(Clone, Debug, Eq, PartialEq)] #[derive(Clone, Debug, Eq, PartialEq)]
pub struct RuntimeDbPath { pub struct RuntimeDbPath {
@@ -136,6 +147,7 @@ pub struct StateRuntime {
pool: Arc<sqlx::SqlitePool>, pool: Arc<sqlx::SqlitePool>,
logs_pool: Arc<sqlx::SqlitePool>, logs_pool: Arc<sqlx::SqlitePool>,
thread_goals: GoalStore, thread_goals: GoalStore,
memories: MemoryStore,
thread_updated_at_millis: Arc<AtomicI64>, thread_updated_at_millis: Arc<AtomicI64>,
} }
@@ -172,9 +184,11 @@ impl StateRuntime {
let state_migrator = runtime_state_migrator(); let state_migrator = runtime_state_migrator();
let logs_migrator = runtime_logs_migrator(); let logs_migrator = runtime_logs_migrator();
let goals_migrator = runtime_goals_migrator(); let goals_migrator = runtime_goals_migrator();
let memories_migrator = runtime_memories_migrator();
let state_path = STATE_DB.path(codex_home.as_path()); let state_path = STATE_DB.path(codex_home.as_path());
let logs_path = LOGS_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 goals_path = GOALS_DB.path(codex_home.as_path());
let memories_path = MEMORIES_DB.path(codex_home.as_path());
let pool = match open_state_sqlite(&state_path, &state_migrator, telemetry_override).await { let pool = match open_state_sqlite(&state_path, &state_migrator, telemetry_override).await {
Ok(db) => Arc::new(db), Ok(db) => Arc::new(db),
Err(err) => { Err(err) => {
@@ -198,6 +212,22 @@ impl StateRuntime {
return Err(err); return Err(err);
} }
}; };
let memories_pool = match open_memories_sqlite(
&memories_path,
&memories_migrator,
telemetry_override,
)
.await
{
Ok(db) => Arc::new(db),
Err(err) => {
warn!(
"failed to open memories db at {}: {err}",
memories_path.display()
);
return Err(err);
}
};
let started = Instant::now(); let started = Instant::now();
let backfill_state_result = ensure_backfill_state_row_in_pool(pool.as_ref()).await; let backfill_state_result = ensure_backfill_state_row_in_pool(pool.as_ref()).await;
crate::telemetry::record_init_result( crate::telemetry::record_init_result(
@@ -225,6 +255,7 @@ impl StateRuntime {
let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0); let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0);
let runtime = Arc::new(Self { let runtime = Arc::new(Self {
thread_goals: GoalStore::new(Arc::clone(&goals_pool)), thread_goals: GoalStore::new(Arc::clone(&goals_pool)),
memories: MemoryStore::new(Arc::clone(&memories_pool), Arc::clone(&pool)),
pool, pool,
logs_pool, logs_pool,
codex_home, codex_home,
@@ -248,6 +279,28 @@ impl StateRuntime {
pub fn thread_goals(&self) -> &GoalStore { pub fn thread_goals(&self) -> &GoalStore {
&self.thread_goals &self.thread_goals
} }
pub fn memories(&self) -> &MemoryStore {
&self.memories
}
pub async fn clear_memory_data_in_sqlite_home(sqlite_home: &Path) -> anyhow::Result<bool> {
let memories_path = MEMORIES_DB.path(sqlite_home);
if !tokio::fs::try_exists(&memories_path).await? {
return Ok(false);
}
let memories_migrator = runtime_memories_migrator();
let pool = open_memories_sqlite(
&memories_path,
&memories_migrator,
/*telemetry_override*/ None,
)
.await?;
memories::clear_memory_data_in_pool(&pool).await?;
pool.close().await;
Ok(true)
}
} }
fn base_sqlite_options(path: &Path) -> SqliteConnectOptions { fn base_sqlite_options(path: &Path) -> SqliteConnectOptions {
@@ -287,6 +340,14 @@ async fn open_goals_sqlite(
open_sqlite(path, migrator, GOALS_DB, telemetry_override).await open_sqlite(path, migrator, GOALS_DB, telemetry_override).await
} }
async fn open_memories_sqlite(
path: &Path,
migrator: &Migrator,
telemetry_override: Option<&dyn DbTelemetry>,
) -> anyhow::Result<SqlitePool> {
open_sqlite(path, migrator, MEMORIES_DB, telemetry_override).await
}
async fn open_sqlite( async fn open_sqlite(
path: &Path, path: &Path,
migrator: &Migrator, migrator: &Migrator,
@@ -363,6 +424,14 @@ pub fn goals_db_path(codex_home: &Path) -> PathBuf {
GOALS_DB.path(codex_home) GOALS_DB.path(codex_home)
} }
pub fn memories_db_filename() -> String {
MEMORIES_DB.filename.to_string()
}
pub fn memories_db_path(codex_home: &Path) -> PathBuf {
MEMORIES_DB.path(codex_home)
}
pub fn runtime_db_paths(codex_home: &Path) -> Vec<RuntimeDbPath> { pub fn runtime_db_paths(codex_home: &Path) -> Vec<RuntimeDbPath> {
RUNTIME_DBS RUNTIME_DBS
.iter() .iter()
@@ -579,6 +648,8 @@ mod tests {
"migrate_logs", "migrate_logs",
"open_goals", "open_goals",
"migrate_goals", "migrate_goals",
"open_memories",
"migrate_memories",
"ensure_backfill_state", "ensure_backfill_state",
"post_init_query", "post_init_query",
] ]
File diff suppressed because it is too large. Load diff
+1
View File
@@ -971,6 +971,7 @@ ON CONFLICT(thread_id, position) DO NOTHING
.execute(self.pool.as_ref()) .execute(self.pool.as_ref())
.await?; .await?;
let rows_affected = result.rows_affected(); let rows_affected = result.rows_affected();
self.memories.delete_thread_memory(thread_id).await?;
if rows_affected > 0 { if rows_affected > 0 {
self.thread_goals.delete_thread_goal(thread_id).await?; self.thread_goals.delete_thread_goal(thread_id).await?;
} }
+2
View File
@@ -40,6 +40,7 @@ pub(crate) enum DbKind {
State, State,
Logs, Logs,
Goals, Goals,
Memories,
} }
impl DbKind { impl DbKind {
@@ -48,6 +49,7 @@ impl DbKind {
Self::State => "state", Self::State => "state",
Self::Logs => "logs", Self::Logs => "logs",
Self::Goals => "goals", Self::Goals => "goals",
Self::Memories => "memories",
} }
} }
} }