diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index b4b271ead..3277e5650 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -5997,7 +5997,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 581869a60..d808219e4 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -9126,7 +9126,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, @@ -14480,7 +14481,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v1/GetConversationSummaryResponse.json b/codex-rs/app-server-protocol/schema/json/v1/GetConversationSummaryResponse.json index 954ac28ed..a54b7ddac 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/GetConversationSummaryResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/GetConversationSummaryResponse.json @@ -113,7 +113,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v1/ListConversationsResponse.json b/codex-rs/app-server-protocol/schema/json/v1/ListConversationsResponse.json index b7e3b8f8f..bb5a0f486 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/ListConversationsResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/ListConversationsResponse.json @@ -113,7 +113,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json index 10b9d5000..de1823e9d 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json @@ -666,7 +666,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json index 6fe793cc1..96e376197 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json @@ -472,7 +472,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json index 444db8943..e0d4cc4f7 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json @@ -472,7 +472,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json index 9b2fb7400..4385565f9 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json @@ -666,7 +666,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json index b6ec97be1..9f81c5ea7 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json @@ -472,7 +472,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json index 0d22e0881..fd75b56a2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json @@ -666,7 +666,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json index 6a9324409..3ad16408e 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json @@ -472,7 +472,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json index 2c1c9950f..3f8bd6686 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json @@ -472,7 +472,8 @@ { "enum": [ "review", - "compact" + "compact", + "memory_consolidation" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/typescript/SubAgentSource.ts b/codex-rs/app-server-protocol/schema/typescript/SubAgentSource.ts index d6da7a466..269a411be 100644 --- a/codex-rs/app-server-protocol/schema/typescript/SubAgentSource.ts +++ b/codex-rs/app-server-protocol/schema/typescript/SubAgentSource.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ThreadId } from "./ThreadId"; -export type SubAgentSource = "review" | "compact" | { "thread_spawn": { parent_thread_id: ThreadId, depth: number, } } | { "other": string }; +export type SubAgentSource = "review" | "compact" | { "thread_spawn": { parent_thread_id: ThreadId, depth: number, } } | "memory_consolidation" | { "other": string }; diff --git a/codex-rs/codex-api/src/requests/headers.rs b/codex-rs/codex-api/src/requests/headers.rs index 02f8c61c3..d1ab83410 100644 --- a/codex-rs/codex-api/src/requests/headers.rs +++ b/codex-rs/codex-api/src/requests/headers.rs @@ -17,6 +17,9 @@ pub(crate) fn subagent_header(source: &Option) -> Option match sub { codex_protocol::protocol::SubAgentSource::Review => Some("review".to_string()), codex_protocol::protocol::SubAgentSource::Compact => Some("compact".to_string()), + codex_protocol::protocol::SubAgentSource::MemoryConsolidation => { + Some("memory_consolidation".to_string()) + } codex_protocol::protocol::SubAgentSource::ThreadSpawn { .. } => { Some("collab_spawn".to_string()) } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index bbf9b26c9..ec0d1e54f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -325,6 +325,9 @@ impl ModelClient { let subagent = match sub { crate::protocol::SubAgentSource::Review => "review".to_string(), crate::protocol::SubAgentSource::Compact => "compact".to_string(), + crate::protocol::SubAgentSource::MemoryConsolidation => { + "memory_consolidation".to_string() + } crate::protocol::SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(), crate::protocol::SubAgentSource::Other(label) => label.clone(), }; diff --git a/codex-rs/core/src/memories/dispatch.rs b/codex-rs/core/src/memories/dispatch.rs deleted file mode 100644 index bfa6daa5a..000000000 --- a/codex-rs/core/src/memories/dispatch.rs +++ /dev/null @@ -1,688 +0,0 @@ -use crate::codex::Session; -use crate::config::Config; -use crate::config::Constrained; -use crate::memories::memory_root; -use crate::memories::metrics; -use crate::memories::phase_two; -use crate::memories::phase2::spawn_phase2_completion_task; -use crate::memories::prompts::build_consolidation_prompt; -use crate::memories::storage::rebuild_raw_memories_file_from_memories; -use crate::memories::storage::sync_rollout_summaries_from_memories; -use codex_protocol::protocol::AskForApproval; -use codex_protocol::protocol::SandboxPolicy; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; -use codex_protocol::user_input::UserInput; -use codex_utils_absolute_path::AbsolutePathBuf; -use std::sync::Arc; -use tracing::debug; -use tracing::info; -use tracing::warn; - -//TODO(jif) clean. - -fn completion_watermark( - claimed_watermark: i64, - latest_memories: &[codex_state::Stage1Output], -) -> i64 { - latest_memories - .iter() - .map(|memory| memory.source_updated_at.timestamp()) - .max() - .unwrap_or(claimed_watermark) - .max(claimed_watermark) -} - -pub(in crate::memories) async fn run_global_memory_consolidation( - session: &Arc, - config: Arc, -) -> bool { - let otel_manager = &session.services.otel_manager; - let Some(state_db) = session.services.state_db.as_deref() else { - warn!("state db unavailable; skipping global memory consolidation"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "skipped_state_db_unavailable")], - ); - return false; - }; - - let claim = match state_db - .try_claim_global_phase2_job(session.conversation_id, phase_two::JOB_LEASE_SECONDS) - .await - { - Ok(claim) => claim, - Err(err) => { - warn!("state db try_claim_global_phase2_job failed during memories startup: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_claim")], - ); - return false; - } - }; - let (ownership_token, claimed_watermark) = match claim { - codex_state::Phase2JobClaimOutcome::Claimed { - ownership_token, - input_watermark, - } => { - otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "claimed")]); - (ownership_token, input_watermark) - } - codex_state::Phase2JobClaimOutcome::SkippedNotDirty => { - debug!("memory phase-2 global lock is up-to-date; skipping consolidation"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "skipped_not_dirty")], - ); - return false; - } - codex_state::Phase2JobClaimOutcome::SkippedRunning => { - debug!("memory phase-2 global consolidation already running; skipping"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "skipped_running")], - ); - return false; - } - }; - - let root = memory_root(&config.codex_home); - let consolidation_config = { - let mut consolidation_config = config.as_ref().clone(); - consolidation_config.cwd = root.clone(); - consolidation_config.permissions.approval_policy = - Constrained::allow_only(AskForApproval::Never); - let mut writable_roots = Vec::new(); - match AbsolutePathBuf::from_absolute_path(consolidation_config.codex_home.clone()) { - Ok(codex_home) => writable_roots.push(codex_home), - Err(err) => warn!( - "memory phase-2 consolidation could not add codex_home writable root {}: {err}", - consolidation_config.codex_home.display() - ), - } - let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite { - writable_roots, - read_only_access: Default::default(), - network_access: false, - exclude_tmpdir_env_var: false, - exclude_slash_tmp: false, - }; - if let Err(err) = consolidation_config - .permissions - .sandbox_policy - .set(consolidation_sandbox_policy) - { - warn!("memory phase-2 consolidation sandbox policy was rejected by constraints: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_sandbox_policy")], - ); - let _ = state_db - .mark_global_phase2_job_failed( - &ownership_token, - "consolidation sandbox policy was rejected by constraints", - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await; - return false; - } - consolidation_config - }; - - let latest_memories = match state_db - .list_stage1_outputs_for_global(phase_two::MAX_RAW_MEMORIES_FOR_GLOBAL) - .await - { - Ok(memories) => memories, - Err(err) => { - warn!("state db list_stage1_outputs_for_global failed during consolidation: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_load_stage1_outputs")], - ); - let _ = state_db - .mark_global_phase2_job_failed( - &ownership_token, - "failed to read stage-1 outputs before global consolidation", - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await; - return false; - } - }; - if !latest_memories.is_empty() { - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_INPUT, - latest_memories.len() as i64, - &[], - ); - } - let completion_watermark = completion_watermark(claimed_watermark, &latest_memories); - if let Err(err) = sync_rollout_summaries_from_memories(&root, &latest_memories).await { - warn!("failed syncing local memory artifacts for global consolidation: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_sync_artifacts")], - ); - let _ = state_db - .mark_global_phase2_job_failed( - &ownership_token, - "failed syncing local memory artifacts", - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await; - return false; - } - - if let Err(err) = rebuild_raw_memories_file_from_memories(&root, &latest_memories).await { - warn!("failed rebuilding raw memories aggregate for global consolidation: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_rebuild_raw_memories")], - ); - let _ = state_db - .mark_global_phase2_job_failed( - &ownership_token, - "failed rebuilding raw memories aggregate", - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await; - return false; - } - if latest_memories.is_empty() { - debug!("memory phase-2 has no stage-1 outputs; finalized local memory artifacts"); - let _ = state_db - .mark_global_phase2_job_succeeded(&ownership_token, completion_watermark) - .await; - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "succeeded_no_input")], - ); - return false; - } - - let prompt = build_consolidation_prompt(&root); - let input = vec![UserInput::Text { - text: prompt, - text_elements: vec![], - }]; - let source = SessionSource::SubAgent(SubAgentSource::Other( - phase_two::MEMORY_CONSOLIDATION_SUBAGENT_LABEL.to_string(), - )); - - match session - .services - .agent_control - .spawn_agent(consolidation_config, input, Some(source)) - .await - { - Ok(consolidation_agent_id) => { - info!( - "memory phase-2 global consolidation agent started: agent_id={consolidation_agent_id}" - ); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "agent_spawned")], - ); - spawn_phase2_completion_task( - session.as_ref(), - ownership_token, - completion_watermark, - consolidation_agent_id, - ); - true - } - Err(err) => { - warn!("failed to spawn global memory consolidation agent: {err}"); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_spawn_agent")], - ); - let _ = state_db - .mark_global_phase2_job_failed( - &ownership_token, - "failed to spawn consolidation agent", - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await; - false - } - } -} - -#[cfg(test)] -mod tests { - use super::completion_watermark; - use super::run_global_memory_consolidation; - use crate::CodexAuth; - use crate::ThreadManager; - use crate::agent::control::AgentControl; - use crate::codex::Session; - use crate::codex::make_session_and_context; - use crate::config::Config; - use crate::config::test_config; - use crate::memories::memory_root; - use crate::memories::raw_memories_file; - use crate::memories::rollout_summaries_dir; - use chrono::Utc; - use codex_protocol::ThreadId; - use codex_protocol::protocol::AskForApproval; - use codex_protocol::protocol::Op; - use codex_protocol::protocol::SandboxPolicy; - use codex_protocol::protocol::SessionSource; - use codex_state::Phase2JobClaimOutcome; - use codex_state::Stage1Output; - use codex_state::ThreadMetadataBuilder; - use pretty_assertions::assert_eq; - use std::path::PathBuf; - use std::sync::Arc; - use tempfile::TempDir; - - struct DispatchHarness { - _codex_home: TempDir, - config: Arc, - session: Arc, - manager: ThreadManager, - state_db: Arc, - } - - impl DispatchHarness { - async fn new() -> Self { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let mut config = test_config(); - config.codex_home = codex_home.path().to_path_buf(); - config.cwd = config.codex_home.clone(); - let config = Arc::new(config); - - let state_db = codex_state::StateRuntime::init( - config.codex_home.clone(), - config.model_provider_id.clone(), - None, - ) - .await - .expect("initialize state db"); - - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - ); - let (mut session, _turn_context) = make_session_and_context().await; - session.services.state_db = Some(Arc::clone(&state_db)); - session.services.agent_control = manager.agent_control(); - - Self { - _codex_home: codex_home, - config, - session: Arc::new(session), - manager, - state_db, - } - } - - async fn seed_stage1_output(&self, source_updated_at: i64) { - let thread_id = ThreadId::new(); - let mut metadata_builder = ThreadMetadataBuilder::new( - thread_id, - self.config - .codex_home - .join(format!("rollout-{thread_id}.jsonl")), - Utc::now(), - SessionSource::Cli, - ); - metadata_builder.cwd = self.config.cwd.clone(); - metadata_builder.model_provider = Some(self.config.model_provider_id.clone()); - let metadata = metadata_builder.build(&self.config.model_provider_id); - - self.state_db - .upsert_thread(&metadata) - .await - .expect("upsert thread metadata"); - - let claim = self - .state_db - .try_claim_stage1_job( - thread_id, - self.session.conversation_id, - source_updated_at, - 3_600, - 64, - ) - .await - .expect("claim stage-1 job"); - let ownership_token = match claim { - codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, - other => panic!("unexpected stage-1 claim outcome: {other:?}"), - }; - assert!( - self.state_db - .mark_stage1_job_succeeded( - thread_id, - &ownership_token, - source_updated_at, - "raw memory", - "rollout summary", - ) - .await - .expect("mark stage-1 success"), - "stage-1 success should enqueue global consolidation" - ); - } - - async fn shutdown_threads(&self) { - self.manager - .remove_and_close_all_threads() - .await - .expect("shutdown spawned threads"); - } - - fn user_input_ops_count(&self) -> usize { - self.manager - .captured_ops() - .into_iter() - .filter(|(_, op)| matches!(op, Op::UserInput { .. })) - .count() - } - } - - #[test] - fn completion_watermark_never_regresses_below_claimed_input_watermark() { - let stage1_output = Stage1Output { - thread_id: ThreadId::new(), - source_updated_at: chrono::DateTime::::from_timestamp(123, 0) - .expect("valid source_updated_at timestamp"), - raw_memory: "raw memory".to_string(), - rollout_summary: "rollout summary".to_string(), - cwd: PathBuf::from("/tmp/workspace"), - generated_at: chrono::DateTime::::from_timestamp(124, 0) - .expect("valid generated_at timestamp"), - }; - - let completion = completion_watermark(1_000, &[stage1_output]); - assert_eq!(completion, 1_000); - } - - #[tokio::test] - async fn dispatch_reclaims_stale_global_lock_and_starts_consolidation() { - let harness = DispatchHarness::new().await; - harness.seed_stage1_output(100).await; - - let stale_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), 0) - .await - .expect("claim stale global lock"); - assert!( - matches!(stale_claim, Phase2JobClaimOutcome::Claimed { .. }), - "stale lock precondition should be claimed" - ); - - let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; - assert!( - scheduled, - "dispatch should reclaim stale lock and spawn one agent" - ); - - let running_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim while running"); - assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning); - - let user_input_ops = harness.user_input_ops_count(); - assert_eq!(user_input_ops, 1); - let thread_ids = harness.manager.list_thread_ids().await; - assert_eq!(thread_ids.len(), 1); - let subagent = harness - .manager - .get_thread(thread_ids[0]) - .await - .expect("get consolidation thread"); - let config_snapshot = subagent.config_snapshot().await; - assert_eq!(config_snapshot.approval_policy, AskForApproval::Never); - assert_eq!(config_snapshot.cwd, memory_root(&harness.config.codex_home)); - match config_snapshot.sandbox_policy { - SandboxPolicy::WorkspaceWrite { writable_roots, .. } => { - assert!( - writable_roots - .iter() - .any(|root| root.as_path() == harness.config.codex_home.as_path()), - "consolidation subagent should have codex_home as writable root" - ); - } - other => panic!("unexpected sandbox policy: {other:?}"), - } - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_schedules_only_one_agent_while_lock_is_running() { - let harness = DispatchHarness::new().await; - harness.seed_stage1_output(200).await; - - let first_run = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; - let second_run = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; - - assert!(first_run, "first dispatch should schedule consolidation"); - assert!( - !second_run, - "second dispatch should skip while the global lock is running" - ); - - let user_input_ops = harness.user_input_ops_count(); - assert_eq!(user_input_ops, 1); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_with_dirty_job_and_no_stage1_outputs_skips_spawn_and_clears_dirty_flag() { - let harness = DispatchHarness::new().await; - harness - .state_db - .enqueue_global_consolidation(999) - .await - .expect("enqueue global consolidation"); - - let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; - assert!( - !scheduled, - "dispatch should not spawn when no stage-1 outputs are available" - ); - assert_eq!(harness.user_input_ops_count(), 0); - - let claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim global job after empty dispatch"); - assert_eq!( - claim, - Phase2JobClaimOutcome::SkippedNotDirty, - "empty dispatch should finalize global job as up-to-date" - ); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_with_empty_stage1_outputs_rebuilds_local_artifacts() { - let harness = DispatchHarness::new().await; - let root = memory_root(&harness.config.codex_home); - let summaries_dir = rollout_summaries_dir(&root); - tokio::fs::create_dir_all(&summaries_dir) - .await - .expect("create rollout summaries dir"); - - let stale_summary_path = summaries_dir.join(format!("{}.md", ThreadId::new())); - tokio::fs::write(&stale_summary_path, "stale summary\n") - .await - .expect("write stale rollout summary"); - let raw_memories_path = raw_memories_file(&root); - tokio::fs::write(&raw_memories_path, "stale raw memories\n") - .await - .expect("write stale raw memories"); - let memory_index_path = root.join("MEMORY.md"); - tokio::fs::write(&memory_index_path, "stale memory index\n") - .await - .expect("write stale memory index"); - let memory_summary_path = root.join("memory_summary.md"); - tokio::fs::write(&memory_summary_path, "stale memory summary\n") - .await - .expect("write stale memory summary"); - let stale_skill_file = root.join("skills/demo/SKILL.md"); - tokio::fs::create_dir_all( - stale_skill_file - .parent() - .expect("skills subdirectory parent should exist"), - ) - .await - .expect("create stale skills dir"); - tokio::fs::write(&stale_skill_file, "stale skill\n") - .await - .expect("write stale skill"); - - harness - .state_db - .enqueue_global_consolidation(999) - .await - .expect("enqueue global consolidation"); - - let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; - assert!( - !scheduled, - "dispatch should skip subagent spawn when no stage-1 outputs are available" - ); - - assert!( - !tokio::fs::try_exists(&stale_summary_path) - .await - .expect("check stale summary existence"), - "empty consolidation should prune stale rollout summary files" - ); - let raw_memories = tokio::fs::read_to_string(&raw_memories_path) - .await - .expect("read rebuilt raw memories"); - assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n"); - assert!( - !tokio::fs::try_exists(&memory_index_path) - .await - .expect("check memory index existence"), - "empty consolidation should remove stale MEMORY.md" - ); - assert!( - !tokio::fs::try_exists(&memory_summary_path) - .await - .expect("check memory summary existence"), - "empty consolidation should remove stale memory_summary.md" - ); - assert!( - !tokio::fs::try_exists(&stale_skill_file) - .await - .expect("check stale skill existence"), - "empty consolidation should remove stale skills artifacts" - ); - assert!( - !tokio::fs::try_exists(root.join("skills")) - .await - .expect("check skills dir existence"), - "empty consolidation should remove stale skills directory" - ); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_marks_job_for_retry_when_spawn_agent_fails() { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let mut config = test_config(); - config.codex_home = codex_home.path().to_path_buf(); - config.cwd = config.codex_home.clone(); - let config = Arc::new(config); - - let state_db = codex_state::StateRuntime::init( - config.codex_home.clone(), - config.model_provider_id.clone(), - None, - ) - .await - .expect("initialize state db"); - - let (mut session, _turn_context) = make_session_and_context().await; - session.services.state_db = Some(Arc::clone(&state_db)); - session.services.agent_control = AgentControl::default(); - let session = Arc::new(session); - - let thread_id = ThreadId::new(); - let mut metadata_builder = ThreadMetadataBuilder::new( - thread_id, - config.codex_home.join(format!("rollout-{thread_id}.jsonl")), - Utc::now(), - SessionSource::Cli, - ); - metadata_builder.cwd = config.cwd.clone(); - metadata_builder.model_provider = Some(config.model_provider_id.clone()); - let metadata = metadata_builder.build(&config.model_provider_id); - state_db - .upsert_thread(&metadata) - .await - .expect("upsert thread metadata"); - - let claim = state_db - .try_claim_stage1_job(thread_id, session.conversation_id, 100, 3_600, 64) - .await - .expect("claim stage-1 job"); - let ownership_token = match claim { - codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, - other => panic!("unexpected stage-1 claim outcome: {other:?}"), - }; - assert!( - state_db - .mark_stage1_job_succeeded( - thread_id, - &ownership_token, - 100, - "raw memory", - "rollout summary", - ) - .await - .expect("mark stage-1 success"), - "stage-1 success should enqueue global consolidation" - ); - - let scheduled = run_global_memory_consolidation(&session, Arc::clone(&config)).await; - assert!( - !scheduled, - "dispatch should return false when consolidation subagent cannot be spawned" - ); - - let retry_claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim global job after spawn failure"); - assert_eq!( - retry_claim, - Phase2JobClaimOutcome::SkippedNotDirty, - "spawn failures should leave the job in retry backoff instead of running" - ); - } -} diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs index 179c90cbb..2892f1faa 100644 --- a/codex-rs/core/src/memories/mod.rs +++ b/codex-rs/core/src/memories/mod.rs @@ -4,7 +4,6 @@ //! - Phase 1: select rollouts, extract stage-1 raw memories, persist stage-1 outputs, and enqueue consolidation. //! - Phase 2: claim a global consolidation lock, materialize consolidation inputs, and dispatch one consolidation agent. -mod dispatch; mod phase1; mod phase2; pub(crate) mod prompts; @@ -58,8 +57,6 @@ mod phase_one { /// Phase 2 (aka `Consolidation`). mod phase_two { - /// Subagent source label used to identify consolidation tasks. - pub(super) const MEMORY_CONSOLIDATION_SUBAGENT_LABEL: &str = "memory_consolidation"; /// Maximum number of recent raw memories retained for global consolidation. pub(super) const MAX_RAW_MEMORIES_FOR_GLOBAL: usize = 1_024; /// Lease duration (seconds) for phase-2 consolidation job ownership. diff --git a/codex-rs/core/src/memories/phase2.rs b/codex-rs/core/src/memories/phase2.rs index ef0f4642b..a43094a72 100644 --- a/codex-rs/core/src/memories/phase2.rs +++ b/codex-rs/core/src/memories/phase2.rs @@ -1,136 +1,351 @@ use crate::agent::AgentStatus; use crate::agent::status::is_final as is_final_agent_status; use crate::codex::Session; +use crate::config::Config; +use crate::memories::memory_root; use crate::memories::metrics; use crate::memories::phase_two; +use crate::memories::prompts::build_consolidation_prompt; +use crate::memories::storage::rebuild_raw_memories_file_from_memories; +use crate::memories::storage::sync_rollout_summaries_from_memories; +use codex_config::Constrained; use codex_protocol::ThreadId; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; +use codex_protocol::user_input::UserInput; +use codex_state::StateRuntime; +use codex_utils_absolute_path::AbsolutePathBuf; use std::sync::Arc; use std::time::Duration; use tokio::sync::watch; -use tracing::debug; -use tracing::info; use tracing::warn; -pub(in crate::memories) fn spawn_phase2_completion_task( - session: &Session, - ownership_token: String, - completion_watermark: i64, - consolidation_agent_id: ThreadId, -) { - let state_db = session.services.state_db.clone(); - let agent_control = session.services.agent_control.clone(); - let otel_manager = session.services.otel_manager.clone(); - - tokio::spawn(async move { - let Some(state_db) = state_db else { - return; - }; - - let status_rx = match agent_control.subscribe_status(consolidation_agent_id).await { - Ok(status_rx) => status_rx, - Err(err) => { - warn!( - "failed to subscribe to global memory consolidation agent {consolidation_agent_id}: {err}" - ); - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_subscribe_status")], - ); - mark_phase2_failed_with_recovery( - state_db.as_ref(), - &ownership_token, - "failed to subscribe to consolidation agent status", - ) - .await; - return; - } - }; - - let final_status = run_phase2_completion_task( - Arc::clone(&state_db), - ownership_token, - completion_watermark, - consolidation_agent_id, - status_rx, - ) - .await; - if matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) { - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "failed_agent_unavailable")], - ); - return; - } - if is_phase2_success(&final_status) { - otel_manager.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - 1, - &[("status", "succeeded")], - ); - } else { - otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "failed")]); - } - - tokio::spawn(async move { - if let Err(err) = agent_control.shutdown_agent(consolidation_agent_id).await { - warn!( - "failed to auto-close global memory consolidation agent {consolidation_agent_id}: {err}" - ); - } - }); - }); +#[derive(Debug, Clone, Default)] +struct Claim { + token: String, + watermark: i64, } -async fn run_phase2_completion_task( - state_db: Arc, - ownership_token: String, - completion_watermark: i64, - consolidation_agent_id: ThreadId, - mut status_rx: watch::Receiver, -) -> AgentStatus { - let final_status = { +#[derive(Debug, Clone, Default)] +struct Counters { + input: i64, +} + +/// Runs memory phase 2 (aka consolidation) in strict order. The method represents the linear +/// flow of the consolidation phase. +pub(super) async fn run(session: &Arc, config: Arc) { + let Some(db) = session.services.state_db.as_deref() else { + // This should not happen. + return; + }; + let root = memory_root(&config.codex_home); + + // 1. Claim the job. + let claim = match job::claim(session, db).await { + Ok(claim) => claim, + Err(e) => { + session.services.otel_manager.counter( + metrics::MEMORY_PHASE_TWO_JOBS, + 1, + &[("status", e)], + ); + return; + } + }; + + // 2. Get the config for the agent + let Some(agent_config) = agent::get_config(config.clone()) else { + // If we can't get the config, we can't consolidate. + tracing::error!("failed to get agent config"); + job::failed(session, db, &claim, "failed_sandbox_policy").await; + return; + }; + + // 3. Query the memories + let raw_memories = match db + .list_stage1_outputs_for_global(phase_two::MAX_RAW_MEMORIES_FOR_GLOBAL) + .await + { + Ok(memories) => memories, + Err(err) => { + tracing::error!("failed to list stage1 outputs from global: {}", err); + job::failed(session, db, &claim, "failed_load_stage1_outputs").await; + return; + } + }; + let new_watermark = get_watermark(claim.watermark, &raw_memories); + + // 4. Update the file system by syncing the raw memories with the one extracted from DB at + // step 3 + // [`rollout_summaries/`] + if let Err(err) = sync_rollout_summaries_from_memories(&root, &raw_memories).await { + tracing::error!("failed syncing local memory artifacts for global consolidation: {err}"); + job::failed(session, db, &claim, "failed_sync_artifacts").await; + return; + } + // [`raw_memories.md`] + if let Err(err) = rebuild_raw_memories_file_from_memories(&root, &raw_memories).await { + tracing::error!("failed syncing local memory artifacts for global consolidation: {err}"); + job::failed(session, db, &claim, "failed_rebuild_raw_memories").await; + return; + } + if raw_memories.is_empty() { + // We check only after sync of the file system. + job::succeed(session, db, &claim, new_watermark, "succeeded_no_input").await; + return; + } + + // 5. Spawn the agent + let prompt = agent::get_prompt(config); + let source = SessionSource::SubAgent(SubAgentSource::MemoryConsolidation); + let thread_id = match session + .services + .agent_control + .spawn_agent(agent_config, prompt, Some(source)) + .await + { + Ok(thread_id) => thread_id, + Err(err) => { + tracing::error!("failed to spawn global memory consolidation agent: {err}"); + job::failed(session, db, &claim, "failed_spawn_agent").await; + return; + } + }; + + // 6. Spawn the agent handler. + agent::handle(session, claim, new_watermark, thread_id); + + // 7. Metrics and logs. + let counters = Counters { + input: raw_memories.len() as i64, + }; + emit_metrics(session, counters); +} + +mod job { + use super::*; + + pub(super) async fn claim( + session: &Arc, + db: &StateRuntime, + ) -> Result { + let otel_manager = &session.services.otel_manager; + let claim = db + .try_claim_global_phase2_job(session.conversation_id, phase_two::JOB_LEASE_SECONDS) + .await + .map_err(|e| { + tracing::error!("failed to claim job: {}", e); + "failed_claim" + })?; + let (token, watermark) = match claim { + codex_state::Phase2JobClaimOutcome::Claimed { + ownership_token, + input_watermark, + } => { + otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "claimed")]); + (ownership_token, input_watermark) + } + codex_state::Phase2JobClaimOutcome::SkippedNotDirty => return Err("skipped_not_dirty"), + codex_state::Phase2JobClaimOutcome::SkippedRunning => return Err("skipped_running"), + }; + + Ok(Claim { token, watermark }) + } + + pub(super) async fn failed( + session: &Arc, + db: &StateRuntime, + claim: &Claim, + reason: &'static str, + ) { + session.services.otel_manager.counter( + metrics::MEMORY_PHASE_TWO_JOBS, + 1, + &[("status", reason)], + ); + if matches!( + db.mark_global_phase2_job_failed( + &claim.token, + reason, + phase_two::JOB_RETRY_DELAY_SECONDS, + ) + .await, + Ok(false) + ) { + let _ = db + .mark_global_phase2_job_failed_if_unowned( + &claim.token, + reason, + phase_two::JOB_RETRY_DELAY_SECONDS, + ) + .await; + } + } + + pub(super) async fn succeed( + session: &Arc, + db: &StateRuntime, + claim: &Claim, + completion_watermark: i64, + reason: &'static str, + ) { + session.services.otel_manager.counter( + metrics::MEMORY_PHASE_TWO_JOBS, + 1, + &[("status", reason)], + ); + let _ = db + .mark_global_phase2_job_succeeded(&claim.token, completion_watermark) + .await; + } +} + +mod agent { + use super::*; + + pub(super) fn get_config(config: Arc) -> Option { + let root = memory_root(&config.codex_home); + let mut consolidation_config = config.as_ref().clone(); + + consolidation_config.cwd = root; + // Approval policy + consolidation_config.permissions.approval_policy = + Constrained::allow_only(AskForApproval::Never); + + // Sandbox policy + let mut writable_roots = Vec::new(); + match AbsolutePathBuf::from_absolute_path(consolidation_config.codex_home.clone()) { + Ok(codex_home) => writable_roots.push(codex_home), + Err(err) => warn!( + "memory phase-2 consolidation could not add codex_home writable root {}: {err}", + consolidation_config.codex_home.display() + ), + } + // The consolidation agent only needs local codex_home write access and no network. + let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite { + writable_roots, + read_only_access: Default::default(), + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }; + consolidation_config + .permissions + .sandbox_policy + .set(consolidation_sandbox_policy) + .ok()?; + + Some(consolidation_config) + } + + pub(super) fn get_prompt(config: Arc) -> Vec { + let root = memory_root(&config.codex_home); + let prompt = build_consolidation_prompt(&root); + vec![UserInput::Text { + text: prompt, + text_elements: vec![], + }] + } + + /// Handle the agent while it is running. + pub(super) fn handle( + session: &Arc, + claim: Claim, + new_watermark: i64, + thread_id: ThreadId, + ) { + let Some(db) = session.services.state_db.clone() else { + return; + }; + let session = session.clone(); + + tokio::spawn(async move { + let agent_control = session.services.agent_control.clone(); + + // TODO(jif) we might have a very small race here. + let rx = match agent_control.subscribe_status(thread_id).await { + Ok(rx) => rx, + Err(err) => { + tracing::error!("agent_control.subscribe_status failed: {err:?}"); + job::failed(&session, &db, &claim, "failed_subscribe_status").await; + return; + } + }; + + // Loop the agent until we have the final status. + let final_status = loop_agent( + db.clone(), + claim.token.clone(), + new_watermark, + thread_id, + rx, + ) + .await; + + if matches!(final_status, AgentStatus::Completed(_)) { + job::succeed(&session, &db, &claim, new_watermark, "succeeded").await; + } else { + job::failed(&session, &db, &claim, "failed_agent").await; + } + + // Fire and forget close of the agent. + if !matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) { + tokio::spawn(async move { + if let Err(err) = agent_control.shutdown_agent(thread_id).await { + warn!( + "failed to auto-close global memory consolidation agent {thread_id}: {err}" + ); + } + }); + } else { + tracing::warn!("The agent was already gone"); + } + }); + } + + async fn loop_agent( + db: Arc, + token: String, + _new_watermark: i64, + thread_id: ThreadId, + mut rx: watch::Receiver, + ) -> AgentStatus { let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(phase_two::JOB_HEARTBEAT_SECONDS)); heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { - let status = status_rx.borrow().clone(); + let status = rx.borrow().clone(); if is_final_agent_status(&status) { break status; } tokio::select! { - changed = status_rx.changed() => { - if changed.is_err() { - warn!( - "lost status updates for global memory consolidation agent {consolidation_agent_id}" + update = rx.changed() => { + if update.is_err() { + tracing::warn!( + "lost status updates for global memory consolidation agent {thread_id}" ); break status; } } _ = heartbeat_interval.tick() => { - match state_db + match db .heartbeat_global_phase2_job( - &ownership_token, + &token, phase_two::JOB_LEASE_SECONDS, ) .await { Ok(true) => {} Ok(false) => { - warn!( - "memory phase-2 heartbeat lost global ownership; finalizing as failure" - ); break AgentStatus::Errored( "lost global phase-2 ownership during heartbeat".to_string(), ); } Err(err) => { - warn!( - "state db heartbeat_global_phase2_job failed during memories startup: {err}" - ); break AgentStatus::Errored(format!( "phase-2 heartbeat update failed: {err}" )); @@ -139,281 +354,30 @@ async fn run_phase2_completion_task( } } } - }; + } +} - let phase2_success = is_phase2_success(&final_status); - info!( - "memory phase-2 global consolidation complete: agent_id={consolidation_agent_id} success={phase2_success} final_status={final_status:?}" +pub(super) fn get_watermark( + claimed_watermark: i64, + latest_memories: &[codex_state::Stage1Output], +) -> i64 { + latest_memories + .iter() + .map(|memory| memory.source_updated_at.timestamp()) + .max() + .unwrap_or(claimed_watermark) + .max(claimed_watermark) // todo double check the claimed here. +} + +fn emit_metrics(session: &Arc, counters: Counters) { + let otel = session.services.otel_manager.clone(); + if counters.input > 0 { + otel.counter(metrics::MEMORY_PHASE_TWO_INPUT, counters.input, &[]); + } + + otel.counter( + metrics::MEMORY_PHASE_TWO_JOBS, + 1, + &[("status", "agent_spawned")], ); - - if phase2_success { - match state_db - .mark_global_phase2_job_succeeded(&ownership_token, completion_watermark) - .await - { - Ok(true) => {} - Ok(false) => { - debug!( - "memory phase-2 success finalization skipped after global ownership changed" - ); - } - Err(err) => { - warn!( - "state db mark_global_phase2_job_succeeded failed during memories startup: {err}" - ); - } - } - return final_status; - } - - let failure_reason = phase2_failure_reason(&final_status); - mark_phase2_failed_with_recovery(state_db.as_ref(), &ownership_token, &failure_reason).await; - warn!( - "memory phase-2 global consolidation agent finished with non-success status: agent_id={consolidation_agent_id} final_status={final_status:?}" - ); - final_status -} - -async fn mark_phase2_failed_with_recovery( - state_db: &codex_state::StateRuntime, - ownership_token: &str, - failure_reason: &str, -) { - match state_db - .mark_global_phase2_job_failed( - ownership_token, - failure_reason, - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await - { - Ok(true) => {} - Ok(false) => match state_db - .mark_global_phase2_job_failed_if_unowned( - ownership_token, - failure_reason, - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await - { - Ok(true) => { - debug!( - "memory phase-2 failure finalization applied fallback update for unowned running job" - ); - } - Ok(false) => { - debug!( - "memory phase-2 failure finalization skipped after global ownership changed" - ); - } - Err(err) => { - warn!( - "state db mark_global_phase2_job_failed_if_unowned failed during memories startup: {err}" - ); - } - }, - Err(err) => { - warn!("state db mark_global_phase2_job_failed failed during memories startup: {err}"); - } - } -} - -fn is_phase2_success(final_status: &AgentStatus) -> bool { - matches!(final_status, AgentStatus::Completed(_)) -} - -fn phase2_failure_reason(final_status: &AgentStatus) -> String { - format!("consolidation agent finished with status {final_status:?}") -} - -#[cfg(test)] -mod tests { - use super::is_phase2_success; - use super::phase2_failure_reason; - use super::run_phase2_completion_task; - use crate::agent::AgentStatus; - use codex_protocol::ThreadId; - use codex_state::Phase2JobClaimOutcome; - use pretty_assertions::assert_eq; - use std::sync::Arc; - - #[test] - fn phase2_success_only_for_completed_status() { - assert!(is_phase2_success(&AgentStatus::Completed(None))); - assert!(!is_phase2_success(&AgentStatus::Running)); - assert!(!is_phase2_success(&AgentStatus::Errored( - "oops".to_string() - ))); - } - - #[test] - fn phase2_failure_reason_includes_status() { - let status = AgentStatus::Errored("boom".to_string()); - let reason = phase2_failure_reason(&status); - assert!(reason.contains("consolidation agent finished with status")); - assert!(reason.contains("boom")); - } - - #[tokio::test] - async fn phase2_completion_marks_succeeded_for_completed_status() { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let state_db = Arc::new( - codex_state::StateRuntime::init( - codex_home.path().to_path_buf(), - "test-provider".to_string(), - None, - ) - .await - .expect("initialize state runtime"), - ); - let owner = ThreadId::new(); - state_db - .enqueue_global_consolidation(123) - .await - .expect("enqueue global consolidation"); - let claim = state_db - .try_claim_global_phase2_job(owner, 3_600) - .await - .expect("claim global phase-2 job"); - let ownership_token = match claim { - Phase2JobClaimOutcome::Claimed { - ownership_token, .. - } => ownership_token, - other => panic!("unexpected phase-2 claim outcome: {other:?}"), - }; - - let (_status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Completed(None)); - run_phase2_completion_task( - Arc::clone(&state_db), - ownership_token.clone(), - 123, - ThreadId::new(), - status_rx, - ) - .await; - - let up_to_date_claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim up-to-date global job"); - assert_eq!(up_to_date_claim, Phase2JobClaimOutcome::SkippedNotDirty); - - state_db - .enqueue_global_consolidation(124) - .await - .expect("enqueue advanced consolidation watermark"); - let rerun_claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim rerun global job"); - assert!( - matches!(rerun_claim, Phase2JobClaimOutcome::Claimed { .. }), - "advanced watermark should be claimable after success finalization" - ); - } - - #[tokio::test] - async fn phase2_completion_marks_failed_when_status_updates_are_lost() { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let state_db = Arc::new( - codex_state::StateRuntime::init( - codex_home.path().to_path_buf(), - "test-provider".to_string(), - None, - ) - .await - .expect("initialize state runtime"), - ); - state_db - .enqueue_global_consolidation(456) - .await - .expect("enqueue global consolidation"); - let claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim global phase-2 job"); - let ownership_token = match claim { - Phase2JobClaimOutcome::Claimed { - ownership_token, .. - } => ownership_token, - other => panic!("unexpected phase-2 claim outcome: {other:?}"), - }; - - let (status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Running); - drop(status_tx); - run_phase2_completion_task( - Arc::clone(&state_db), - ownership_token, - 456, - ThreadId::new(), - status_rx, - ) - .await; - - let claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim after failure finalization"); - assert_eq!( - claim, - Phase2JobClaimOutcome::SkippedNotDirty, - "failure finalization should leave global job in retry-backoff, not running ownership" - ); - } - - #[tokio::test] - async fn phase2_completion_heartbeat_loss_does_not_steal_active_other_owner() { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let state_db = Arc::new( - codex_state::StateRuntime::init( - codex_home.path().to_path_buf(), - "test-provider".to_string(), - None, - ) - .await - .expect("initialize state runtime"), - ); - state_db - .enqueue_global_consolidation(789) - .await - .expect("enqueue global consolidation"); - let claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim global phase-2 job"); - let claimed_token = match claim { - Phase2JobClaimOutcome::Claimed { - ownership_token, .. - } => ownership_token, - other => panic!("unexpected phase-2 claim outcome: {other:?}"), - }; - - let (_status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Running); - run_phase2_completion_task( - Arc::clone(&state_db), - "non-owner-token".to_string(), - 789, - ThreadId::new(), - status_rx, - ) - .await; - - let claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), 3_600) - .await - .expect("claim after heartbeat ownership loss"); - assert_eq!( - claim, - Phase2JobClaimOutcome::SkippedRunning, - "heartbeat ownership-loss handling should not steal a live owner lease" - ); - assert_eq!( - state_db - .mark_global_phase2_job_succeeded(claimed_token.as_str(), 789) - .await - .expect("mark original owner success"), - true, - "the original owner should still be able to finalize" - ); - } } diff --git a/codex-rs/core/src/memories/start.rs b/codex-rs/core/src/memories/start.rs index 5b47b6c2b..3ac5fae85 100644 --- a/codex-rs/core/src/memories/start.rs +++ b/codex-rs/core/src/memories/start.rs @@ -2,6 +2,7 @@ use crate::codex::Session; use crate::config::Config; use crate::features::Feature; use crate::memories::phase1; +use crate::memories::phase2; use codex_protocol::protocol::SessionSource; use std::sync::Arc; use tracing::warn; @@ -36,6 +37,6 @@ pub(crate) fn start_memories_startup_task( // Run phase 1. phase1::run(&session).await; // Run phase 2. - crate::memories::dispatch::run_global_memory_consolidation(&session, config).await; + phase2::run(&session, config).await; }); } diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs index 1baea576c..82e4517a5 100644 --- a/codex-rs/core/src/memories/tests.rs +++ b/codex-rs/core/src/memories/tests.rs @@ -87,3 +87,498 @@ async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only assert!(raw_memories.contains(&keep_id)); assert!(raw_memories.contains("cwd: /tmp/workspace")); } + +mod phase2 { + use crate::CodexAuth; + use crate::ThreadManager; + use crate::agent::AgentControl; + use crate::codex::Session; + use crate::codex::make_session_and_context; + use crate::config::Config; + use crate::config::test_config; + use crate::memories::memory_root; + use crate::memories::phase2; + use crate::memories::raw_memories_file; + use crate::memories::rollout_summaries_dir; + use chrono::Utc; + use codex_config::Constrained; + use codex_protocol::ThreadId; + use codex_protocol::protocol::AskForApproval; + use codex_protocol::protocol::Op; + use codex_protocol::protocol::SandboxPolicy; + use codex_protocol::protocol::SessionSource; + use codex_state::Phase2JobClaimOutcome; + use codex_state::Stage1Output; + use codex_state::ThreadMetadataBuilder; + use std::path::PathBuf; + use std::sync::Arc; + use tempfile::TempDir; + + fn stage1_output_with_source_updated_at(source_updated_at: i64) -> Stage1Output { + Stage1Output { + thread_id: ThreadId::new(), + source_updated_at: chrono::DateTime::::from_timestamp(source_updated_at, 0) + .expect("valid source_updated_at timestamp"), + raw_memory: "raw memory".to_string(), + rollout_summary: "rollout summary".to_string(), + cwd: PathBuf::from("/tmp/workspace"), + generated_at: chrono::DateTime::::from_timestamp(source_updated_at + 1, 0) + .expect("valid generated_at timestamp"), + } + } + + struct DispatchHarness { + _codex_home: TempDir, + config: Arc, + session: Arc, + manager: ThreadManager, + state_db: Arc, + } + + impl DispatchHarness { + async fn new() -> Self { + let codex_home = tempfile::tempdir().expect("create temp codex home"); + let mut config = test_config(); + config.codex_home = codex_home.path().to_path_buf(); + config.cwd = config.codex_home.clone(); + let config = Arc::new(config); + + let state_db = codex_state::StateRuntime::init( + config.codex_home.clone(), + config.model_provider_id.clone(), + None, + ) + .await + .expect("initialize state db"); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone(), + ); + let (mut session, _turn_context) = make_session_and_context().await; + session.services.state_db = Some(Arc::clone(&state_db)); + session.services.agent_control = manager.agent_control(); + + Self { + _codex_home: codex_home, + config, + session: Arc::new(session), + manager, + state_db, + } + } + + async fn seed_stage1_output(&self, source_updated_at: i64) { + let thread_id = ThreadId::new(); + let mut metadata_builder = ThreadMetadataBuilder::new( + thread_id, + self.config + .codex_home + .join(format!("rollout-{thread_id}.jsonl")), + Utc::now(), + SessionSource::Cli, + ); + metadata_builder.cwd = self.config.cwd.clone(); + metadata_builder.model_provider = Some(self.config.model_provider_id.clone()); + let metadata = metadata_builder.build(&self.config.model_provider_id); + + self.state_db + .upsert_thread(&metadata) + .await + .expect("upsert thread metadata"); + + let claim = self + .state_db + .try_claim_stage1_job( + thread_id, + self.session.conversation_id, + source_updated_at, + 3_600, + 64, + ) + .await + .expect("claim stage-1 job"); + let ownership_token = match claim { + codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, + other => panic!("unexpected stage-1 claim outcome: {other:?}"), + }; + assert!( + self.state_db + .mark_stage1_job_succeeded( + thread_id, + &ownership_token, + source_updated_at, + "raw memory", + "rollout summary", + ) + .await + .expect("mark stage-1 success"), + "stage-1 success should enqueue global consolidation" + ); + } + + async fn shutdown_threads(&self) { + self.manager + .remove_and_close_all_threads() + .await + .expect("shutdown spawned threads"); + } + + fn user_input_ops_count(&self) -> usize { + self.manager + .captured_ops() + .into_iter() + .filter(|(_, op)| matches!(op, Op::UserInput { .. })) + .count() + } + } + + #[test] + fn completion_watermark_never_regresses_below_claimed_input_watermark() { + let stage1_output = stage1_output_with_source_updated_at(123); + + let completion = phase2::get_watermark(1_000, &[stage1_output]); + pretty_assertions::assert_eq!(completion, 1_000); + } + + #[test] + fn completion_watermark_uses_claimed_watermark_when_there_are_no_memories() { + let completion = phase2::get_watermark(777, &[]); + pretty_assertions::assert_eq!(completion, 777); + } + + #[test] + fn completion_watermark_uses_latest_memory_timestamp_when_it_is_newer() { + let older = stage1_output_with_source_updated_at(123); + let newer = stage1_output_with_source_updated_at(456); + + let completion = phase2::get_watermark(200, &[older, newer]); + pretty_assertions::assert_eq!(completion, 456); + } + + #[tokio::test] + async fn dispatch_skips_when_global_job_is_not_dirty() { + let harness = DispatchHarness::new().await; + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + } + + #[tokio::test] + async fn dispatch_skips_when_global_job_is_already_running() { + let harness = DispatchHarness::new().await; + harness + .state_db + .enqueue_global_consolidation(123) + .await + .expect("enqueue global consolidation"); + let claimed = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim running global lock"); + assert!( + matches!(claimed, Phase2JobClaimOutcome::Claimed { .. }), + "precondition should claim the running lock" + ); + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + let running_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim while lock is still running"); + pretty_assertions::assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning); + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + } + + #[tokio::test] + async fn dispatch_reclaims_stale_global_lock_and_starts_consolidation() { + let harness = DispatchHarness::new().await; + harness.seed_stage1_output(100).await; + + let stale_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 0) + .await + .expect("claim stale global lock"); + assert!( + matches!(stale_claim, Phase2JobClaimOutcome::Claimed { .. }), + "stale lock precondition should be claimed" + ); + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + let running_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim while running"); + pretty_assertions::assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning); + + let user_input_ops = harness.user_input_ops_count(); + pretty_assertions::assert_eq!(user_input_ops, 1); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 1); + let subagent = harness + .manager + .get_thread(thread_ids[0]) + .await + .expect("get consolidation thread"); + let config_snapshot = subagent.config_snapshot().await; + pretty_assertions::assert_eq!(config_snapshot.approval_policy, AskForApproval::Never); + pretty_assertions::assert_eq!(config_snapshot.cwd, memory_root(&harness.config.codex_home)); + match config_snapshot.sandbox_policy { + SandboxPolicy::WorkspaceWrite { writable_roots, .. } => { + assert!( + writable_roots + .iter() + .any(|root| root.as_path() == harness.config.codex_home.as_path()), + "consolidation subagent should have codex_home as writable root" + ); + } + other => panic!("unexpected sandbox policy: {other:?}"), + } + + harness.shutdown_threads().await; + } + + #[tokio::test] + async fn dispatch_with_empty_stage1_outputs_rebuilds_local_artifacts() { + let harness = DispatchHarness::new().await; + let root = memory_root(&harness.config.codex_home); + let summaries_dir = rollout_summaries_dir(&root); + tokio::fs::create_dir_all(&summaries_dir) + .await + .expect("create rollout summaries dir"); + + let stale_summary_path = summaries_dir.join(format!("{}.md", ThreadId::new())); + tokio::fs::write(&stale_summary_path, "stale summary\n") + .await + .expect("write stale rollout summary"); + let raw_memories_path = raw_memories_file(&root); + tokio::fs::write(&raw_memories_path, "stale raw memories\n") + .await + .expect("write stale raw memories"); + let memory_index_path = root.join("MEMORY.md"); + tokio::fs::write(&memory_index_path, "stale memory index\n") + .await + .expect("write stale memory index"); + let memory_summary_path = root.join("memory_summary.md"); + tokio::fs::write(&memory_summary_path, "stale memory summary\n") + .await + .expect("write stale memory summary"); + let stale_skill_file = root.join("skills/demo/SKILL.md"); + tokio::fs::create_dir_all( + stale_skill_file + .parent() + .expect("skills subdirectory parent should exist"), + ) + .await + .expect("create stale skills dir"); + tokio::fs::write(&stale_skill_file, "stale skill\n") + .await + .expect("write stale skill"); + + harness + .state_db + .enqueue_global_consolidation(999) + .await + .expect("enqueue global consolidation"); + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + assert!( + !tokio::fs::try_exists(&stale_summary_path) + .await + .expect("check stale summary existence"), + "empty consolidation should prune stale rollout summary files" + ); + let raw_memories = tokio::fs::read_to_string(&raw_memories_path) + .await + .expect("read rebuilt raw memories"); + pretty_assertions::assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n"); + assert!( + !tokio::fs::try_exists(&memory_index_path) + .await + .expect("check memory index existence"), + "empty consolidation should remove stale MEMORY.md" + ); + assert!( + !tokio::fs::try_exists(&memory_summary_path) + .await + .expect("check memory summary existence"), + "empty consolidation should remove stale memory_summary.md" + ); + assert!( + !tokio::fs::try_exists(&stale_skill_file) + .await + .expect("check stale skill existence"), + "empty consolidation should remove stale skills artifacts" + ); + assert!( + !tokio::fs::try_exists(root.join("skills")) + .await + .expect("check skills dir existence"), + "empty consolidation should remove stale skills directory" + ); + let next_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim global job after empty consolidation success"); + pretty_assertions::assert_eq!(next_claim, Phase2JobClaimOutcome::SkippedNotDirty); + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + + harness.shutdown_threads().await; + } + + #[tokio::test] + async fn dispatch_marks_job_for_retry_when_sandbox_policy_cannot_be_overridden() { + let harness = DispatchHarness::new().await; + harness + .state_db + .enqueue_global_consolidation(99) + .await + .expect("enqueue global consolidation"); + let mut constrained_config = harness.config.as_ref().clone(); + constrained_config.permissions.sandbox_policy = + Constrained::allow_only(SandboxPolicy::DangerFullAccess); + + phase2::run(&harness.session, Arc::new(constrained_config)).await; + + let retry_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim global job after sandbox policy failure"); + pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty); + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + } + + #[tokio::test] + async fn dispatch_marks_job_for_retry_when_syncing_artifacts_fails() { + let harness = DispatchHarness::new().await; + harness.seed_stage1_output(100).await; + let root = memory_root(&harness.config.codex_home); + tokio::fs::write(&root, "not a directory") + .await + .expect("create file at memory root"); + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + let retry_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim global job after sync failure"); + pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty); + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + } + + #[tokio::test] + async fn dispatch_marks_job_for_retry_when_rebuilding_raw_memories_fails() { + let harness = DispatchHarness::new().await; + harness.seed_stage1_output(100).await; + let root = memory_root(&harness.config.codex_home); + tokio::fs::create_dir_all(raw_memories_file(&root)) + .await + .expect("create raw_memories.md as a directory"); + + phase2::run(&harness.session, Arc::clone(&harness.config)).await; + + let retry_claim = harness + .state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim global job after rebuild failure"); + pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty); + pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); + let thread_ids = harness.manager.list_thread_ids().await; + pretty_assertions::assert_eq!(thread_ids.len(), 0); + } + + #[tokio::test] + async fn dispatch_marks_job_for_retry_when_spawn_agent_fails() { + let codex_home = tempfile::tempdir().expect("create temp codex home"); + let mut config = test_config(); + config.codex_home = codex_home.path().to_path_buf(); + config.cwd = config.codex_home.clone(); + let config = Arc::new(config); + + let state_db = codex_state::StateRuntime::init( + config.codex_home.clone(), + config.model_provider_id.clone(), + None, + ) + .await + .expect("initialize state db"); + + let (mut session, _turn_context) = make_session_and_context().await; + session.services.state_db = Some(Arc::clone(&state_db)); + session.services.agent_control = AgentControl::default(); + let session = Arc::new(session); + + let thread_id = ThreadId::new(); + let mut metadata_builder = ThreadMetadataBuilder::new( + thread_id, + config.codex_home.join(format!("rollout-{thread_id}.jsonl")), + Utc::now(), + SessionSource::Cli, + ); + metadata_builder.cwd = config.cwd.clone(); + metadata_builder.model_provider = Some(config.model_provider_id.clone()); + let metadata = metadata_builder.build(&config.model_provider_id); + state_db + .upsert_thread(&metadata) + .await + .expect("upsert thread metadata"); + + let claim = state_db + .try_claim_stage1_job(thread_id, session.conversation_id, 100, 3_600, 64) + .await + .expect("claim stage-1 job"); + let ownership_token = match claim { + codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, + other => panic!("unexpected stage-1 claim outcome: {other:?}"), + }; + assert!( + state_db + .mark_stage1_job_succeeded( + thread_id, + &ownership_token, + 100, + "raw memory", + "rollout summary", + ) + .await + .expect("mark stage-1 success"), + "stage-1 success should enqueue global consolidation" + ); + + phase2::run(&session, Arc::clone(&config)).await; + + let retry_claim = state_db + .try_claim_global_phase2_job(ThreadId::new(), 3_600) + .await + .expect("claim global job after spawn failure"); + pretty_assertions::assert_eq!( + retry_claim, + Phase2JobClaimOutcome::SkippedNotDirty, + "spawn failures should leave the job in retry backoff instead of running" + ); + } +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 3b65aa05c..e44bf4c1d 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1825,6 +1825,7 @@ pub enum SubAgentSource { parent_thread_id: ThreadId, depth: i32, }, + MemoryConsolidation, Other(String), } @@ -1846,6 +1847,7 @@ impl fmt::Display for SubAgentSource { match self { SubAgentSource::Review => f.write_str("review"), SubAgentSource::Compact => f.write_str("compact"), + SubAgentSource::MemoryConsolidation => f.write_str("memory_consolidation"), SubAgentSource::ThreadSpawn { parent_thread_id, depth,