Codex Handbook
state/src/runtime/memories.rs 5268 lines
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(&current)            .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(&current)            .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(&current)            .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(&current)            .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;    }}