openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

codex: wire startup memory extraction and consolidation

Browse Source
This commit is contained in:
jif-oai
2026-02-09 16:16:03 +00:00
parent 8c15c5c60f
commit 22b8626e43
2 changed files with 515 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

514
codex-rs/core/src/codex.rs
View File

@@ -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<Self>, config: Arc<Config>, 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<Self>,
config: Arc<Config>,
) -> 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::<HashSet<PathBuf>>()
.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::<Vec<_>>()
.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<Self>,
config: Arc<Config>,
candidate: memories::RolloutCandidate,
model_info: ModelInfo,
otel_manager: OtelManager,
reasoning_effort: Option<ReasoningEffortConfig>,
reasoning_summary: ReasoningSummaryConfig,
turn_metadata_header: Option<String>,
) -> Option<PathBuf> {
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<Self>, config: Arc<Config>, 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<String> {
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<Arc<AuthManager>>,
@@ -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)
}

2
codex-rs/core/src/features.rs
View File

@@ -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,