diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3cd7fc8099..84feec0cef 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1326,6 +1326,7 @@ impl Session { &sess, Arc::clone(&config), &session_configuration.session_source, + Some(sess.conversation_id.to_string()), ); Ok(sess) @@ -3639,10 +3640,15 @@ mod handlers { state.session_configuration.session_source.clone() }; - crate::memories::start_memories_startup_task(sess, Arc::clone(config), &session_source); + crate::memories::start_memories_startup_task( + sess, + Arc::clone(config), + &session_source, + Some(sub_id.clone()), + ); sess.send_event_raw(Event { - id: sub_id.clone(), + id: sub_id, msg: EventMsg::Warning(WarningEvent { message: "Memory update triggered.".to_string(), }), diff --git a/codex-rs/core/src/memories/startup/dispatch.rs b/codex-rs/core/src/memories/startup/dispatch.rs index aaf21fed8a..9f0bb586c4 100644 --- a/codex-rs/core/src/memories/startup/dispatch.rs +++ b/codex-rs/core/src/memories/startup/dispatch.rs @@ -20,6 +20,7 @@ use super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS; use super::super::prompts::build_consolidation_prompt; use super::super::storage::rebuild_raw_memories_file_from_memories; use super::super::storage::sync_rollout_summaries_from_memories; +use super::emit_memory_progress; use super::phase2::spawn_phase2_completion_task; fn completion_watermark( @@ -37,9 +38,16 @@ fn completion_watermark( pub(super) async fn run_global_memory_consolidation( session: &Arc, config: Arc, + progress_sub_id: &Option, ) -> bool { let Some(state_db) = session.services.state_db.as_deref() else { warn!("state db unavailable; skipping global memory consolidation"); + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 skipped (state db unavailable)", + ) + .await; return false; }; @@ -50,6 +58,12 @@ pub(super) async fn run_global_memory_consolidation( Ok(claim) => claim, Err(err) => { warn!("state db try_claim_global_phase2_job failed during memories startup: {err}"); + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed to claim global job", + ) + .await; return false; } }; @@ -60,10 +74,13 @@ pub(super) async fn run_global_memory_consolidation( } => (ownership_token, input_watermark), codex_state::Phase2JobClaimOutcome::SkippedNotDirty => { debug!("memory phase-2 global lock is up-to-date; skipping consolidation"); + emit_memory_progress(session.as_ref(), progress_sub_id, "phase 2 up to date").await; return false; } codex_state::Phase2JobClaimOutcome::SkippedRunning => { debug!("memory phase-2 global consolidation already running; skipping"); + emit_memory_progress(session.as_ref(), progress_sub_id, "phase 2 already running") + .await; return false; } }; @@ -100,6 +117,12 @@ pub(super) async fn run_global_memory_consolidation( PHASE_TWO_JOB_RETRY_DELAY_SECONDS, ) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed (sandbox policy rejected)", + ) + .await; return false; } consolidation_config @@ -119,6 +142,12 @@ pub(super) async fn run_global_memory_consolidation( PHASE_TWO_JOB_RETRY_DELAY_SECONDS, ) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed (could not load stage-1 outputs)", + ) + .await; return false; } }; @@ -132,6 +161,12 @@ pub(super) async fn run_global_memory_consolidation( PHASE_TWO_JOB_RETRY_DELAY_SECONDS, ) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed (could not sync local artifacts)", + ) + .await; return false; } @@ -144,6 +179,12 @@ pub(super) async fn run_global_memory_consolidation( PHASE_TWO_JOB_RETRY_DELAY_SECONDS, ) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed (could not rebuild raw memories)", + ) + .await; return false; } if latest_memories.is_empty() { @@ -151,6 +192,12 @@ pub(super) async fn run_global_memory_consolidation( let _ = state_db .mark_global_phase2_job_succeeded(&ownership_token, completion_watermark) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 complete (no stage-1 outputs)", + ) + .await; return false; } @@ -173,11 +220,13 @@ pub(super) async fn run_global_memory_consolidation( info!( "memory phase-2 global consolidation agent started: agent_id={consolidation_agent_id}" ); + emit_memory_progress(session.as_ref(), progress_sub_id, "phase 2 running").await; spawn_phase2_completion_task( - session.as_ref(), + Arc::clone(session), ownership_token, completion_watermark, consolidation_agent_id, + progress_sub_id.clone(), ); true } @@ -190,6 +239,12 @@ pub(super) async fn run_global_memory_consolidation( PHASE_TWO_JOB_RETRY_DELAY_SECONDS, ) .await; + emit_memory_progress( + session.as_ref(), + progress_sub_id, + "phase 2 failed (could not spawn consolidation agent)", + ) + .await; false } } @@ -361,7 +416,8 @@ mod tests { ); let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; + run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config), &None) + .await; assert!( scheduled, "dispatch should reclaim stale lock and spawn one agent" @@ -407,9 +463,11 @@ mod tests { harness.seed_stage1_output(200).await; let first_run = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; + run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config), &None) + .await; let second_run = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; + run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config), &None) + .await; assert!(first_run, "first dispatch should schedule consolidation"); assert!( @@ -433,7 +491,8 @@ mod tests { .expect("enqueue global consolidation"); let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; + run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config), &None) + .await; assert!( !scheduled, "dispatch should not spawn when no stage-1 outputs are available" @@ -498,7 +557,8 @@ mod tests { .expect("enqueue global consolidation"); let scheduled = - run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await; + run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config), &None) + .await; assert!( !scheduled, "dispatch should skip subagent spawn when no stage-1 outputs are available" @@ -600,7 +660,7 @@ mod tests { "stage-1 success should enqueue global consolidation" ); - let scheduled = run_global_memory_consolidation(&session, Arc::clone(&config)).await; + let scheduled = run_global_memory_consolidation(&session, Arc::clone(&config), &None).await; assert!( !scheduled, "dispatch should return false when consolidation subagent cannot be spawned" diff --git a/codex-rs/core/src/memories/startup/mod.rs b/codex-rs/core/src/memories/startup/mod.rs index 476853e0cb..37cb41912c 100644 --- a/codex-rs/core/src/memories/startup/mod.rs +++ b/codex-rs/core/src/memories/startup/mod.rs @@ -12,9 +12,14 @@ 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::BackgroundEventEvent; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; use futures::StreamExt; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tracing::info; use tracing::warn; @@ -49,6 +54,7 @@ pub(crate) fn start_memories_startup_task( session: &Arc, config: Arc, source: &SessionSource, + progress_sub_id: Option, ) { if config.ephemeral || !config.features.enabled(Feature::MemoryTool) @@ -62,7 +68,7 @@ pub(crate) fn start_memories_startup_task( let Some(session) = weak_session.upgrade() else { return; }; - if let Err(err) = run_memories_startup_pipeline(&session, config).await { + if let Err(err) = run_memories_startup_pipeline(&session, config, progress_sub_id).await { warn!("memories startup pipeline failed: {err}"); } }); @@ -77,12 +83,26 @@ pub(crate) fn start_memories_startup_task( pub(super) async fn run_memories_startup_pipeline( session: &Arc, config: Arc, + progress_sub_id: Option, ) -> CodexResult<()> { let Some(state_db) = session.services.state_db.as_deref() else { warn!("state db unavailable for memories startup pipeline; skipping"); + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + "phase 1 skipped (state db unavailable)", + ) + .await; return Ok(()); }; + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + "phase 1 scanning candidates", + ) + .await; + let allowed_sources = INTERACTIVE_SESSION_SOURCES .iter() .map(ToString::to_string) @@ -110,6 +130,12 @@ pub(super) async fn run_memories_startup_pipeline( }; let claimed_count = claimed_candidates.len(); + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + format!("phase 1 running (0/{claimed_count} done)"), + ) + .await; let mut succeeded_count = 0; if claimed_count > 0 { let turn_context = session.new_default_turn().await; @@ -117,14 +143,17 @@ pub(super) async fn run_memories_startup_pipeline( turn_context.as_ref(), turn_context.resolve_turn_metadata_header().await, ); + let completed_count = Arc::new(AtomicUsize::new(0)); succeeded_count = futures::stream::iter(claimed_candidates.into_iter()) .map(|claim| { let session = Arc::clone(session); let stage_one_context = stage_one_context.clone(); + let progress_sub_id = progress_sub_id.clone(); + let completed_count = Arc::clone(&completed_count); async move { let thread = claim.thread; - let stage_one_output = match extract::extract_stage_one_output( + let job_succeeded = match extract::extract_stage_one_output( session.as_ref(), &thread.rollout_path, &thread.cwd, @@ -132,7 +161,6 @@ pub(super) async fn run_memories_startup_pipeline( ) .await { - Ok(output) => output, Err(reason) => { if let Some(state_db) = session.services.state_db.as_deref() { let _ = state_db @@ -144,33 +172,46 @@ pub(super) async fn run_memories_startup_pipeline( ) .await; } - return false; + false + } + Ok(stage_one_output) => { + if let Some(state_db) = session.services.state_db.as_deref() { + if stage_one_output.raw_memory.is_empty() + && stage_one_output.rollout_summary.is_empty() + { + state_db + .mark_stage1_job_succeeded_no_output( + thread.id, + &claim.ownership_token, + ) + .await + .unwrap_or(false) + } else { + state_db + .mark_stage1_job_succeeded( + thread.id, + &claim.ownership_token, + thread.updated_at.timestamp(), + &stage_one_output.raw_memory, + &stage_one_output.rollout_summary, + ) + .await + .unwrap_or(false) + } + } else { + false + } } }; - let Some(state_db) = session.services.state_db.as_deref() else { - return false; - }; - - if stage_one_output.raw_memory.is_empty() - && stage_one_output.rollout_summary.is_empty() - { - return state_db - .mark_stage1_job_succeeded_no_output(thread.id, &claim.ownership_token) - .await - .unwrap_or(false); - } - - state_db - .mark_stage1_job_succeeded( - thread.id, - &claim.ownership_token, - thread.updated_at.timestamp(), - &stage_one_output.raw_memory, - &stage_one_output.rollout_summary, - ) - .await - .unwrap_or(false) + let done = completed_count.fetch_add(1, Ordering::Relaxed) + 1; + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + format!("phase 1 running ({done}/{claimed_count} done)"), + ) + .await; + job_succeeded } }) .buffer_unordered(super::PHASE_ONE_CONCURRENCY_LIMIT) @@ -185,9 +226,16 @@ pub(super) async fn run_memories_startup_pipeline( "memory stage-1 extraction complete: {} job(s) claimed, {} succeeded", claimed_count, succeeded_count ); + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + format!("phase 1 complete ({succeeded_count}/{claimed_count} succeeded)"), + ) + .await; - let consolidation_job_count = - usize::from(dispatch::run_global_memory_consolidation(session, config).await); + let consolidation_job_count = usize::from( + dispatch::run_global_memory_consolidation(session, config, &progress_sub_id).await, + ); info!( "memory consolidation dispatch complete: {} job(s) scheduled", consolidation_job_count @@ -196,6 +244,24 @@ pub(super) async fn run_memories_startup_pipeline( Ok(()) } +pub(super) async fn emit_memory_progress( + session: &Session, + progress_sub_id: &Option, + message: impl Into, +) { + let Some(sub_id) = progress_sub_id.as_ref() else { + return; + }; + session + .send_event_raw(Event { + id: sub_id.clone(), + msg: EventMsg::BackgroundEvent(BackgroundEventEvent { + message: format!("memory startup: {}", message.into()), + }), + }) + .await; +} + #[cfg(test)] mod tests { use super::run_memories_startup_pipeline; @@ -208,7 +274,7 @@ mod tests { let (session, _turn_context) = make_session_and_context().await; let session = Arc::new(session); let config = Arc::new(test_config()); - run_memories_startup_pipeline(&session, config) + run_memories_startup_pipeline(&session, config, None) .await .expect("startup pipeline should skip cleanly without state db"); } diff --git a/codex-rs/core/src/memories/startup/phase2.rs b/codex-rs/core/src/memories/startup/phase2.rs index 137add7830..abc2e54315 100644 --- a/codex-rs/core/src/memories/startup/phase2.rs +++ b/codex-rs/core/src/memories/startup/phase2.rs @@ -12,18 +12,26 @@ use tracing::warn; 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::emit_memory_progress; pub(super) fn spawn_phase2_completion_task( - session: &Session, + session: Arc, ownership_token: String, completion_watermark: i64, consolidation_agent_id: ThreadId, + progress_sub_id: Option, ) { let state_db = session.services.state_db.clone(); let agent_control = session.services.agent_control.clone(); tokio::spawn(async move { let Some(state_db) = state_db else { + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + "phase 2 failed (state db unavailable)", + ) + .await; return; }; @@ -39,6 +47,12 @@ pub(super) fn spawn_phase2_completion_task( "failed to subscribe to consolidation agent status", ) .await; + emit_memory_progress( + session.as_ref(), + &progress_sub_id, + "phase 2 failed (status subscription unavailable)", + ) + .await; return; } }; @@ -51,6 +65,13 @@ pub(super) fn spawn_phase2_completion_task( status_rx, ) .await; + let progress = match &final_status { + AgentStatus::Completed(_) => "phase 2 complete", + AgentStatus::Errored(_) | AgentStatus::NotFound => "phase 2 failed", + AgentStatus::Shutdown => "phase 2 cancelled", + AgentStatus::PendingInit | AgentStatus::Running => "phase 2 failed", + }; + emit_memory_progress(session.as_ref(), &progress_sub_id, progress).await; if matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) { return; } diff --git a/codex-rs/tui/src/bottom_pane/status_line_setup.rs b/codex-rs/tui/src/bottom_pane/status_line_setup.rs index 29bcf7b9c9..7c80ca876c 100644 --- a/codex-rs/tui/src/bottom_pane/status_line_setup.rs +++ b/codex-rs/tui/src/bottom_pane/status_line_setup.rs @@ -91,6 +91,9 @@ pub(crate) enum StatusLineItem { /// Full session UUID. SessionId, + + /// Most recent memory startup pipeline progress update. + MemoryProgress, } impl StatusLineItem { @@ -124,6 +127,9 @@ impl StatusLineItem { StatusLineItem::SessionId => { "Current session identifier (omitted until session starts)" } + StatusLineItem::MemoryProgress => { + "Memory startup phase progress (omitted until available)" + } } } @@ -148,6 +154,7 @@ impl StatusLineItem { StatusLineItem::TotalInputTokens => "17,588 in", StatusLineItem::TotalOutputTokens => "265 out", StatusLineItem::SessionId => "019c19bd-ceb6-73b0-adc8-8ec0397b85cf", + StatusLineItem::MemoryProgress => "memory p2 running", } } } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 0ae056f836..bca2eb1855 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -611,6 +611,8 @@ pub(crate) struct ChatWidget { session_network_proxy: Option, // Shared latch so we only warn once about invalid status-line item IDs. status_line_invalid_items_warned: Arc, + // Most recent memory startup progress message (without prefix). + memory_progress_status: Option, // Cached git branch name for the status line (None if unknown). status_line_branch: Option, // CWD used to resolve the cached branch; change resets branch state. @@ -2119,11 +2121,28 @@ impl ChatWidget { fn on_background_event(&mut self, message: String) { debug!("BackgroundEvent: {message}"); + if let Some(progress) = Self::memory_progress_from_background_event(&message) { + self.memory_progress_status = Some(progress); + self.refresh_status_line(); + } self.bottom_pane.ensure_status_indicator(); self.bottom_pane.set_interrupt_hint_visible(true); self.set_status_header(message); } + fn memory_progress_from_background_event(message: &str) -> Option { + let (prefix, progress) = message.split_once(':')?; + if !prefix.trim().eq_ignore_ascii_case("memory startup") { + return None; + } + let progress = progress.trim(); + if progress.is_empty() { + None + } else { + Some(progress.to_string()) + } + } + fn on_undo_started(&mut self, event: UndoStartedEvent) { self.bottom_pane.ensure_status_indicator(); self.bottom_pane.set_interrupt_hint_visible(false); @@ -2664,6 +2683,7 @@ impl ChatWidget { current_cwd, session_network_proxy: None, status_line_invalid_items_warned, + memory_progress_status: None, status_line_branch: None, status_line_branch_cwd: None, status_line_branch_pending: false, @@ -2829,6 +2849,7 @@ impl ChatWidget { current_cwd, session_network_proxy: None, status_line_invalid_items_warned, + memory_progress_status: None, status_line_branch: None, status_line_branch_cwd: None, status_line_branch_pending: false, @@ -2983,6 +3004,7 @@ impl ChatWidget { current_cwd, session_network_proxy: None, status_line_invalid_items_warned, + memory_progress_status: None, status_line_branch: None, status_line_branch_cwd: None, status_line_branch_pending: false, @@ -4431,6 +4453,7 @@ impl ChatWidget { format_tokens_compact(self.status_line_total_usage().output_tokens) )), StatusLineItem::SessionId => self.thread_id.map(|id| id.to_string()), + StatusLineItem::MemoryProgress => self.memory_progress_status.clone(), } } diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index a126c5f984..a703eefb7f 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -1123,6 +1123,7 @@ async fn make_chatwidget_manual( current_cwd: None, session_network_proxy: None, status_line_invalid_items_warned: Arc::new(AtomicBool::new(false)), + memory_progress_status: None, status_line_branch: None, status_line_branch_cwd: None, status_line_branch_pending: false, @@ -5820,6 +5821,40 @@ async fn status_line_branch_refreshes_after_interrupt() { assert!(chat.status_line_branch_pending); } +#[tokio::test] +async fn status_line_memory_progress_updates_from_background_events() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; + chat.config.tui_status_line = Some(vec!["memory-progress".to_string()]); + + chat.refresh_status_line(); + assert_eq!( + chat.status_line_value_for_item(&StatusLineItem::MemoryProgress), + None + ); + + chat.handle_codex_event(Event { + id: "bg-1".into(), + msg: EventMsg::BackgroundEvent(BackgroundEventEvent { + message: "memory startup: phase 1 running".to_string(), + }), + }); + assert_eq!( + chat.status_line_value_for_item(&StatusLineItem::MemoryProgress), + Some("phase 1 running".to_string()) + ); + + chat.handle_codex_event(Event { + id: "bg-2".into(), + msg: EventMsg::BackgroundEvent(BackgroundEventEvent { + message: "Waiting for `vim`".to_string(), + }), + }); + assert_eq!( + chat.status_line_value_for_item(&StatusLineItem::MemoryProgress), + Some("phase 1 running".to_string()) + ); +} + #[tokio::test] async fn stream_recovery_restores_previous_status_header() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await;