From 3a419ba624ef6b561b29fe25aafd625696e5648a Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 10 Feb 2026 19:30:41 +0000 Subject: [PATCH] nit fixes --- codex-rs/state/src/runtime.rs | 152 +++++++++++++++++++++++++++ codex-rs/state/src/runtime/memory.rs | 56 +++++++++- 2 files changed, 206 insertions(+), 2 deletions(-) diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 102ab50ac7..c301c648d3 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -1365,6 +1365,158 @@ mod tests { let _ = tokio::fs::remove_dir_all(codex_home).await; } + #[tokio::test] + async fn list_stage1_outputs_for_cwd_scope_matches_canonical_equivalent_paths() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("initialize runtime"); + + let workspace = codex_home.join("workspace"); + tokio::fs::create_dir_all(&workspace) + .await + .expect("create workspace"); + let non_normalized_cwd = workspace.join("..").join("workspace"); + let canonical_scope_key = workspace + .canonicalize() + .expect("canonicalize workspace") + .display() + .to_string(); + + 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, + non_normalized_cwd, + )) + .await + .expect("upsert thread"); + + let claim = runtime + .try_claim_stage1_job(thread_id, owner, 100, 3600) + .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(), + 100, + "raw memory", + "summary", + ) + .await + .expect("mark stage1 succeeded"), + "stage1 success should persist output" + ); + + let outputs = runtime + .list_stage1_outputs_for_scope("cwd", canonical_scope_key.as_str(), 10) + .await + .expect("list stage1 outputs for canonical cwd scope"); + assert_eq!(outputs.len(), 1); + assert_eq!(outputs[0].thread_id, thread_id); + assert_eq!(outputs[0].summary, "summary"); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + + #[tokio::test] + async fn list_pending_scope_consolidations_omits_unclaimable_jobs() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("initialize runtime"); + let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); + + runtime + .enqueue_scope_consolidation("cwd", "scope-running", 200) + .await + .expect("enqueue running scope"); + runtime + .enqueue_scope_consolidation("cwd", "scope-backoff", 199) + .await + .expect("enqueue backoff scope"); + runtime + .enqueue_scope_consolidation("cwd", "scope-exhausted", 198) + .await + .expect("enqueue exhausted scope"); + runtime + .enqueue_scope_consolidation("cwd", "scope-claimable-a", 90) + .await + .expect("enqueue claimable scope a"); + runtime + .enqueue_scope_consolidation("cwd", "scope-claimable-b", 89) + .await + .expect("enqueue claimable scope b"); + + let running_claim = runtime + .try_claim_phase2_job("cwd", "scope-running", owner, 3600) + .await + .expect("claim running scope"); + assert!( + matches!(running_claim, Phase2JobClaimOutcome::Claimed { .. }), + "scope-running should be claimed" + ); + + let backoff_claim = runtime + .try_claim_phase2_job("cwd", "scope-backoff", owner, 3600) + .await + .expect("claim backoff scope"); + let backoff_token = match backoff_claim { + Phase2JobClaimOutcome::Claimed { + ownership_token, .. + } => ownership_token, + other => panic!("unexpected backoff claim outcome: {other:?}"), + }; + assert!( + runtime + .mark_phase2_job_failed( + "cwd", + "scope-backoff", + backoff_token.as_str(), + "temporary failure", + 3600, + ) + .await + .expect("mark backoff scope failed"), + "backoff scope should transition to retry backoff" + ); + + sqlx::query("UPDATE jobs SET retry_remaining = 0 WHERE kind = ? AND job_key = ?") + .bind("memory_consolidate_cwd") + .bind("scope-exhausted") + .execute(runtime.pool.as_ref()) + .await + .expect("set exhausted scope retries to zero"); + + let pending = runtime + .list_pending_scope_consolidations(2) + .await + .expect("list pending scopes"); + assert_eq!( + pending, + vec![ + PendingScopeConsolidation { + scope_kind: "cwd".to_string(), + scope_key: "scope-claimable-a".to_string(), + }, + PendingScopeConsolidation { + scope_kind: "cwd".to_string(), + scope_key: "scope-claimable-b".to_string(), + }, + ] + ); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + fn test_thread_metadata( codex_home: &Path, thread_id: ThreadId, diff --git a/codex-rs/state/src/runtime/memory.rs b/codex-rs/state/src/runtime/memory.rs index 16e2f6baaf..dc5b94a92d 100644 --- a/codex-rs/state/src/runtime/memory.rs +++ b/codex-rs/state/src/runtime/memory.rs @@ -4,6 +4,9 @@ use crate::model::Stage1OutputRow; use chrono::Duration; use sqlx::Executor; use sqlx::Sqlite; +use std::collections::HashSet; +use std::path::Path; +use std::path::PathBuf; const JOB_KIND_MEMORY_STAGE1: &str = "memory_stage1"; const JOB_KIND_MEMORY_CONSOLIDATE_CWD: &str = "memory_consolidate_cwd"; @@ -27,6 +30,10 @@ fn scope_kind_for_job_kind(job_kind: &str) -> Option<&'static str> { } } +fn normalize_cwd_for_scope_matching(cwd: &str) -> Option { + Path::new(cwd).canonicalize().ok() +} + impl StateRuntime { pub async fn claim_stage1_jobs_for_startup( &self, @@ -116,7 +123,7 @@ WHERE thread_id = ? let rows = match scope_kind { MEMORY_SCOPE_KIND_CWD => { - sqlx::query( + let mut rows = sqlx::query( r#" SELECT so.thread_id, so.source_updated_at, so.raw_memory, so.summary, so.generated_at FROM stage1_outputs AS so @@ -129,7 +136,46 @@ LIMIT ? .bind(scope_key) .bind(n as i64) .fetch_all(self.pool.as_ref()) - .await? + .await?; + + if rows.len() < n + && let Some(normalized_scope_key) = normalize_cwd_for_scope_matching(scope_key) + { + let mut selected_thread_ids = rows + .iter() + .map(|row| row.try_get::("thread_id")) + .collect::, _>>()?; + let candidate_rows = sqlx::query( + r#" +SELECT so.thread_id, so.source_updated_at, so.raw_memory, so.summary, so.generated_at, t.cwd AS thread_cwd +FROM stage1_outputs AS so +JOIN threads AS t ON t.id = so.thread_id +ORDER BY so.source_updated_at DESC, so.thread_id DESC + "#, + ) + .fetch_all(self.pool.as_ref()) + .await?; + + for row in candidate_rows { + if rows.len() >= n { + break; + } + let thread_id: String = row.try_get("thread_id")?; + if selected_thread_ids.contains(&thread_id) { + continue; + } + let thread_cwd: String = row.try_get("thread_cwd")?; + if let Some(normalized_thread_cwd) = + normalize_cwd_for_scope_matching(&thread_cwd) + && normalized_thread_cwd == normalized_scope_key + { + selected_thread_ids.insert(thread_id); + rows.push(row); + } + } + } + + rows } MEMORY_SCOPE_KIND_USER => { sqlx::query( @@ -446,6 +492,7 @@ WHERE kind = ? AND job_key = ? if limit == 0 { return Ok(Vec::new()); } + let now = Utc::now().timestamp(); let rows = sqlx::query( r#" @@ -454,12 +501,17 @@ FROM jobs WHERE kind IN (?, ?) AND input_watermark IS NOT NULL AND input_watermark > COALESCE(last_success_watermark, 0) + AND retry_remaining > 0 + AND (retry_at IS NULL OR retry_at <= ?) + AND (status != 'running' OR lease_until IS NULL OR lease_until <= ?) ORDER BY input_watermark DESC, kind ASC, job_key ASC LIMIT ? "#, ) .bind(JOB_KIND_MEMORY_CONSOLIDATE_CWD) .bind(JOB_KIND_MEMORY_CONSOLIDATE_USER) + .bind(now) + .bind(now) .bind(limit as i64) .fetch_all(self.pool.as_ref()) .await?;