use super::threads::ThreadFilterOptions;use super::threads::push_thread_filters;use super::*;use crate::SortDirection;use crate::model::Phase2JobClaimOutcome;use crate::model::Stage1JobClaim;use crate::model::Stage1JobClaimOutcome;use crate::model::Stage1Output;use crate::model::Stage1StartupClaimParams;use crate::model::ThreadRow;use chrono::DateTime;use chrono::Duration;use sqlx::Executor;use sqlx::QueryBuilder;use sqlx::Sqlite;use uuid::Uuid;const JOB_KIND_MEMORY_STAGE1: &str = "memory_stage1";const JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL: &str = "memory_consolidate_global";const MEMORY_CONSOLIDATION_JOB_KEY: &str = "global";const PHASE2_SUCCESS_COOLDOWN_SECONDS: i64 = 6 * 60 * 60;const PHASE2_INPUT_SELECTION_PAGE_SIZE: usize = 512;const DEFAULT_RETRY_REMAINING: i64 = 3;/// Store for generated memory state and memory extraction/consolidation jobs.#[derive(Clone)]pub struct MemoryStore { pool: Arc<SqlitePool>, state_pool: Arc<SqlitePool>,}impl MemoryStore { pub(crate) fn new(pool: Arc<SqlitePool>, state_pool: Arc<SqlitePool>) -> Self { Self { pool, state_pool } } pub(crate) async fn close(&self) { self.pool.close().await; } /// Deletes all persisted memory state in one transaction. /// /// This removes every `stage1_outputs` row and all `jobs` rows for the /// stage-1 (`memory_stage1`) and phase-2 (`memory_consolidate_global`) /// memory pipelines. pub async fn clear_memory_data(&self) -> anyhow::Result<()> { clear_memory_data_in_pool(self.pool.as_ref()).await } /// Record usage for cited stage-1 outputs. /// /// Each thread id increments `usage_count` by one and sets `last_usage` to /// the current Unix timestamp. Missing rows are ignored. pub async fn record_stage1_output_usage( &self, thread_ids: &[ThreadId], ) -> anyhow::Result<usize> { if thread_ids.is_empty() { return Ok(0); } let now = Utc::now().timestamp(); let mut tx = self.pool.begin().await?; let mut updated_rows = 0; for thread_id in thread_ids { updated_rows += sqlx::query( r#"UPDATE stage1_outputsSET usage_count = COALESCE(usage_count, 0) + 1, last_usage = ?WHERE thread_id = ? "#, ) .bind(now) .bind(thread_id.to_string()) .execute(&mut *tx) .await? .rows_affected() as usize; } tx.commit().await?; Ok(updated_rows) } async fn stage1_source_needs_update( &self, thread_id: ThreadId, source_updated_at: i64, ) -> anyhow::Result<bool> { let thread_id = thread_id.to_string(); let existing_output = sqlx::query( r#"SELECT source_updated_atFROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .fetch_optional(self.pool.as_ref()) .await?; if let Some(existing_output) = existing_output { let existing_source_updated_at: i64 = existing_output.try_get("source_updated_at")?; if existing_source_updated_at >= source_updated_at { return Ok(false); } } let existing_job = sqlx::query( r#"SELECT last_success_watermarkFROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .fetch_optional(self.pool.as_ref()) .await?; if let Some(existing_job) = existing_job { let last_success_watermark = existing_job.try_get::<Option<i64>, _>("last_success_watermark")?; if last_success_watermark.is_some_and(|watermark| watermark >= source_updated_at) { return Ok(false); } } Ok(true) } /// Selects and claims stage-1 startup jobs for stale threads. /// /// Query behavior: /// - starts from `threads` filtered to active threads and allowed sources /// (`push_thread_filters`) /// - excludes threads with `memory_mode != 'enabled'` /// - excludes the current thread id /// - keeps only threads whose millisecond `updated_at` is in the age window /// - checks memory staleness against the memories DB /// - orders by `updated_at_ms DESC` and applies `scan_limit` to bound /// state-DB work before probing the memories DB /// /// For each selected thread, this function calls [`Self::try_claim_stage1_job`] /// with `source_updated_at = thread.updated_at.timestamp()` and returns up to /// `max_claimed` successful claims. pub async fn claim_stage1_jobs_for_startup( &self, current_thread_id: ThreadId, 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 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_millis(); let idle_cutoff = (Utc::now() - Duration::hours(min_rollout_idle_hours.max(0))).timestamp_millis(); let mut builder = QueryBuilder::<Sqlite>::new( r#"SELECT threads.id, threads.rollout_path, threads.created_at_ms AS created_at, threads.updated_at_ms AS updated_at, threads.source, threads.thread_source, threads.agent_path, threads.agent_nickname, threads.agent_role, threads.model_provider, threads.model, threads.reasoning_effort, threads.cwd, threads.cli_version, threads.title, threads.preview, threads.sandbox_policy, threads.approval_mode, threads.tokens_used, threads.first_user_message, threads.archived_at, threads.git_sha, threads.git_branch, threads.git_origin_urlFROM threads "#, ); push_thread_filters( &mut builder, ThreadFilterOptions { archived_only: false, allowed_sources, model_providers: None, cwd_filters: None, anchor: None, sort_key: SortKey::UpdatedAt, sort_direction: SortDirection::Desc, search_term: None, }, ); builder.push(" AND threads.memory_mode = 'enabled'"); builder .push(" AND threads.id != ") .push_bind(current_thread_id.as_str()); builder .push(" AND ") .push("threads.updated_at_ms") .push(" >= ") .push_bind(max_age_cutoff); builder .push(" AND ") .push("threads.updated_at_ms") .push(" <= ") .push_bind(idle_cutoff); let scan_limit_i64 = i64::try_from(scan_limit).unwrap_or(i64::MAX); builder.push(" ORDER BY threads.updated_at_ms DESC LIMIT "); builder.push_bind(scan_limit_i64); let items = builder .build() .fetch_all(self.state_pool.as_ref()) .await? .into_iter() .map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) .collect::<Result<Vec<_>, _>>()?; let mut claimed = Vec::new(); for item in items { if claimed.len() >= max_claimed { break; } if !self .stage1_source_needs_update(item.id, item.updated_at.timestamp()) .await? { continue; } if let Stage1JobClaimOutcome::Claimed { ownership_token } = self .try_claim_stage1_job( item.id, worker_id, item.updated_at.timestamp(), lease_seconds, max_claimed, ) .await? { claimed.push(Stage1JobClaim { thread: item, ownership_token, }); } } Ok(claimed) } pub(super) async fn delete_thread_memory(&self, thread_id: ThreadId) -> anyhow::Result<()> { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); let mut tx = self.pool.begin().await?; let existing_output = sqlx::query( r#"SELECT selected_for_phase2FROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .fetch_optional(&mut *tx) .await?; let was_selected_for_phase2 = existing_output .map(|row| row.try_get::<i64, _>("selected_for_phase2")) .transpose()? .is_some_and(|selected| selected != 0); let deleted_rows = sqlx::query( r#"DELETE FROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .execute(&mut *tx) .await? .rows_affected(); sqlx::query( r#"DELETE FROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .execute(&mut *tx) .await?; if deleted_rows > 0 && was_selected_for_phase2 { enqueue_global_consolidation_with_executor(&mut *tx, now).await?; } tx.commit().await?; Ok(()) } /// Lists the most recent non-empty stage-1 outputs for global consolidation. /// /// Query behavior: /// - filters out rows where both `raw_memory` and `rollout_summary` are blank /// - hydrates thread `cwd`, `rollout_path`, and `git_branch` from the state DB /// - filters out missing or non-enabled threads /// - orders by `source_updated_at DESC, thread_id DESC` /// - returns the first `n` visible outputs pub async fn list_stage1_outputs_for_global( &self, n: usize, ) -> anyhow::Result<Vec<Stage1Output>> { if n == 0 { return Ok(Vec::new()); } let rows = sqlx::query( r#"SELECT so.thread_id, so.source_updated_at, so.raw_memory, so.rollout_summary, so.rollout_slug, so.generated_atFROM stage1_outputs AS soWHERE length(trim(so.raw_memory)) > 0 OR length(trim(so.rollout_summary)) > 0ORDER BY so.source_updated_at DESC, so.thread_id DESC "#, ) .fetch_all(self.pool.as_ref()) .await?; let mut outputs = Vec::new(); for row in rows { if let Some(output) = self.stage1_output_from_row_if_thread_enabled(&row).await? { outputs.push(output); if outputs.len() >= n { break; } } } Ok(outputs) } /// Prunes stale stage-1 outputs while preserving the latest phase-2 /// baseline and stage-1 job watermarks. /// /// Query behavior: /// - considers only rows with `selected_for_phase2 = 0` /// - keeps recency as `COALESCE(last_usage, source_updated_at)` /// - removes rows older than `max_unused_days` /// - prunes at most `limit` rows ordered from stalest to newest pub async fn prune_stage1_outputs_for_retention( &self, max_unused_days: i64, limit: usize, ) -> anyhow::Result<usize> { if limit == 0 { return Ok(0); } let cutoff = (Utc::now() - Duration::days(max_unused_days.max(0))).timestamp(); let rows_affected = sqlx::query( r#"DELETE FROM stage1_outputsWHERE thread_id IN ( SELECT thread_id FROM stage1_outputs WHERE selected_for_phase2 = 0 AND COALESCE(last_usage, source_updated_at) < ? ORDER BY COALESCE(last_usage, source_updated_at) ASC, source_updated_at ASC, thread_id ASC LIMIT ?) "#, ) .bind(cutoff) .bind(limit as i64) .execute(self.pool.as_ref()) .await? .rows_affected(); Ok(rows_affected as usize) } /// Returns the current phase-2 input set. /// /// Query behavior: /// - current selection keeps only non-empty stage-1 outputs whose /// `last_usage` is within `max_unused_days`, or whose /// `source_updated_at` is within that window when the memory has never /// been used /// - eligible rows are ranked by `usage_count DESC`, /// `COALESCE(last_usage, source_updated_at) DESC`, `source_updated_at DESC`, /// `thread_id DESC` /// - the selected top-N rows are returned in stable `thread_id ASC` order /// /// The returned rows are the complete Phase 2 filesystem input. Phase 2 /// syncs these rows directly; deletions are represented by the workspace /// diff against the previous successful memory baseline. pub async fn get_phase2_input_selection( &self, n: usize, max_unused_days: i64, ) -> anyhow::Result<Vec<Stage1Output>> { if n == 0 { return Ok(Vec::new()); } let cutoff = (Utc::now() - Duration::days(max_unused_days.max(0))).timestamp(); let page_size = n.clamp(1, PHASE2_INPUT_SELECTION_PAGE_SIZE); let page_size_i64 = i64::try_from(page_size).unwrap_or(i64::MAX); let mut offset = 0_i64; let mut selected_keys = Vec::with_capacity(n); while selected_keys.len() < n { let candidate_rows = sqlx::query( r#"SELECT so.thread_id, so.source_updated_atFROM stage1_outputs AS soWHERE (length(trim(so.raw_memory)) > 0 OR length(trim(so.rollout_summary)) > 0) AND ( (so.last_usage IS NOT NULL AND so.last_usage >= ?) OR (so.last_usage IS NULL AND so.source_updated_at >= ?) )ORDER BY COALESCE(so.usage_count, 0) DESC, COALESCE(so.last_usage, so.source_updated_at) DESC, so.source_updated_at DESC, so.thread_id DESCLIMIT ? OFFSET ? "#, ) .bind(cutoff) .bind(cutoff) .bind(page_size_i64) .bind(offset) .fetch_all(self.pool.as_ref()) .await?; if candidate_rows.is_empty() { break; } let candidate_count = i64::try_from(candidate_rows.len()).unwrap_or(i64::MAX); for row in candidate_rows { let thread_id: String = row.try_get("thread_id")?; let source_updated_at: i64 = row.try_get("source_updated_at")?; if self .enabled_thread_metadata(ThreadId::try_from(thread_id.as_str())?) .await? .is_some() { selected_keys.push((thread_id, source_updated_at)); if selected_keys.len() >= n { break; } } } offset = offset.saturating_add(candidate_count); } let mut selected = Vec::with_capacity(selected_keys.len()); for (thread_id, source_updated_at) in selected_keys { let Some(row) = sqlx::query( r#"SELECT so.thread_id, so.source_updated_at, so.raw_memory, so.rollout_summary, so.rollout_slug, so.generated_atFROM stage1_outputs AS soWHERE so.thread_id = ? AND so.source_updated_at = ? "#, ) .bind(thread_id.as_str()) .bind(source_updated_at) .fetch_optional(self.pool.as_ref()) .await? else { continue; }; if let Some(output) = self.stage1_output_from_row_if_thread_enabled(&row).await? { selected.push(output); } } selected.sort_by_key(|entry| entry.thread_id.to_string()); Ok(selected) } async fn stage1_output_from_row_if_thread_enabled( &self, row: &sqlx::sqlite::SqliteRow, ) -> anyhow::Result<Option<Stage1Output>> { let thread_id: String = row.try_get("thread_id")?; let Some(thread) = self .enabled_thread_metadata(ThreadId::try_from(thread_id.as_str())?) .await? else { return Ok(None); }; Ok(Some(stage1_output_from_row_and_thread(row, thread)?)) } async fn enabled_thread_metadata( &self, thread_id: ThreadId, ) -> anyhow::Result<Option<ThreadMetadata>> { let row = sqlx::query( r#"SELECT threads.id, threads.rollout_path, threads.created_at_ms AS created_at, threads.updated_at_ms AS updated_at, threads.source, threads.thread_source, threads.agent_nickname, threads.agent_role, threads.agent_path, threads.model_provider, threads.model, threads.reasoning_effort, threads.cwd, threads.cli_version, threads.title, threads.preview, threads.sandbox_policy, threads.approval_mode, threads.tokens_used, threads.first_user_message, threads.archived_at, threads.git_sha, threads.git_branch, threads.git_origin_urlFROM threadsWHERE threads.id = ? AND threads.memory_mode = 'enabled' "#, ) .bind(thread_id.to_string()) .fetch_optional(self.state_pool.as_ref()) .await?; row.map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) .transpose() } /// Marks a thread as polluted and enqueues phase-2 forgetting when the /// thread participated in the last successful phase-2 baseline. pub async fn mark_thread_memory_mode_polluted( &self, thread_id: ThreadId, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); let selected_for_phase2 = sqlx::query_scalar::<_, i64>( r#"SELECT selected_for_phase2FROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .fetch_optional(self.pool.as_ref()) .await? .unwrap_or(0); let rows_affected = sqlx::query( r#"UPDATE threadsSET memory_mode = 'polluted'WHERE id = ? AND memory_mode != 'polluted' "#, ) .bind(thread_id.as_str()) .execute(self.state_pool.as_ref()) .await? .rows_affected(); if selected_for_phase2 != 0 { self.enqueue_global_consolidation(now).await?; } Ok(rows_affected > 0) } /// Attempts to claim a stage-1 job for a thread at `source_updated_at`. /// /// Claim semantics: /// - skips as up-to-date when either: /// - `stage1_outputs.source_updated_at >= source_updated_at`, or /// - `jobs.last_success_watermark >= source_updated_at` /// - inserts or updates a `jobs` row to `running` only when: /// - global running job count for `memory_stage1` is below `max_running_jobs` /// - existing row is not actively running with a valid lease /// - retry backoff (if present) has elapsed, or `source_updated_at` advanced /// - retries remain, or `source_updated_at` advanced (which resets retries) /// /// The update path refreshes ownership token, lease, and `input_watermark`. /// If claiming fails, a follow-up read maps current row state to a precise /// skip outcome (`SkippedRunning`, `SkippedRetryBackoff`, or /// `SkippedRetryExhausted`). pub async fn try_claim_stage1_job( &self, thread_id: ThreadId, 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(); let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; let existing_output = sqlx::query( r#"SELECT source_updated_atFROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .fetch_optional(&mut *tx) .await?; if let Some(existing_output) = existing_output { let existing_source_updated_at: i64 = existing_output.try_get("source_updated_at")?; if existing_source_updated_at >= source_updated_at { tx.commit().await?; return Ok(Stage1JobClaimOutcome::SkippedUpToDate); } } let existing_job = sqlx::query( r#"SELECT last_success_watermarkFROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .fetch_optional(&mut *tx) .await?; if let Some(existing_job) = existing_job { let last_success_watermark = existing_job.try_get::<Option<i64>, _>("last_success_watermark")?; if last_success_watermark.is_some_and(|watermark| watermark >= source_updated_at) { tx.commit().await?; return Ok(Stage1JobClaimOutcome::SkippedUpToDate); } } let rows_affected = 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)SELECT ?, ?, 'running', ?, ?, ?, NULL, ?, NULL, ?, NULL, ?, NULLWHERE ( SELECT COUNT(*) FROM jobs WHERE kind = ? AND status = 'running' AND lease_until IS NOT NULL AND lease_until > ?) < ?ON CONFLICT(kind, job_key) DO UPDATE SET status = 'running', worker_id = excluded.worker_id, ownership_token = excluded.ownership_token, started_at = excluded.started_at, finished_at = NULL, lease_until = excluded.lease_until, retry_at = NULL, retry_remaining = CASE WHEN excluded.input_watermark > COALESCE(jobs.input_watermark, -1) THEN ? ELSE jobs.retry_remaining END, last_error = NULL, input_watermark = excluded.input_watermarkWHERE (jobs.status != 'running' OR jobs.lease_until IS NULL OR jobs.lease_until <= excluded.started_at) AND ( jobs.retry_at IS NULL OR jobs.retry_at <= excluded.started_at OR excluded.input_watermark > COALESCE(jobs.input_watermark, -1) ) AND ( jobs.retry_remaining > 0 OR excluded.input_watermark > COALESCE(jobs.input_watermark, -1) ) AND ( SELECT COUNT(*) FROM jobs AS running_jobs WHERE running_jobs.kind = excluded.kind AND running_jobs.status = 'running' AND running_jobs.lease_until IS NOT NULL AND running_jobs.lease_until > excluded.started_at AND running_jobs.job_key != excluded.job_key ) < ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(worker_id.as_str()) .bind(ownership_token.as_str()) .bind(now) .bind(lease_until) .bind(DEFAULT_RETRY_REMAINING) .bind(source_updated_at) .bind(JOB_KIND_MEMORY_STAGE1) .bind(now) .bind(max_running_jobs) .bind(DEFAULT_RETRY_REMAINING) .bind(max_running_jobs) .execute(&mut *tx) .await? .rows_affected(); if rows_affected > 0 { tx.commit().await?; return Ok(Stage1JobClaimOutcome::Claimed { ownership_token }); } let existing_job = sqlx::query( r#"SELECT status, lease_until, retry_at, retry_remainingFROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .fetch_optional(&mut *tx) .await?; tx.commit().await?; 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 { return Ok(Stage1JobClaimOutcome::SkippedRetryExhausted); } if retry_at.is_some_and(|retry_at| retry_at > now) { return Ok(Stage1JobClaimOutcome::SkippedRetryBackoff); } if status == "running" && existing_lease_until.is_some_and(|lease_until| lease_until > now) { return Ok(Stage1JobClaimOutcome::SkippedRunning); } } Ok(Stage1JobClaimOutcome::SkippedRunning) } /// Marks a claimed stage-1 job successful and upserts generated output. /// /// Transaction behavior: /// - updates `jobs` only for the currently owned running row /// - sets `status='done'` and `last_success_watermark = input_watermark` /// - upserts `stage1_outputs` for the thread, replacing existing output only /// when `source_updated_at` is newer or equal /// - preserves any existing `selected_for_phase2` baseline until the next /// successful phase-2 run rewrites the baseline selection, including the /// snapshot timestamp chosen during that run /// - persists optional `rollout_slug` for rollout summary artifact naming /// - enqueues/advances the global phase-2 job watermark using /// `source_updated_at` pub async fn mark_stage1_job_succeeded( &self, thread_id: ThreadId, ownership_token: &str, source_updated_at: i64, raw_memory: &str, rollout_summary: &str, rollout_slug: Option<&str>, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); let mut tx = self.pool.begin().await?; let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'done', finished_at = ?, lease_until = NULL, last_error = NULL, last_success_watermark = input_watermarkWHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(now) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) .execute(&mut *tx) .await? .rows_affected(); if rows_affected == 0 { tx.commit().await?; return Ok(false); } sqlx::query( r#"INSERT INTO stage1_outputs ( thread_id, source_updated_at, raw_memory, rollout_summary, rollout_slug, generated_at) VALUES (?, ?, ?, ?, ?, ?)ON CONFLICT(thread_id) DO UPDATE SET source_updated_at = excluded.source_updated_at, raw_memory = excluded.raw_memory, rollout_summary = excluded.rollout_summary, rollout_slug = excluded.rollout_slug, generated_at = excluded.generated_atWHERE excluded.source_updated_at >= stage1_outputs.source_updated_at "#, ) .bind(thread_id.as_str()) .bind(source_updated_at) .bind(raw_memory) .bind(rollout_summary) .bind(rollout_slug) .bind(now) .execute(&mut *tx) .await?; enqueue_global_consolidation_with_executor(&mut *tx, source_updated_at).await?; tx.commit().await?; Ok(true) } /// Marks a claimed stage-1 job successful when extraction produced no output. /// /// Transaction behavior: /// - updates `jobs` only for the currently owned running row /// - sets `status='done'` and `last_success_watermark = input_watermark` /// - deletes any existing `stage1_outputs` row for the thread /// - enqueues/advances the global phase-2 job watermark using the claimed /// `input_watermark` only when deleting an existing `stage1_outputs` row pub async fn mark_stage1_job_succeeded_no_output( &self, thread_id: ThreadId, ownership_token: &str, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); let mut tx = self.pool.begin().await?; let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'done', finished_at = ?, lease_until = NULL, last_error = NULL, last_success_watermark = input_watermarkWHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(now) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) .execute(&mut *tx) .await? .rows_affected(); if rows_affected == 0 { tx.commit().await?; return Ok(false); } let source_updated_at = sqlx::query( r#"SELECT input_watermarkFROM jobsWHERE kind = ? AND job_key = ? AND ownership_token = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) .fetch_one(&mut *tx) .await? .try_get::<i64, _>("input_watermark")?; let deleted_rows = sqlx::query( r#"DELETE FROM stage1_outputsWHERE thread_id = ? "#, ) .bind(thread_id.as_str()) .execute(&mut *tx) .await? .rows_affected(); if deleted_rows > 0 { enqueue_global_consolidation_with_executor(&mut *tx, source_updated_at).await?; } tx.commit().await?; Ok(true) } /// Marks a claimed stage-1 job as failed and schedules retry backoff. /// /// Query behavior: /// - updates only the owned running row for `(kind='memory_stage1', job_key)` /// - sets `status='error'`, clears lease, writes `last_error` /// - decrements `retry_remaining` /// - sets `retry_at = now + retry_delay_seconds` pub async fn mark_stage1_job_failed( &self, thread_id: ThreadId, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let retry_at = now.saturating_add(retry_delay_seconds.max(0)); let thread_id = thread_id.to_string(); let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?, retry_remaining = retry_remaining - 1, last_error = ?WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(now) .bind(retry_at) .bind(failure_reason) .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) .execute(self.pool.as_ref()) .await? .rows_affected(); Ok(rows_affected > 0) } /// Enqueues or advances the global phase-2 consolidation job watermark. /// /// The underlying upsert keeps the job `running` when already running, resets /// `pending/error` jobs to `pending`, and advances `input_watermark` as /// bookkeeping even when `source_updated_at` is older than prior maxima. /// Phase 2 does not use this watermark as a dirty check; git workspace diffing /// decides whether consolidation work exists after the lock is claimed. pub async fn enqueue_global_consolidation(&self, input_watermark: i64) -> anyhow::Result<()> { enqueue_global_consolidation_with_executor(self.pool.as_ref(), input_watermark).await } /// Attempts to claim the global phase-2 consolidation lock. /// /// Claim semantics: /// - reads the singleton global job row (`kind='memory_consolidate_global'`) /// - creates and claims the singleton row when it does not exist yet /// - does not use DB watermarks to decide whether Phase 2 has work; git workspace /// dirtiness is the source of truth after the caller materializes inputs /// - returns `SkippedRetryUnavailable` when retry backoff is active /// - returns `SkippedRunning` when an active running lease exists /// - returns `SkippedCooldown` when the latest successful run finished /// within the phase-2 success cooldown /// - otherwise updates the row to `running`, sets ownership + lease, and /// returns `Claimed` pub async fn try_claim_global_phase2_job( &self, worker_id: ThreadId, lease_seconds: i64, ) -> anyhow::Result<Phase2JobClaimOutcome> { let now = Utc::now().timestamp(); let lease_until = now.saturating_add(lease_seconds.max(0)); let cooldown_cutoff = now.saturating_sub(PHASE2_SUCCESS_COOLDOWN_SECONDS); let ownership_token = Uuid::new_v4().to_string(); let worker_id = worker_id.to_string(); let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; let existing_job = sqlx::query( r#"SELECT status, lease_until, retry_at, input_watermark, finished_at, last_errorFROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .fetch_optional(&mut *tx) .await?; let Some(existing_job) = existing_job else { let rows_affected = 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, 0, 0) "#, ) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(worker_id.as_str()) .bind(ownership_token.as_str()) .bind(now) .bind(lease_until) .bind(DEFAULT_RETRY_REMAINING) .execute(&mut *tx) .await? .rows_affected(); tx.commit().await?; return if rows_affected == 0 { Ok(Phase2JobClaimOutcome::SkippedRunning) } else { Ok(Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark: 0, }) }; }; let input_watermark: Option<i64> = existing_job.try_get("input_watermark")?; let input_watermark_value = input_watermark.unwrap_or(0); 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 finished_at: Option<i64> = existing_job.try_get("finished_at")?; let last_error: Option<String> = existing_job.try_get("last_error")?; if retry_at.is_some_and(|retry_at| retry_at > now) { tx.commit().await?; return Ok(Phase2JobClaimOutcome::SkippedRetryUnavailable); } if status == "running" && existing_lease_until.is_some_and(|lease_until| lease_until > now) { tx.commit().await?; return Ok(Phase2JobClaimOutcome::SkippedRunning); } if last_error.is_none() && finished_at.is_some_and(|finished_at| finished_at > cooldown_cutoff) { tx.commit().await?; return Ok(Phase2JobClaimOutcome::SkippedCooldown); } let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'running', worker_id = ?, ownership_token = ?, started_at = ?, finished_at = NULL, lease_until = ?, retry_at = NULL, last_error = NULLWHERE kind = ? AND job_key = ? AND (status != 'running' OR lease_until IS NULL OR lease_until <= ?) AND (retry_at IS NULL OR retry_at <= ?) AND (last_error IS NOT NULL OR finished_at IS NULL OR finished_at <= ?) "#, ) .bind(worker_id.as_str()) .bind(ownership_token.as_str()) .bind(now) .bind(lease_until) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(now) .bind(now) .bind(cooldown_cutoff) .execute(&mut *tx) .await? .rows_affected(); tx.commit().await?; if rows_affected == 0 { Ok(Phase2JobClaimOutcome::SkippedRunning) } else { Ok(Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark: input_watermark_value, }) } } /// Extends the lease for an owned running phase-2 global job. /// /// Query behavior: /// - `UPDATE jobs SET lease_until = ?` for the singleton global row /// - requires `status='running'` and matching `ownership_token` pub async fn heartbeat_global_phase2_job( &self, ownership_token: &str, lease_seconds: i64, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let lease_until = now.saturating_add(lease_seconds.max(0)); let rows_affected = sqlx::query( r#"UPDATE jobsSET lease_until = ?WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(lease_until) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) .execute(self.pool.as_ref()) .await? .rows_affected(); Ok(rows_affected > 0) } /// Marks the owned running global phase-2 job as succeeded. /// /// Query behavior: /// - updates only the owned running singleton global row /// - sets `status='done'`, clears lease/errors /// - advances `last_success_watermark` to /// `max(existing_last_success_watermark, completed_watermark)` /// - rewrites `selected_for_phase2` so only the exact selected stage-1 /// snapshots remain marked as part of the latest successful phase-2 /// selection, and persists each selected snapshot's `source_updated_at` pub async fn mark_global_phase2_job_succeeded( &self, ownership_token: &str, completed_watermark: i64, selected_outputs: &[Stage1Output], ) -> anyhow::Result<bool> { let mut tx = self.pool.begin().await?; let rows_affected = mark_global_phase2_job_succeeded_row(&mut *tx, ownership_token, completed_watermark) .await?; if rows_affected == 0 { tx.commit().await?; return Ok(false); } sqlx::query( r#"UPDATE stage1_outputsSET selected_for_phase2 = 0, selected_for_phase2_source_updated_at = NULLWHERE selected_for_phase2 != 0 OR selected_for_phase2_source_updated_at IS NOT NULL "#, ) .execute(&mut *tx) .await?; for output in selected_outputs { sqlx::query( r#"UPDATE stage1_outputsSET selected_for_phase2 = 1, selected_for_phase2_source_updated_at = ?WHERE thread_id = ? AND source_updated_at = ? "#, ) .bind(output.source_updated_at.timestamp()) .bind(output.thread_id.to_string()) .bind(output.source_updated_at.timestamp()) .execute(&mut *tx) .await?; } tx.commit().await?; Ok(true) } /// Marks the owned running global phase-2 job as failed and schedules retry. /// /// Query behavior: /// - updates only the owned running singleton global row /// - sets `status='error'`, clears lease /// - writes failure reason and retry time /// - decrements `retry_remaining` without going below zero pub async fn mark_global_phase2_job_failed( &self, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let retry_at = now.saturating_add(retry_delay_seconds.max(0)); let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?, retry_remaining = max(retry_remaining - 1, 0), last_error = ?WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(now) .bind(retry_at) .bind(failure_reason) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) .execute(self.pool.as_ref()) .await? .rows_affected(); Ok(rows_affected > 0) } /// Fallback failure finalization when ownership may have been lost. /// /// Query behavior: /// - same state transition as [`Self::mark_global_phase2_job_failed`] /// - matches rows where `ownership_token = ? OR ownership_token IS NULL` /// - allows recovering a stuck unowned running row pub async fn mark_global_phase2_job_failed_if_unowned( &self, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let retry_at = now.saturating_add(retry_delay_seconds.max(0)); let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?, retry_remaining = max(retry_remaining - 1, 0), last_error = ?WHERE kind = ? AND job_key = ? AND status = 'running' AND (ownership_token = ? OR ownership_token IS NULL) "#, ) .bind(now) .bind(retry_at) .bind(failure_reason) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) .execute(self.pool.as_ref()) .await? .rows_affected(); Ok(rows_affected > 0) }}async fn mark_global_phase2_job_succeeded_row<'e, E>( executor: E, ownership_token: &str, completed_watermark: i64,) -> anyhow::Result<u64>where E: Executor<'e, Database = Sqlite>,{ let now = Utc::now().timestamp(); let rows_affected = sqlx::query( r#"UPDATE jobsSET status = 'done', finished_at = ?, lease_until = NULL, last_error = NULL, last_success_watermark = max(COALESCE(last_success_watermark, 0), ?)WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ? "#, ) .bind(now) .bind(completed_watermark) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) .execute(executor) .await? .rows_affected(); Ok(rows_affected)}pub(super) async fn clear_memory_data_in_pool(pool: &SqlitePool) -> anyhow::Result<()> { let mut tx = pool.begin().await?; sqlx::query( r#"DELETE FROM stage1_outputs "#, ) .execute(&mut *tx) .await?; sqlx::query( r#"DELETE FROM jobsWHERE kind = ? OR kind = ? "#, ) .bind(JOB_KIND_MEMORY_STAGE1) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .execute(&mut *tx) .await?; tx.commit().await?; Ok(())}fn stage1_output_from_row_and_thread( row: &sqlx::sqlite::SqliteRow, thread: ThreadMetadata,) -> anyhow::Result<Stage1Output> { let source_updated_at: i64 = row.try_get("source_updated_at")?; let generated_at: i64 = row.try_get("generated_at")?; let source_updated_at = datetime_from_epoch_seconds(source_updated_at)?; let generated_at = datetime_from_epoch_seconds(generated_at)?; Ok(Stage1Output { thread_id: thread.id, rollout_path: thread.rollout_path, 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: thread.cwd, git_branch: thread.git_branch, generated_at, })}fn datetime_from_epoch_seconds(secs: i64) -> anyhow::Result<DateTime<Utc>> { DateTime::<Utc>::from_timestamp(secs, 0) .ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}"))}async fn enqueue_global_consolidation_with_executor<'e, E>( executor: E, input_watermark: i64,) -> anyhow::Result<()>where E: Executor<'e, Database = Sqlite>,{ 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 (?, ?, 'pending', NULL, NULL, NULL, NULL, NULL, NULL, ?, NULL, ?, 0)ON CONFLICT(kind, job_key) DO UPDATE SET status = CASE WHEN jobs.status = 'running' THEN 'running' ELSE 'pending' END, retry_at = CASE WHEN jobs.status = 'running' THEN jobs.retry_at ELSE NULL END, retry_remaining = max(jobs.retry_remaining, excluded.retry_remaining), input_watermark = CASE WHEN excluded.input_watermark > COALESCE(jobs.input_watermark, 0) THEN excluded.input_watermark ELSE COALESCE(jobs.input_watermark, 0) + 1 END "#, ) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(DEFAULT_RETRY_REMAINING) .bind(input_watermark) .execute(executor) .await?; Ok(())}#[cfg(test)]impl StateRuntime { async fn clear_memory_data(&self) -> anyhow::Result<()> { self.memories.clear_memory_data().await } async fn record_stage1_output_usage(&self, thread_ids: &[ThreadId]) -> anyhow::Result<usize> { self.memories.record_stage1_output_usage(thread_ids).await } async fn claim_stage1_jobs_for_startup( &self, current_thread_id: ThreadId, params: Stage1StartupClaimParams<'_>, ) -> anyhow::Result<Vec<Stage1JobClaim>> { self.memories .claim_stage1_jobs_for_startup(current_thread_id, params) .await } async fn list_stage1_outputs_for_global(&self, n: usize) -> anyhow::Result<Vec<Stage1Output>> { self.memories.list_stage1_outputs_for_global(n).await } async fn prune_stage1_outputs_for_retention( &self, max_unused_days: i64, limit: usize, ) -> anyhow::Result<usize> { self.memories .prune_stage1_outputs_for_retention(max_unused_days, limit) .await } async fn get_phase2_input_selection( &self, n: usize, max_unused_days: i64, ) -> anyhow::Result<Vec<Stage1Output>> { self.memories .get_phase2_input_selection(n, max_unused_days) .await } async fn mark_thread_memory_mode_polluted(&self, thread_id: ThreadId) -> anyhow::Result<bool> { self.memories .mark_thread_memory_mode_polluted(thread_id) .await } async fn try_claim_stage1_job( &self, thread_id: ThreadId, worker_id: ThreadId, source_updated_at: i64, lease_seconds: i64, max_running_jobs: usize, ) -> anyhow::Result<Stage1JobClaimOutcome> { self.memories .try_claim_stage1_job( thread_id, worker_id, source_updated_at, lease_seconds, max_running_jobs, ) .await } async fn mark_stage1_job_succeeded( &self, thread_id: ThreadId, ownership_token: &str, source_updated_at: i64, raw_memory: &str, rollout_summary: &str, rollout_slug: Option<&str>, ) -> anyhow::Result<bool> { self.memories .mark_stage1_job_succeeded( thread_id, ownership_token, source_updated_at, raw_memory, rollout_summary, rollout_slug, ) .await } async fn mark_stage1_job_succeeded_no_output( &self, thread_id: ThreadId, ownership_token: &str, ) -> anyhow::Result<bool> { self.memories .mark_stage1_job_succeeded_no_output(thread_id, ownership_token) .await } async fn mark_stage1_job_failed( &self, thread_id: ThreadId, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { self.memories .mark_stage1_job_failed( thread_id, ownership_token, failure_reason, retry_delay_seconds, ) .await } async fn enqueue_global_consolidation(&self, input_watermark: i64) -> anyhow::Result<()> { self.memories .enqueue_global_consolidation(input_watermark) .await } async fn try_claim_global_phase2_job( &self, worker_id: ThreadId, lease_seconds: i64, ) -> anyhow::Result<Phase2JobClaimOutcome> { self.memories .try_claim_global_phase2_job(worker_id, lease_seconds) .await } async fn mark_global_phase2_job_succeeded( &self, ownership_token: &str, completed_watermark: i64, selected_outputs: &[Stage1Output], ) -> anyhow::Result<bool> { self.memories .mark_global_phase2_job_succeeded( ownership_token, completed_watermark, selected_outputs, ) .await } async fn mark_global_phase2_job_failed( &self, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { self.memories .mark_global_phase2_job_failed(ownership_token, failure_reason, retry_delay_seconds) .await } async fn mark_global_phase2_job_failed_if_unowned( &self, ownership_token: &str, failure_reason: &str, retry_delay_seconds: i64, ) -> anyhow::Result<bool> { self.memories .mark_global_phase2_job_failed_if_unowned( ownership_token, failure_reason, retry_delay_seconds, ) .await }}#[cfg(test)]mod tests { use super::JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL; use super::JOB_KIND_MEMORY_STAGE1; use super::MEMORY_CONSOLIDATION_JOB_KEY; use super::PHASE2_SUCCESS_COOLDOWN_SECONDS; use super::StateRuntime; use super::test_support::test_thread_metadata; use super::test_support::unique_temp_dir; use crate::model::Phase2JobClaimOutcome; use crate::model::Stage1JobClaimOutcome; use crate::model::Stage1StartupClaimParams; use chrono::Duration; use chrono::Utc; use codex_protocol::ThreadId; use pretty_assertions::assert_eq; use sqlx::Row; use std::sync::Arc; use uuid::Uuid; fn stable_thread_id(value: &str) -> ThreadId { ThreadId::from_string(value).expect("thread id") } fn memory_pool(runtime: &StateRuntime) -> &sqlx::SqlitePool { runtime.memories().pool.as_ref() } async fn age_phase2_success_beyond_cooldown(runtime: &StateRuntime) { sqlx::query("UPDATE jobs SET finished_at = ? WHERE kind = ? AND job_key = ?") .bind(Utc::now().timestamp() - PHASE2_SUCCESS_COOLDOWN_SECONDS - 1) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .execute(memory_pool(runtime)) .await .expect("age phase2 success beyond cooldown"); } #[tokio::test] async fn stage1_claim_skips_when_up_to_date() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.join("a")); runtime .upsert_thread(&metadata) .await .expect("upsert thread"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); 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, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 job"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), /*source_updated_at*/ 100, "raw", "sum", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "stage1 success should finalize for current token" ); let up_to_date = runtime .try_claim_stage1_job( thread_id, owner_b, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 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, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 newer source"); assert!( matches!(needs_rerun, Stage1JobClaimOutcome::Claimed { .. }), "newer source_updated_at should be claimable" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn stage1_running_stale_can_be_stolen_but_fresh_running_is_skipped() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let cwd = codex_home.join("workspace"); runtime .upsert_thread(&test_thread_metadata(&codex_home, thread_id, cwd)) .await .expect("upsert thread"); let claim_a = runtime .try_claim_stage1_job( thread_id, owner_a, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim a"); assert!(matches!(claim_a, Stage1JobClaimOutcome::Claimed { .. })); let claim_b_fresh = runtime .try_claim_stage1_job( thread_id, owner_b, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim b fresh"); assert_eq!(claim_b_fresh, Stage1JobClaimOutcome::SkippedRunning); sqlx::query("UPDATE jobs SET lease_until = 0 WHERE kind = 'memory_stage1' AND job_key = ?") .bind(thread_id.to_string()) .execute(memory_pool(&runtime)) .await .expect("force stale lease"); let claim_b_stale = runtime .try_claim_stage1_job( thread_id, owner_b, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim b stale"); assert!(matches!( claim_b_stale, Stage1JobClaimOutcome::Claimed { .. } )); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn stage1_concurrent_claim_for_same_thread_is_conflict_safe() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let thread_id_a = thread_id; let thread_id_b = thread_id; let runtime_a = Arc::clone(&runtime); let runtime_b = Arc::clone(&runtime); let claim_with_retry = |runtime: Arc<StateRuntime>, thread_id: ThreadId, owner: ThreadId| async move { for attempt in 0..5 { match runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 64, ) .await { Ok(outcome) => return outcome, Err(err) if err.to_string().contains("database is locked") && attempt < 4 => { tokio::time::sleep(std::time::Duration::from_millis(10)).await; } Err(err) => panic!("claim stage1 should not fail: {err}"), } } panic!("claim stage1 should have returned within retry budget") }; let (claim_a, claim_b) = tokio::join!( claim_with_retry(runtime_a, thread_id_a, owner_a), claim_with_retry(runtime_b, thread_id_b, owner_b), ); let claim_outcomes = vec![claim_a, claim_b]; let claimed_count = claim_outcomes .iter() .filter(|outcome| matches!(outcome, Stage1JobClaimOutcome::Claimed { .. })) .count(); assert_eq!(claimed_count, 1); assert!( claim_outcomes.iter().all(|outcome| { matches!( outcome, Stage1JobClaimOutcome::Claimed { .. } | Stage1JobClaimOutcome::SkippedRunning ) }), "unexpected claim outcomes: {claim_outcomes:?}" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn stage1_concurrent_claims_respect_running_cap() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_a, codex_home.join("workspace-a"), )) .await .expect("upsert thread a"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_b, codex_home.join("workspace-b"), )) .await .expect("upsert thread b"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let runtime_a = Arc::clone(&runtime); let runtime_b = Arc::clone(&runtime); let (claim_a, claim_b) = tokio::join!( async move { runtime_a .try_claim_stage1_job( thread_a, owner_a, /*source_updated_at*/ 100, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 1, ) .await .expect("claim stage1 thread a") }, async move { runtime_b .try_claim_stage1_job( thread_b, owner_b, /*source_updated_at*/ 101, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 1, ) .await .expect("claim stage1 thread b") }, ); let claim_outcomes = vec![claim_a, claim_b]; let claimed_count = claim_outcomes .iter() .filter(|outcome| matches!(outcome, Stage1JobClaimOutcome::Claimed { .. })) .count(); assert_eq!(claimed_count, 1); assert!( claim_outcomes .iter() .any(|outcome| { matches!(outcome, Stage1JobClaimOutcome::SkippedRunning) }), "one concurrent claim should be throttled by running cap: {claim_outcomes:?}" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] 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()) .await .expect("initialize runtime"); let now = Utc::now(); 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 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"); let mut current = test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current")); current.created_at = now; current.updated_at = now; runtime .upsert_thread(¤t) .await .expect("upsert current"); 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; old.updated_at = old_at; runtime.upsert_thread(&old).await.expect("upsert old"); let allowed_sources = vec!["cli".to_string()]; let claims = runtime .claim_stage1_jobs_for_startup( current_thread_id, 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, eligible_idle_thread_id); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn claim_stage1_jobs_bounds_state_scan_before_memory_probes() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let now = Utc::now(); let eligible_newer_at = now - Duration::hours(13); let eligible_older_at = now - Duration::hours(14); let current_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); let up_to_date_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("up-to-date thread id"); let stale_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("stale thread id"); let worker_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("worker id"); let mut current = test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current")); current.created_at = now; current.updated_at = now; runtime .upsert_thread(¤t) .await .expect("upsert current thread"); let mut up_to_date = test_thread_metadata( &codex_home, up_to_date_thread_id, codex_home.join("up-to-date"), ); up_to_date.created_at = eligible_newer_at; up_to_date.updated_at = eligible_newer_at; runtime .upsert_thread(&up_to_date) .await .expect("upsert up-to-date thread"); let up_to_date_claim = runtime .try_claim_stage1_job( up_to_date_thread_id, worker_id, up_to_date.updated_at.timestamp(), /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim up-to-date thread for seed"); let up_to_date_token = match up_to_date_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected seed claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( up_to_date_thread_id, up_to_date_token.as_str(), up_to_date.updated_at.timestamp(), "raw", "summary", /*rollout_slug*/ None, ) .await .expect("mark up-to-date thread succeeded"), "seed stage1 success should complete for up-to-date thread" ); let mut stale = test_thread_metadata(&codex_home, stale_thread_id, codex_home.join("stale")); stale.created_at = eligible_older_at; stale.updated_at = eligible_older_at; runtime .upsert_thread(&stale) .await .expect("upsert stale thread"); let allowed_sources = vec!["cli".to_string()]; let claims_with_one_scanned_thread = runtime .claim_stage1_jobs_for_startup( current_thread_id, Stage1StartupClaimParams { scan_limit: 1, max_claimed: 1, max_age_days: 30, min_rollout_idle_hours: 12, allowed_sources: allowed_sources.as_slice(), lease_seconds: 3600, }, ) .await .expect("claim stage1 startup jobs"); assert_eq!(claims_with_one_scanned_thread.len(), 0); let claims = runtime .claim_stage1_jobs_for_startup( current_thread_id, Stage1StartupClaimParams { scan_limit: 2, max_claimed: 1, max_age_days: 30, min_rollout_idle_hours: 12, allowed_sources: allowed_sources.as_slice(), lease_seconds: 3600, }, ) .await .expect("claim stage1 startup jobs with wider scan"); assert_eq!(claims.len(), 1); assert_eq!(claims[0].thread.id, stale_thread_id); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn claim_stage1_jobs_skips_threads_with_disabled_memory_mode() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let now = Utc::now(); let eligible_at = now - Duration::hours(13); let current_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); let disabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("disabled thread id"); let enabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("enabled thread id"); let mut current = test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current")); current.created_at = now; current.updated_at = now; runtime .upsert_thread(¤t) .await .expect("upsert current thread"); let mut disabled = test_thread_metadata(&codex_home, disabled_thread_id, codex_home.join("disabled")); disabled.created_at = eligible_at; disabled.updated_at = eligible_at; runtime .upsert_thread(&disabled) .await .expect("upsert disabled thread"); sqlx::query("UPDATE threads SET memory_mode = 'disabled' WHERE id = ?") .bind(disabled_thread_id.to_string()) .execute(runtime.pool.as_ref()) .await .expect("disable thread memory mode"); let mut enabled = test_thread_metadata(&codex_home, enabled_thread_id, codex_home.join("enabled")); enabled.created_at = eligible_at; enabled.updated_at = eligible_at; runtime .upsert_thread(&enabled) .await .expect("upsert enabled thread"); let allowed_sources = vec!["cli".to_string()]; let claims = runtime .claim_stage1_jobs_for_startup( current_thread_id, Stage1StartupClaimParams { scan_limit: 10, max_claimed: 10, max_age_days: 30, min_rollout_idle_hours: 12, allowed_sources: allowed_sources.as_slice(), lease_seconds: 3600, }, ) .await .expect("claim stage1 startup jobs"); assert_eq!(claims.len(), 1); assert_eq!(claims[0].thread.id, enabled_thread_id); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn clear_memory_data_clears_rows_and_preserves_thread_memory_modes() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let now = Utc::now() - Duration::hours(13); let worker_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("worker id"); let enabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("enabled thread id"); let disabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("disabled thread id"); let mut enabled = test_thread_metadata(&codex_home, enabled_thread_id, codex_home.join("enabled")); enabled.created_at = now; enabled.updated_at = now; runtime .upsert_thread(&enabled) .await .expect("upsert enabled thread"); let claim = runtime .try_claim_stage1_job( enabled_thread_id, worker_id, enabled.updated_at.timestamp(), /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim enabled thread"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( enabled_thread_id, ownership_token.as_str(), enabled.updated_at.timestamp(), "raw", "summary", /*rollout_slug*/ None, ) .await .expect("mark enabled thread succeeded"), "stage1 success should be recorded" ); runtime .enqueue_global_consolidation(enabled.updated_at.timestamp()) .await .expect("enqueue global consolidation"); let mut disabled = test_thread_metadata(&codex_home, disabled_thread_id, codex_home.join("disabled")); disabled.created_at = now; disabled.updated_at = now; runtime .upsert_thread(&disabled) .await .expect("upsert disabled thread"); sqlx::query("UPDATE threads SET memory_mode = 'disabled' WHERE id = ?") .bind(disabled_thread_id.to_string()) .execute(runtime.pool.as_ref()) .await .expect("disable existing thread"); runtime .clear_memory_data() .await .expect("clear memory data"); let stage1_outputs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs") .fetch_one(memory_pool(&runtime)) .await .expect("count stage1 outputs"); assert_eq!(stage1_outputs_count, 0); let memory_jobs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM jobs WHERE kind = ? OR kind = ?") .bind(JOB_KIND_MEMORY_STAGE1) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .fetch_one(memory_pool(&runtime)) .await .expect("count memory jobs"); assert_eq!(memory_jobs_count, 0); let enabled_memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(enabled_thread_id.to_string()) .fetch_one(runtime.pool.as_ref()) .await .expect("read enabled thread memory mode"); assert_eq!(enabled_memory_mode, "enabled"); let disabled_memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(disabled_thread_id.to_string()) .fetch_one(runtime.pool.as_ref()) .await .expect("read disabled thread memory mode"); assert_eq!(disabled_memory_mode, "disabled"); 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()) .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(memory_pool(&runtime)) .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 countFROM jobsWHERE kind = 'memory_stage1' AND status = 'running' AND lease_until IS NOT NULL AND lease_until > ? "#, ) .bind(Utc::now().timestamp()) .fetch_one(memory_pool(&runtime)) .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; } #[tokio::test] async fn claim_stage1_jobs_processes_two_full_batches_across_startup_passes() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let current_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); let mut current = test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current")); current.created_at = Utc::now(); current.updated_at = Utc::now(); runtime .upsert_thread(¤t) .await .expect("upsert current"); let eligible_at = Utc::now() - Duration::hours(13); for idx in 0..200 { 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 eligible thread"); } let allowed_sources = vec!["cli".to_string()]; let first_claims = runtime .claim_stage1_jobs_for_startup( current_thread_id, Stage1StartupClaimParams { scan_limit: 5_000, max_claimed: 64, max_age_days: 30, min_rollout_idle_hours: 12, allowed_sources: allowed_sources.as_slice(), lease_seconds: 3_600, }, ) .await .expect("first stage1 startup claim"); assert_eq!(first_claims.len(), 64); for claim in first_claims { assert!( runtime .mark_stage1_job_succeeded( claim.thread.id, claim.ownership_token.as_str(), claim.thread.updated_at.timestamp(), "raw", "summary", /*rollout_slug*/ None, ) .await .expect("mark first-batch stage1 success"), "first batch stage1 completion should succeed" ); } let second_claims = runtime .claim_stage1_jobs_for_startup( current_thread_id, Stage1StartupClaimParams { scan_limit: 5_000, max_claimed: 64, max_age_days: 30, min_rollout_idle_hours: 12, allowed_sources: allowed_sources.as_slice(), lease_seconds: 3_600, }, ) .await .expect("second stage1 startup claim"); assert_eq!(second_claims.len(), 64); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn delete_thread_removes_stage1_output_and_enqueues_phase2_when_selected() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let cwd = codex_home.join("workspace"); runtime .upsert_thread(&test_thread_metadata(&codex_home, thread_id, cwd)) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), /*source_updated_at*/ 100, "raw", "sum", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "mark stage1 succeeded should write stage1_outputs" ); let count_before = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("count before delete") .try_get::<i64, _>("count") .expect("count value"); assert_eq!(count_before, 1); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs"); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 succeeded"), "phase2 success should mark selected stage1 output" ); let before_delete = Utc::now().timestamp(); assert_eq!( runtime .delete_thread(thread_id) .await .expect("delete thread"), 1 ); let count_after = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("count after delete") .try_get::<i64, _>("count") .expect("count value"); assert_eq!(count_after, 0); let phase2_job = sqlx::query( r#"SELECT status, input_watermarkFROM jobsWHERE kind = ? AND job_key = ? "#, ) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .fetch_one(memory_pool(&runtime)) .await .expect("load phase2 job after delete"); let status: String = phase2_job.try_get("status").expect("status"); let input_watermark: i64 = phase2_job .try_get("input_watermark") .expect("input watermark"); assert_eq!(status, "pending"); assert!(input_watermark >= before_delete); let visible_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs after thread delete"); assert_eq!(visible_outputs.len(), 0); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_stage1_job_succeeded_no_output_skips_phase2_when_output_was_already_absent() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded_no_output(thread_id, ownership_token.as_str()) .await .expect("mark stage1 succeeded without output"), "stage1 no-output success should complete the job" ); let output_row_count = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load stage1 output count") .try_get::<i64, _>("count") .expect("stage1 output count"); assert_eq!( output_row_count, 0, "stage1 no-output success should not persist empty stage1 outputs" ); let up_to_date = runtime .try_claim_stage1_job( thread_id, owner_b, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 up-to-date"); assert_eq!(up_to_date, Stage1JobClaimOutcome::SkippedUpToDate); let global_job_row_count = sqlx::query("SELECT COUNT(*) AS count FROM jobs WHERE kind = ?") .bind("memory_consolidate_global") .fetch_one(memory_pool(&runtime)) .await .expect("load phase2 job row count") .try_get::<i64, _>("count") .expect("phase2 job row count"); assert_eq!( global_job_row_count, 0, "no-output without an existing stage1 output should not enqueue phase2" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_stage1_job_succeeded_no_output_enqueues_phase2_when_deleting_output() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let first_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim initial stage1"); let first_token = match first_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected initial stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, first_token.as_str(), /*source_updated_at*/ 100, "raw", "sum", /*rollout_slug*/ None ) .await .expect("mark initial stage1 succeeded"), "initial stage1 success should create stage1 output" ); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after initial output"); let (phase2_token, phase2_input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim after initial output: {other:?}"), }; assert_eq!(phase2_input_watermark, 100); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), phase2_input_watermark, &[], ) .await .expect("mark initial phase2 succeeded"), "initial phase2 success should finalize the global job" ); let no_output_claim = runtime .try_claim_stage1_job( thread_id, owner_b, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 for no-output delete"); let no_output_token = match no_output_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected no-output stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded_no_output(thread_id, no_output_token.as_str()) .await .expect("mark stage1 no-output after existing output"), "no-output should succeed when deleting an existing stage1 output" ); let output_row_count = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load stage1 output count after delete") .try_get::<i64, _>("count") .expect("stage1 output count"); assert_eq!(output_row_count, 0); age_phase2_success_beyond_cooldown(&runtime).await; let claim_phase2 = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after no-output deletion"); let (phase2_token, phase2_input_watermark) = match claim_phase2 { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim after no-output deletion: {other:?}"), }; assert_eq!(phase2_input_watermark, 101); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), phase2_input_watermark, &[], ) .await .expect("mark phase2 succeeded after no-output delete") ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn stage1_retry_exhaustion_does_not_block_newer_watermark() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); for attempt in 0..3 { let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 for retry exhaustion"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!( "attempt {} should claim stage1 before retries are exhausted: {other:?}", attempt + 1 ), }; assert!( runtime .mark_stage1_job_failed( thread_id, ownership_token.as_str(), "boom", /*retry_delay_seconds*/ 0 ) .await .expect("mark stage1 failed"), "attempt {} should decrement retry budget", attempt + 1 ); } let exhausted_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 after retry exhaustion"); assert_eq!( exhausted_claim, Stage1JobClaimOutcome::SkippedRetryExhausted ); let newer_source_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3_600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 with newer source watermark"); assert!( matches!(newer_source_claim, Stage1JobClaimOutcome::Claimed { .. }), "newer source watermark should reset retry budget and be claimable" ); let job_row = sqlx::query( "SELECT retry_remaining, input_watermark FROM jobs WHERE kind = ? AND job_key = ?", ) .bind("memory_stage1") .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load stage1 job row after newer-source claim"); assert_eq!( job_row .try_get::<i64, _>("retry_remaining") .expect("retry_remaining"), 3 ); assert_eq!( job_row .try_get::<i64, _>("input_watermark") .expect("input_watermark"), 101 ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_global_lock_respects_success_cooldown() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .enqueue_global_consolidation(/*input_watermark*/ 100) .await .expect("enqueue global consolidation"); let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (ownership_token, input_watermark) = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; assert!( runtime .mark_global_phase2_job_succeeded(ownership_token.as_str(), input_watermark, &[],) .await .expect("mark phase2 succeeded"), "phase2 success should finalize for current token" ); let claim_after_success = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after success"); assert_eq!(claim_after_success, Phase2JobClaimOutcome::SkippedCooldown); runtime .enqueue_global_consolidation(/*input_watermark*/ 101) .await .expect("enqueue global consolidation after success"); let claim_after_enqueue = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after enqueue"); assert_eq!(claim_after_enqueue, Phase2JobClaimOutcome::SkippedCooldown); age_phase2_success_beyond_cooldown(&runtime).await; let claim_after_cooldown = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after cooldown"); assert!(matches!( claim_after_cooldown, Phase2JobClaimOutcome::Claimed { .. } )); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_global_lock_can_be_claimed_after_retry_budget_is_exhausted() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 100) .await .expect("enqueue global consolidation"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); for attempt in 0..3 { let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3_600) .await .expect("claim phase2 before retry exhaustion"); let ownership_token = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, .. } => ownership_token, other => panic!( "attempt {} should claim phase2 before retries are exhausted: {other:?}", attempt + 1 ), }; assert!( runtime .mark_global_phase2_job_failed( ownership_token.as_str(), "boom", /*retry_delay_seconds*/ 0, ) .await .expect("mark phase2 failed"), "attempt {} should decrement retry budget", attempt + 1 ); } let job_row = sqlx::query("SELECT retry_remaining FROM jobs WHERE kind = ? AND job_key = ?") .bind("memory_consolidate_global") .bind("global") .fetch_one(memory_pool(&runtime)) .await .expect("load phase2 job row after retry exhaustion"); assert_eq!( job_row .try_get::<i64, _>("retry_remaining") .expect("retry_remaining"), 0 ); let claim_after_exhaustion = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3_600) .await .expect("claim phase2 after retry exhaustion"); assert!( matches!( claim_after_exhaustion, Phase2JobClaimOutcome::Claimed { .. } ), "phase2 claim should only lock; workspace diffing decides whether there is work" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn list_stage1_outputs_for_global_returns_latest_outputs() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_id_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id_a, codex_home.join("workspace-a"), )) .await .expect("upsert thread a"); let mut metadata_b = test_thread_metadata(&codex_home, thread_id_b, codex_home.join("workspace-b")); metadata_b.git_branch = Some("feature/stage1-b".to_string()); runtime .upsert_thread(&metadata_b) .await .expect("upsert thread b"); let claim = runtime .try_claim_stage1_job( thread_id_a, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 a"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id_a, ownership_token.as_str(), /*source_updated_at*/ 100, "raw memory a", "summary a", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded a"), "stage1 success should persist output a" ); let claim = runtime .try_claim_stage1_job( thread_id_b, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 b"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id_b, ownership_token.as_str(), /*source_updated_at*/ 101, "raw memory b", "summary b", Some("rollout-b"), ) .await .expect("mark stage1 succeeded b"), "stage1 success should persist output b" ); let outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs for global"); assert_eq!(outputs.len(), 2); assert_eq!(outputs[0].thread_id, thread_id_b); assert_eq!(outputs[0].rollout_summary, "summary b"); assert_eq!(outputs[0].rollout_slug.as_deref(), Some("rollout-b")); assert_eq!(outputs[0].cwd, codex_home.join("workspace-b")); assert_eq!(outputs[0].git_branch.as_deref(), Some("feature/stage1-b")); assert_eq!(outputs[1].thread_id, thread_id_a); assert_eq!(outputs[1].rollout_summary, "summary a"); assert_eq!(outputs[1].rollout_slug, None); assert_eq!(outputs[1].cwd, codex_home.join("workspace-a")); assert_eq!(outputs[1].git_branch, None); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn list_stage1_outputs_for_global_skips_empty_payloads() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_non_empty = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_id_empty = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id_non_empty, codex_home.join("workspace-non-empty"), )) .await .expect("upsert non-empty thread"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id_empty, codex_home.join("workspace-empty"), )) .await .expect("upsert empty thread"); sqlx::query( r#"INSERT INTO stage1_outputs (thread_id, source_updated_at, raw_memory, rollout_summary, generated_at)VALUES (?, ?, ?, ?, ?) "#, ) .bind(thread_id_non_empty.to_string()) .bind(100_i64) .bind("raw memory") .bind("summary") .bind(100_i64) .execute(memory_pool(&runtime)) .await .expect("insert non-empty stage1 output"); sqlx::query( r#"INSERT INTO stage1_outputs (thread_id, source_updated_at, raw_memory, rollout_summary, generated_at)VALUES (?, ?, ?, ?, ?) "#, ) .bind(thread_id_empty.to_string()) .bind(101_i64) .bind("") .bind("") .bind(101_i64) .execute(memory_pool(&runtime)) .await .expect("insert empty stage1 output"); let outputs = runtime .list_stage1_outputs_for_global(/*n*/ 1) .await .expect("list stage1 outputs for global"); assert_eq!(outputs.len(), 1); assert_eq!(outputs[0].thread_id, thread_id_non_empty); assert_eq!(outputs[0].rollout_summary, "summary"); assert_eq!(outputs[0].cwd, codex_home.join("workspace-non-empty")); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn list_stage1_outputs_for_global_skips_polluted_threads() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_enabled = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_id_polluted = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); for (thread_id, workspace) in [ (thread_id_enabled, "workspace-enabled"), (thread_id_polluted, "workspace-polluted"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), /*source_updated_at*/ 100, "raw memory", "summary", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); } runtime .set_thread_memory_mode(thread_id_polluted, "polluted") .await .expect("mark thread polluted"); let outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs for global"); assert_eq!(outputs.len(), 1); assert_eq!(outputs[0].thread_id, thread_id_enabled); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_returns_current_selected_rows() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_id_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); let thread_id_c = stable_thread_id("00000000-0000-4000-8000-000000000003"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); for (thread_id, workspace) in [ (thread_id_a, "workspace-a"), (thread_id_b, "workspace-b"), (thread_id_c, "workspace-c"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } for (thread_id, updated_at, slug) in [ (thread_id_a, 100, Some("rollout-a")), (thread_id_b, 101, Some("rollout-b")), (thread_id_c, 102, Some("rollout-c")), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), updated_at, &format!("raw-{updated_at}"), &format!("summary-{updated_at}"), slug, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); } let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (ownership_token, input_watermark) = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; assert_eq!(input_watermark, 102); let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs for global") .into_iter() .filter(|output| output.thread_id == thread_id_c || output.thread_id == thread_id_a) .collect::<Vec<_>>(); assert!( runtime .mark_global_phase2_job_succeeded( ownership_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success with selection"), "phase2 success should persist selected rows" ); let selection = runtime .get_phase2_input_selection(/*n*/ 2, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!(selection.len(), 2); assert_eq!( selection .iter() .map(|output| output.thread_id) .collect::<Vec<_>>(), vec![thread_id_b, thread_id_c] ); let selected_c = selection .iter() .find(|output| output.thread_id == thread_id_c) .expect("thread c should be selected"); assert_eq!( selected_c.rollout_path, codex_home.join(format!("rollout-{thread_id_c}.jsonl")) ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_excludes_polluted_previous_selection() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_enabled = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_id_polluted = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); for (thread_id, updated_at) in [(thread_id_enabled, 100), (thread_id_polluted, 101)] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(thread_id.to_string()), )) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), updated_at, &format!("raw-{updated_at}"), &format!("summary-{updated_at}"), /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); } let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (ownership_token, input_watermark) = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs for global"); assert!( runtime .mark_global_phase2_job_succeeded( ownership_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should persist selected rows" ); runtime .set_thread_memory_mode(thread_id_polluted, "polluted") .await .expect("mark thread polluted"); let selection = runtime .get_phase2_input_selection(/*n*/ 2, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!(selection.len(), 1); assert_eq!(selection[0].thread_id, thread_id_enabled); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_thread_memory_mode_polluted_enqueues_phase2_for_selected_threads() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), /*source_updated_at*/ 100, "raw", "summary", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs"); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should persist selected rows" ); assert!( runtime .mark_thread_memory_mode_polluted(thread_id) .await .expect("mark thread polluted"), "thread should transition to polluted" ); age_phase2_success_beyond_cooldown(&runtime).await; let next_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after pollution"); assert!(matches!(next_claim, Phase2JobClaimOutcome::Claimed { .. })); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_thread_memory_mode_polluted_enqueues_phase2_when_already_polluted() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), /*source_updated_at*/ 100, "raw", "summary", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 10) .await .expect("list stage1 outputs"); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should persist selected rows" ); sqlx::query("UPDATE threads SET memory_mode = 'polluted' WHERE id = ?") .bind(thread_id.to_string()) .execute(runtime.pool.as_ref()) .await .expect("mark thread polluted before memory enqueue"); assert!( !runtime .mark_thread_memory_mode_polluted(thread_id) .await .expect("mark already polluted thread"), "already polluted thread should not report a state transition" ); age_phase2_success_beyond_cooldown(&runtime).await; let next_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2 after already-polluted enqueue"); assert!(matches!(next_claim, Phase2JobClaimOutcome::Claimed { .. })); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_returns_regenerated_selected_rows() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let first_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim initial stage1"); let first_token = match first_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, first_token.as_str(), /*source_updated_at*/ 100, "raw-100", "summary-100", Some("rollout-100"), ) .await .expect("mark initial stage1 success"), "initial stage1 success should persist output" ); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 1) .await .expect("list selected outputs"); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should persist selected rows" ); let refreshed_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim refreshed stage1"); let refreshed_token = match refreshed_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, refreshed_token.as_str(), /*source_updated_at*/ 101, "raw-101", "summary-101", Some("rollout-101"), ) .await .expect("mark refreshed stage1 success"), "refreshed stage1 success should persist output" ); let selection = runtime .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!(selection.len(), 1); assert_eq!(selection[0].thread_id, thread_id); assert_eq!(selection[0].source_updated_at.timestamp(), 101); let (selected_for_phase2, selected_for_phase2_source_updated_at) = sqlx::query_as::<_, (i64, Option<i64>)>( "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load selected_for_phase2"); assert_eq!(selected_for_phase2, 1); assert_eq!(selected_for_phase2_source_updated_at, Some(100)); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_uses_current_ranking_after_refreshes() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_id_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); let thread_id_c = stable_thread_id("00000000-0000-4000-8000-000000000003"); let thread_id_d = stable_thread_id("00000000-0000-4000-8000-000000000004"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); for (thread_id, workspace) in [ (thread_id_a, "workspace-a"), (thread_id_b, "workspace-b"), (thread_id_c, "workspace-c"), (thread_id_d, "workspace-d"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } for (thread_id, updated_at, slug) in [ (thread_id_a, 100, Some("rollout-a-100")), (thread_id_b, 101, Some("rollout-b-101")), (thread_id_c, 99, Some("rollout-c-99")), (thread_id_d, 98, Some("rollout-d-98")), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim initial stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), updated_at, &format!("raw-{updated_at}"), &format!("summary-{updated_at}"), slug, ) .await .expect("mark stage1 succeeded"), "stage1 success should persist output" ); } let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 2) .await .expect("list selected outputs"); assert_eq!( selected_outputs .iter() .map(|output| output.thread_id) .collect::<Vec<_>>(), vec![thread_id_b, thread_id_a] ); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should persist selected rows" ); for (thread_id, updated_at, slug) in [ (thread_id_a, 102, Some("rollout-a-102")), (thread_id_c, 103, Some("rollout-c-103")), (thread_id_d, 104, Some("rollout-d-104")), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim refreshed stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), updated_at, &format!("raw-{updated_at}"), &format!("summary-{updated_at}"), slug, ) .await .expect("mark refreshed stage1 success"), "refreshed stage1 success should persist output" ); } let selection = runtime .get_phase2_input_selection(/*n*/ 2, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!( selection .iter() .map(|output| output.thread_id) .collect::<Vec<_>>(), vec![thread_id_c, thread_id_d] ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_global_phase2_job_succeeded_updates_selected_snapshot_timestamp() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let initial_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim initial stage1"); let initial_token = match initial_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, initial_token.as_str(), /*source_updated_at*/ 100, "raw-100", "summary-100", Some("rollout-100"), ) .await .expect("mark initial stage1 success"), "initial stage1 success should persist output" ); let first_phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim first phase2"); let (first_phase2_token, first_input_watermark) = match first_phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected first phase2 claim outcome: {other:?}"), }; let first_selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 1) .await .expect("list first selected outputs"); assert!( runtime .mark_global_phase2_job_succeeded( first_phase2_token.as_str(), first_input_watermark, &first_selected_outputs, ) .await .expect("mark first phase2 success"), "first phase2 success should persist selected rows" ); let refreshed_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim refreshed stage1"); let refreshed_token = match refreshed_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected refreshed stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, refreshed_token.as_str(), /*source_updated_at*/ 101, "raw-101", "summary-101", Some("rollout-101"), ) .await .expect("mark refreshed stage1 success"), "refreshed stage1 success should persist output" ); age_phase2_success_beyond_cooldown(&runtime).await; let second_phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim second phase2"); let (second_phase2_token, second_input_watermark) = match second_phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected second phase2 claim outcome: {other:?}"), }; let second_selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 1) .await .expect("list second selected outputs"); assert_eq!( second_selected_outputs[0].source_updated_at.timestamp(), 101 ); assert!( runtime .mark_global_phase2_job_succeeded( second_phase2_token.as_str(), second_input_watermark, &second_selected_outputs, ) .await .expect("mark second phase2 success"), "second phase2 success should persist selected rows" ); let selection = runtime .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection after refresh"); assert_eq!(selection.len(), 1); assert_eq!(selection[0].thread_id, thread_id); let (selected_for_phase2, selected_for_phase2_source_updated_at) = sqlx::query_as::<_, (i64, Option<i64>)>( "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load selected snapshot after phase2"); assert_eq!(selected_for_phase2, 1); assert_eq!(selected_for_phase2_source_updated_at, Some(101)); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_global_phase2_job_succeeded_only_marks_exact_selected_snapshots() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join("workspace"), )) .await .expect("upsert thread"); let initial_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim initial stage1"); let initial_token = match initial_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, initial_token.as_str(), /*source_updated_at*/ 100, "raw-100", "summary-100", Some("rollout-100"), ) .await .expect("mark initial stage1 success"), "initial stage1 success should persist output" ); let phase2_claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim phase2"); let (phase2_token, input_watermark) = match phase2_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected phase2 claim outcome: {other:?}"), }; let selected_outputs = runtime .list_stage1_outputs_for_global(/*n*/ 1) .await .expect("list selected outputs"); assert_eq!(selected_outputs[0].source_updated_at.timestamp(), 100); let refreshed_claim = runtime .try_claim_stage1_job( thread_id, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim refreshed stage1"); let refreshed_token = match refreshed_claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, refreshed_token.as_str(), /*source_updated_at*/ 101, "raw-101", "summary-101", Some("rollout-101"), ) .await .expect("mark refreshed stage1 success"), "refreshed stage1 success should persist output" ); assert!( runtime .mark_global_phase2_job_succeeded( phase2_token.as_str(), input_watermark, &selected_outputs, ) .await .expect("mark phase2 success"), "phase2 success should still complete" ); let (selected_for_phase2, selected_for_phase2_source_updated_at) = sqlx::query_as::<_, (i64, Option<i64>)>( "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load selected_for_phase2"); assert_eq!(selected_for_phase2, 0); assert_eq!(selected_for_phase2_source_updated_at, None); let selection = runtime .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!(selection.len(), 1); assert_eq!(selection[0].source_updated_at.timestamp(), 101); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn record_stage1_output_usage_updates_usage_metadata() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id a"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id b"); let missing = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("missing id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_a, codex_home.join("workspace-a"), )) .await .expect("upsert thread a"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_b, codex_home.join("workspace-b"), )) .await .expect("upsert thread b"); let claim_a = runtime .try_claim_stage1_job( thread_a, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 a"); let token_a = match claim_a { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome for a: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_a, token_a.as_str(), /*source_updated_at*/ 100, "raw a", "sum a", /*rollout_slug*/ None ) .await .expect("mark stage1 succeeded a") ); let claim_b = runtime .try_claim_stage1_job( thread_b, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 b"); let token_b = match claim_b { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome for b: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_b, token_b.as_str(), /*source_updated_at*/ 101, "raw b", "sum b", /*rollout_slug*/ None ) .await .expect("mark stage1 succeeded b") ); let updated_rows = runtime .record_stage1_output_usage(&[thread_a, thread_a, thread_b, missing]) .await .expect("record stage1 output usage"); assert_eq!(updated_rows, 3); let row_a = sqlx::query("SELECT usage_count, last_usage FROM stage1_outputs WHERE thread_id = ?") .bind(thread_a.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load stage1 usage row a"); let row_b = sqlx::query("SELECT usage_count, last_usage FROM stage1_outputs WHERE thread_id = ?") .bind(thread_b.to_string()) .fetch_one(memory_pool(&runtime)) .await .expect("load stage1 usage row b"); assert_eq!( row_a .try_get::<i64, _>("usage_count") .expect("usage_count a"), 2 ); assert_eq!( row_b .try_get::<i64, _>("usage_count") .expect("usage_count b"), 1 ); let last_usage_a = row_a.try_get::<i64, _>("last_usage").expect("last_usage a"); let last_usage_b = row_b.try_get::<i64, _>("last_usage").expect("last_usage b"); assert_eq!(last_usage_a, last_usage_b); assert!(last_usage_a > 0); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_prioritizes_usage_count_then_recent_usage() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let now = Utc::now(); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let thread_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); let thread_c = stable_thread_id("00000000-0000-4000-8000-000000000003"); for (thread_id, workspace) in [ (thread_a, "workspace-a"), (thread_b, "workspace-b"), (thread_c, "workspace-c"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } for (thread_id, generated_at, summary) in [ (thread_a, now - Duration::days(3), "summary-a"), (thread_b, now - Duration::days(2), "summary-b"), (thread_c, now - Duration::days(1), "summary-c"), ] { let source_updated_at = generated_at.timestamp(); let claim = runtime .try_claim_stage1_job( thread_id, owner, source_updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), source_updated_at, &format!("raw-{summary}"), summary, /*rollout_slug*/ None, ) .await .expect("mark stage1 success"), "stage1 success should persist output" ); } for (thread_id, usage_count, last_usage) in [ (thread_a, 5_i64, now - Duration::days(10)), (thread_b, 5_i64, now - Duration::days(1)), (thread_c, 1_i64, now - Duration::hours(1)), ] { sqlx::query( "UPDATE stage1_outputs SET usage_count = ?, last_usage = ? WHERE thread_id = ?", ) .bind(usage_count) .bind(last_usage.timestamp()) .bind(thread_id.to_string()) .execute(memory_pool(&runtime)) .await .expect("update usage metadata"); } let selection = runtime .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 30) .await .expect("load phase2 input selection"); assert_eq!( selection .iter() .map(|output| output.thread_id) .collect::<Vec<_>>(), vec![thread_b] ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_excludes_stale_used_memories_but_keeps_fresh_never_used() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let now = Utc::now(); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let thread_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); let thread_c = stable_thread_id("00000000-0000-4000-8000-000000000003"); for (thread_id, workspace) in [ (thread_a, "workspace-a"), (thread_b, "workspace-b"), (thread_c, "workspace-c"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } for (thread_id, generated_at, summary) in [ (thread_a, now - Duration::days(40), "summary-a"), (thread_b, now - Duration::days(2), "summary-b"), (thread_c, now - Duration::days(50), "summary-c"), ] { let source_updated_at = generated_at.timestamp(); let claim = runtime .try_claim_stage1_job( thread_id, owner, source_updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), source_updated_at, &format!("raw-{summary}"), summary, /*rollout_slug*/ None, ) .await .expect("mark stage1 success"), "stage1 success should persist output" ); } for (thread_id, usage_count, last_usage) in [ (thread_a, Some(9_i64), Some(now - Duration::days(31))), (thread_b, None, None), (thread_c, Some(1_i64), Some(now - Duration::days(1))), ] { sqlx::query( "UPDATE stage1_outputs SET usage_count = ?, last_usage = ? WHERE thread_id = ?", ) .bind(usage_count) .bind(last_usage.map(|value| value.timestamp())) .bind(thread_id.to_string()) .execute(memory_pool(&runtime)) .await .expect("update usage metadata"); } let selection = runtime .get_phase2_input_selection(/*n*/ 3, /*max_unused_days*/ 30) .await .expect("load phase2 input selection"); assert_eq!( selection .iter() .map(|output| output.thread_id) .collect::<Vec<_>>(), vec![thread_b, thread_c] ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn get_phase2_input_selection_prefers_recent_thread_updates_over_recent_generation() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let older_thread = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("older thread id"); let newer_thread = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("newer thread id"); for (thread_id, workspace) in [ (older_thread, "workspace-older"), (newer_thread, "workspace-newer"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } for (thread_id, source_updated_at, summary) in [ (older_thread, 100_i64, "summary-older"), (newer_thread, 200_i64, "summary-newer"), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, source_updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), source_updated_at, &format!("raw-{summary}"), summary, /*rollout_slug*/ None, ) .await .expect("mark stage1 success"), "stage1 success should persist output" ); } sqlx::query("UPDATE stage1_outputs SET generated_at = ? WHERE thread_id = ?") .bind(300_i64) .bind(older_thread.to_string()) .execute(memory_pool(&runtime)) .await .expect("update older generated_at"); sqlx::query("UPDATE stage1_outputs SET generated_at = ? WHERE thread_id = ?") .bind(150_i64) .bind(newer_thread.to_string()) .execute(memory_pool(&runtime)) .await .expect("update newer generated_at"); let selection = runtime .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 36_500) .await .expect("load phase2 input selection"); assert_eq!(selection.len(), 1); assert_eq!(selection[0].thread_id, newer_thread); assert_eq!(selection[0].source_updated_at.timestamp(), 200); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn prune_stage1_outputs_for_retention_prunes_stale_unselected_rows_only() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let stale_unused = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("stale unused"); let stale_used = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("stale used"); let stale_selected = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("stale selected"); let fresh_used = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("fresh used"); for (thread_id, workspace) in [ (stale_unused, "workspace-stale-unused"), (stale_used, "workspace-stale-used"), (stale_selected, "workspace-stale-selected"), (fresh_used, "workspace-fresh-used"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } let now = Utc::now().timestamp(); for (thread_id, source_updated_at, summary) in [ ( stale_unused, now - Duration::days(60).num_seconds(), "stale-unused", ), ( stale_used, now - Duration::days(50).num_seconds(), "stale-used", ), ( stale_selected, now - Duration::days(45).num_seconds(), "stale-selected", ), ( fresh_used, now - Duration::days(10).num_seconds(), "fresh-used", ), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, source_updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), source_updated_at, &format!("raw-{summary}"), summary, /*rollout_slug*/ None, ) .await .expect("mark stage1 success"), "stage1 success should persist output" ); } sqlx::query( "UPDATE stage1_outputs SET usage_count = ?, last_usage = ? WHERE thread_id = ?", ) .bind(3_i64) .bind(now - Duration::days(40).num_seconds()) .bind(stale_used.to_string()) .execute(memory_pool(&runtime)) .await .expect("set stale used metadata"); sqlx::query( "UPDATE stage1_outputs SET selected_for_phase2 = 1, selected_for_phase2_source_updated_at = source_updated_at WHERE thread_id = ?", ) .bind(stale_selected.to_string()) .execute(memory_pool(&runtime)) .await .expect("mark selected for phase2"); sqlx::query( "UPDATE stage1_outputs SET usage_count = ?, last_usage = ? WHERE thread_id = ?", ) .bind(8_i64) .bind(now - Duration::days(2).num_seconds()) .bind(fresh_used.to_string()) .execute(memory_pool(&runtime)) .await .expect("set fresh used metadata"); let before_jobs_count = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM jobs WHERE kind = 'memory_stage1'") .fetch_one(memory_pool(&runtime)) .await .expect("count stage1 jobs before prune"); let pruned = runtime .prune_stage1_outputs_for_retention(/*max_unused_days*/ 30, /*limit*/ 100) .await .expect("prune stage1 outputs"); assert_eq!(pruned, 2); let remaining = sqlx::query_scalar::<_, String>( "SELECT thread_id FROM stage1_outputs ORDER BY thread_id", ) .fetch_all(memory_pool(&runtime)) .await .expect("load remaining stage1 outputs"); let mut expected_remaining = vec![fresh_used.to_string(), stale_selected.to_string()]; expected_remaining.sort(); assert_eq!(remaining, expected_remaining); let after_jobs_count = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM jobs WHERE kind = 'memory_stage1'") .fetch_one(memory_pool(&runtime)) .await .expect("count stage1 jobs after prune"); assert_eq!(after_jobs_count, before_jobs_count); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn prune_stage1_outputs_for_retention_respects_batch_limit() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread a"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread b"); let thread_c = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread c"); for (thread_id, workspace) in [ (thread_a, "workspace-a"), (thread_b, "workspace-b"), (thread_c, "workspace-c"), ] { runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_id, codex_home.join(workspace), )) .await .expect("upsert thread"); } let now = Utc::now().timestamp(); for (thread_id, source_updated_at, summary) in [ (thread_a, now - Duration::days(60).num_seconds(), "stale-a"), (thread_b, now - Duration::days(50).num_seconds(), "stale-b"), (thread_c, now - Duration::days(40).num_seconds(), "stale-c"), ] { let claim = runtime .try_claim_stage1_job( thread_id, owner, source_updated_at, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1"); let ownership_token = match claim { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_id, ownership_token.as_str(), source_updated_at, &format!("raw-{summary}"), summary, /*rollout_slug*/ None, ) .await .expect("mark stage1 success"), "stage1 success should persist output" ); } let pruned = runtime .prune_stage1_outputs_for_retention(/*max_unused_days*/ 30, /*limit*/ 2) .await .expect("prune stage1 outputs with limit"); assert_eq!(pruned, 2); let remaining_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs") .fetch_one(memory_pool(&runtime)) .await .expect("count remaining stage1 outputs"); assert_eq!(remaining_count, 1); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn mark_stage1_job_succeeded_enqueues_global_consolidation() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id a"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id b"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_a, codex_home.join("workspace-a"), )) .await .expect("upsert thread a"); runtime .upsert_thread(&test_thread_metadata( &codex_home, thread_b, codex_home.join("workspace-b"), )) .await .expect("upsert thread b"); let claim_a = runtime .try_claim_stage1_job( thread_a, owner, /*source_updated_at*/ 100, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 a"); let token_a = match claim_a { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome for thread a: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_a, token_a.as_str(), /*source_updated_at*/ 100, "raw-a", "summary-a", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded a"), "stage1 success should persist output for thread a" ); let claim_b = runtime .try_claim_stage1_job( thread_b, owner, /*source_updated_at*/ 101, /*lease_seconds*/ 3600, /*max_running_jobs*/ 64, ) .await .expect("claim stage1 b"); let token_b = match claim_b { Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, other => panic!("unexpected stage1 claim outcome for thread b: {other:?}"), }; assert!( runtime .mark_stage1_job_succeeded( thread_b, token_b.as_str(), /*source_updated_at*/ 101, "raw-b", "summary-b", /*rollout_slug*/ None, ) .await .expect("mark stage1 succeeded b"), "stage1 success should persist output for thread b" ); let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3600) .await .expect("claim global consolidation"); let input_watermark = match claim { Phase2JobClaimOutcome::Claimed { input_watermark, .. } => input_watermark, other => panic!("unexpected global consolidation claim outcome: {other:?}"), }; assert_eq!(input_watermark, 101); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_global_lock_allows_only_one_fresh_runner() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 200) .await .expect("enqueue global consolidation"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b"); let running_claim = runtime .try_claim_global_phase2_job(owner_a, /*lease_seconds*/ 3600) .await .expect("claim global lock"); assert!( matches!(running_claim, Phase2JobClaimOutcome::Claimed { .. }), "first owner should claim global lock" ); let second_claim = runtime .try_claim_global_phase2_job(owner_b, /*lease_seconds*/ 3600) .await .expect("claim global lock from second owner"); assert_eq!(second_claim, Phase2JobClaimOutcome::SkippedRunning); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_global_lock_creates_missing_job_row() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b"); let claim = runtime .try_claim_global_phase2_job(owner_a, /*lease_seconds*/ 3_600) .await .expect("claim global phase2 lock"); let ownership_token = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => { assert_eq!(input_watermark, 0); ownership_token } other => panic!("unexpected phase2 lock claim outcome: {other:?}"), }; let second_claim = runtime .try_claim_global_phase2_job(owner_b, /*lease_seconds*/ 3_600) .await .expect("claim global phase2 lock from second owner"); assert_eq!(second_claim, Phase2JobClaimOutcome::SkippedRunning); assert!( runtime .mark_global_phase2_job_succeeded( ownership_token.as_str(), /*completed_watermark*/ 0, &[] ) .await .expect("mark phase2 lock success") ); let claim_after_success = runtime .try_claim_global_phase2_job(owner_b, /*lease_seconds*/ 3_600) .await .expect("claim global phase2 lock after success"); assert_eq!(claim_after_success, Phase2JobClaimOutcome::SkippedCooldown); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_global_lock_stale_lease_allows_takeover() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 300) .await .expect("enqueue global consolidation"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b"); let initial_claim = runtime .try_claim_global_phase2_job(owner_a, /*lease_seconds*/ 3600) .await .expect("claim initial global lock"); let token_a = match initial_claim { Phase2JobClaimOutcome::Claimed { ownership_token, .. } => ownership_token, other => panic!("unexpected initial claim outcome: {other:?}"), }; sqlx::query("UPDATE jobs SET lease_until = ? WHERE kind = ? AND job_key = ?") .bind(Utc::now().timestamp() - 1) .bind("memory_consolidate_global") .bind("global") .execute(memory_pool(&runtime)) .await .expect("expire global consolidation lease"); let takeover_claim = runtime .try_claim_global_phase2_job(owner_b, /*lease_seconds*/ 3600) .await .expect("claim stale global lock"); let (token_b, input_watermark) = match takeover_claim { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => (ownership_token, input_watermark), other => panic!("unexpected takeover claim outcome: {other:?}"), }; assert_ne!(token_a, token_b); assert_eq!(input_watermark, 300); assert_eq!( runtime .mark_global_phase2_job_succeeded( token_a.as_str(), /*completed_watermark*/ 300, &[] ) .await .expect("mark stale owner success result"), false, "stale owner should lose finalization ownership after takeover" ); assert!( runtime .mark_global_phase2_job_succeeded( token_b.as_str(), /*completed_watermark*/ 300, &[] ) .await .expect("mark takeover owner success"), "takeover owner should finalize consolidation" ); let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn enqueue_global_consolidation_keeps_phase2_input_watermark_monotonic() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 500) .await .expect("enqueue initial consolidation"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a"); let claim_a = runtime .try_claim_global_phase2_job(owner_a, /*lease_seconds*/ 3_600) .await .expect("claim initial consolidation"); let token_a = match claim_a { Phase2JobClaimOutcome::Claimed { ownership_token, input_watermark, } => { assert_eq!(input_watermark, 500); ownership_token } other => panic!("unexpected initial phase2 claim outcome: {other:?}"), }; assert!( runtime .mark_global_phase2_job_succeeded( token_a.as_str(), /*completed_watermark*/ 500, &[] ) .await .expect("mark initial phase2 success"), "initial phase2 success should finalize" ); runtime .enqueue_global_consolidation(/*input_watermark*/ 400) .await .expect("enqueue lower-watermark consolidation"); age_phase2_success_beyond_cooldown(&runtime).await; let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b"); let claim_b = runtime .try_claim_global_phase2_job(owner_b, /*lease_seconds*/ 3_600) .await .expect("claim lower-watermark consolidation"); match claim_b { Phase2JobClaimOutcome::Claimed { input_watermark, .. } => { assert!( input_watermark > 500, "lower-watermark enqueue should still advance the bookkeeping watermark" ); } other => panic!("unexpected lower-watermark phase2 claim outcome: {other:?}"), } let _ = tokio::fs::remove_dir_all(codex_home).await; } #[tokio::test] async fn phase2_failure_fallback_updates_unowned_running_job() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 400) .await .expect("enqueue global consolidation"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner"); let claim = runtime .try_claim_global_phase2_job(owner, /*lease_seconds*/ 3_600) .await .expect("claim global consolidation"); let ownership_token = match claim { Phase2JobClaimOutcome::Claimed { ownership_token, .. } => ownership_token, other => panic!("unexpected claim outcome: {other:?}"), }; sqlx::query("UPDATE jobs SET ownership_token = NULL WHERE kind = ? AND job_key = ?") .bind("memory_consolidate_global") .bind("global") .execute(memory_pool(&runtime)) .await .expect("clear ownership token"); assert_eq!( runtime .mark_global_phase2_job_failed( ownership_token.as_str(), "lost", /*retry_delay_seconds*/ 3_600 ) .await .expect("mark phase2 failed with strict ownership"), false, "strict failure update should not match unowned running job" ); assert!( runtime .mark_global_phase2_job_failed_if_unowned( ownership_token.as_str(), "lost", /*retry_delay_seconds*/ 3_600 ) .await .expect("fallback failure update should match unowned running job"), "fallback failure update should transition the unowned running job" ); let claim = runtime .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) .await .expect("claim after fallback failure"); assert_eq!(claim, Phase2JobClaimOutcome::SkippedRetryUnavailable); let _ = tokio::fs::remove_dir_all(codex_home).await; }}