diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs index 71ea14afa..7867f9064 100644 --- a/codex-rs/core/src/memories/mod.rs +++ b/codex-rs/core/src/memories/mod.rs @@ -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. diff --git a/codex-rs/core/src/memories/startup/mod.rs b/codex-rs/core/src/memories/startup/mod.rs index df908bac9..443641196 100644 --- a/codex-rs/core/src/memories/startup/mod.rs +++ b/codex-rs/core/src/memories/startup/mod.rs @@ -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 { diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index a521d9997..4179cf095 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -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; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index e29b5aafb..d8424ccbd 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -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::("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 { diff --git a/codex-rs/state/src/runtime/memory.rs b/codex-rs/state/src/runtime/memory.rs index 0f706d926..1906944f8 100644 --- a/codex-rs/state/src/runtime/memory.rs +++ b/codex-rs/state/src/runtime/memory.rs @@ -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> { + 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::::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::, _>>()?; - 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 { 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 = existing_job.try_get("lease_until")?; + let retry_at: Option = 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::("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 = existing_job.try_get("lease_until")?; - let retry_at: Option = 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(