From ba3a423c3dc2cdf5225648b385cf367bb8df0f6d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 10 Feb 2026 17:50:41 +0000 Subject: [PATCH] further cleanup --- codex-rs/core/src/memories/layout.rs | 7 - codex-rs/core/src/memories/mod.rs | 1 - codex-rs/core/src/memories/selection.rs | 47 ---- .../startup/{phase_two => }/dispatch.rs | 18 +- .../startup/{phase_one => }/extract.rs | 20 +- codex-rs/core/src/memories/startup/mod.rs | 199 ++++++++++++-- .../src/memories/startup/phase_one/mod.rs | 247 ------------------ .../src/memories/startup/phase_two/mod.rs | 60 ----- .../memories/startup/{phase_two => }/watch.rs | 8 +- codex-rs/core/src/memories/tests.rs | 108 -------- codex-rs/core/src/memories/types.rs | 15 -- codex-rs/state/src/lib.rs | 1 + codex-rs/state/src/runtime.rs | 66 +++++ codex-rs/state/src/runtime/memory.rs | 190 +++++++++++--- 14 files changed, 410 insertions(+), 577 deletions(-) delete mode 100644 codex-rs/core/src/memories/selection.rs rename codex-rs/core/src/memories/startup/{phase_two => }/dispatch.rs (92%) rename codex-rs/core/src/memories/startup/{phase_one => }/extract.rs (89%) delete mode 100644 codex-rs/core/src/memories/startup/phase_one/mod.rs delete mode 100644 codex-rs/core/src/memories/startup/phase_two/mod.rs rename codex-rs/core/src/memories/startup/{phase_two => }/watch.rs (96%) diff --git a/codex-rs/core/src/memories/layout.rs b/codex-rs/core/src/memories/layout.rs index 1efd1cba95..8df5d2340a 100644 --- a/codex-rs/core/src/memories/layout.rs +++ b/codex-rs/core/src/memories/layout.rs @@ -24,13 +24,6 @@ pub(super) fn memory_root_for_cwd(codex_home: &Path, cwd: &Path) -> PathBuf { codex_home.join("memories").join(bucket).join(MEMORY_SUBDIR) } -/// Returns the DB scope key for a cwd-scoped memory entry. -/// -/// This uses the same normalization/fallback behavior as cwd bucket derivation. -pub(super) fn memory_scope_key_for_cwd(cwd: &Path) -> String { - normalize_cwd_for_memory(cwd).display().to_string() -} - /// Returns the on-disk user-shared memory root directory. pub(super) fn memory_root_for_user(codex_home: &Path) -> PathBuf { codex_home diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs index ec47d0179f..71ea14afa1 100644 --- a/codex-rs/core/src/memories/mod.rs +++ b/codex-rs/core/src/memories/mod.rs @@ -8,7 +8,6 @@ mod layout; mod prompts; mod rollout; mod scope; -mod selection; mod stage_one; mod startup; mod storage; diff --git a/codex-rs/core/src/memories/selection.rs b/codex-rs/core/src/memories/selection.rs deleted file mode 100644 index 4afc91916a..0000000000 --- a/codex-rs/core/src/memories/selection.rs +++ /dev/null @@ -1,47 +0,0 @@ -use chrono::Duration; -use chrono::Utc; -use codex_protocol::ThreadId; -use codex_state::ThreadMetadata; - -use super::types::RolloutCandidate; - -/// Selects rollout candidates that need stage-1 memory extraction. -/// -/// A rollout is selected when it is not the active thread and was updated -/// within the configured max age window. -pub(super) fn select_rollout_candidates_from_db( - items: &[ThreadMetadata], - current_thread_id: ThreadId, - max_items: usize, - max_age_days: i64, -) -> Vec { - if max_items == 0 { - return Vec::new(); - } - - let cutoff = Utc::now() - Duration::days(max_age_days.max(0)); - - let mut candidates = Vec::new(); - - for item in items { - if item.id == current_thread_id { - continue; - } - if item.updated_at < cutoff { - continue; - } - - candidates.push(RolloutCandidate { - thread_id: item.id, - rollout_path: item.rollout_path.clone(), - cwd: item.cwd.clone(), - source_updated_at: item.updated_at.timestamp(), - }); - - if candidates.len() >= max_items { - break; - } - } - - candidates -} diff --git a/codex-rs/core/src/memories/startup/phase_two/dispatch.rs b/codex-rs/core/src/memories/startup/dispatch.rs similarity index 92% rename from codex-rs/core/src/memories/startup/phase_two/dispatch.rs rename to codex-rs/core/src/memories/startup/dispatch.rs index 389cb7c519..c97bae80a1 100644 --- a/codex-rs/core/src/memories/startup/phase_two/dispatch.rs +++ b/codex-rs/core/src/memories/startup/dispatch.rs @@ -8,15 +8,15 @@ use tracing::debug; use tracing::info; use tracing::warn; -use super::super::super::MAX_RAW_MEMORIES_PER_SCOPE; -use super::super::super::MEMORY_CONSOLIDATION_SUBAGENT_LABEL; -use super::super::super::PHASE_TWO_JOB_LEASE_SECONDS; -use super::super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS; -use super::super::super::prompts::build_consolidation_prompt; -use super::super::super::storage::rebuild_memory_summary_from_memories; -use super::super::super::storage::sync_raw_memories_from_memories; -use super::super::super::storage::wipe_consolidation_outputs; -use super::super::MemoryScopeTarget; +use super::super::MAX_RAW_MEMORIES_PER_SCOPE; +use super::super::MEMORY_CONSOLIDATION_SUBAGENT_LABEL; +use super::super::PHASE_TWO_JOB_LEASE_SECONDS; +use super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS; +use super::super::prompts::build_consolidation_prompt; +use super::super::storage::rebuild_memory_summary_from_memories; +use super::super::storage::sync_raw_memories_from_memories; +use super::super::storage::wipe_consolidation_outputs; +use super::MemoryScopeTarget; use super::watch::spawn_phase2_completion_task; pub(super) async fn run_memory_consolidation_for_scope( diff --git a/codex-rs/core/src/memories/startup/phase_one/extract.rs b/codex-rs/core/src/memories/startup/extract.rs similarity index 89% rename from codex-rs/core/src/memories/startup/phase_one/extract.rs rename to codex-rs/core/src/memories/startup/extract.rs index decf2cf252..dd5f25854a 100644 --- a/codex-rs/core/src/memories/startup/phase_one/extract.rs +++ b/codex-rs/core/src/memories/startup/extract.rs @@ -18,21 +18,21 @@ use crate::memories::rollout::serialize_filtered_rollout_response_items; use crate::memories::stage_one::RAW_MEMORY_PROMPT; use crate::memories::stage_one::parse_stage_one_output; use crate::memories::stage_one::stage_one_output_schema; -use crate::memories::types::RolloutCandidate; use crate::memories::types::StageOneOutput; +use std::path::Path; pub(super) async fn extract_stage_one_output( session: &Session, - candidate: &RolloutCandidate, + rollout_path: &Path, stage_one_context: &StageOneRequestContext, ) -> Result { let (rollout_items, _thread_id, parse_errors) = - match RolloutRecorder::load_rollout_items(&candidate.rollout_path).await { + match RolloutRecorder::load_rollout_items(rollout_path).await { Ok(result) => result, Err(err) => { warn!( "failed to load rollout {} for memories: {err}", - candidate.rollout_path.display() + rollout_path.display() ); return Err("failed to load rollout"); } @@ -40,7 +40,7 @@ pub(super) async fn extract_stage_one_output( if parse_errors > 0 { warn!( "rollout {} had {parse_errors} parse errors while preparing stage-1 memory input", - candidate.rollout_path.display() + rollout_path.display() ); } @@ -52,7 +52,7 @@ pub(super) async fn extract_stage_one_output( Err(err) => { warn!( "failed to prepare filtered rollout payload {} for memories: {err}", - candidate.rollout_path.display() + rollout_path.display() ); return Err("failed to serialize filtered rollout"); } @@ -63,7 +63,7 @@ pub(super) async fn extract_stage_one_output( id: None, role: "user".to_string(), content: vec![ContentItem::InputText { - text: build_stage_one_input_message(&candidate.rollout_path, &rollout_contents), + text: build_stage_one_input_message(rollout_path, &rollout_contents), }], end_turn: None, phase: None, @@ -93,7 +93,7 @@ pub(super) async fn extract_stage_one_output( Err(err) => { warn!( "stage-1 memory request failed for rollout {}: {err}", - candidate.rollout_path.display() + rollout_path.display() ); return Err("stage-1 memory request failed"); } @@ -104,7 +104,7 @@ pub(super) async fn extract_stage_one_output( Err(err) => { warn!( "failed while waiting for stage-1 memory response for rollout {}: {err}", - candidate.rollout_path.display() + rollout_path.display() ); return Err("stage-1 memory response stream failed"); } @@ -115,7 +115,7 @@ pub(super) async fn extract_stage_one_output( Err(err) => { warn!( "invalid stage-1 memory payload for rollout {}: {err}", - candidate.rollout_path.display() + rollout_path.display() ); Err("invalid stage-1 memory payload") } diff --git a/codex-rs/core/src/memories/startup/mod.rs b/codex-rs/core/src/memories/startup/mod.rs index 4a92407778..df908bac95 100644 --- a/codex-rs/core/src/memories/startup/mod.rs +++ b/codex-rs/core/src/memories/startup/mod.rs @@ -1,7 +1,9 @@ -mod phase_one; -mod phase_two; +mod dispatch; +mod extract; +mod watch; use crate::codex::Session; +use crate::codex::TurnContext; use crate::config::Config; use crate::error::Result as CodexResult; use crate::features::Feature; @@ -11,9 +13,13 @@ use crate::memories::scope::MEMORY_SCOPE_KEY_USER; use crate::memories::scope::MEMORY_SCOPE_KIND_CWD; use crate::memories::scope::MEMORY_SCOPE_KIND_USER; use crate::rollout::INTERACTIVE_SESSION_SOURCES; -use crate::rollout::list::ThreadSortKey; -use crate::state_db; +use codex_otel::OtelManager; +use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::SessionSource; +use futures::StreamExt; +use serde_json::Value; use std::path::PathBuf; use std::sync::Arc; use tracing::info; @@ -21,6 +27,27 @@ use tracing::warn; pub(super) const PHASE_ONE_THREAD_SCAN_LIMIT: usize = 5_000; +#[derive(Clone)] +struct StageOneRequestContext { + model_info: ModelInfo, + otel_manager: OtelManager, + reasoning_effort: Option, + reasoning_summary: ReasoningSummaryConfig, + turn_metadata_header: Option, +} + +impl StageOneRequestContext { + fn from_turn_context(turn_context: &TurnContext, turn_metadata_header: Option) -> Self { + Self { + model_info: turn_context.model_info.clone(), + otel_manager: turn_context.otel_manager.clone(), + reasoning_effort: turn_context.reasoning_effort, + reasoning_summary: turn_context.reasoning_summary, + turn_metadata_header, + } + } +} + /// Canonical memory scope metadata used by both startup phases. #[derive(Clone, Debug, PartialEq, Eq)] pub(super) struct MemoryScopeTarget { @@ -112,42 +139,162 @@ pub(super) async fn run_memories_startup_pipeline( session: &Arc, config: Arc, ) -> CodexResult<()> { - let Some(page) = state_db::list_threads_db( - session.services.state_db.as_deref(), - &config.codex_home, - PHASE_ONE_THREAD_SCAN_LIMIT, - None, - ThreadSortKey::UpdatedAt, - INTERACTIVE_SESSION_SOURCES, - None, - false, - ) - .await - else { + let Some(state_db) = session.services.state_db.as_deref() else { warn!("state db unavailable for memories startup pipeline; skipping"); return Ok(()); }; - let phase_one = phase_one::run_phase_one(session, &page.items).await; + let allowed_sources = INTERACTIVE_SESSION_SOURCES + .iter() + .map(|value| match serde_json::to_value(value) { + Ok(Value::String(s)) => s, + Ok(other) => other.to_string(), + Err(_) => String::new(), + }) + .collect::>(); + + let claimed_candidates = match state_db + .claim_stage1_jobs_for_startup( + session.conversation_id, + PHASE_ONE_THREAD_SCAN_LIMIT, + super::MAX_ROLLOUTS_PER_STARTUP, + super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS, + allowed_sources.as_slice(), + super::PHASE_ONE_JOB_LEASE_SECONDS, + ) + .await + { + Ok(claims) => claims, + Err(err) => { + warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}"); + Vec::new() + } + }; + + let claimed_count = claimed_candidates.len(); + let mut succeeded_count = 0; + if claimed_count > 0 { + let turn_context = session.new_default_turn().await; + let stage_one_context = StageOneRequestContext::from_turn_context( + turn_context.as_ref(), + turn_context.resolve_turn_metadata_header().await, + ); + + succeeded_count = futures::stream::iter(claimed_candidates.into_iter()) + .map(|claim| { + let session = Arc::clone(session); + let stage_one_context = stage_one_context.clone(); + async move { + let thread = claim.thread; + let stage_one_output = match extract::extract_stage_one_output( + session.as_ref(), + &thread.rollout_path, + &stage_one_context, + ) + .await + { + Ok(output) => output, + Err(reason) => { + if let Some(state_db) = session.services.state_db.as_deref() { + let _ = state_db + .mark_stage1_job_failed( + thread.id, + &claim.ownership_token, + reason, + super::PHASE_ONE_JOB_RETRY_DELAY_SECONDS, + ) + .await; + } + return false; + } + }; + + let Some(state_db) = session.services.state_db.as_deref() else { + return false; + }; + + state_db + .mark_stage1_job_succeeded( + thread.id, + &claim.ownership_token, + thread.updated_at.timestamp(), + &stage_one_output.raw_memory, + &stage_one_output.summary, + ) + .await + .unwrap_or(false) + } + }) + .buffer_unordered(super::PHASE_ONE_CONCURRENCY_LIMIT) + .collect::>() + .await + .into_iter() + .filter(|ok| *ok) + .count(); + } + info!( - "memory phase-1 candidate selection complete: {} claimed candidate(s) from {} indexed thread(s)", - phase_one.claimed_candidate_count, - page.items.len() - ); - info!( - "memory phase-1 extraction complete: {} scope(s) touched", - phase_one.touched_scope_count + "memory stage-1 extraction complete: {} job(s) claimed, {} succeeded", + claimed_count, succeeded_count ); - let consolidation_scope_count = phase_two::run_phase_two(session, config).await; + let consolidation_scope_count = run_consolidation_dispatch(session, config).await; info!( - "memory phase-2 consolidation dispatch complete: {} scope(s) scheduled", + "memory consolidation dispatch complete: {} scope(s) scheduled", consolidation_scope_count ); Ok(()) } +async fn run_consolidation_dispatch(session: &Arc, config: Arc) -> usize { + let scopes = list_consolidation_scopes( + session.as_ref(), + config.as_ref(), + super::MAX_ROLLOUTS_PER_STARTUP, + ) + .await; + let consolidation_scope_count = scopes.len(); + + futures::stream::iter(scopes.into_iter()) + .map(|scope| { + let session = Arc::clone(session); + let config = Arc::clone(&config); + async move { + dispatch::run_memory_consolidation_for_scope(session, config, scope).await; + } + }) + .buffer_unordered(super::PHASE_TWO_CONCURRENCY_LIMIT) + .collect::>() + .await; + + consolidation_scope_count +} + +async fn list_consolidation_scopes( + session: &Session, + config: &Config, + limit: usize, +) -> Vec { + if limit == 0 { + return Vec::new(); + } + + let Some(state_db) = session.services.state_db.as_deref() else { + return Vec::new(); + }; + + let pending_scopes = match state_db.list_pending_scope_consolidations(limit).await { + Ok(scopes) => scopes, + Err(_) => return Vec::new(), + }; + + pending_scopes + .into_iter() + .filter_map(|scope| memory_scope_target_for_pending_scope(config, scope)) + .collect() +} + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/core/src/memories/startup/phase_one/mod.rs b/codex-rs/core/src/memories/startup/phase_one/mod.rs deleted file mode 100644 index 0d6118298a..0000000000 --- a/codex-rs/core/src/memories/startup/phase_one/mod.rs +++ /dev/null @@ -1,247 +0,0 @@ -mod extract; - -use codex_otel::OtelManager; -use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; -use codex_protocol::openai_models::ModelInfo; -use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; -use extract::extract_stage_one_output; -use futures::StreamExt; -use tracing::warn; - -use super::super::MAX_ROLLOUTS_PER_STARTUP; -use super::super::PHASE_ONE_CONCURRENCY_LIMIT; -use super::super::PHASE_ONE_JOB_LEASE_SECONDS; -use super::super::PHASE_ONE_JOB_RETRY_DELAY_SECONDS; -use super::super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS; -use super::super::selection::select_rollout_candidates_from_db; -use super::super::types::RolloutCandidate; -use crate::codex::Session; -use crate::codex::TurnContext; -use crate::memories::layout::memory_scope_key_for_cwd; -use crate::memories::scope::MEMORY_SCOPE_KEY_USER; -use crate::memories::scope::MEMORY_SCOPE_KIND_CWD; -use crate::memories::scope::MEMORY_SCOPE_KIND_USER; -use std::sync::Arc; - -/// Result counters for startup phase-1 extraction. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct PhaseOneRunResult { - /// Number of rollout candidates that were successfully claimed. - pub(super) claimed_candidate_count: usize, - /// Number of scope refresh/enqueue operations performed. - pub(super) touched_scope_count: usize, -} - -#[derive(Clone, Debug)] -pub(super) struct ClaimedStageOneCandidate { - pub(super) candidate: RolloutCandidate, - pub(super) ownership_token: String, -} - -#[derive(Clone)] -struct StageOneRequestContext { - model_info: ModelInfo, - otel_manager: OtelManager, - reasoning_effort: Option, - reasoning_summary: ReasoningSummaryConfig, - turn_metadata_header: Option, -} - -impl StageOneRequestContext { - fn from_turn_context(turn_context: &TurnContext, turn_metadata_header: Option) -> Self { - Self { - model_info: turn_context.model_info.clone(), - otel_manager: turn_context.otel_manager.clone(), - reasoning_effort: turn_context.reasoning_effort, - reasoning_summary: turn_context.reasoning_summary, - turn_metadata_header, - } - } -} - -/// Runs startup phase 1: -/// -/// 1. Select rollout candidates from thread metadata. -/// 2. Claim stage-1 jobs per thread. -/// 3. Execute stage-1 extraction requests in parallel. -/// 4. Persist stage-1 outputs and enqueue consolidation jobs. -pub(super) async fn run_phase_one( - session: &Arc, - thread_items: &[codex_state::ThreadMetadata], -) -> PhaseOneRunResult { - let selection_candidates = select_rollout_candidates_from_db( - thread_items, - session.conversation_id, - super::PHASE_ONE_THREAD_SCAN_LIMIT, - PHASE_ONE_MAX_ROLLOUT_AGE_DAYS, - ); - let claimed_candidates = claim_stage_one_candidates( - session.as_ref(), - selection_candidates, - MAX_ROLLOUTS_PER_STARTUP, - ) - .await; - - if claimed_candidates.is_empty() { - return PhaseOneRunResult { - claimed_candidate_count: 0, - touched_scope_count: 0, - }; - } - - let turn_context = session.new_default_turn().await; - let stage_one_context = StageOneRequestContext::from_turn_context( - turn_context.as_ref(), - turn_context.resolve_turn_metadata_header().await, - ); - - let touched_scope_count = - futures::stream::iter(claimed_candidates.iter().cloned()) - .map(|claimed_candidate| { - let session = Arc::clone(session); - let stage_one_context = stage_one_context.clone(); - async move { - process_memory_candidate(session, claimed_candidate, stage_one_context).await - } - }) - .buffer_unordered(PHASE_ONE_CONCURRENCY_LIMIT) - .collect::>() - .await - .into_iter() - .sum::(); - - PhaseOneRunResult { - claimed_candidate_count: claimed_candidates.len(), - touched_scope_count, - } -} - -async fn claim_stage_one_candidates( - session: &Session, - candidates: Vec, - max_claimed_candidates: usize, -) -> Vec { - if max_claimed_candidates == 0 { - return Vec::new(); - } - - let Some(state_db) = session.services.state_db.as_deref() else { - return Vec::new(); - }; - - let mut claimed = Vec::new(); - for candidate in candidates { - if claimed.len() >= max_claimed_candidates { - break; - } - - let claim = match state_db - .try_claim_stage1_job( - candidate.thread_id, - session.conversation_id, - candidate.source_updated_at, - PHASE_ONE_JOB_LEASE_SECONDS, - ) - .await - { - Ok(claim) => claim, - Err(err) => { - warn!( - "state db try_claim_stage1_job failed for rollout {}: {err}", - candidate.rollout_path.display() - ); - continue; - } - }; - - if let codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } = claim { - claimed.push(ClaimedStageOneCandidate { - candidate, - ownership_token, - }); - } - } - - claimed -} - -async fn process_memory_candidate( - session: Arc, - claimed_candidate: ClaimedStageOneCandidate, - stage_one_context: StageOneRequestContext, -) -> usize { - let candidate = claimed_candidate.candidate; - - let stage_one_output = - match extract_stage_one_output(session.as_ref(), &candidate, &stage_one_context).await { - Ok(output) => output, - Err(reason) => { - if let Some(state_db) = session.services.state_db.as_deref() { - let _ = state_db - .mark_stage1_job_failed( - candidate.thread_id, - &claimed_candidate.ownership_token, - reason, - PHASE_ONE_JOB_RETRY_DELAY_SECONDS, - ) - .await; - } - return 0; - } - }; - - let Some(state_db) = session.services.state_db.as_deref() else { - return 0; - }; - - if !state_db - .mark_stage1_job_succeeded( - candidate.thread_id, - &claimed_candidate.ownership_token, - candidate.source_updated_at, - &stage_one_output.raw_memory, - &stage_one_output.summary, - ) - .await - .unwrap_or(false) - { - return 0; - } - - let mut touched_scope_count = 0; - let cwd_scope_key = memory_scope_key_for_cwd(&candidate.cwd); - - if let Err(err) = state_db - .enqueue_scope_consolidation( - MEMORY_SCOPE_KIND_CWD, - &cwd_scope_key, - candidate.source_updated_at, - ) - .await - { - warn!( - "failed enqueueing scope consolidation for scope {}:{}: {err}", - MEMORY_SCOPE_KIND_CWD, cwd_scope_key - ); - } else { - touched_scope_count += 1; - } - - if let Err(err) = state_db - .enqueue_scope_consolidation( - MEMORY_SCOPE_KIND_USER, - MEMORY_SCOPE_KEY_USER, - candidate.source_updated_at, - ) - .await - { - warn!( - "failed enqueueing scope consolidation for scope {}:{}: {err}", - MEMORY_SCOPE_KIND_USER, MEMORY_SCOPE_KEY_USER - ); - } else { - touched_scope_count += 1; - } - - touched_scope_count -} diff --git a/codex-rs/core/src/memories/startup/phase_two/mod.rs b/codex-rs/core/src/memories/startup/phase_two/mod.rs deleted file mode 100644 index 95b78e42aa..0000000000 --- a/codex-rs/core/src/memories/startup/phase_two/mod.rs +++ /dev/null @@ -1,60 +0,0 @@ -mod dispatch; -mod watch; - -use super::super::MAX_ROLLOUTS_PER_STARTUP; -use super::super::PHASE_TWO_CONCURRENCY_LIMIT; -use super::MemoryScopeTarget; -use super::memory_scope_target_for_pending_scope; -use crate::codex::Session; -use crate::config::Config; -use futures::StreamExt; -use std::sync::Arc; - -/// Runs startup phase 2: -/// -/// 1. Load scopes pending consolidation from the DB. -/// 2. Claim scope jobs. -/// 3. Spawn consolidation agents for owned scopes. -pub(super) async fn run_phase_two(session: &Arc, config: Arc) -> usize { - let scopes = - list_phase2_scopes(session.as_ref(), config.as_ref(), MAX_ROLLOUTS_PER_STARTUP).await; - let consolidation_scope_count = scopes.len(); - - futures::stream::iter(scopes.into_iter()) - .map(|scope| { - let session = Arc::clone(session); - let config = Arc::clone(&config); - async move { - dispatch::run_memory_consolidation_for_scope(session, config, scope).await; - } - }) - .buffer_unordered(PHASE_TWO_CONCURRENCY_LIMIT) - .collect::>() - .await; - - consolidation_scope_count -} - -async fn list_phase2_scopes( - session: &Session, - config: &Config, - limit: usize, -) -> Vec { - if limit == 0 { - return Vec::new(); - } - - let Some(state_db) = session.services.state_db.as_deref() else { - return Vec::new(); - }; - - let pending_scopes = match state_db.list_pending_scope_consolidations(limit).await { - Ok(scopes) => scopes, - Err(_) => return Vec::new(), - }; - - pending_scopes - .into_iter() - .filter_map(|scope| memory_scope_target_for_pending_scope(config, scope)) - .collect() -} diff --git a/codex-rs/core/src/memories/startup/phase_two/watch.rs b/codex-rs/core/src/memories/startup/watch.rs similarity index 96% rename from codex-rs/core/src/memories/startup/phase_two/watch.rs rename to codex-rs/core/src/memories/startup/watch.rs index c4a0483ecf..ea4341e11b 100644 --- a/codex-rs/core/src/memories/startup/phase_two/watch.rs +++ b/codex-rs/core/src/memories/startup/watch.rs @@ -7,10 +7,10 @@ use tracing::debug; use tracing::info; use tracing::warn; -use super::super::super::PHASE_TWO_JOB_HEARTBEAT_SECONDS; -use super::super::super::PHASE_TWO_JOB_LEASE_SECONDS; -use super::super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS; -use super::super::MemoryScopeTarget; +use super::super::PHASE_TWO_JOB_HEARTBEAT_SECONDS; +use super::super::PHASE_TWO_JOB_LEASE_SECONDS; +use super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS; +use super::MemoryScopeTarget; pub(super) fn spawn_phase2_completion_task( session: &Session, diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs index e4b6f442c0..337b553174 100644 --- a/codex-rs/core/src/memories/tests.rs +++ b/codex-rs/core/src/memories/tests.rs @@ -1,15 +1,12 @@ use super::rollout::StageOneResponseItemKinds; use super::rollout::StageOneRolloutFilter; use super::rollout::serialize_filtered_rollout_response_items; -use super::selection::select_rollout_candidates_from_db; use super::stage_one::parse_stage_one_output; use super::storage::rebuild_memory_summary_from_memories; use super::storage::sync_raw_memories_from_memories; use super::storage::wipe_consolidation_outputs; -use crate::memories::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS; use crate::memories::layout::ensure_layout; use crate::memories::layout::memory_root_for_cwd; -use crate::memories::layout::memory_scope_key_for_cwd; use crate::memories::layout::memory_summary_file; use crate::memories::layout::raw_memories_dir; use chrono::TimeZone; @@ -20,43 +17,9 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::RolloutItem; use codex_state::Stage1Output; -use codex_state::ThreadMetadata; use pretty_assertions::assert_eq; -use std::path::PathBuf; use tempfile::tempdir; -fn thread_metadata( - thread_id: ThreadId, - path: PathBuf, - cwd: PathBuf, - title: &str, - updated_at_secs: i64, -) -> ThreadMetadata { - let updated_at = Utc - .timestamp_opt(updated_at_secs, 0) - .single() - .expect("timestamp"); - ThreadMetadata { - id: thread_id, - rollout_path: path, - created_at: updated_at, - updated_at, - source: "cli".to_string(), - model_provider: "openai".to_string(), - cwd, - cli_version: "test".to_string(), - title: title.to_string(), - sandbox_policy: "read_only".to_string(), - approval_mode: "on_request".to_string(), - tokens_used: 0, - first_user_message: None, - archived_at: None, - git_branch: None, - git_sha: None, - git_origin_url: None, - } -} - #[test] fn memory_root_varies_by_cwd() { let dir = tempdir().expect("tempdir"); @@ -100,22 +63,6 @@ fn memory_root_encoding_avoids_component_collisions() { assert!(!root_hash.display().to_string().contains("workspace")); } -#[test] -fn memory_scope_key_uses_normalized_cwd() { - let dir = tempdir().expect("tempdir"); - let workspace = dir.path().join("workspace"); - std::fs::create_dir_all(&workspace).expect("mkdir workspace"); - std::fs::create_dir_all(workspace.join("nested")).expect("mkdir nested"); - - let alias = workspace.join("nested").join(".."); - let normalized = workspace - .canonicalize() - .expect("canonical workspace path should resolve"); - let alias_key = memory_scope_key_for_cwd(&alias); - let normalized_key = memory_scope_key_for_cwd(&normalized); - assert_eq!(alias_key, normalized_key); -} - #[test] fn parse_stage_one_output_accepts_fenced_json() { let raw = "```json\n{\"rawMemory\":\"abc\",\"summary\":\"short\"}\n```"; @@ -224,61 +171,6 @@ fn serialize_filtered_rollout_response_items_filters_by_response_item_kind() { assert!(matches!(parsed[0], ResponseItem::Message { .. })); } -#[test] -fn select_rollout_candidates_filters_by_age_window() { - let dir = tempdir().expect("tempdir"); - let cwd_a = dir.path().join("workspace-a"); - let cwd_b = dir.path().join("workspace-b"); - std::fs::create_dir_all(&cwd_a).expect("mkdir cwd a"); - std::fs::create_dir_all(&cwd_b).expect("mkdir cwd b"); - - let now = Utc::now().timestamp(); - let current_thread_id = ThreadId::default(); - let recent_thread_id = ThreadId::default(); - let old_thread_id = ThreadId::default(); - let recent_two_thread_id = ThreadId::default(); - - let current = thread_metadata( - current_thread_id, - dir.path().join("current.jsonl"), - cwd_a.clone(), - "current", - now, - ); - let recent = thread_metadata( - recent_thread_id, - dir.path().join("recent.jsonl"), - cwd_a, - "recent", - now - 10, - ); - let old = thread_metadata( - old_thread_id, - dir.path().join("old.jsonl"), - cwd_b.clone(), - "old", - now - (PHASE_ONE_MAX_ROLLOUT_AGE_DAYS + 1) * 24 * 60 * 60, - ); - let recent_two = thread_metadata( - recent_two_thread_id, - dir.path().join("recent-two.jsonl"), - cwd_b, - "recent-two", - now - 20, - ); - - let candidates = select_rollout_candidates_from_db( - &[current, recent, old, recent_two], - current_thread_id, - 5, - PHASE_ONE_MAX_ROLLOUT_AGE_DAYS, - ); - - assert_eq!(candidates.len(), 2); - assert_eq!(candidates[0].thread_id, recent_thread_id); - assert_eq!(candidates[1].thread_id, recent_two_thread_id); -} - #[tokio::test] async fn prune_and_rebuild_summary_keeps_latest_memories_only() { let dir = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/memories/types.rs b/codex-rs/core/src/memories/types.rs index e962abfd50..5054ff5410 100644 --- a/codex-rs/core/src/memories/types.rs +++ b/codex-rs/core/src/memories/types.rs @@ -1,19 +1,4 @@ -use codex_protocol::ThreadId; use serde::Deserialize; -use std::path::PathBuf; - -/// A rollout selected for stage-1 memory extraction during startup. -#[derive(Debug, Clone)] -pub(super) struct RolloutCandidate { - /// Source thread identifier for this rollout. - pub(super) thread_id: ThreadId, - /// Absolute path to the rollout file to summarize. - pub(super) rollout_path: PathBuf, - /// Thread working directory used for per-project memory bucketing. - pub(super) cwd: PathBuf, - /// Thread update timestamp (unix seconds) used for stage-1 staleness checks. - pub(super) source_updated_at: i64, -} /// Parsed stage-1 model output payload. #[derive(Debug, Clone, Deserialize)] diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index f0f9bae1a9..a521d99979 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -35,6 +35,7 @@ pub use runtime::PendingScopeConsolidation; pub use runtime::Phase2JobClaimOutcome; pub use runtime::STATE_DB_FILENAME; pub use runtime::STATE_DB_VERSION; +pub use runtime::Stage1JobClaim; pub use runtime::Stage1JobClaimOutcome; pub use runtime::state_db_filename; pub use runtime::state_db_path; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index dea29ae256..102ab50ac7 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -41,6 +41,7 @@ pub const STATE_DB_VERSION: u32 = 4; const MEMORY_SCOPE_KIND_CWD: &str = "cwd"; const MEMORY_SCOPE_KIND_USER: &str = "user"; +const MEMORY_SCOPE_KEY_USER: &str = "user"; const METRIC_DB_INIT: &str = "codex.db.init"; @@ -69,6 +70,13 @@ pub enum Stage1JobClaimOutcome { SkippedRetryExhausted, } +/// Claimed stage-1 job with thread metadata. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Stage1JobClaim { + pub thread: ThreadMetadata, + pub ownership_token: String, +} + /// Scope row used to queue phase-2 consolidation work. #[derive(Debug, Clone, PartialEq, Eq)] pub struct PendingScopeConsolidation { @@ -918,6 +926,7 @@ mod tests { use super::ThreadMetadata; use super::state_db_filename; use chrono::DateTime; + use chrono::Duration; use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::protocol::AskForApproval; @@ -1166,6 +1175,63 @@ mod tests { let _ = tokio::fs::remove_dir_all(codex_home).await; } + #[tokio::test] + async fn claim_stage1_jobs_filters_by_age_and_current_thread() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("initialize runtime"); + + let now = Utc::now(); + let recent_at = now - Duration::seconds(10); + let old_at = now - Duration::days(31); + + let current_thread_id = + ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); + let recent_thread_id = + ThreadId::from_string(&Uuid::new_v4().to_string()).expect("recent thread id"); + let old_thread_id = + ThreadId::from_string(&Uuid::new_v4().to_string()).expect("old thread id"); + + let mut current = + test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current")); + current.created_at = now; + current.updated_at = now; + runtime + .upsert_thread(¤t) + .await + .expect("upsert current"); + + let mut recent = + test_thread_metadata(&codex_home, recent_thread_id, codex_home.join("recent")); + recent.created_at = recent_at; + recent.updated_at = recent_at; + runtime.upsert_thread(&recent).await.expect("upsert recent"); + + let mut 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, + 10, + 5, + 30, + allowed_sources.as_slice(), + 3600, + ) + .await + .expect("claim stage1 jobs"); + + assert_eq!(claims.len(), 1); + assert_eq!(claims[0].thread.id, recent_thread_id); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + #[tokio::test] async fn stage1_output_cascades_on_thread_delete() { let codex_home = unique_temp_dir(); diff --git a/codex-rs/state/src/runtime/memory.rs b/codex-rs/state/src/runtime/memory.rs index bb04199653..16e2f6baaf 100644 --- a/codex-rs/state/src/runtime/memory.rs +++ b/codex-rs/state/src/runtime/memory.rs @@ -1,6 +1,9 @@ use super::*; use crate::Stage1Output; use crate::model::Stage1OutputRow; +use chrono::Duration; +use sqlx::Executor; +use sqlx::Sqlite; const JOB_KIND_MEMORY_STAGE1: &str = "memory_stage1"; const JOB_KIND_MEMORY_CONSOLIDATE_CWD: &str = "memory_consolidate_cwd"; @@ -25,6 +28,63 @@ fn scope_kind_for_job_kind(job_kind: &str) -> Option<&'static str> { } impl StateRuntime { + pub async fn claim_stage1_jobs_for_startup( + &self, + current_thread_id: ThreadId, + scan_limit: usize, + max_claimed: usize, + max_age_days: i64, + allowed_sources: &[String], + lease_seconds: i64, + ) -> anyhow::Result> { + if scan_limit == 0 || max_claimed == 0 { + return Ok(Vec::new()); + } + + let page = self + .list_threads( + scan_limit, + None, + SortKey::UpdatedAt, + allowed_sources, + None, + false, + ) + .await?; + + let cutoff = Utc::now() - Duration::days(max_age_days.max(0)); + let mut claimed = Vec::new(); + + for item in page.items { + if claimed.len() >= max_claimed { + break; + } + if item.id == current_thread_id { + continue; + } + if item.updated_at < cutoff { + continue; + } + + if let Stage1JobClaimOutcome::Claimed { ownership_token } = self + .try_claim_stage1_job( + item.id, + current_thread_id, + item.updated_at.timestamp(), + lease_seconds, + ) + .await? + { + claimed.push(Stage1JobClaim { + thread: item, + ownership_token, + }); + } + } + + Ok(claimed) + } + pub async fn get_stage1_output( &self, thread_id: ThreadId, @@ -42,7 +102,6 @@ WHERE thread_id = ? row.map(|row| Stage1OutputRow::try_from_row(&row).and_then(Stage1Output::try_from)) .transpose() - .map_err(Into::into) } pub async fn list_stage1_outputs_for_scope( @@ -92,7 +151,6 @@ LIMIT ? rows.into_iter() .map(|row| Stage1OutputRow::try_from_row(&row).and_then(Stage1Output::try_from)) .collect::, _>>() - .map_err(Into::into) } pub async fn try_claim_stage1_job( @@ -296,6 +354,34 @@ WHERE excluded.source_updated_at >= stage1_outputs.source_updated_at .execute(&mut *tx) .await?; + if let Some(thread_row) = sqlx::query( + r#" +SELECT cwd +FROM threads +WHERE id = ? + "#, + ) + .bind(thread_id.as_str()) + .fetch_optional(&mut *tx) + .await? + { + let cwd: String = thread_row.try_get("cwd")?; + enqueue_scope_consolidation_with_executor( + &mut *tx, + MEMORY_SCOPE_KIND_CWD, + &cwd, + source_updated_at, + ) + .await?; + enqueue_scope_consolidation_with_executor( + &mut *tx, + MEMORY_SCOPE_KIND_USER, + MEMORY_SCOPE_KEY_USER, + source_updated_at, + ) + .await?; + } + tx.commit().await?; Ok(true) } @@ -344,48 +430,13 @@ WHERE kind = ? AND job_key = ? scope_key: &str, input_watermark: i64, ) -> anyhow::Result<()> { - let Some(job_kind) = job_kind_for_scope(scope_kind) else { - return Ok(()); - }; - - 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 = max(COALESCE(jobs.input_watermark, 0), excluded.input_watermark) - "#, + enqueue_scope_consolidation_with_executor( + self.pool.as_ref(), + scope_kind, + scope_key, + input_watermark, ) - .bind(job_kind) - .bind(scope_key) - .bind(DEFAULT_RETRY_REMAINING) - .bind(input_watermark) - .execute(self.pool.as_ref()) - .await?; - - Ok(()) + .await } pub async fn list_pending_scope_consolidations( @@ -641,3 +692,56 @@ WHERE kind = ? AND job_key = ? Ok(rows_affected > 0) } } + +async fn enqueue_scope_consolidation_with_executor<'e, E>( + executor: E, + scope_kind: &str, + scope_key: &str, + input_watermark: i64, +) -> anyhow::Result<()> +where + E: Executor<'e, Database = Sqlite>, +{ + let Some(job_kind) = job_kind_for_scope(scope_kind) else { + return Ok(()); + }; + + 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 = max(COALESCE(jobs.input_watermark, 0), excluded.input_watermark) + "#, + ) + .bind(job_kind) + .bind(scope_key) + .bind(DEFAULT_RETRY_REMAINING) + .bind(input_watermark) + .execute(executor) + .await?; + + Ok(()) +}