From 22b8626e43445d1c822331f0a6cbabf23d91c5a7 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 9 Feb 2026 16:16:03 +0000 Subject: [PATCH] codex: wire startup memory extraction and consolidation --- codex-rs/core/src/codex.rs | 514 ++++++++++++++++++++++++++++++++++ codex-rs/core/src/features.rs | 2 +- 2 files changed, 515 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 5d0c8c239f..1c1d4d4489 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -12,6 +12,7 @@ use crate::agent::AgentControl; use crate::agent::AgentStatus; use crate::agent::MAX_THREAD_SPAWN_DEPTH; use crate::agent::agent_status_from_event; +use crate::agent::status::is_final as is_final_agent_status; use crate::analytics_client::AnalyticsEventsClient; use crate::analytics_client::build_track_events_context; use crate::compact; @@ -104,6 +105,7 @@ use crate::client::ModelClient; use crate::client::ModelClientSession; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; +use crate::client_common::ResponseStream; use crate::codex_thread::ThreadConfigSnapshot; use crate::compact::collect_user_messages; use crate::config::Config; @@ -138,6 +140,7 @@ use crate::mcp::effective_mcp_servers; use crate::mcp::maybe_prompt_and_install_mcp_dependencies; use crate::mcp::with_codex_apps_mcp; use crate::mcp_connection_manager::McpConnectionManager; +use crate::memories; use crate::mentions::build_connector_slug_counts; use crate::mentions::build_skill_name_counts; use crate::mentions::collect_explicit_app_paths; @@ -178,8 +181,10 @@ use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::protocol::TurnDiffEvent; use crate::protocol::WarningEvent; +use crate::rollout::INTERACTIVE_SESSION_SOURCES; use crate::rollout::RolloutRecorder; use crate::rollout::RolloutRecorderParams; +use crate::rollout::list::ThreadSortKey; use crate::rollout::map_session_init_error; use crate::rollout::metadata; use crate::shell; @@ -781,6 +786,510 @@ impl Session { }); } + fn start_memories_startup_task(self: &Arc, config: Arc, source: &SessionSource) { + if config.ephemeral + || !config.features.enabled(Feature::MemoryTool) + || matches!(source, SessionSource::SubAgent(_)) + { + return; + } + + let weak_sess = Arc::downgrade(self); + tokio::spawn(async move { + let Some(sess) = weak_sess.upgrade() else { + return; + }; + if let Err(err) = sess.run_memories_startup_pipeline(config).await { + warn!("memories startup pipeline failed: {err}"); + } + }); + } + + async fn run_memories_startup_pipeline( + self: &Arc, + config: Arc, + ) -> CodexResult<()> { + let turn_context = self.new_default_turn().await; + + let Some(page) = state_db::list_threads_db( + self.services.state_db.as_deref(), + &config.codex_home, + 200, + None, + ThreadSortKey::UpdatedAt, + INTERACTIVE_SESSION_SOURCES, + None, + false, + ) + .await + else { + warn!("state db unavailable for memories startup pipeline; skipping"); + return Ok(()); + }; + + let mut existing_memories = Vec::new(); + for item in &page.items { + if let Some(memory) = state_db::get_thread_memory( + self.services.state_db.as_deref(), + item.id, + "run_memories_startup_pipeline", + ) + .await + { + existing_memories.push(memory); + } + } + + let candidates = memories::select_rollout_candidates_from_db( + &page.items, + self.conversation_id, + &existing_memories, + memories::MAX_ROLLOUTS_PER_STARTUP, + ); + info!( + "memory phase-1 candidate selection complete: {} candidate(s) from {} indexed thread(s)", + candidates.len(), + page.items.len() + ); + + if candidates.is_empty() { + return Ok(()); + } + + let turn_metadata_header = turn_context.resolve_turn_metadata_header().await; + let model_info = turn_context.model_info.clone(); + let otel_manager = turn_context.otel_manager.clone(); + let reasoning_effort = turn_context.reasoning_effort; + let reasoning_summary = turn_context.reasoning_summary; + let touched_cwds = futures::stream::iter(candidates.into_iter()) + .map(|candidate| { + let session = Arc::clone(self); + let config = Arc::clone(&config); + let model_info = model_info.clone(); + let otel_manager = otel_manager.clone(); + let turn_metadata_header = turn_metadata_header.clone(); + async move { + session + .process_memory_trace_candidate( + config, + candidate, + model_info, + otel_manager, + reasoning_effort, + reasoning_summary, + turn_metadata_header, + ) + .await + } + }) + .buffer_unordered(memories::PHASE_ONE_CONCURRENCY_LIMIT) + .filter_map(futures::future::ready) + .collect::>() + .await; + info!( + "memory phase-1 extraction complete: {} cwd(s) touched", + touched_cwds.len() + ); + + if touched_cwds.is_empty() { + return Ok(()); + } + + let consolidation_cwd_count = touched_cwds.len(); + futures::stream::iter(touched_cwds.into_iter()) + .map(|cwd| { + let session = Arc::clone(self); + let config = Arc::clone(&config); + async move { + session.run_memory_consolidation_for_cwd(config, cwd).await; + } + }) + .buffer_unordered(memories::PHASE_ONE_CONCURRENCY_LIMIT) + .collect::>() + .await; + info!( + "memory phase-2 consolidation dispatch complete: {} cwd(s) scheduled", + consolidation_cwd_count + ); + + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn process_memory_trace_candidate( + self: Arc, + config: Arc, + candidate: memories::RolloutCandidate, + model_info: ModelInfo, + otel_manager: OtelManager, + reasoning_effort: Option, + reasoning_summary: ReasoningSummaryConfig, + turn_metadata_header: Option, + ) -> Option { + let memory_root = memories::memory_root_for_cwd(&config.codex_home, &candidate.cwd); + if let Err(err) = memories::ensure_layout(&memory_root).await { + warn!( + "failed to create memory layout for cwd {}: {err}", + candidate.cwd.display() + ); + return None; + } + + let (rollout_items, _thread_id, parse_errors) = + match RolloutRecorder::load_rollout_items(&candidate.rollout_path).await { + Ok(result) => result, + Err(err) => { + warn!( + "failed to load rollout {} for memories: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + if parse_errors > 0 { + warn!( + "rollout {} had {parse_errors} parse errors while preparing stage-1 memory input", + candidate.rollout_path.display() + ); + } + + let rollout_contents = match memories::serialize_filtered_rollout_response_items( + &rollout_items, + memories::StageOneRolloutFilter::default(), + ) { + Ok(contents) => contents, + Err(err) => { + warn!( + "failed to prepare filtered rollout payload {} for memories: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + + let prompt = Prompt { + input: vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: memories::build_stage_one_input_message( + &candidate.rollout_path, + &rollout_contents, + ), + }], + end_turn: None, + phase: None, + }], + tools: Vec::new(), + parallel_tool_calls: false, + base_instructions: BaseInstructions { + text: memories::TRACE_MEMORY_PROMPT.to_string(), + }, + personality: None, + output_schema: Some(memories::stage_one_output_schema()), + }; + + let mut client_session = self.services.model_client.new_session(); + let mut stream = match client_session + .stream( + &prompt, + &model_info, + &otel_manager, + reasoning_effort, + reasoning_summary, + turn_metadata_header.as_deref(), + ) + .await + { + Ok(stream) => stream, + Err(err) => { + warn!( + "stage-1 memory request failed for rollout {}: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + + let output_text = match Self::collect_response_text_until_completed(&mut stream).await { + Ok(text) => text, + Err(err) => { + warn!( + "failed while waiting for stage-1 memory response for rollout {}: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + + let stage_one_output = match memories::parse_stage_one_output(&output_text) { + Ok(output) => output, + Err(err) => { + warn!( + "invalid stage-1 memory payload for rollout {}: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + + let trace_summary_path = match memories::write_trace_memory( + &memory_root, + &candidate, + &stage_one_output.trace_memory, + ) + .await + { + Ok(path) => path, + Err(err) => { + warn!( + "failed to write trace summary for rollout {}: {err}", + candidate.rollout_path.display() + ); + return None; + } + }; + + if state_db::upsert_thread_memory( + self.services.state_db.as_deref(), + candidate.thread_id, + &stage_one_output.trace_memory, + &stage_one_output.summary, + "run_memories_startup_pipeline", + ) + .await + .is_none() + { + warn!( + "failed to upsert thread memory for rollout {}; removing {}", + candidate.rollout_path.display(), + trace_summary_path.display() + ); + if let Err(err) = tokio::fs::remove_file(&trace_summary_path).await + && err.kind() != std::io::ErrorKind::NotFound + { + warn!( + "failed to remove orphaned trace summary {}: {err}", + trace_summary_path.display() + ); + } + return None; + } + info!( + "memory phase-1 trace persisted: rollout={} cwd={} trace_path={}", + candidate.rollout_path.display(), + candidate.cwd.display(), + trace_summary_path.display() + ); + + Some(candidate.cwd) + } + + async fn run_memory_consolidation_for_cwd(self: Arc, config: Arc, cwd: PathBuf) { + let lock_owner = self.conversation_id; + let Some(lock_acquired) = state_db::try_acquire_memory_consolidation_lock( + self.services.state_db.as_deref(), + &cwd, + lock_owner, + memories::CONSOLIDATION_LOCK_LEASE_SECONDS, + "run_memories_startup_pipeline", + ) + .await + else { + warn!( + "failed to acquire memory consolidation lock for cwd {}; skipping consolidation", + cwd.display() + ); + return; + }; + if !lock_acquired { + debug!( + "memory consolidation lock already held for cwd {}; skipping", + cwd.display() + ); + return; + } + + let Some(latest_memories) = state_db::get_last_n_thread_memories_for_cwd( + self.services.state_db.as_deref(), + &cwd, + memories::MAX_TRACES_PER_CWD, + "run_memories_startup_pipeline", + ) + .await + else { + warn!( + "failed to read recent thread memories for cwd {}; skipping consolidation", + cwd.display() + ); + let _ = state_db::release_memory_consolidation_lock( + self.services.state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + return; + }; + + let memory_root = memories::memory_root_for_cwd(&config.codex_home, &cwd); + if let Err(err) = + memories::prune_to_recent_traces_and_rebuild_summary(&memory_root, &latest_memories) + .await + { + warn!( + "failed to refresh phase-1 memory outputs for cwd {}: {err}", + cwd.display() + ); + let _ = state_db::release_memory_consolidation_lock( + self.services.state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + return; + } + + if let Err(err) = memories::wipe_consolidation_outputs(&memory_root).await { + warn!( + "failed to wipe previous consolidation outputs for cwd {}: {err}", + cwd.display() + ); + let _ = state_db::release_memory_consolidation_lock( + self.services.state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + return; + } + + let prompt = memories::build_consolidation_prompt(&memory_root); + let mut consolidation_config = config.as_ref().clone(); + consolidation_config.cwd = memory_root.clone(); + let source = SessionSource::SubAgent(SubAgentSource::Other( + memories::MEMORY_CONSOLIDATION_SUBAGENT_LABEL.to_string(), + )); + match self + .services + .agent_control + .spawn_agent(consolidation_config, prompt, Some(source)) + .await + { + Ok(consolidation_agent_id) => { + info!( + "memory phase-2 consolidation agent started: cwd={} agent_id={}", + cwd.display(), + consolidation_agent_id + ); + self.spawn_memory_lock_release_task(cwd, lock_owner, consolidation_agent_id); + } + Err(err) => { + warn!( + "failed to spawn memory consolidation agent for cwd {}: {err}", + cwd.display() + ); + let _ = state_db::release_memory_consolidation_lock( + self.services.state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + } + } + } + + fn spawn_memory_lock_release_task( + &self, + cwd: PathBuf, + lock_owner: ThreadId, + consolidation_agent_id: ThreadId, + ) { + let state_db = self.services.state_db.clone(); + let agent_control = self.services.agent_control.clone(); + tokio::spawn(async move { + let mut status_rx = match agent_control.subscribe_status(consolidation_agent_id).await { + Ok(status_rx) => status_rx, + Err(err) => { + warn!( + "failed to subscribe to memory consolidation agent {} for cwd {}: {err}", + consolidation_agent_id, + cwd.display() + ); + let _ = state_db::release_memory_consolidation_lock( + state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + return; + } + }; + + let final_status = loop { + let status = status_rx.borrow().clone(); + if is_final_agent_status(&status) { + break Some(status); + } + if status_rx.changed().await.is_err() { + warn!( + "lost status updates for memory consolidation agent {} in cwd {}; releasing lock", + consolidation_agent_id, + cwd.display() + ); + break Some(status); + } + }; + + let _ = state_db::release_memory_consolidation_lock( + state_db.as_deref(), + &cwd, + lock_owner, + "run_memories_startup_pipeline", + ) + .await; + info!( + "memory phase-2 consolidation agent finished: cwd={} agent_id={} final_status={:?}", + cwd.display(), + consolidation_agent_id, + final_status + ); + }); + } + + async fn collect_response_text_until_completed( + stream: &mut ResponseStream, + ) -> CodexResult { + let mut output_text = String::new(); + + loop { + let Some(event) = stream.next().await else { + return Err(CodexErr::Stream( + "stream closed before response.completed".to_string(), + None, + )); + }; + + match event? { + ResponseEvent::OutputTextDelta(delta) => output_text.push_str(&delta), + ResponseEvent::OutputItemDone(item) => { + if output_text.is_empty() + && let ResponseItem::Message { content, .. } = item + && let Some(text) = crate::compact::content_items_to_text(&content) + { + output_text.push_str(&text); + } + } + ResponseEvent::Completed { .. } => return Ok(output_text), + _ => {} + } + } + } + #[allow(clippy::too_many_arguments)] fn make_turn_context( auth_manager: Option>, @@ -1187,6 +1696,11 @@ impl Session { // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. sess.record_initial_history(initial_history).await; + sess.start_memories_startup_task( + Arc::clone(&config), + &session_configuration.session_source, + ); + Ok(sess) } diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index 201430c9bd..9c4867be01 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -103,7 +103,7 @@ pub enum Feature { RuntimeMetrics, /// Persist rollout metadata to a local SQLite database. Sqlite, - /// Enable the get_memory tool backed by SQLite thread memories. + /// Enable startup memory extraction and file-backed memory consolidation. MemoryTool, /// Append additional AGENTS.md guidance to user instructions. ChildAgentsMd,