openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-16 12:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: mem v2 - PR1

Browse Source
This commit is contained in:
jif-oai
2026-02-10 21:15:08 +00:00
parent 0a72832eb9
commit 9f902100a9
5 changed files with 322 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/core/src/memories/mod.rs
View File

@@ -29,6 +29,8 @@ const PHASE_TWO_CONCURRENCY_LIMIT: usize = MAX_ROLLOUTS_PER_STARTUP;
const MAX_RAW_MEMORIES_PER_SCOPE: usize = 64;
/// Maximum rollout age considered for phase-1 extraction.
const PHASE_ONE_MAX_ROLLOUT_AGE_DAYS: i64 = 30;
/// Minimum rollout idle time required before phase-1 extraction.
const PHASE_ONE_MIN_ROLLOUT_IDLE_HOURS: i64 = 12;
/// Lease duration (seconds) for phase-1 job ownership.
const PHASE_ONE_JOB_LEASE_SECONDS: i64 = 3_600;
/// Backoff delay (seconds) before retrying a failed stage-1 extraction job.

13
codex-rs/core/src/memories/startup/mod.rs
View File

@@ -156,11 +156,14 @@ pub(super) async fn run_memories_startup_pipeline(
let claimed_candidates = match state_db
.claim_stage1_jobs_for_startup(
session.conversation_id,
PHASE_ONE_THREAD_SCAN_LIMIT,
super::MAX_ROLLOUTS_PER_STARTUP,
super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
allowed_sources.as_slice(),
super::PHASE_ONE_JOB_LEASE_SECONDS,
codex_state::Stage1StartupClaimParams {
scan_limit: PHASE_ONE_THREAD_SCAN_LIMIT,
max_claimed: super::MAX_ROLLOUTS_PER_STARTUP,
max_age_days: super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
min_rollout_idle_hours: super::PHASE_ONE_MIN_ROLLOUT_IDLE_HOURS,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: super::PHASE_ONE_JOB_LEASE_SECONDS,
},
)
.await
{

1
codex-rs/state/src/lib.rs
View File

@@ -37,6 +37,7 @@ pub use runtime::STATE_DB_FILENAME;
pub use runtime::STATE_DB_VERSION;
pub use runtime::Stage1JobClaim;
pub use runtime::Stage1JobClaimOutcome;
pub use runtime::Stage1StartupClaimParams;
pub use runtime::state_db_filename;
pub use runtime::state_db_path;

221
codex-rs/state/src/runtime.rs
View File

@@ -77,6 +77,16 @@ pub struct Stage1JobClaim {
pub ownership_token: String,
}
#[derive(Debug, Clone, Copy)]
pub struct Stage1StartupClaimParams<'a> {
pub scan_limit: usize,
pub max_claimed: usize,
pub max_age_days: i64,
pub min_rollout_idle_hours: i64,
pub allowed_sources: &'a [String],
pub lease_seconds: i64,
}
/// Scope row used to queue phase-2 consolidation work.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PendingScopeConsolidation {
@@ -922,6 +932,7 @@ mod tests {
use super::STATE_DB_FILENAME;
use super::STATE_DB_VERSION;
use super::Stage1JobClaimOutcome;
use super::Stage1StartupClaimParams;
use super::StateRuntime;
use super::ThreadMetadata;
use super::state_db_filename;
@@ -1095,7 +1106,7 @@ mod tests {
let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
let claim = runtime
.try_claim_stage1_job(thread_id, owner_a, 100, 3600)
.try_claim_stage1_job(thread_id, owner_a, 100, 3600, 64)
.await
.expect("claim stage1 job");
let ownership_token = match claim {
@@ -1112,13 +1123,13 @@ mod tests {
);
let up_to_date = runtime
.try_claim_stage1_job(thread_id, owner_b, 100, 3600)
.try_claim_stage1_job(thread_id, owner_b, 100, 3600, 64)
.await
.expect("claim stage1 up-to-date");
assert_eq!(up_to_date, Stage1JobClaimOutcome::SkippedUpToDate);
let needs_rerun = runtime
.try_claim_stage1_job(thread_id, owner_b, 101, 3600)
.try_claim_stage1_job(thread_id, owner_b, 101, 3600, 64)
.await
.expect("claim stage1 newer source");
assert!(
@@ -1146,13 +1157,13 @@ mod tests {
.expect("upsert thread");
let claim_a = runtime
.try_claim_stage1_job(thread_id, owner_a, 100, 3600)
.try_claim_stage1_job(thread_id, owner_a, 100, 3600, 64)
.await
.expect("claim a");
assert!(matches!(claim_a, Stage1JobClaimOutcome::Claimed { .. }));
let claim_b_fresh = runtime
.try_claim_stage1_job(thread_id, owner_b, 100, 3600)
.try_claim_stage1_job(thread_id, owner_b, 100, 3600, 64)
.await
.expect("claim b fresh");
assert_eq!(claim_b_fresh, Stage1JobClaimOutcome::SkippedRunning);
@@ -1164,7 +1175,7 @@ mod tests {
.expect("force stale lease");
let claim_b_stale = runtime
.try_claim_stage1_job(thread_id, owner_b, 100, 3600)
.try_claim_stage1_job(thread_id, owner_b, 100, 3600, 64)
.await
.expect("claim b stale");
assert!(matches!(
@@ -1176,20 +1187,26 @@ mod tests {
}
#[tokio::test]
async fn claim_stage1_jobs_filters_by_age_and_current_thread() {
async fn claim_stage1_jobs_filters_by_age_idle_and_current_thread() {
let codex_home = unique_temp_dir();
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
.await
.expect("initialize runtime");
let now = Utc::now();
let recent_at = now - Duration::seconds(10);
let fresh_at = now - Duration::hours(1);
let just_under_idle_at = now - Duration::hours(12) + Duration::minutes(1);
let eligible_idle_at = now - Duration::hours(12) - Duration::minutes(1);
let old_at = now - Duration::days(31);
let current_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id");
let recent_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("recent thread id");
let fresh_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("fresh thread id");
let just_under_idle_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("just under idle thread id");
let eligible_idle_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("eligible idle thread id");
let old_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("old thread id");
@@ -1202,11 +1219,35 @@ mod tests {
.await
.expect("upsert current");
let mut recent =
test_thread_metadata(&codex_home, recent_thread_id, codex_home.join("recent"));
recent.created_at = recent_at;
recent.updated_at = recent_at;
runtime.upsert_thread(&recent).await.expect("upsert recent");
let mut fresh =
test_thread_metadata(&codex_home, fresh_thread_id, codex_home.join("fresh"));
fresh.created_at = fresh_at;
fresh.updated_at = fresh_at;
runtime.upsert_thread(&fresh).await.expect("upsert fresh");
let mut just_under_idle = test_thread_metadata(
&codex_home,
just_under_idle_thread_id,
codex_home.join("just-under-idle"),
);
just_under_idle.created_at = just_under_idle_at;
just_under_idle.updated_at = just_under_idle_at;
runtime
.upsert_thread(&just_under_idle)
.await
.expect("upsert just-under-idle");
let mut eligible_idle = test_thread_metadata(
&codex_home,
eligible_idle_thread_id,
codex_home.join("eligible-idle"),
);
eligible_idle.created_at = eligible_idle_at;
eligible_idle.updated_at = eligible_idle_at;
runtime
.upsert_thread(&eligible_idle)
.await
.expect("upsert eligible-idle");
let mut old = test_thread_metadata(&codex_home, old_thread_id, codex_home.join("old"));
old.created_at = old_at;
@@ -1217,17 +1258,147 @@ mod tests {
let claims = runtime
.claim_stage1_jobs_for_startup(
current_thread_id,
10,
5,
30,
allowed_sources.as_slice(),
3600,
Stage1StartupClaimParams {
scan_limit: 1,
max_claimed: 5,
max_age_days: 30,
min_rollout_idle_hours: 12,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: 3600,
},
)
.await
.expect("claim stage1 jobs");
assert_eq!(claims.len(), 1);
assert_eq!(claims[0].thread.id, recent_thread_id);
assert_eq!(claims[0].thread.id, eligible_idle_thread_id);
let _ = tokio::fs::remove_dir_all(codex_home).await;
}
#[tokio::test]
async fn claim_stage1_jobs_enforces_global_running_cap() {
let codex_home = unique_temp_dir();
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
.await
.expect("initialize runtime");
let current_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id");
runtime
.upsert_thread(&test_thread_metadata(
&codex_home,
current_thread_id,
codex_home.join("current"),
))
.await
.expect("upsert current");
let now = Utc::now();
let started_at = now.timestamp();
let lease_until = started_at + 3600;
let eligible_at = now - Duration::hours(13);
let existing_running = 10usize;
let total_candidates = 80usize;
for idx in 0..total_candidates {
let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id");
let mut metadata = test_thread_metadata(
&codex_home,
thread_id,
codex_home.join(format!("thread-{idx}")),
);
metadata.created_at = eligible_at - Duration::seconds(idx as i64);
metadata.updated_at = eligible_at - Duration::seconds(idx as i64);
runtime
.upsert_thread(&metadata)
.await
.expect("upsert thread");
if idx < existing_running {
sqlx::query(
r#"
INSERT INTO jobs (
kind,
job_key,
status,
worker_id,
ownership_token,
started_at,
finished_at,
lease_until,
retry_at,
retry_remaining,
last_error,
input_watermark,
last_success_watermark
) VALUES (?, ?, 'running', ?, ?, ?, NULL, ?, NULL, ?, NULL, ?, NULL)
"#,
)
.bind("memory_stage1")
.bind(thread_id.to_string())
.bind(current_thread_id.to_string())
.bind(Uuid::new_v4().to_string())
.bind(started_at)
.bind(lease_until)
.bind(3)
.bind(metadata.updated_at.timestamp())
.execute(runtime.pool.as_ref())
.await
.expect("seed running stage1 job");
}
}
let allowed_sources = vec!["cli".to_string()];
let claims = runtime
.claim_stage1_jobs_for_startup(
current_thread_id,
Stage1StartupClaimParams {
scan_limit: 200,
max_claimed: 64,
max_age_days: 30,
min_rollout_idle_hours: 12,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: 3600,
},
)
.await
.expect("claim stage1 jobs");
assert_eq!(claims.len(), 54);
let running_count = sqlx::query(
r#"
SELECT COUNT(*) AS count
FROM jobs
WHERE kind = 'memory_stage1'
AND status = 'running'
AND lease_until IS NOT NULL
AND lease_until > ?
"#,
)
.bind(Utc::now().timestamp())
.fetch_one(runtime.pool.as_ref())
.await
.expect("count running stage1 jobs")
.try_get::<i64, _>("count")
.expect("running count value");
assert_eq!(running_count, 64);
let more_claims = runtime
.claim_stage1_jobs_for_startup(
current_thread_id,
Stage1StartupClaimParams {
scan_limit: 200,
max_claimed: 64,
max_age_days: 30,
min_rollout_idle_hours: 12,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: 3600,
},
)
.await
.expect("claim stage1 jobs with cap reached");
assert_eq!(more_claims.len(), 0);
let _ = tokio::fs::remove_dir_all(codex_home).await;
}
@@ -1248,7 +1419,7 @@ mod tests {
.expect("upsert thread");
let claim = runtime
.try_claim_stage1_job(thread_id, owner, 100, 3600)
.try_claim_stage1_job(thread_id, owner, 100, 3600, 64)
.await
.expect("claim stage1");
let ownership_token = match claim {
@@ -1395,7 +1566,7 @@ mod tests {
.expect("upsert thread");
let claim = runtime
.try_claim_stage1_job(thread_id, owner, 100, 3600)
.try_claim_stage1_job(thread_id, owner, 100, 3600, 64)
.await
.expect("claim stage1");
let ownership_token = match claim {
@@ -1459,7 +1630,7 @@ mod tests {
.expect("upsert thread b");
let claim_a = runtime
.try_claim_stage1_job(thread_a, owner, 100, 3600)
.try_claim_stage1_job(thread_a, owner, 100, 3600, 64)
.await
.expect("claim stage1 a");
let token_a = match claim_a {
@@ -1475,7 +1646,7 @@ mod tests {
);
let claim_b = runtime
.try_claim_stage1_job(thread_b, owner, 101, 3600)
.try_claim_stage1_job(thread_b, owner, 101, 3600, 64)
.await
.expect("claim stage1 b");
let token_b = match claim_b {

159
codex-rs/state/src/runtime/memory.rs
View File

@@ -1,8 +1,10 @@
use super::*;
use crate::Stage1Output;
use crate::model::Stage1OutputRow;
use crate::model::ThreadRow;
use chrono::Duration;
use sqlx::Executor;
use sqlx::QueryBuilder;
use sqlx::Sqlite;
use std::collections::HashSet;
use std::path::Path;
@@ -38,47 +40,87 @@ impl StateRuntime {
pub async fn claim_stage1_jobs_for_startup(
&self,
current_thread_id: ThreadId,
scan_limit: usize,
max_claimed: usize,
max_age_days: i64,
allowed_sources: &[String],
lease_seconds: i64,
params: Stage1StartupClaimParams<'_>,
) -> anyhow::Result<Vec<Stage1JobClaim>> {
let Stage1StartupClaimParams {
scan_limit,
max_claimed,
max_age_days,
min_rollout_idle_hours,
allowed_sources,
lease_seconds,
} = params;
if scan_limit == 0 || max_claimed == 0 {
return Ok(Vec::new());
}
let page = self
.list_threads(
scan_limit,
None,
SortKey::UpdatedAt,
allowed_sources,
None,
false,
)
.await?;
let worker_id = current_thread_id;
let current_thread_id = worker_id.to_string();
let max_age_cutoff = (Utc::now() - Duration::days(max_age_days.max(0))).timestamp();
let idle_cutoff = (Utc::now() - Duration::hours(min_rollout_idle_hours.max(0))).timestamp();
let mut builder = QueryBuilder::<Sqlite>::new(
r#"
SELECT
id,
rollout_path,
created_at,
updated_at,
source,
model_provider,
cwd,
cli_version,
title,
sandbox_policy,
approval_mode,
tokens_used,
first_user_message,
archived_at,
git_sha,
git_branch,
git_origin_url
FROM threads
"#,
);
push_thread_filters(
&mut builder,
false,
allowed_sources,
None,
None,
SortKey::UpdatedAt,
);
builder
.push(" AND id != ")
.push_bind(current_thread_id.as_str());
builder
.push(" AND updated_at >= ")
.push_bind(max_age_cutoff);
builder.push(" AND updated_at <= ").push_bind(idle_cutoff);
push_thread_order_and_limit(&mut builder, SortKey::UpdatedAt, scan_limit);
let items = builder
.build()
.fetch_all(self.pool.as_ref())
.await?
.into_iter()
.map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from))
.collect::<Result<Vec<_>, _>>()?;
let cutoff = Utc::now() - Duration::days(max_age_days.max(0));
let mut claimed = Vec::new();
for item in page.items {
for item in items {
if claimed.len() >= max_claimed {
break;
}
if item.id == current_thread_id {
continue;
}
if item.updated_at < cutoff {
continue;
}
if let Stage1JobClaimOutcome::Claimed { ownership_token } = self
.try_claim_stage1_job(
item.id,
current_thread_id,
worker_id,
item.updated_at.timestamp(),
lease_seconds,
max_claimed,
)
.await?
{
@@ -202,9 +244,11 @@ LIMIT ?
worker_id: ThreadId,
source_updated_at: i64,
lease_seconds: i64,
max_running_jobs: usize,
) -> anyhow::Result<Stage1JobClaimOutcome> {
let now = Utc::now().timestamp();
let lease_until = now.saturating_add(lease_seconds.max(0));
let max_running_jobs = max_running_jobs as i64;
let ownership_token = Uuid::new_v4().to_string();
let thread_id = thread_id.to_string();
let worker_id = worker_id.to_string();
@@ -241,7 +285,53 @@ WHERE kind = ? AND job_key = ?
.fetch_optional(&mut *tx)
.await?;
let Some(existing_job) = existing_job else {
let should_insert = if let Some(existing_job) = existing_job {
let status: String = existing_job.try_get("status")?;
let existing_lease_until: Option<i64> = existing_job.try_get("lease_until")?;
let retry_at: Option<i64> = existing_job.try_get("retry_at")?;
let retry_remaining: i64 = existing_job.try_get("retry_remaining")?;
if retry_remaining <= 0 {
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRetryExhausted);
}
if retry_at.is_some_and(|retry_at| retry_at > now) {
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRetryBackoff);
}
if status == "running"
&& existing_lease_until.is_some_and(|lease_until| lease_until > now)
{
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRunning);
}
false
} else {
true
};
let fresh_running_jobs = sqlx::query(
r#"
SELECT COUNT(*) AS count
FROM jobs
WHERE kind = ?
AND status = 'running'
AND lease_until IS NOT NULL
AND lease_until > ?
"#,
)
.bind(JOB_KIND_MEMORY_STAGE1)
.bind(now)
.fetch_one(&mut *tx)
.await?
.try_get::<i64, _>("count")?;
if fresh_running_jobs >= max_running_jobs {
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRunning);
}
if should_insert {
sqlx::query(
r#"
INSERT INTO jobs (
@@ -273,25 +363,6 @@ INSERT INTO jobs (
.await?;
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::Claimed { ownership_token });
};
let status: String = existing_job.try_get("status")?;
let existing_lease_until: Option<i64> = existing_job.try_get("lease_until")?;
let retry_at: Option<i64> = existing_job.try_get("retry_at")?;
let retry_remaining: i64 = existing_job.try_get("retry_remaining")?;
if retry_remaining <= 0 {
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRetryExhausted);
}
if retry_at.is_some_and(|retry_at| retry_at > now) {
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRetryBackoff);
}
if status == "running" && existing_lease_until.is_some_and(|lease_until| lease_until > now)
{
tx.commit().await?;
return Ok(Stage1JobClaimOutcome::SkippedRunning);
}
let rows_affected = sqlx::query(