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

further cleanup

Browse Source
This commit is contained in:
jif-oai
2026-02-10 17:50:41 +00:00
parent e8ee5ec201
commit ba3a423c3d
14 changed files with 410 additions and 577 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
codex-rs/core/src/memories/layout.rs
View File

@@ -24,13 +24,6 @@ pub(super) fn memory_root_for_cwd(codex_home: &Path, cwd: &Path) -> PathBuf {
codex_home.join("memories").join(bucket).join(MEMORY_SUBDIR)
}
/// Returns the DB scope key for a cwd-scoped memory entry.
///
/// This uses the same normalization/fallback behavior as cwd bucket derivation.
pub(super) fn memory_scope_key_for_cwd(cwd: &Path) -> String {
normalize_cwd_for_memory(cwd).display().to_string()
}
/// Returns the on-disk user-shared memory root directory.
pub(super) fn memory_root_for_user(codex_home: &Path) -> PathBuf {
codex_home

1
codex-rs/core/src/memories/mod.rs
View File

@@ -8,7 +8,6 @@ mod layout;
mod prompts;
mod rollout;
mod scope;
mod selection;
mod stage_one;
mod startup;
mod storage;

47
codex-rs/core/src/memories/selection.rs
View File

@@ -1,47 +0,0 @@
use chrono::Duration;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_state::ThreadMetadata;
use super::types::RolloutCandidate;
/// Selects rollout candidates that need stage-1 memory extraction.
///
/// A rollout is selected when it is not the active thread and was updated
/// within the configured max age window.
pub(super) fn select_rollout_candidates_from_db(
items: &[ThreadMetadata],
current_thread_id: ThreadId,
max_items: usize,
max_age_days: i64,
) -> Vec<RolloutCandidate> {
if max_items == 0 {
return Vec::new();
}
let cutoff = Utc::now() - Duration::days(max_age_days.max(0));
let mut candidates = Vec::new();
for item in items {
if item.id == current_thread_id {
continue;
}
if item.updated_at < cutoff {
continue;
}
candidates.push(RolloutCandidate {
thread_id: item.id,
rollout_path: item.rollout_path.clone(),
cwd: item.cwd.clone(),
source_updated_at: item.updated_at.timestamp(),
});
if candidates.len() >= max_items {
break;
}
}
candidates
}

18
codex-rs/core/src/memories/startup/phase_two/dispatch.rs → codex-rs/core/src/memories/startup/dispatch.rs
View File

@@ -8,15 +8,15 @@ use tracing::debug;
use tracing::info;
use tracing::warn;
use super::super::super::MAX_RAW_MEMORIES_PER_SCOPE;
use super::super::super::MEMORY_CONSOLIDATION_SUBAGENT_LABEL;
use super::super::super::PHASE_TWO_JOB_LEASE_SECONDS;
use super::super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS;
use super::super::super::prompts::build_consolidation_prompt;
use super::super::super::storage::rebuild_memory_summary_from_memories;
use super::super::super::storage::sync_raw_memories_from_memories;
use super::super::super::storage::wipe_consolidation_outputs;
use super::super::MemoryScopeTarget;
use super::super::MAX_RAW_MEMORIES_PER_SCOPE;
use super::super::MEMORY_CONSOLIDATION_SUBAGENT_LABEL;
use super::super::PHASE_TWO_JOB_LEASE_SECONDS;
use super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS;
use super::super::prompts::build_consolidation_prompt;
use super::super::storage::rebuild_memory_summary_from_memories;
use super::super::storage::sync_raw_memories_from_memories;
use super::super::storage::wipe_consolidation_outputs;
use super::MemoryScopeTarget;
use super::watch::spawn_phase2_completion_task;
pub(super) async fn run_memory_consolidation_for_scope(

20
codex-rs/core/src/memories/startup/phase_one/extract.rs → codex-rs/core/src/memories/startup/extract.rs
View File

@@ -18,21 +18,21 @@ use crate::memories::rollout::serialize_filtered_rollout_response_items;
use crate::memories::stage_one::RAW_MEMORY_PROMPT;
use crate::memories::stage_one::parse_stage_one_output;
use crate::memories::stage_one::stage_one_output_schema;
use crate::memories::types::RolloutCandidate;
use crate::memories::types::StageOneOutput;
use std::path::Path;
pub(super) async fn extract_stage_one_output(
session: &Session,
candidate: &RolloutCandidate,
rollout_path: &Path,
stage_one_context: &StageOneRequestContext,
) -> Result<StageOneOutput, &'static str> {
let (rollout_items, _thread_id, parse_errors) =
match RolloutRecorder::load_rollout_items(&candidate.rollout_path).await {
match RolloutRecorder::load_rollout_items(rollout_path).await {
Ok(result) => result,
Err(err) => {
warn!(
"failed to load rollout {} for memories: {err}",
candidate.rollout_path.display()
rollout_path.display()
);
return Err("failed to load rollout");
}
@@ -40,7 +40,7 @@ pub(super) async fn extract_stage_one_output(
if parse_errors > 0 {
warn!(
"rollout {} had {parse_errors} parse errors while preparing stage-1 memory input",
candidate.rollout_path.display()
rollout_path.display()
);
}
@@ -52,7 +52,7 @@ pub(super) async fn extract_stage_one_output(
Err(err) => {
warn!(
"failed to prepare filtered rollout payload {} for memories: {err}",
candidate.rollout_path.display()
rollout_path.display()
);
return Err("failed to serialize filtered rollout");
}
@@ -63,7 +63,7 @@ pub(super) async fn extract_stage_one_output(
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: build_stage_one_input_message(&candidate.rollout_path, &rollout_contents),
text: build_stage_one_input_message(rollout_path, &rollout_contents),
}],
end_turn: None,
phase: None,
@@ -93,7 +93,7 @@ pub(super) async fn extract_stage_one_output(
Err(err) => {
warn!(
"stage-1 memory request failed for rollout {}: {err}",
candidate.rollout_path.display()
rollout_path.display()
);
return Err("stage-1 memory request failed");
}
@@ -104,7 +104,7 @@ pub(super) async fn extract_stage_one_output(
Err(err) => {
warn!(
"failed while waiting for stage-1 memory response for rollout {}: {err}",
candidate.rollout_path.display()
rollout_path.display()
);
return Err("stage-1 memory response stream failed");
}
@@ -115,7 +115,7 @@ pub(super) async fn extract_stage_one_output(
Err(err) => {
warn!(
"invalid stage-1 memory payload for rollout {}: {err}",
candidate.rollout_path.display()
rollout_path.display()
);
Err("invalid stage-1 memory payload")
}

199
codex-rs/core/src/memories/startup/mod.rs
View File

@@ -1,7 +1,9 @@
mod phase_one;
mod phase_two;
mod dispatch;
mod extract;
mod watch;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::config::Config;
use crate::error::Result as CodexResult;
use crate::features::Feature;
@@ -11,9 +13,13 @@ use crate::memories::scope::MEMORY_SCOPE_KEY_USER;
use crate::memories::scope::MEMORY_SCOPE_KIND_CWD;
use crate::memories::scope::MEMORY_SCOPE_KIND_USER;
use crate::rollout::INTERACTIVE_SESSION_SOURCES;
use crate::rollout::list::ThreadSortKey;
use crate::state_db;
use codex_otel::OtelManager;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::SessionSource;
use futures::StreamExt;
use serde_json::Value;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::info;
@@ -21,6 +27,27 @@ use tracing::warn;
pub(super) const PHASE_ONE_THREAD_SCAN_LIMIT: usize = 5_000;
#[derive(Clone)]
struct StageOneRequestContext {
model_info: ModelInfo,
otel_manager: OtelManager,
reasoning_effort: Option<ReasoningEffortConfig>,
reasoning_summary: ReasoningSummaryConfig,
turn_metadata_header: Option<String>,
}
impl StageOneRequestContext {
fn from_turn_context(turn_context: &TurnContext, turn_metadata_header: Option<String>) -> Self {
Self {
model_info: turn_context.model_info.clone(),
otel_manager: turn_context.otel_manager.clone(),
reasoning_effort: turn_context.reasoning_effort,
reasoning_summary: turn_context.reasoning_summary,
turn_metadata_header,
}
}
}
/// Canonical memory scope metadata used by both startup phases.
#[derive(Clone, Debug, PartialEq, Eq)]
pub(super) struct MemoryScopeTarget {
@@ -112,42 +139,162 @@ pub(super) async fn run_memories_startup_pipeline(
session: &Arc<Session>,
config: Arc<Config>,
) -> CodexResult<()> {
let Some(page) = state_db::list_threads_db(
session.services.state_db.as_deref(),
&config.codex_home,
PHASE_ONE_THREAD_SCAN_LIMIT,
None,
ThreadSortKey::UpdatedAt,
INTERACTIVE_SESSION_SOURCES,
None,
false,
)
.await
else {
let Some(state_db) = session.services.state_db.as_deref() else {
warn!("state db unavailable for memories startup pipeline; skipping");
return Ok(());
};
let phase_one = phase_one::run_phase_one(session, &page.items).await;
let allowed_sources = INTERACTIVE_SESSION_SOURCES
.iter()
.map(|value| match serde_json::to_value(value) {
Ok(Value::String(s)) => s,
Ok(other) => other.to_string(),
Err(_) => String::new(),
})
.collect::<Vec<_>>();
let claimed_candidates = match state_db
.claim_stage1_jobs_for_startup(
session.conversation_id,
PHASE_ONE_THREAD_SCAN_LIMIT,
super::MAX_ROLLOUTS_PER_STARTUP,
super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
allowed_sources.as_slice(),
super::PHASE_ONE_JOB_LEASE_SECONDS,
)
.await
{
Ok(claims) => claims,
Err(err) => {
warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}");
Vec::new()
}
};
let claimed_count = claimed_candidates.len();
let mut succeeded_count = 0;
if claimed_count > 0 {
let turn_context = session.new_default_turn().await;
let stage_one_context = StageOneRequestContext::from_turn_context(
turn_context.as_ref(),
turn_context.resolve_turn_metadata_header().await,
);
succeeded_count = futures::stream::iter(claimed_candidates.into_iter())
.map(|claim| {
let session = Arc::clone(session);
let stage_one_context = stage_one_context.clone();
async move {
let thread = claim.thread;
let stage_one_output = match extract::extract_stage_one_output(
session.as_ref(),
&thread.rollout_path,
&stage_one_context,
)
.await
{
Ok(output) => output,
Err(reason) => {
if let Some(state_db) = session.services.state_db.as_deref() {
let _ = state_db
.mark_stage1_job_failed(
thread.id,
&claim.ownership_token,
reason,
super::PHASE_ONE_JOB_RETRY_DELAY_SECONDS,
)
.await;
}
return false;
}
};
let Some(state_db) = session.services.state_db.as_deref() else {
return false;
};
state_db
.mark_stage1_job_succeeded(
thread.id,
&claim.ownership_token,
thread.updated_at.timestamp(),
&stage_one_output.raw_memory,
&stage_one_output.summary,
)
.await
.unwrap_or(false)
}
})
.buffer_unordered(super::PHASE_ONE_CONCURRENCY_LIMIT)
.collect::<Vec<bool>>()
.await
.into_iter()
.filter(|ok| *ok)
.count();
}
info!(
"memory phase-1 candidate selection complete: {} claimed candidate(s) from {} indexed thread(s)",
phase_one.claimed_candidate_count,
page.items.len()
);
info!(
"memory phase-1 extraction complete: {} scope(s) touched",
phase_one.touched_scope_count
"memory stage-1 extraction complete: {} job(s) claimed, {} succeeded",
claimed_count, succeeded_count
);
let consolidation_scope_count = phase_two::run_phase_two(session, config).await;
let consolidation_scope_count = run_consolidation_dispatch(session, config).await;
info!(
"memory phase-2 consolidation dispatch complete: {} scope(s) scheduled",
"memory consolidation dispatch complete: {} scope(s) scheduled",
consolidation_scope_count
);
Ok(())
}
async fn run_consolidation_dispatch(session: &Arc<Session>, config: Arc<Config>) -> usize {
let scopes = list_consolidation_scopes(
session.as_ref(),
config.as_ref(),
super::MAX_ROLLOUTS_PER_STARTUP,
)
.await;
let consolidation_scope_count = scopes.len();
futures::stream::iter(scopes.into_iter())
.map(|scope| {
let session = Arc::clone(session);
let config = Arc::clone(&config);
async move {
dispatch::run_memory_consolidation_for_scope(session, config, scope).await;
}
})
.buffer_unordered(super::PHASE_TWO_CONCURRENCY_LIMIT)
.collect::<Vec<_>>()
.await;
consolidation_scope_count
}
async fn list_consolidation_scopes(
session: &Session,
config: &Config,
limit: usize,
) -> Vec<MemoryScopeTarget> {
if limit == 0 {
return Vec::new();
}
let Some(state_db) = session.services.state_db.as_deref() else {
return Vec::new();
};
let pending_scopes = match state_db.list_pending_scope_consolidations(limit).await {
Ok(scopes) => scopes,
Err(_) => return Vec::new(),
};
pending_scopes
.into_iter()
.filter_map(|scope| memory_scope_target_for_pending_scope(config, scope))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;

247
codex-rs/core/src/memories/startup/phase_one/mod.rs
View File

@@ -1,247 +0,0 @@
mod extract;
use codex_otel::OtelManager;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use extract::extract_stage_one_output;
use futures::StreamExt;
use tracing::warn;
use super::super::MAX_ROLLOUTS_PER_STARTUP;
use super::super::PHASE_ONE_CONCURRENCY_LIMIT;
use super::super::PHASE_ONE_JOB_LEASE_SECONDS;
use super::super::PHASE_ONE_JOB_RETRY_DELAY_SECONDS;
use super::super::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS;
use super::super::selection::select_rollout_candidates_from_db;
use super::super::types::RolloutCandidate;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::memories::layout::memory_scope_key_for_cwd;
use crate::memories::scope::MEMORY_SCOPE_KEY_USER;
use crate::memories::scope::MEMORY_SCOPE_KIND_CWD;
use crate::memories::scope::MEMORY_SCOPE_KIND_USER;
use std::sync::Arc;
/// Result counters for startup phase-1 extraction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct PhaseOneRunResult {
/// Number of rollout candidates that were successfully claimed.
pub(super) claimed_candidate_count: usize,
/// Number of scope refresh/enqueue operations performed.
pub(super) touched_scope_count: usize,
}
#[derive(Clone, Debug)]
pub(super) struct ClaimedStageOneCandidate {
pub(super) candidate: RolloutCandidate,
pub(super) ownership_token: String,
}
#[derive(Clone)]
struct StageOneRequestContext {
model_info: ModelInfo,
otel_manager: OtelManager,
reasoning_effort: Option<ReasoningEffortConfig>,
reasoning_summary: ReasoningSummaryConfig,
turn_metadata_header: Option<String>,
}
impl StageOneRequestContext {
fn from_turn_context(turn_context: &TurnContext, turn_metadata_header: Option<String>) -> Self {
Self {
model_info: turn_context.model_info.clone(),
otel_manager: turn_context.otel_manager.clone(),
reasoning_effort: turn_context.reasoning_effort,
reasoning_summary: turn_context.reasoning_summary,
turn_metadata_header,
}
}
}
/// Runs startup phase 1:
///
/// 1. Select rollout candidates from thread metadata.
/// 2. Claim stage-1 jobs per thread.
/// 3. Execute stage-1 extraction requests in parallel.
/// 4. Persist stage-1 outputs and enqueue consolidation jobs.
pub(super) async fn run_phase_one(
session: &Arc<Session>,
thread_items: &[codex_state::ThreadMetadata],
) -> PhaseOneRunResult {
let selection_candidates = select_rollout_candidates_from_db(
thread_items,
session.conversation_id,
super::PHASE_ONE_THREAD_SCAN_LIMIT,
PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
);
let claimed_candidates = claim_stage_one_candidates(
session.as_ref(),
selection_candidates,
MAX_ROLLOUTS_PER_STARTUP,
)
.await;
if claimed_candidates.is_empty() {
return PhaseOneRunResult {
claimed_candidate_count: 0,
touched_scope_count: 0,
};
}
let turn_context = session.new_default_turn().await;
let stage_one_context = StageOneRequestContext::from_turn_context(
turn_context.as_ref(),
turn_context.resolve_turn_metadata_header().await,
);
let touched_scope_count =
futures::stream::iter(claimed_candidates.iter().cloned())
.map(|claimed_candidate| {
let session = Arc::clone(session);
let stage_one_context = stage_one_context.clone();
async move {
process_memory_candidate(session, claimed_candidate, stage_one_context).await
}
})
.buffer_unordered(PHASE_ONE_CONCURRENCY_LIMIT)
.collect::<Vec<usize>>()
.await
.into_iter()
.sum::<usize>();
PhaseOneRunResult {
claimed_candidate_count: claimed_candidates.len(),
touched_scope_count,
}
}
async fn claim_stage_one_candidates(
session: &Session,
candidates: Vec<RolloutCandidate>,
max_claimed_candidates: usize,
) -> Vec<ClaimedStageOneCandidate> {
if max_claimed_candidates == 0 {
return Vec::new();
}
let Some(state_db) = session.services.state_db.as_deref() else {
return Vec::new();
};
let mut claimed = Vec::new();
for candidate in candidates {
if claimed.len() >= max_claimed_candidates {
break;
}
let claim = match state_db
.try_claim_stage1_job(
candidate.thread_id,
session.conversation_id,
candidate.source_updated_at,
PHASE_ONE_JOB_LEASE_SECONDS,
)
.await
{
Ok(claim) => claim,
Err(err) => {
warn!(
"state db try_claim_stage1_job failed for rollout {}: {err}",
candidate.rollout_path.display()
);
continue;
}
};
if let codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } = claim {
claimed.push(ClaimedStageOneCandidate {
candidate,
ownership_token,
});
}
}
claimed
}
async fn process_memory_candidate(
session: Arc<Session>,
claimed_candidate: ClaimedStageOneCandidate,
stage_one_context: StageOneRequestContext,
) -> usize {
let candidate = claimed_candidate.candidate;
let stage_one_output =
match extract_stage_one_output(session.as_ref(), &candidate, &stage_one_context).await {
Ok(output) => output,
Err(reason) => {
if let Some(state_db) = session.services.state_db.as_deref() {
let _ = state_db
.mark_stage1_job_failed(
candidate.thread_id,
&claimed_candidate.ownership_token,
reason,
PHASE_ONE_JOB_RETRY_DELAY_SECONDS,
)
.await;
}
return 0;
}
};
let Some(state_db) = session.services.state_db.as_deref() else {
return 0;
};
if !state_db
.mark_stage1_job_succeeded(
candidate.thread_id,
&claimed_candidate.ownership_token,
candidate.source_updated_at,
&stage_one_output.raw_memory,
&stage_one_output.summary,
)
.await
.unwrap_or(false)
{
return 0;
}
let mut touched_scope_count = 0;
let cwd_scope_key = memory_scope_key_for_cwd(&candidate.cwd);
if let Err(err) = state_db
.enqueue_scope_consolidation(
MEMORY_SCOPE_KIND_CWD,
&cwd_scope_key,
candidate.source_updated_at,
)
.await
{
warn!(
"failed enqueueing scope consolidation for scope {}:{}: {err}",
MEMORY_SCOPE_KIND_CWD, cwd_scope_key
);
} else {
touched_scope_count += 1;
}
if let Err(err) = state_db
.enqueue_scope_consolidation(
MEMORY_SCOPE_KIND_USER,
MEMORY_SCOPE_KEY_USER,
candidate.source_updated_at,
)
.await
{
warn!(
"failed enqueueing scope consolidation for scope {}:{}: {err}",
MEMORY_SCOPE_KIND_USER, MEMORY_SCOPE_KEY_USER
);
} else {
touched_scope_count += 1;
}
touched_scope_count
}

60
codex-rs/core/src/memories/startup/phase_two/mod.rs
View File

@@ -1,60 +0,0 @@
mod dispatch;
mod watch;
use super::super::MAX_ROLLOUTS_PER_STARTUP;
use super::super::PHASE_TWO_CONCURRENCY_LIMIT;
use super::MemoryScopeTarget;
use super::memory_scope_target_for_pending_scope;
use crate::codex::Session;
use crate::config::Config;
use futures::StreamExt;
use std::sync::Arc;
/// Runs startup phase 2:
///
/// 1. Load scopes pending consolidation from the DB.
/// 2. Claim scope jobs.
/// 3. Spawn consolidation agents for owned scopes.
pub(super) async fn run_phase_two(session: &Arc<Session>, config: Arc<Config>) -> usize {
let scopes =
list_phase2_scopes(session.as_ref(), config.as_ref(), MAX_ROLLOUTS_PER_STARTUP).await;
let consolidation_scope_count = scopes.len();
futures::stream::iter(scopes.into_iter())
.map(|scope| {
let session = Arc::clone(session);
let config = Arc::clone(&config);
async move {
dispatch::run_memory_consolidation_for_scope(session, config, scope).await;
}
})
.buffer_unordered(PHASE_TWO_CONCURRENCY_LIMIT)
.collect::<Vec<_>>()
.await;
consolidation_scope_count
}
async fn list_phase2_scopes(
session: &Session,
config: &Config,
limit: usize,
) -> Vec<MemoryScopeTarget> {
if limit == 0 {
return Vec::new();
}
let Some(state_db) = session.services.state_db.as_deref() else {
return Vec::new();
};
let pending_scopes = match state_db.list_pending_scope_consolidations(limit).await {
Ok(scopes) => scopes,
Err(_) => return Vec::new(),
};
pending_scopes
.into_iter()
.filter_map(|scope| memory_scope_target_for_pending_scope(config, scope))
.collect()
}

8
codex-rs/core/src/memories/startup/phase_two/watch.rs → codex-rs/core/src/memories/startup/watch.rs
View File

@@ -7,10 +7,10 @@ use tracing::debug;
use tracing::info;
use tracing::warn;
use super::super::super::PHASE_TWO_JOB_HEARTBEAT_SECONDS;
use super::super::super::PHASE_TWO_JOB_LEASE_SECONDS;
use super::super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS;
use super::super::MemoryScopeTarget;
use super::super::PHASE_TWO_JOB_HEARTBEAT_SECONDS;
use super::super::PHASE_TWO_JOB_LEASE_SECONDS;
use super::super::PHASE_TWO_JOB_RETRY_DELAY_SECONDS;
use super::MemoryScopeTarget;
pub(super) fn spawn_phase2_completion_task(
session: &Session,

108
codex-rs/core/src/memories/tests.rs
View File

@@ -1,15 +1,12 @@
use super::rollout::StageOneResponseItemKinds;
use super::rollout::StageOneRolloutFilter;
use super::rollout::serialize_filtered_rollout_response_items;
use super::selection::select_rollout_candidates_from_db;
use super::stage_one::parse_stage_one_output;
use super::storage::rebuild_memory_summary_from_memories;
use super::storage::sync_raw_memories_from_memories;
use super::storage::wipe_consolidation_outputs;
use crate::memories::PHASE_ONE_MAX_ROLLOUT_AGE_DAYS;
use crate::memories::layout::ensure_layout;
use crate::memories::layout::memory_root_for_cwd;
use crate::memories::layout::memory_scope_key_for_cwd;
use crate::memories::layout::memory_summary_file;
use crate::memories::layout::raw_memories_dir;
use chrono::TimeZone;
@@ -20,43 +17,9 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::CompactedItem;
use codex_protocol::protocol::RolloutItem;
use codex_state::Stage1Output;
use codex_state::ThreadMetadata;
use pretty_assertions::assert_eq;
use std::path::PathBuf;
use tempfile::tempdir;
fn thread_metadata(
thread_id: ThreadId,
path: PathBuf,
cwd: PathBuf,
title: &str,
updated_at_secs: i64,
) -> ThreadMetadata {
let updated_at = Utc
.timestamp_opt(updated_at_secs, 0)
.single()
.expect("timestamp");
ThreadMetadata {
id: thread_id,
rollout_path: path,
created_at: updated_at,
updated_at,
source: "cli".to_string(),
model_provider: "openai".to_string(),
cwd,
cli_version: "test".to_string(),
title: title.to_string(),
sandbox_policy: "read_only".to_string(),
approval_mode: "on_request".to_string(),
tokens_used: 0,
first_user_message: None,
archived_at: None,
git_branch: None,
git_sha: None,
git_origin_url: None,
}
}
#[test]
fn memory_root_varies_by_cwd() {
let dir = tempdir().expect("tempdir");
@@ -100,22 +63,6 @@ fn memory_root_encoding_avoids_component_collisions() {
assert!(!root_hash.display().to_string().contains("workspace"));
}
#[test]
fn memory_scope_key_uses_normalized_cwd() {
let dir = tempdir().expect("tempdir");
let workspace = dir.path().join("workspace");
std::fs::create_dir_all(&workspace).expect("mkdir workspace");
std::fs::create_dir_all(workspace.join("nested")).expect("mkdir nested");
let alias = workspace.join("nested").join("..");
let normalized = workspace
.canonicalize()
.expect("canonical workspace path should resolve");
let alias_key = memory_scope_key_for_cwd(&alias);
let normalized_key = memory_scope_key_for_cwd(&normalized);
assert_eq!(alias_key, normalized_key);
}
#[test]
fn parse_stage_one_output_accepts_fenced_json() {
let raw = "```json\n{\"rawMemory\":\"abc\",\"summary\":\"short\"}\n```";
@@ -224,61 +171,6 @@ fn serialize_filtered_rollout_response_items_filters_by_response_item_kind() {
assert!(matches!(parsed[0], ResponseItem::Message { .. }));
}
#[test]
fn select_rollout_candidates_filters_by_age_window() {
let dir = tempdir().expect("tempdir");
let cwd_a = dir.path().join("workspace-a");
let cwd_b = dir.path().join("workspace-b");
std::fs::create_dir_all(&cwd_a).expect("mkdir cwd a");
std::fs::create_dir_all(&cwd_b).expect("mkdir cwd b");
let now = Utc::now().timestamp();
let current_thread_id = ThreadId::default();
let recent_thread_id = ThreadId::default();
let old_thread_id = ThreadId::default();
let recent_two_thread_id = ThreadId::default();
let current = thread_metadata(
current_thread_id,
dir.path().join("current.jsonl"),
cwd_a.clone(),
"current",
now,
);
let recent = thread_metadata(
recent_thread_id,
dir.path().join("recent.jsonl"),
cwd_a,
"recent",
now - 10,
);
let old = thread_metadata(
old_thread_id,
dir.path().join("old.jsonl"),
cwd_b.clone(),
"old",
now - (PHASE_ONE_MAX_ROLLOUT_AGE_DAYS + 1) * 24 * 60 * 60,
);
let recent_two = thread_metadata(
recent_two_thread_id,
dir.path().join("recent-two.jsonl"),
cwd_b,
"recent-two",
now - 20,
);
let candidates = select_rollout_candidates_from_db(
&[current, recent, old, recent_two],
current_thread_id,
5,
PHASE_ONE_MAX_ROLLOUT_AGE_DAYS,
);
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].thread_id, recent_thread_id);
assert_eq!(candidates[1].thread_id, recent_two_thread_id);
}
#[tokio::test]
async fn prune_and_rebuild_summary_keeps_latest_memories_only() {
let dir = tempdir().expect("tempdir");

15
codex-rs/core/src/memories/types.rs
View File

@@ -1,19 +1,4 @@
use codex_protocol::ThreadId;
use serde::Deserialize;
use std::path::PathBuf;
/// A rollout selected for stage-1 memory extraction during startup.
#[derive(Debug, Clone)]
pub(super) struct RolloutCandidate {
/// Source thread identifier for this rollout.
pub(super) thread_id: ThreadId,
/// Absolute path to the rollout file to summarize.
pub(super) rollout_path: PathBuf,
/// Thread working directory used for per-project memory bucketing.
pub(super) cwd: PathBuf,
/// Thread update timestamp (unix seconds) used for stage-1 staleness checks.
pub(super) source_updated_at: i64,
}
/// Parsed stage-1 model output payload.
#[derive(Debug, Clone, Deserialize)]

1
codex-rs/state/src/lib.rs
View File

@@ -35,6 +35,7 @@ pub use runtime::PendingScopeConsolidation;
pub use runtime::Phase2JobClaimOutcome;
pub use runtime::STATE_DB_FILENAME;
pub use runtime::STATE_DB_VERSION;
pub use runtime::Stage1JobClaim;
pub use runtime::Stage1JobClaimOutcome;
pub use runtime::state_db_filename;
pub use runtime::state_db_path;

66
codex-rs/state/src/runtime.rs
View File

@@ -41,6 +41,7 @@ pub const STATE_DB_VERSION: u32 = 4;
const MEMORY_SCOPE_KIND_CWD: &str = "cwd";
const MEMORY_SCOPE_KIND_USER: &str = "user";
const MEMORY_SCOPE_KEY_USER: &str = "user";
const METRIC_DB_INIT: &str = "codex.db.init";
@@ -69,6 +70,13 @@ pub enum Stage1JobClaimOutcome {
SkippedRetryExhausted,
}
/// Claimed stage-1 job with thread metadata.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Stage1JobClaim {
pub thread: ThreadMetadata,
pub ownership_token: String,
}
/// Scope row used to queue phase-2 consolidation work.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PendingScopeConsolidation {
@@ -918,6 +926,7 @@ mod tests {
use super::ThreadMetadata;
use super::state_db_filename;
use chrono::DateTime;
use chrono::Duration;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
@@ -1166,6 +1175,63 @@ mod tests {
let _ = tokio::fs::remove_dir_all(codex_home).await;
}
#[tokio::test]
async fn claim_stage1_jobs_filters_by_age_and_current_thread() {
let codex_home = unique_temp_dir();
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
.await
.expect("initialize runtime");
let now = Utc::now();
let recent_at = now - Duration::seconds(10);
let old_at = now - Duration::days(31);
let current_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id");
let recent_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("recent thread id");
let old_thread_id =
ThreadId::from_string(&Uuid::new_v4().to_string()).expect("old thread id");
let mut current =
test_thread_metadata(&codex_home, current_thread_id, codex_home.join("current"));
current.created_at = now;
current.updated_at = now;
runtime
.upsert_thread(&current)
.await
.expect("upsert current");
let mut recent =
test_thread_metadata(&codex_home, recent_thread_id, codex_home.join("recent"));
recent.created_at = recent_at;
recent.updated_at = recent_at;
runtime.upsert_thread(&recent).await.expect("upsert recent");
let mut old = test_thread_metadata(&codex_home, old_thread_id, codex_home.join("old"));
old.created_at = old_at;
old.updated_at = old_at;
runtime.upsert_thread(&old).await.expect("upsert old");
let allowed_sources = vec!["cli".to_string()];
let claims = runtime
.claim_stage1_jobs_for_startup(
current_thread_id,
10,
5,
30,
allowed_sources.as_slice(),
3600,
)
.await
.expect("claim stage1 jobs");
assert_eq!(claims.len(), 1);
assert_eq!(claims[0].thread.id, recent_thread_id);
let _ = tokio::fs::remove_dir_all(codex_home).await;
}
#[tokio::test]
async fn stage1_output_cascades_on_thread_delete() {
let codex_home = unique_temp_dir();

190
codex-rs/state/src/runtime/memory.rs
View File

@@ -1,6 +1,9 @@
use super::*;
use crate::Stage1Output;
use crate::model::Stage1OutputRow;
use chrono::Duration;
use sqlx::Executor;
use sqlx::Sqlite;
const JOB_KIND_MEMORY_STAGE1: &str = "memory_stage1";
const JOB_KIND_MEMORY_CONSOLIDATE_CWD: &str = "memory_consolidate_cwd";
@@ -25,6 +28,63 @@ fn scope_kind_for_job_kind(job_kind: &str) -> Option<&'static str> {
}
impl StateRuntime {
pub async fn claim_stage1_jobs_for_startup(
&self,
current_thread_id: ThreadId,
scan_limit: usize,
max_claimed: usize,
max_age_days: i64,
allowed_sources: &[String],
lease_seconds: i64,
) -> anyhow::Result<Vec<Stage1JobClaim>> {
if scan_limit == 0 || max_claimed == 0 {
return Ok(Vec::new());
}
let page = self
.list_threads(
scan_limit,
None,
SortKey::UpdatedAt,
allowed_sources,
None,
false,
)
.await?;
let cutoff = Utc::now() - Duration::days(max_age_days.max(0));
let mut claimed = Vec::new();
for item in page.items {
if claimed.len() >= max_claimed {
break;
}
if item.id == current_thread_id {
continue;
}
if item.updated_at < cutoff {
continue;
}
if let Stage1JobClaimOutcome::Claimed { ownership_token } = self
.try_claim_stage1_job(
item.id,
current_thread_id,
item.updated_at.timestamp(),
lease_seconds,
)
.await?
{
claimed.push(Stage1JobClaim {
thread: item,
ownership_token,
});
}
}
Ok(claimed)
}
pub async fn get_stage1_output(
&self,
thread_id: ThreadId,
@@ -42,7 +102,6 @@ WHERE thread_id = ?
row.map(|row| Stage1OutputRow::try_from_row(&row).and_then(Stage1Output::try_from))
.transpose()
.map_err(Into::into)
}
pub async fn list_stage1_outputs_for_scope(
@@ -92,7 +151,6 @@ LIMIT ?
rows.into_iter()
.map(|row| Stage1OutputRow::try_from_row(&row).and_then(Stage1Output::try_from))
.collect::<Result<Vec<_>, _>>()
.map_err(Into::into)
}
pub async fn try_claim_stage1_job(
@@ -296,6 +354,34 @@ WHERE excluded.source_updated_at >= stage1_outputs.source_updated_at
.execute(&mut *tx)
.await?;
if let Some(thread_row) = sqlx::query(
r#"
SELECT cwd
FROM threads
WHERE id = ?
"#,
)
.bind(thread_id.as_str())
.fetch_optional(&mut *tx)
.await?
{
let cwd: String = thread_row.try_get("cwd")?;
enqueue_scope_consolidation_with_executor(
&mut *tx,
MEMORY_SCOPE_KIND_CWD,
&cwd,
source_updated_at,
)
.await?;
enqueue_scope_consolidation_with_executor(
&mut *tx,
MEMORY_SCOPE_KIND_USER,
MEMORY_SCOPE_KEY_USER,
source_updated_at,
)
.await?;
}
tx.commit().await?;
Ok(true)
}
@@ -344,48 +430,13 @@ WHERE kind = ? AND job_key = ?
scope_key: &str,
input_watermark: i64,
) -> anyhow::Result<()> {
let Some(job_kind) = job_kind_for_scope(scope_kind) else {
return Ok(());
};
sqlx::query(
r#"
INSERT INTO jobs (
kind,
job_key,
status,
worker_id,
ownership_token,
started_at,
finished_at,
lease_until,
retry_at,
retry_remaining,
last_error,
input_watermark,
last_success_watermark
) VALUES (?, ?, 'pending', NULL, NULL, NULL, NULL, NULL, NULL, ?, NULL, ?, 0)
ON CONFLICT(kind, job_key) DO UPDATE SET
status = CASE
WHEN jobs.status = 'running' THEN 'running'
ELSE 'pending'
END,
retry_at = CASE
WHEN jobs.status = 'running' THEN jobs.retry_at
ELSE NULL
END,
retry_remaining = max(jobs.retry_remaining, excluded.retry_remaining),
input_watermark = max(COALESCE(jobs.input_watermark, 0), excluded.input_watermark)
"#,
enqueue_scope_consolidation_with_executor(
self.pool.as_ref(),
scope_kind,
scope_key,
input_watermark,
)
.bind(job_kind)
.bind(scope_key)
.bind(DEFAULT_RETRY_REMAINING)
.bind(input_watermark)
.execute(self.pool.as_ref())
.await?;
Ok(())
.await
}
pub async fn list_pending_scope_consolidations(
@@ -641,3 +692,56 @@ WHERE kind = ? AND job_key = ?
Ok(rows_affected > 0)
}
}
async fn enqueue_scope_consolidation_with_executor<'e, E>(
executor: E,
scope_kind: &str,
scope_key: &str,
input_watermark: i64,
) -> anyhow::Result<()>
where
E: Executor<'e, Database = Sqlite>,
{
let Some(job_kind) = job_kind_for_scope(scope_kind) else {
return Ok(());
};
sqlx::query(
r#"
INSERT INTO jobs (
kind,
job_key,
status,
worker_id,
ownership_token,
started_at,
finished_at,
lease_until,
retry_at,
retry_remaining,
last_error,
input_watermark,
last_success_watermark
) VALUES (?, ?, 'pending', NULL, NULL, NULL, NULL, NULL, NULL, ?, NULL, ?, 0)
ON CONFLICT(kind, job_key) DO UPDATE SET
status = CASE
WHEN jobs.status = 'running' THEN 'running'
ELSE 'pending'
END,
retry_at = CASE
WHEN jobs.status = 'running' THEN jobs.retry_at
ELSE NULL
END,
retry_remaining = max(jobs.retry_remaining, excluded.retry_remaining),
input_watermark = max(COALESCE(jobs.input_watermark, 0), excluded.input_watermark)
"#,
)
.bind(job_kind)
.bind(scope_key)
.bind(DEFAULT_RETRY_REMAINING)
.bind(input_watermark)
.execute(executor)
.await?;
Ok(())
}