mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: split memories part 2 (#19860)
Keep extracting memories out of core and moving the write trigger in the app-server This is temporary and it should move at the client level as a follow-up This makes core fully independant from `codex-memories-write` --------- Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
jif-oai
Codex
@@ -42,3 +42,75 @@ pub(crate) async fn clear_memory_root_contents(memory_root: &Path) -> std::io::R
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[tokio::test]
|
||||
async fn clear_memory_root_contents_preserves_root_directory() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let root = dir.path().join("memories");
|
||||
let nested_dir = root.join("rollout_summaries");
|
||||
tokio::fs::create_dir_all(&nested_dir)
|
||||
.await
|
||||
.expect("create rollout summaries dir");
|
||||
tokio::fs::write(root.join("MEMORY.md"), "stale memory index\n")
|
||||
.await
|
||||
.expect("write memory index");
|
||||
tokio::fs::write(nested_dir.join("rollout.md"), "stale rollout\n")
|
||||
.await
|
||||
.expect("write rollout summary");
|
||||
|
||||
clear_memory_root_contents(&root)
|
||||
.await
|
||||
.expect("clear memory root contents");
|
||||
|
||||
assert!(
|
||||
tokio::fs::try_exists(&root)
|
||||
.await
|
||||
.expect("check memory root existence"),
|
||||
"memory root should still exist after clearing contents"
|
||||
);
|
||||
let mut entries = tokio::fs::read_dir(&root)
|
||||
.await
|
||||
.expect("read memory root after clear");
|
||||
assert!(
|
||||
entries
|
||||
.next_entry()
|
||||
.await
|
||||
.expect("read next entry")
|
||||
.is_none(),
|
||||
"memory root should be empty after clearing contents"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn clear_memory_root_contents_rejects_symlinked_root() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let target = dir.path().join("outside");
|
||||
tokio::fs::create_dir_all(&target)
|
||||
.await
|
||||
.expect("create symlink target dir");
|
||||
let target_file = target.join("keep.txt");
|
||||
tokio::fs::write(&target_file, "keep\n")
|
||||
.await
|
||||
.expect("write target file");
|
||||
|
||||
let root = dir.path().join("memories");
|
||||
std::os::unix::fs::symlink(&target, &root).expect("create memory root symlink");
|
||||
|
||||
let err = clear_memory_root_contents(&root)
|
||||
.await
|
||||
.expect_err("symlinked memory root should be rejected");
|
||||
assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
|
||||
assert!(
|
||||
tokio::fs::try_exists(&target_file)
|
||||
.await
|
||||
.expect("check target file existence"),
|
||||
"rejecting a symlinked memory root should not delete the symlink target"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,12 +1,16 @@
|
||||
//! Write-path helpers for Codex memories.
|
||||
//! Write-path implementation for Codex memories.
|
||||
//!
|
||||
//! This crate owns the file-backed memory artifact helpers, Phase 1 and Phase
|
||||
//! 2 prompt rendering, extension pruning, and workspace diffing. Runtime
|
||||
//! orchestration for Phase 1 and Phase 2 remains in `codex-core`.
|
||||
//! This crate owns the startup memory pipeline, file-backed memory artifact
|
||||
//! helpers, Phase 1 and Phase 2 prompt rendering, extension pruning, and
|
||||
//! workspace diffing.
|
||||
|
||||
mod control;
|
||||
mod extensions;
|
||||
mod phase1;
|
||||
mod phase2;
|
||||
mod prompts;
|
||||
mod runtime;
|
||||
mod start;
|
||||
mod storage;
|
||||
pub mod workspace;
|
||||
|
||||
@@ -18,6 +22,7 @@ pub use control::clear_memory_roots_contents;
|
||||
pub use extensions::prune_old_extension_resources;
|
||||
pub use prompts::build_consolidation_prompt;
|
||||
pub use prompts::build_stage_one_input_message;
|
||||
pub use start::start_memories_startup_task;
|
||||
pub use storage::rebuild_raw_memories_file_from_memories;
|
||||
pub use storage::rollout_summary_file_stem;
|
||||
pub use storage::sync_rollout_summaries_from_memories;
|
||||
@@ -36,6 +41,9 @@ pub const DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT: usize = 150_000;
|
||||
/// and model output.
|
||||
pub const STAGE_ONE_CONTEXT_WINDOW_PERCENT: i64 = 70;
|
||||
|
||||
#[cfg(test)]
|
||||
mod startup_tests;
|
||||
|
||||
mod artifacts {
|
||||
pub(super) const EXTENSIONS_SUBDIR: &str = "extensions";
|
||||
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
|
||||
|
||||
@@ -0,0 +1,803 @@
|
||||
use crate::STAGE_ONE_PROMPT;
|
||||
use crate::build_stage_one_input_message;
|
||||
use crate::runtime::MemoryStartupContext;
|
||||
use crate::runtime::StageOneRequestContext;
|
||||
use codex_config::types::MemoriesConfig;
|
||||
use codex_core::Prompt;
|
||||
use codex_core::RolloutRecorder;
|
||||
use codex_core::config::Config;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::models::BaseInstructions;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::openai_models::ReasoningEffort;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_rollout::INTERACTIVE_SESSION_SOURCES;
|
||||
use codex_rollout::should_persist_response_item_for_memories;
|
||||
use codex_secrets::redact_secrets;
|
||||
use futures::StreamExt;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
const MODEL: &str = "gpt-5.4-mini";
|
||||
const REASONING_EFFORT: ReasoningEffort = ReasoningEffort::Low;
|
||||
const CONCURRENCY_LIMIT: usize = 8;
|
||||
const JOB_LEASE_SECONDS: i64 = 3_600;
|
||||
const JOB_RETRY_DELAY_SECONDS: i64 = 3_600;
|
||||
const THREAD_SCAN_LIMIT: usize = 5_000;
|
||||
const PRUNE_BATCH_SIZE: usize = 200;
|
||||
const MEMORY_PHASE_ONE_JOBS: &str = "codex.memory.phase1";
|
||||
const MEMORY_PHASE_ONE_E2E_MS: &str = "codex.memory.phase1.e2e_ms";
|
||||
const MEMORY_PHASE_ONE_OUTPUT: &str = "codex.memory.phase1.output";
|
||||
const MEMORY_PHASE_ONE_TOKEN_USAGE: &str = "codex.memory.phase1.token_usage";
|
||||
|
||||
struct JobResult {
|
||||
outcome: JobOutcome,
|
||||
token_usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
enum JobOutcome {
|
||||
SucceededWithOutput,
|
||||
SucceededNoOutput,
|
||||
Failed,
|
||||
}
|
||||
|
||||
struct Stats {
|
||||
claimed: usize,
|
||||
succeeded_with_output: usize,
|
||||
succeeded_no_output: usize,
|
||||
failed: usize,
|
||||
total_token_usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
/// Phase 1 model output payload.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct StageOneOutput {
|
||||
/// Detailed markdown raw memory for a single rollout.
|
||||
#[serde(rename = "raw_memory")]
|
||||
pub(crate) raw_memory: String,
|
||||
/// Compact summary line used for routing and indexing.
|
||||
#[serde(rename = "rollout_summary")]
|
||||
pub(crate) rollout_summary: String,
|
||||
/// Optional slug used to derive rollout summary artifact filenames.
|
||||
#[serde(default, rename = "rollout_slug")]
|
||||
pub(crate) rollout_slug: Option<String>,
|
||||
}
|
||||
|
||||
/// Runs memory phase 1 in strict step order:
|
||||
/// 1) claim eligible rollout jobs
|
||||
/// 2) build one stage-1 request context
|
||||
/// 3) run stage-1 extraction jobs in parallel
|
||||
/// 4) emit metrics and logs
|
||||
pub async fn run(context: Arc<MemoryStartupContext>, config: Arc<Config>) {
|
||||
let stage_one_context = build_request_context(context.as_ref(), config.as_ref()).await;
|
||||
let _phase_one_e2e_timer = stage_one_context.start_timer(MEMORY_PHASE_ONE_E2E_MS);
|
||||
|
||||
// 1. Claim startup job.
|
||||
let Some(claimed_candidates) = claim_startup_jobs(context.as_ref(), &config.memories).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if claimed_candidates.is_empty() {
|
||||
stage_one_context.counter(
|
||||
MEMORY_PHASE_ONE_JOBS,
|
||||
/*inc*/ 1,
|
||||
&[("status", "skipped_no_candidates")],
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// 3. Run the parallel sampling.
|
||||
let outcomes = run_jobs(
|
||||
context,
|
||||
config,
|
||||
claimed_candidates,
|
||||
stage_one_context.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
// 4. Metrics and logs.
|
||||
let counts = aggregate_stats(outcomes);
|
||||
emit_metrics(&stage_one_context, &counts);
|
||||
info!(
|
||||
"memory stage-1 extraction complete: {} job(s) claimed, {} succeeded ({} with output, {} no output), {} failed",
|
||||
counts.claimed,
|
||||
counts.succeeded_with_output + counts.succeeded_no_output,
|
||||
counts.succeeded_with_output,
|
||||
counts.succeeded_no_output,
|
||||
counts.failed
|
||||
);
|
||||
}
|
||||
|
||||
/// Prune old un-used "dead" raw memories.
|
||||
pub async fn prune(context: &MemoryStartupContext, config: &Config) {
|
||||
if let Some(db) = context.state_db() {
|
||||
let max_unused_days = config.memories.max_unused_days;
|
||||
match db
|
||||
.prune_stage1_outputs_for_retention(max_unused_days, PRUNE_BATCH_SIZE)
|
||||
.await
|
||||
{
|
||||
Ok(pruned) => {
|
||||
if pruned > 0 {
|
||||
info!(
|
||||
"memory startup pruned {pruned} stale stage-1 output row(s) older than {max_unused_days} days"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"state db prune_stage1_outputs_for_retention failed during memories startup: {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON schema used to constrain phase-1 model output.
|
||||
pub fn output_schema() -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"rollout_summary": { "type": "string" },
|
||||
"rollout_slug": { "type": ["string", "null"] },
|
||||
"raw_memory": { "type": "string" }
|
||||
},
|
||||
"required": ["rollout_summary", "rollout_slug", "raw_memory"],
|
||||
"additionalProperties": false
|
||||
})
|
||||
}
|
||||
|
||||
async fn claim_startup_jobs(
|
||||
context: &MemoryStartupContext,
|
||||
memories_config: &MemoriesConfig,
|
||||
) -> Option<Vec<codex_state::Stage1JobClaim>> {
|
||||
let Some(state_db) = context.state_db() else {
|
||||
// This should not happen.
|
||||
warn!("state db unavailable while claiming phase-1 startup jobs; skipping");
|
||||
return None;
|
||||
};
|
||||
|
||||
let allowed_sources = INTERACTIVE_SESSION_SOURCES
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
match state_db
|
||||
.claim_stage1_jobs_for_startup(
|
||||
context.thread_id(),
|
||||
codex_state::Stage1StartupClaimParams {
|
||||
scan_limit: THREAD_SCAN_LIMIT,
|
||||
max_claimed: memories_config.max_rollouts_per_startup,
|
||||
max_age_days: memories_config.max_rollout_age_days,
|
||||
min_rollout_idle_hours: memories_config.min_rollout_idle_hours,
|
||||
allowed_sources: allowed_sources.as_slice(),
|
||||
lease_seconds: JOB_LEASE_SECONDS,
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(claims) => Some(claims),
|
||||
Err(err) => {
|
||||
warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_request_context(
|
||||
context: &MemoryStartupContext,
|
||||
config: &Config,
|
||||
) -> StageOneRequestContext {
|
||||
let model_name = config
|
||||
.memories
|
||||
.extract_model
|
||||
.clone()
|
||||
.unwrap_or(MODEL.to_string());
|
||||
context
|
||||
.stage_one_request_context(config, &model_name, REASONING_EFFORT)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_jobs(
|
||||
context: Arc<MemoryStartupContext>,
|
||||
config: Arc<Config>,
|
||||
claimed_candidates: Vec<codex_state::Stage1JobClaim>,
|
||||
stage_one_context: StageOneRequestContext,
|
||||
) -> Vec<JobResult> {
|
||||
futures::stream::iter(claimed_candidates.into_iter())
|
||||
.map(|claim| {
|
||||
let context = Arc::clone(&context);
|
||||
let config = Arc::clone(&config);
|
||||
let stage_one_context = stage_one_context.clone();
|
||||
async move {
|
||||
job::run(context.as_ref(), config.as_ref(), claim, &stage_one_context).await
|
||||
}
|
||||
})
|
||||
.buffer_unordered(CONCURRENCY_LIMIT)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
}
|
||||
|
||||
mod job {
|
||||
use super::*;
|
||||
|
||||
pub(crate) async fn run(
|
||||
context: &MemoryStartupContext,
|
||||
config: &Config,
|
||||
claim: codex_state::Stage1JobClaim,
|
||||
stage_one_context: &StageOneRequestContext,
|
||||
) -> JobResult {
|
||||
let claimed_thread = claim.thread;
|
||||
let (stage_one_output, token_usage) = match sample(
|
||||
context,
|
||||
config,
|
||||
&claimed_thread.rollout_path,
|
||||
&claimed_thread.cwd,
|
||||
stage_one_context,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(output) => output,
|
||||
Err(reason) => {
|
||||
result::failed(
|
||||
context,
|
||||
claimed_thread.id,
|
||||
&claim.ownership_token,
|
||||
&reason.to_string(),
|
||||
)
|
||||
.await;
|
||||
return JobResult {
|
||||
outcome: JobOutcome::Failed,
|
||||
token_usage: None,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
if stage_one_output.raw_memory.is_empty() || stage_one_output.rollout_summary.is_empty() {
|
||||
return JobResult {
|
||||
outcome: result::no_output(context, claimed_thread.id, &claim.ownership_token)
|
||||
.await,
|
||||
token_usage,
|
||||
};
|
||||
}
|
||||
|
||||
JobResult {
|
||||
outcome: result::success(
|
||||
context,
|
||||
claimed_thread.id,
|
||||
&claim.ownership_token,
|
||||
claimed_thread.updated_at.timestamp(),
|
||||
&stage_one_output.raw_memory,
|
||||
&stage_one_output.rollout_summary,
|
||||
stage_one_output.rollout_slug.as_deref(),
|
||||
)
|
||||
.await,
|
||||
token_usage,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the rollout and perform the actual sampling.
|
||||
async fn sample(
|
||||
context: &MemoryStartupContext,
|
||||
config: &Config,
|
||||
rollout_path: &Path,
|
||||
rollout_cwd: &Path,
|
||||
stage_one_context: &StageOneRequestContext,
|
||||
) -> anyhow::Result<(StageOneOutput, Option<TokenUsage>)> {
|
||||
let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?;
|
||||
let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?;
|
||||
|
||||
let mut prompt = Prompt::default();
|
||||
prompt.input = vec![ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: build_stage_one_input_message(
|
||||
&stage_one_context.model_info,
|
||||
rollout_path,
|
||||
rollout_cwd,
|
||||
&rollout_contents,
|
||||
)?,
|
||||
}],
|
||||
phase: None,
|
||||
}];
|
||||
prompt.base_instructions = BaseInstructions {
|
||||
text: STAGE_ONE_PROMPT.to_string(),
|
||||
};
|
||||
prompt.output_schema = Some(output_schema());
|
||||
prompt.output_schema_strict = true;
|
||||
|
||||
let (result, token_usage) = context
|
||||
.stream_stage_one_prompt(config, &prompt, stage_one_context)
|
||||
.await?;
|
||||
|
||||
let mut output: StageOneOutput = serde_json::from_str(&result)?;
|
||||
output.raw_memory = redact_secrets(output.raw_memory);
|
||||
output.rollout_summary = redact_secrets(output.rollout_summary);
|
||||
output.rollout_slug = output.rollout_slug.map(redact_secrets);
|
||||
|
||||
Ok((output, token_usage))
|
||||
}
|
||||
|
||||
mod result {
|
||||
use super::*;
|
||||
|
||||
pub(crate) async fn failed(
|
||||
context: &MemoryStartupContext,
|
||||
thread_id: codex_protocol::ThreadId,
|
||||
ownership_token: &str,
|
||||
reason: &str,
|
||||
) {
|
||||
tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}");
|
||||
if let Some(state_db) = context.state_db() {
|
||||
let _ = state_db
|
||||
.mark_stage1_job_failed(
|
||||
thread_id,
|
||||
ownership_token,
|
||||
reason,
|
||||
JOB_RETRY_DELAY_SECONDS,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn no_output(
|
||||
context: &MemoryStartupContext,
|
||||
thread_id: codex_protocol::ThreadId,
|
||||
ownership_token: &str,
|
||||
) -> JobOutcome {
|
||||
let Some(state_db) = context.state_db() else {
|
||||
return JobOutcome::Failed;
|
||||
};
|
||||
|
||||
if state_db
|
||||
.mark_stage1_job_succeeded_no_output(thread_id, ownership_token)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
{
|
||||
JobOutcome::SucceededNoOutput
|
||||
} else {
|
||||
JobOutcome::Failed
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn success(
|
||||
context: &MemoryStartupContext,
|
||||
thread_id: codex_protocol::ThreadId,
|
||||
ownership_token: &str,
|
||||
source_updated_at: i64,
|
||||
raw_memory: &str,
|
||||
rollout_summary: &str,
|
||||
rollout_slug: Option<&str>,
|
||||
) -> JobOutcome {
|
||||
let Some(state_db) = context.state_db() else {
|
||||
return JobOutcome::Failed;
|
||||
};
|
||||
|
||||
if state_db
|
||||
.mark_stage1_job_succeeded(
|
||||
thread_id,
|
||||
ownership_token,
|
||||
source_updated_at,
|
||||
raw_memory,
|
||||
rollout_summary,
|
||||
rollout_slug,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
{
|
||||
JobOutcome::SucceededWithOutput
|
||||
} else {
|
||||
JobOutcome::Failed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializes filtered stage-1 memory items for prompt inclusion.
|
||||
pub(super) fn serialize_filtered_rollout_response_items(
|
||||
items: &[RolloutItem],
|
||||
) -> codex_protocol::error::Result<String> {
|
||||
let filtered = items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
if let RolloutItem::ResponseItem(item) = item {
|
||||
sanitize_response_item_for_memories(item)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let serialized = serde_json::to_string(&filtered).map_err(|err| {
|
||||
CodexErr::InvalidRequest(format!("failed to serialize rollout memory: {err}"))
|
||||
})?;
|
||||
Ok(redact_secrets(serialized))
|
||||
}
|
||||
|
||||
fn sanitize_response_item_for_memories(item: &ResponseItem) -> Option<ResponseItem> {
|
||||
let ResponseItem::Message {
|
||||
id,
|
||||
role,
|
||||
content,
|
||||
phase,
|
||||
} = item
|
||||
else {
|
||||
return should_persist_response_item_for_memories(item).then(|| item.clone());
|
||||
};
|
||||
|
||||
if role == "developer" {
|
||||
return None;
|
||||
}
|
||||
|
||||
if role != "user" {
|
||||
return Some(item.clone());
|
||||
}
|
||||
|
||||
let content = content
|
||||
.iter()
|
||||
.filter(|content_item| !is_memory_excluded_contextual_user_fragment(content_item))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if content.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(ResponseItem::Message {
|
||||
id: id.clone(),
|
||||
role: role.clone(),
|
||||
content,
|
||||
phase: phase.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_memory_excluded_contextual_user_fragment(content_item: &ContentItem) -> bool {
|
||||
let ContentItem::InputText { text } = content_item else {
|
||||
return false;
|
||||
};
|
||||
|
||||
matches_marked_fragment(text, "# AGENTS.md instructions for ", "</INSTRUCTIONS>")
|
||||
|| matches_marked_fragment(text, "<skill>", "</skill>")
|
||||
}
|
||||
|
||||
fn matches_marked_fragment(text: &str, start_marker: &str, end_marker: &str) -> bool {
|
||||
let trimmed = text.trim_start();
|
||||
let starts_with_marker = trimmed
|
||||
.get(..start_marker.len())
|
||||
.is_some_and(|candidate| candidate.eq_ignore_ascii_case(start_marker));
|
||||
let trimmed = trimmed.trim_end();
|
||||
let ends_with_marker = trimmed
|
||||
.get(trimmed.len().saturating_sub(end_marker.len())..)
|
||||
.is_some_and(|candidate| candidate.eq_ignore_ascii_case(end_marker));
|
||||
starts_with_marker && ends_with_marker
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn classifies_memory_excluded_fragments() {
|
||||
let cases = [
|
||||
(
|
||||
"# AGENTS.md instructions for /tmp\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
|
||||
true,
|
||||
),
|
||||
(
|
||||
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
|
||||
true,
|
||||
),
|
||||
(
|
||||
"<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>",
|
||||
false,
|
||||
),
|
||||
];
|
||||
|
||||
for (text, expected) in cases {
|
||||
assert_eq!(
|
||||
is_memory_excluded_contextual_user_fragment(&ContentItem::InputText {
|
||||
text: text.to_string(),
|
||||
}),
|
||||
expected,
|
||||
"{text}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_schema_requires_rollout_slug_and_keeps_it_nullable() {
|
||||
let schema = output_schema();
|
||||
let properties = schema
|
||||
.get("properties")
|
||||
.and_then(Value::as_object)
|
||||
.expect("properties object");
|
||||
let required = schema
|
||||
.get("required")
|
||||
.and_then(Value::as_array)
|
||||
.expect("required array");
|
||||
|
||||
let mut required_keys = required
|
||||
.iter()
|
||||
.map(|key| key.as_str().expect("required key string"))
|
||||
.collect::<Vec<_>>();
|
||||
required_keys.sort_unstable();
|
||||
|
||||
assert!(
|
||||
properties.contains_key("rollout_slug"),
|
||||
"schema should declare rollout_slug"
|
||||
);
|
||||
|
||||
let rollout_slug_type = properties
|
||||
.get("rollout_slug")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|entry| entry.get("type"))
|
||||
.and_then(Value::as_array)
|
||||
.expect("rollout_slug type array");
|
||||
let mut rollout_slug_types = rollout_slug_type
|
||||
.iter()
|
||||
.map(|entry| entry.as_str().expect("type entry string"))
|
||||
.collect::<Vec<_>>();
|
||||
rollout_slug_types.sort_unstable();
|
||||
|
||||
assert_eq!(
|
||||
required_keys,
|
||||
vec!["raw_memory", "rollout_slug", "rollout_summary"]
|
||||
);
|
||||
assert_eq!(rollout_slug_types, vec!["null", "string"]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn aggregate_stats(outcomes: Vec<JobResult>) -> Stats {
|
||||
let claimed = outcomes.len();
|
||||
let mut succeeded_with_output = 0;
|
||||
let mut succeeded_no_output = 0;
|
||||
let mut failed = 0;
|
||||
let mut total_token_usage = TokenUsage::default();
|
||||
let mut has_token_usage = false;
|
||||
|
||||
for outcome in outcomes {
|
||||
match outcome.outcome {
|
||||
JobOutcome::SucceededWithOutput => succeeded_with_output += 1,
|
||||
JobOutcome::SucceededNoOutput => succeeded_no_output += 1,
|
||||
JobOutcome::Failed => failed += 1,
|
||||
}
|
||||
|
||||
if let Some(token_usage) = outcome.token_usage {
|
||||
total_token_usage.add_assign(&token_usage);
|
||||
has_token_usage = true;
|
||||
}
|
||||
}
|
||||
|
||||
Stats {
|
||||
claimed,
|
||||
succeeded_with_output,
|
||||
succeeded_no_output,
|
||||
failed,
|
||||
total_token_usage: has_token_usage.then_some(total_token_usage),
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_metrics(context: &StageOneRequestContext, counts: &Stats) {
|
||||
if counts.claimed > 0 {
|
||||
context.counter(
|
||||
MEMORY_PHASE_ONE_JOBS,
|
||||
counts.claimed as i64,
|
||||
&[("status", "claimed")],
|
||||
);
|
||||
}
|
||||
if counts.succeeded_with_output > 0 {
|
||||
context.counter(
|
||||
MEMORY_PHASE_ONE_JOBS,
|
||||
counts.succeeded_with_output as i64,
|
||||
&[("status", "succeeded")],
|
||||
);
|
||||
context.counter(
|
||||
MEMORY_PHASE_ONE_OUTPUT,
|
||||
counts.succeeded_with_output as i64,
|
||||
&[],
|
||||
);
|
||||
}
|
||||
if counts.succeeded_no_output > 0 {
|
||||
context.counter(
|
||||
MEMORY_PHASE_ONE_JOBS,
|
||||
counts.succeeded_no_output as i64,
|
||||
&[("status", "succeeded_no_output")],
|
||||
);
|
||||
}
|
||||
if counts.failed > 0 {
|
||||
context.counter(
|
||||
MEMORY_PHASE_ONE_JOBS,
|
||||
counts.failed as i64,
|
||||
&[("status", "failed")],
|
||||
);
|
||||
}
|
||||
if let Some(token_usage) = counts.total_token_usage.as_ref() {
|
||||
context.histogram(
|
||||
MEMORY_PHASE_ONE_TOKEN_USAGE,
|
||||
token_usage.total_tokens.max(0),
|
||||
&[("token_type", "total")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_ONE_TOKEN_USAGE,
|
||||
token_usage.input_tokens.max(0),
|
||||
&[("token_type", "input")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_ONE_TOKEN_USAGE,
|
||||
token_usage.cached_input(),
|
||||
&[("token_type", "cached_input")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_ONE_TOKEN_USAGE,
|
||||
token_usage.output_tokens.max(0),
|
||||
&[("token_type", "output")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_ONE_TOKEN_USAGE,
|
||||
token_usage.reasoning_output_tokens.max(0),
|
||||
&[("token_type", "reasoning_output")],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn serializes_memory_rollout_with_agents_removed_but_environment_kept() {
|
||||
let mixed_contextual_message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![
|
||||
ContentItem::InputText {
|
||||
text:
|
||||
"# AGENTS.md instructions for /tmp\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
|
||||
.to_string(),
|
||||
},
|
||||
ContentItem::InputText {
|
||||
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>"
|
||||
.to_string(),
|
||||
},
|
||||
],
|
||||
phase: None,
|
||||
};
|
||||
let skill_message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text:
|
||||
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>"
|
||||
.to_string(),
|
||||
}],
|
||||
phase: None,
|
||||
};
|
||||
let subagent_message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: "<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>"
|
||||
.to_string(),
|
||||
}],
|
||||
phase: None,
|
||||
};
|
||||
|
||||
let serialized = job::serialize_filtered_rollout_response_items(&[
|
||||
RolloutItem::ResponseItem(mixed_contextual_message),
|
||||
RolloutItem::ResponseItem(skill_message),
|
||||
RolloutItem::ResponseItem(subagent_message.clone()),
|
||||
])
|
||||
.expect("serialize");
|
||||
let parsed: Vec<ResponseItem> = serde_json::from_str(&serialized).expect("parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
vec![
|
||||
ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>"
|
||||
.to_string(),
|
||||
}],
|
||||
phase: None,
|
||||
},
|
||||
subagent_message,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_memory_rollout_redacts_secrets_before_prompt_upload() {
|
||||
let serialized =
|
||||
job::serialize_filtered_rollout_response_items(&[RolloutItem::ResponseItem(
|
||||
ResponseItem::FunctionCallOutput {
|
||||
call_id: "call_123".to_string(),
|
||||
output: codex_protocol::models::FunctionCallOutputPayload {
|
||||
body: codex_protocol::models::FunctionCallOutputBody::Text(
|
||||
r#"{"token":"sk-abcdefghijklmnopqrstuvwxyz123456"}"#.to_string(),
|
||||
),
|
||||
success: Some(true),
|
||||
},
|
||||
},
|
||||
)])
|
||||
.expect("serialize");
|
||||
|
||||
assert!(!serialized.contains("sk-abcdefghijklmnopqrstuvwxyz123456"));
|
||||
assert!(serialized.contains("[REDACTED_SECRET]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_outcomes_sums_token_usage_across_all_jobs() {
|
||||
let counts = aggregate_stats(vec![
|
||||
JobResult {
|
||||
outcome: JobOutcome::SucceededWithOutput,
|
||||
token_usage: Some(TokenUsage {
|
||||
input_tokens: 10,
|
||||
cached_input_tokens: 2,
|
||||
output_tokens: 3,
|
||||
reasoning_output_tokens: 1,
|
||||
total_tokens: 13,
|
||||
}),
|
||||
},
|
||||
JobResult {
|
||||
outcome: JobOutcome::SucceededNoOutput,
|
||||
token_usage: Some(TokenUsage {
|
||||
input_tokens: 7,
|
||||
cached_input_tokens: 1,
|
||||
output_tokens: 2,
|
||||
reasoning_output_tokens: 0,
|
||||
total_tokens: 9,
|
||||
}),
|
||||
},
|
||||
JobResult {
|
||||
outcome: JobOutcome::Failed,
|
||||
token_usage: None,
|
||||
},
|
||||
]);
|
||||
|
||||
assert_eq!(counts.claimed, 3);
|
||||
assert_eq!(counts.succeeded_with_output, 1);
|
||||
assert_eq!(counts.succeeded_no_output, 1);
|
||||
assert_eq!(counts.failed, 1);
|
||||
assert_eq!(
|
||||
counts.total_token_usage,
|
||||
Some(TokenUsage {
|
||||
input_tokens: 17,
|
||||
cached_input_tokens: 3,
|
||||
output_tokens: 5,
|
||||
reasoning_output_tokens: 1,
|
||||
total_tokens: 22,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_outcomes_keeps_usage_empty_when_no_job_reports_it() {
|
||||
let counts = aggregate_stats(vec![
|
||||
JobResult {
|
||||
outcome: JobOutcome::SucceededWithOutput,
|
||||
token_usage: None,
|
||||
},
|
||||
JobResult {
|
||||
outcome: JobOutcome::Failed,
|
||||
token_usage: None,
|
||||
},
|
||||
]);
|
||||
|
||||
assert_eq!(counts.claimed, 2);
|
||||
assert_eq!(counts.total_token_usage, None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,565 @@
|
||||
use crate::build_consolidation_prompt;
|
||||
use crate::memory_root;
|
||||
use crate::prune_old_extension_resources;
|
||||
use crate::rebuild_raw_memories_file_from_memories;
|
||||
use crate::runtime::MemoryStartupContext;
|
||||
use crate::runtime::SpawnedConsolidationAgent;
|
||||
use crate::sync_rollout_summaries_from_memories;
|
||||
use crate::workspace::memory_workspace_diff;
|
||||
use crate::workspace::prepare_memory_workspace;
|
||||
use crate::workspace::reset_memory_workspace_baseline;
|
||||
use crate::workspace::write_workspace_diff;
|
||||
use codex_config::Constrained;
|
||||
use codex_core::config::Config;
|
||||
use codex_features::Feature;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::AgentStatus;
|
||||
use codex_protocol::protocol::AskForApproval;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_state::Stage1Output;
|
||||
use codex_state::StateRuntime;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
const MODEL: &str = "gpt-5.4";
|
||||
const REASONING_EFFORT: codex_protocol::openai_models::ReasoningEffort =
|
||||
codex_protocol::openai_models::ReasoningEffort::Medium;
|
||||
const JOB_LEASE_SECONDS: i64 = 3_600;
|
||||
const JOB_RETRY_DELAY_SECONDS: i64 = 3_600;
|
||||
const JOB_HEARTBEAT_SECONDS: u64 = 90;
|
||||
const MEMORY_PHASE_TWO_JOBS: &str = "codex.memory.phase2";
|
||||
const MEMORY_PHASE_TWO_E2E_MS: &str = "codex.memory.phase2.e2e_ms";
|
||||
const MEMORY_PHASE_TWO_INPUT: &str = "codex.memory.phase2.input";
|
||||
const MEMORY_PHASE_TWO_TOKEN_USAGE: &str = "codex.memory.phase2.token_usage";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct Claim {
|
||||
token: String,
|
||||
watermark: i64,
|
||||
}
|
||||
|
||||
#[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 async fn run(context: Arc<MemoryStartupContext>, config: Arc<Config>) {
|
||||
let phase_two_e2e_timer = context.start_timer(MEMORY_PHASE_TWO_E2E_MS);
|
||||
|
||||
let Some(db) = context.state_db() else {
|
||||
// This should not happen.
|
||||
return;
|
||||
};
|
||||
let root = memory_root(&config.codex_home);
|
||||
let max_raw_memories = config.memories.max_raw_memories_for_consolidation;
|
||||
let max_unused_days = config.memories.max_unused_days;
|
||||
|
||||
// 1. Claim the global Phase 2 lock before touching the memory workspace.
|
||||
let claim = match job::claim(context.as_ref(), db.as_ref()).await {
|
||||
Ok(claim) => claim,
|
||||
Err(e) => {
|
||||
context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", e)]);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 2. Ensure the memories root has a git baseline repository.
|
||||
if let Err(err) = prepare_memory_workspace(&root).await {
|
||||
tracing::error!("failed preparing memory workspace: {err}");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_prepare_workspace",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
// 3. Build the locked-down config used by the consolidation agent.
|
||||
let Some(agent_config) = agent::get_config(config.as_ref()) else {
|
||||
// If we can't get the config, we can't consolidate.
|
||||
tracing::error!("failed to get agent config");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_sandbox_policy",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
|
||||
// 4. Load current DB-backed Phase 2 inputs.
|
||||
let raw_memories = match db
|
||||
.get_phase2_input_selection(max_raw_memories, max_unused_days)
|
||||
.await
|
||||
{
|
||||
Ok(raw_memories) => raw_memories,
|
||||
Err(err) => {
|
||||
tracing::error!("failed to list stage1 outputs from global: {err}");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_load_stage1_outputs",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let raw_memory_count = raw_memories.len();
|
||||
let new_watermark = get_watermark(claim.watermark, &raw_memories);
|
||||
|
||||
// 5. Sync the current inputs into the memory workspace.
|
||||
if let Err(err) = sync_phase2_workspace_inputs(&root, &raw_memories).await {
|
||||
tracing::error!("failed syncing phase2 workspace inputs: {err}");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_sync_workspace_inputs",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
// 6. Use git to decide whether the synced workspace actually changed.
|
||||
let workspace_diff = match memory_workspace_diff(&root).await {
|
||||
Ok(diff) => diff,
|
||||
Err(err) => {
|
||||
tracing::error!("failed checking memory workspace changes: {err}");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_workspace_status",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
if !workspace_diff.has_changes() {
|
||||
tracing::error!("Phase 2 no changes");
|
||||
// We check only after sync of the file system.
|
||||
job::succeed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
new_watermark,
|
||||
&raw_memories,
|
||||
"succeeded_no_workspace_changes",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
// 7. Persist the diff for the consolidation agent to inspect.
|
||||
if let Err(err) = write_workspace_diff(&root, &workspace_diff).await {
|
||||
tracing::error!("failed writing memory workspace diff file: {err}");
|
||||
job::failed(
|
||||
context.as_ref(),
|
||||
db.as_ref(),
|
||||
&claim,
|
||||
"failed_workspace_diff_file",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
// 8. Spawn the consolidation agent.
|
||||
let prompt = agent::get_prompt(&root);
|
||||
let agent = match context
|
||||
.spawn_consolidation_agent(agent_config, prompt)
|
||||
.await
|
||||
{
|
||||
Ok(agent) => agent,
|
||||
Err(err) => {
|
||||
tracing::error!("failed to spawn global memory consolidation agent: {err}");
|
||||
job::failed(context.as_ref(), db.as_ref(), &claim, "failed_spawn_agent").await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 9. Hand off completion handling, heartbeats, and baseline reset.
|
||||
agent::handle(
|
||||
Arc::clone(&context),
|
||||
claim,
|
||||
new_watermark,
|
||||
raw_memories.clone(),
|
||||
root,
|
||||
agent,
|
||||
phase_two_e2e_timer,
|
||||
);
|
||||
|
||||
// 10. Emit dispatch metrics.
|
||||
let counters = Counters {
|
||||
input: raw_memory_count as i64,
|
||||
};
|
||||
emit_metrics(context.as_ref(), counters);
|
||||
}
|
||||
|
||||
async fn sync_phase2_workspace_inputs(
|
||||
root: &Path,
|
||||
raw_memories: &[Stage1Output],
|
||||
) -> std::io::Result<()> {
|
||||
let raw_memory_count = raw_memories.len();
|
||||
sync_rollout_summaries_from_memories(root, raw_memories, raw_memory_count).await?;
|
||||
rebuild_raw_memories_file_from_memories(root, raw_memories, raw_memory_count).await?;
|
||||
prune_old_extension_resources(root).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
mod job {
|
||||
use super::*;
|
||||
|
||||
pub(super) async fn claim(
|
||||
context: &MemoryStartupContext,
|
||||
db: &StateRuntime,
|
||||
) -> Result<Claim, &'static str> {
|
||||
let claim = db
|
||||
.try_claim_global_phase2_job(context.thread_id(), 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,
|
||||
} => {
|
||||
context.counter(
|
||||
MEMORY_PHASE_TWO_JOBS,
|
||||
/*inc*/ 1,
|
||||
&[("status", "claimed")],
|
||||
);
|
||||
(ownership_token, input_watermark)
|
||||
}
|
||||
codex_state::Phase2JobClaimOutcome::SkippedRetryUnavailable => {
|
||||
return Err("skipped_retry_unavailable");
|
||||
}
|
||||
codex_state::Phase2JobClaimOutcome::SkippedRunning => return Err("skipped_running"),
|
||||
};
|
||||
|
||||
Ok(Claim { token, watermark })
|
||||
}
|
||||
|
||||
pub(super) async fn failed(
|
||||
context: &MemoryStartupContext,
|
||||
db: &StateRuntime,
|
||||
claim: &Claim,
|
||||
reason: &'static str,
|
||||
) {
|
||||
context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]);
|
||||
if matches!(
|
||||
db.mark_global_phase2_job_failed(&claim.token, reason, JOB_RETRY_DELAY_SECONDS,)
|
||||
.await,
|
||||
Ok(false)
|
||||
) {
|
||||
let _ = db
|
||||
.mark_global_phase2_job_failed_if_unowned(
|
||||
&claim.token,
|
||||
reason,
|
||||
JOB_RETRY_DELAY_SECONDS,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn succeed(
|
||||
context: &MemoryStartupContext,
|
||||
db: &StateRuntime,
|
||||
claim: &Claim,
|
||||
completion_watermark: i64,
|
||||
selected_outputs: &[codex_state::Stage1Output],
|
||||
reason: &'static str,
|
||||
) -> bool {
|
||||
context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]);
|
||||
db.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
mod agent {
|
||||
use super::*;
|
||||
use tracing::warn;
|
||||
|
||||
pub(super) fn get_config(config: &Config) -> Option<Config> {
|
||||
let root = memory_root(&config.codex_home);
|
||||
let mut agent_config = config.clone();
|
||||
|
||||
agent_config.cwd = root.clone();
|
||||
// Consolidation threads must never feed back into phase-1 memory generation.
|
||||
agent_config.ephemeral = true;
|
||||
agent_config.memories.generate_memories = false;
|
||||
agent_config.memories.use_memories = false;
|
||||
agent_config.include_apps_instructions = false;
|
||||
agent_config.mcp_servers = Constrained::allow_only(HashMap::new());
|
||||
// Approval policy
|
||||
agent_config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never);
|
||||
// Consolidation runs as an internal worker and must not recursively delegate.
|
||||
let _ = agent_config.features.disable(Feature::SpawnCsv);
|
||||
let _ = agent_config.features.disable(Feature::Collab);
|
||||
let _ = agent_config.features.disable(Feature::MemoryTool);
|
||||
let _ = agent_config.features.disable(Feature::Apps);
|
||||
let _ = agent_config.features.disable(Feature::Plugins);
|
||||
let _ = agent_config
|
||||
.features
|
||||
.disable(Feature::SkillMcpDependencyInstall);
|
||||
|
||||
// Sandbox policy
|
||||
let writable_roots = vec![root];
|
||||
// The consolidation agent only needs local memory-root write access and no network.
|
||||
let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite {
|
||||
writable_roots,
|
||||
network_access: false,
|
||||
exclude_tmpdir_env_var: true,
|
||||
exclude_slash_tmp: true,
|
||||
};
|
||||
agent_config
|
||||
.permissions
|
||||
.set_legacy_sandbox_policy(consolidation_sandbox_policy, agent_config.cwd.as_path())
|
||||
.ok()?;
|
||||
|
||||
agent_config.model = Some(
|
||||
config
|
||||
.memories
|
||||
.consolidation_model
|
||||
.clone()
|
||||
.unwrap_or(MODEL.to_string()),
|
||||
);
|
||||
agent_config.model_reasoning_effort = Some(REASONING_EFFORT);
|
||||
|
||||
Some(agent_config)
|
||||
}
|
||||
|
||||
pub(super) fn get_prompt(root: &Path) -> Vec<UserInput> {
|
||||
let prompt = build_consolidation_prompt(root);
|
||||
vec![UserInput::Text {
|
||||
text: prompt,
|
||||
text_elements: vec![],
|
||||
}]
|
||||
}
|
||||
|
||||
/// Handle the agent while it is running.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn handle(
|
||||
context: Arc<MemoryStartupContext>,
|
||||
claim: Claim,
|
||||
new_watermark: i64,
|
||||
selected_outputs: Vec<codex_state::Stage1Output>,
|
||||
memory_root: codex_utils_absolute_path::AbsolutePathBuf,
|
||||
agent: SpawnedConsolidationAgent,
|
||||
phase_two_e2e_timer: Option<codex_otel::Timer>,
|
||||
) {
|
||||
let Some(db) = context.state_db() else {
|
||||
return;
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _phase_two_e2e_timer = phase_two_e2e_timer;
|
||||
let SpawnedConsolidationAgent { thread_id, thread } = agent;
|
||||
|
||||
// Loop the agent until we have the final status.
|
||||
let final_status =
|
||||
loop_agent(db.clone(), claim.token.clone(), thread_id, &thread).await;
|
||||
|
||||
if matches!(final_status, AgentStatus::Completed(_)) {
|
||||
if let Some(token_usage) = thread
|
||||
.token_usage_info()
|
||||
.await
|
||||
.map(|info| info.total_token_usage)
|
||||
{
|
||||
emit_token_usage_metrics(context.as_ref(), &token_usage);
|
||||
}
|
||||
// Do not reset the workspace baseline if we lost the lock.
|
||||
let still_owns_lock = match db
|
||||
.heartbeat_global_phase2_job(&claim.token, JOB_LEASE_SECONDS)
|
||||
.await
|
||||
.inspect_err(|err| {
|
||||
tracing::error!(
|
||||
"failed confirming global memory consolidation ownership before resetting workspace baseline: {err}"
|
||||
);
|
||||
}) {
|
||||
Ok(true) => true,
|
||||
Ok(false) => {
|
||||
tracing::error!(
|
||||
"lost global memory consolidation ownership before resetting workspace baseline"
|
||||
);
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
job::failed(context.as_ref(), &db, &claim, "failed_confirm_ownership")
|
||||
.await;
|
||||
false
|
||||
}
|
||||
};
|
||||
if still_owns_lock {
|
||||
if let Err(err) = reset_memory_workspace_baseline(&memory_root).await {
|
||||
tracing::error!("failed resetting memory workspace baseline: {err}");
|
||||
job::failed(context.as_ref(), &db, &claim, "failed_workspace_commit").await;
|
||||
} else if !job::succeed(
|
||||
context.as_ref(),
|
||||
&db,
|
||||
&claim,
|
||||
new_watermark,
|
||||
&selected_outputs,
|
||||
"succeeded",
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
"failed marking global memory consolidation job succeeded after resetting workspace baseline"
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
job::failed(context.as_ref(), &db, &claim, "failed_agent").await;
|
||||
}
|
||||
|
||||
let cleanup_context = Arc::clone(&context);
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = cleanup_context
|
||||
.shutdown_consolidation_agent(SpawnedConsolidationAgent { thread_id, thread })
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"failed to auto-close global memory consolidation agent {thread_id}: {err}"
|
||||
);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async fn loop_agent(
|
||||
db: Arc<StateRuntime>,
|
||||
token: String,
|
||||
thread_id: ThreadId,
|
||||
thread: &codex_core::CodexThread,
|
||||
) -> AgentStatus {
|
||||
let mut heartbeat_interval =
|
||||
tokio::time::interval(Duration::from_secs(JOB_HEARTBEAT_SECONDS));
|
||||
heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
let mut status_poll_interval = tokio::time::interval(Duration::from_secs(1));
|
||||
status_poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
let session_termination = thread.wait_until_terminated();
|
||||
tokio::pin!(session_termination);
|
||||
|
||||
loop {
|
||||
let status = thread.agent_status().await;
|
||||
if is_final_agent_status(&status) {
|
||||
break status;
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut session_termination => {
|
||||
let status = thread.agent_status().await;
|
||||
if is_final_agent_status(&status) {
|
||||
break status;
|
||||
}
|
||||
tracing::warn!(
|
||||
"memory consolidation agent {thread_id} exited before final status; last status was {status:?}"
|
||||
);
|
||||
break AgentStatus::Errored(format!(
|
||||
"memory consolidation agent exited before final status: {status:?}"
|
||||
));
|
||||
}
|
||||
_ = status_poll_interval.tick() => {
|
||||
}
|
||||
_ = heartbeat_interval.tick() => {
|
||||
match db
|
||||
.heartbeat_global_phase2_job(
|
||||
&token,
|
||||
JOB_LEASE_SECONDS,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
tracing::warn!(
|
||||
"lost global phase-2 ownership during heartbeat for memory consolidation agent {thread_id}"
|
||||
);
|
||||
break AgentStatus::Errored(
|
||||
"lost global phase-2 ownership during heartbeat".to_string(),
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
"phase-2 heartbeat update failed for memory consolidation agent {thread_id}: {err}"
|
||||
);
|
||||
break AgentStatus::Errored(format!(
|
||||
"phase-2 heartbeat update failed: {err}"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
fn is_final_agent_status(status: &AgentStatus) -> bool {
|
||||
!matches!(
|
||||
status,
|
||||
AgentStatus::PendingInit | AgentStatus::Running | AgentStatus::Interrupted
|
||||
)
|
||||
}
|
||||
|
||||
fn emit_metrics(context: &MemoryStartupContext, counters: Counters) {
|
||||
if counters.input > 0 {
|
||||
context.counter(MEMORY_PHASE_TWO_INPUT, counters.input, &[]);
|
||||
}
|
||||
|
||||
context.counter(
|
||||
MEMORY_PHASE_TWO_JOBS,
|
||||
/*inc*/ 1,
|
||||
&[("status", "agent_spawned")],
|
||||
);
|
||||
}
|
||||
|
||||
fn emit_token_usage_metrics(context: &MemoryStartupContext, token_usage: &TokenUsage) {
|
||||
context.histogram(
|
||||
MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.total_tokens.max(0),
|
||||
&[("token_type", "total")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.input_tokens.max(0),
|
||||
&[("token_type", "input")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.cached_input(),
|
||||
&[("token_type", "cached_input")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.output_tokens.max(0),
|
||||
&[("token_type", "output")],
|
||||
);
|
||||
context.histogram(
|
||||
MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.reasoning_output_tokens.max(0),
|
||||
&[("token_type", "reasoning_output")],
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,293 @@
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ModelClient;
|
||||
use codex_core::NewThread;
|
||||
use codex_core::Prompt;
|
||||
use codex_core::ResponseEvent;
|
||||
use codex_core::StartThreadOptions;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::content_items_to_text;
|
||||
use codex_core::resolve_installation_id;
|
||||
use codex_features::Feature;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
|
||||
use codex_login::default_client::originator;
|
||||
use codex_otel::SessionTelemetry;
|
||||
use codex_otel::TelemetryAuthMode;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::config_types::ReasoningSummary;
|
||||
use codex_protocol::config_types::ServiceTier;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
use codex_protocol::openai_models::ReasoningEffort;
|
||||
use codex_protocol::protocol::InitialHistory;
|
||||
use codex_protocol::protocol::InternalSessionSource;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_rollout_trace::InferenceTraceContext;
|
||||
use codex_state::StateRuntime;
|
||||
use codex_terminal_detection::user_agent;
|
||||
use futures::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
pub(crate) struct SpawnedConsolidationAgent {
|
||||
pub(crate) thread_id: ThreadId,
|
||||
pub(crate) thread: Arc<CodexThread>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct StageOneRequestContext {
|
||||
pub(crate) model_info: ModelInfo,
|
||||
pub(crate) session_telemetry: SessionTelemetry,
|
||||
pub(crate) reasoning_effort: Option<ReasoningEffort>,
|
||||
pub(crate) reasoning_summary: ReasoningSummary,
|
||||
pub(crate) service_tier: Option<ServiceTier>,
|
||||
pub(crate) turn_metadata_header: Option<String>,
|
||||
}
|
||||
|
||||
impl StageOneRequestContext {
|
||||
pub(crate) fn start_timer(&self, name: &str) -> Option<codex_otel::Timer> {
|
||||
self.session_telemetry.start_timer(name, &[]).ok()
|
||||
}
|
||||
|
||||
pub(crate) fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) {
|
||||
self.session_telemetry.counter(name, inc, tags);
|
||||
}
|
||||
|
||||
pub(crate) fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) {
|
||||
self.session_telemetry.histogram(name, value, tags);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct MemoryStartupContext {
|
||||
thread_id: ThreadId,
|
||||
thread: Arc<CodexThread>,
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
session_telemetry: SessionTelemetry,
|
||||
}
|
||||
|
||||
impl MemoryStartupContext {
|
||||
pub(crate) fn new(
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
thread_id: ThreadId,
|
||||
thread: Arc<CodexThread>,
|
||||
config: &Config,
|
||||
source: SessionSource,
|
||||
) -> Self {
|
||||
let auth = auth_manager.auth_cached();
|
||||
let auth = auth.as_ref();
|
||||
let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from);
|
||||
let account_id = auth.and_then(CodexAuth::get_account_id);
|
||||
let account_email = auth.and_then(CodexAuth::get_account_email);
|
||||
let model = config.model.as_deref().unwrap_or("unknown");
|
||||
let auth_env_telemetry = collect_auth_env_telemetry(
|
||||
&config.model_provider,
|
||||
auth_manager.codex_api_key_env_enabled(),
|
||||
);
|
||||
let session_telemetry = SessionTelemetry::new(
|
||||
thread_id,
|
||||
model,
|
||||
model,
|
||||
account_id,
|
||||
account_email,
|
||||
auth_mode,
|
||||
originator().value,
|
||||
config.otel.log_user_prompt,
|
||||
user_agent(),
|
||||
source,
|
||||
)
|
||||
.with_auth_env(auth_env_telemetry.to_otel_metadata());
|
||||
|
||||
Self {
|
||||
thread_id,
|
||||
thread,
|
||||
thread_manager,
|
||||
auth_manager,
|
||||
session_telemetry,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn thread_id(&self) -> ThreadId {
|
||||
self.thread_id
|
||||
}
|
||||
|
||||
pub(crate) fn state_db(&self) -> Option<Arc<StateRuntime>> {
|
||||
self.thread.state_db()
|
||||
}
|
||||
|
||||
pub(crate) fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) {
|
||||
self.session_telemetry.counter(name, inc, tags);
|
||||
}
|
||||
|
||||
pub(crate) fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) {
|
||||
self.session_telemetry.histogram(name, value, tags);
|
||||
}
|
||||
|
||||
pub(crate) fn start_timer(&self, name: &str) -> Option<codex_otel::Timer> {
|
||||
self.session_telemetry.start_timer(name, &[]).ok()
|
||||
}
|
||||
|
||||
pub(crate) async fn stage_one_request_context(
|
||||
&self,
|
||||
config: &Config,
|
||||
model_name: &str,
|
||||
reasoning_effort: ReasoningEffort,
|
||||
) -> StageOneRequestContext {
|
||||
let config_snapshot = self.thread.config_snapshot().await;
|
||||
let model_info = self
|
||||
.thread_manager
|
||||
.get_models_manager()
|
||||
.get_model_info(model_name, &config.to_models_manager_config())
|
||||
.await;
|
||||
let turn_metadata_header =
|
||||
codex_core::build_turn_metadata_header(&config.cwd, /*sandbox*/ None).await;
|
||||
let reasoning_summary = config
|
||||
.model_reasoning_summary
|
||||
.unwrap_or(model_info.default_reasoning_summary);
|
||||
|
||||
StageOneRequestContext {
|
||||
model_info,
|
||||
turn_metadata_header,
|
||||
session_telemetry: self
|
||||
.session_telemetry
|
||||
.clone()
|
||||
.with_model(model_name, model_name),
|
||||
reasoning_effort: Some(reasoning_effort),
|
||||
reasoning_summary,
|
||||
service_tier: config_snapshot.service_tier,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn stream_stage_one_prompt(
|
||||
&self,
|
||||
config: &Config,
|
||||
prompt: &Prompt,
|
||||
context: &StageOneRequestContext,
|
||||
) -> anyhow::Result<(String, Option<TokenUsage>)> {
|
||||
let installation_id = resolve_installation_id(&config.codex_home).await?;
|
||||
let session_source = self.thread.config_snapshot().await.session_source;
|
||||
let model_client = ModelClient::new(
|
||||
Some(Arc::clone(&self.auth_manager)),
|
||||
self.thread_id,
|
||||
installation_id,
|
||||
config.model_provider.clone(),
|
||||
session_source,
|
||||
config.model_verbosity,
|
||||
config.features.enabled(Feature::EnableRequestCompression),
|
||||
config.features.enabled(Feature::RuntimeMetrics),
|
||||
/*beta_features_header*/ None,
|
||||
);
|
||||
|
||||
let mut client_session = model_client.new_session();
|
||||
let mut stream = client_session
|
||||
.stream(
|
||||
prompt,
|
||||
&context.model_info,
|
||||
&context.session_telemetry,
|
||||
context.reasoning_effort,
|
||||
context.reasoning_summary,
|
||||
context.service_tier,
|
||||
context.turn_metadata_header.as_deref(),
|
||||
&InferenceTraceContext::disabled(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut result = String::new();
|
||||
let mut token_usage = None;
|
||||
while let Some(message) = stream.next().await.transpose()? {
|
||||
match message {
|
||||
ResponseEvent::OutputTextDelta(delta) => result.push_str(&delta),
|
||||
ResponseEvent::OutputItemDone(item) => {
|
||||
if result.is_empty()
|
||||
&& let codex_protocol::models::ResponseItem::Message { content, .. } = item
|
||||
&& let Some(text) = content_items_to_text(&content)
|
||||
{
|
||||
result.push_str(&text);
|
||||
}
|
||||
}
|
||||
ResponseEvent::Completed {
|
||||
token_usage: usage, ..
|
||||
} => {
|
||||
token_usage = usage;
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((result, token_usage))
|
||||
}
|
||||
|
||||
pub(crate) async fn spawn_consolidation_agent(
|
||||
&self,
|
||||
config: Config,
|
||||
prompt: Vec<UserInput>,
|
||||
) -> anyhow::Result<SpawnedConsolidationAgent> {
|
||||
let environments = self
|
||||
.thread_manager
|
||||
.default_environment_selections(&config.cwd);
|
||||
let NewThread {
|
||||
thread_id, thread, ..
|
||||
} = self
|
||||
.thread_manager
|
||||
.start_thread_with_options(StartThreadOptions {
|
||||
config,
|
||||
initial_history: InitialHistory::New,
|
||||
session_source: Some(SessionSource::Internal(
|
||||
InternalSessionSource::MemoryConsolidation,
|
||||
)),
|
||||
dynamic_tools: Vec::new(),
|
||||
persist_extended_history: false,
|
||||
metrics_service_name: None,
|
||||
parent_trace: None,
|
||||
environments,
|
||||
})
|
||||
.await?;
|
||||
|
||||
let agent = SpawnedConsolidationAgent { thread_id, thread };
|
||||
if let Err(err) = agent
|
||||
.thread
|
||||
.submit(Op::UserInput {
|
||||
items: prompt,
|
||||
environments: None,
|
||||
final_output_json_schema: None,
|
||||
responsesapi_client_metadata: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
if let Err(shutdown_err) = self.shutdown_consolidation_agent(agent).await {
|
||||
tracing::warn!(
|
||||
"failed to shut down consolidation agent after submit error: {shutdown_err}"
|
||||
);
|
||||
}
|
||||
return Err(err.into());
|
||||
}
|
||||
|
||||
Ok(agent)
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown_consolidation_agent(
|
||||
&self,
|
||||
agent: SpawnedConsolidationAgent,
|
||||
) -> anyhow::Result<()> {
|
||||
let SpawnedConsolidationAgent { thread_id, thread } = agent;
|
||||
let thread = self
|
||||
.thread_manager
|
||||
.remove_thread(&thread_id)
|
||||
.await
|
||||
.unwrap_or(thread);
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(10), thread.shutdown_and_wait())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!("memory consolidation agent {thread_id} shutdown timed out")
|
||||
})??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
use crate::phase1;
|
||||
use crate::phase2;
|
||||
use crate::runtime::MemoryStartupContext;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_features::Feature;
|
||||
use codex_login::AuthManager;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use std::sync::Arc;
|
||||
use tracing::warn;
|
||||
|
||||
/// Starts the asynchronous startup memory pipeline for an eligible root session.
|
||||
///
|
||||
/// The pipeline is skipped for ephemeral sessions, disabled feature flags, and
|
||||
/// subagent sessions.
|
||||
pub fn start_memories_startup_task(
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
thread_id: ThreadId,
|
||||
thread: Arc<CodexThread>,
|
||||
config: Arc<Config>,
|
||||
source: &SessionSource,
|
||||
) {
|
||||
if config.ephemeral
|
||||
|| !config.features.enabled(Feature::MemoryTool)
|
||||
|| source.is_non_root_agent()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let context = Arc::new(MemoryStartupContext::new(
|
||||
thread_manager,
|
||||
auth_manager,
|
||||
thread_id,
|
||||
thread,
|
||||
config.as_ref(),
|
||||
source.clone(),
|
||||
));
|
||||
|
||||
if context.state_db().is_none() {
|
||||
warn!("state db unavailable for memories startup pipeline; skipping");
|
||||
return;
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Clean memories to make preserve DB size
|
||||
phase1::prune(context.as_ref(), &config).await;
|
||||
// Run phase 1.
|
||||
phase1::run(Arc::clone(&context), Arc::clone(&config)).await;
|
||||
// Run phase 2.
|
||||
phase2::run(context, config).await;
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
use crate::start_memories_startup_task;
|
||||
use codex_features::Feature;
|
||||
use codex_git_utils::diff_since_latest_init;
|
||||
use codex_git_utils::reset_git_repository;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::config_types::ServiceTier;
|
||||
use codex_protocol::openai_models::ReasoningEffort;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use core_test_support::responses::ResponseMock;
|
||||
use core_test_support::responses::ResponsesRequest;
|
||||
use core_test_support::responses::ev_assistant_message;
|
||||
use core_test_support::responses::ev_completed;
|
||||
use core_test_support::responses::ev_response_created;
|
||||
use core_test_support::responses::mount_sse_once;
|
||||
use core_test_support::responses::sse;
|
||||
use core_test_support::responses::start_mock_server;
|
||||
use core_test_support::test_codex::TestCodex;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use core_test_support::wait_for_event;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tempfile::TempDir;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::Instant;
|
||||
|
||||
#[tokio::test]
|
||||
async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> anyhow::Result<()> {
|
||||
let server = start_mock_server().await;
|
||||
let home = Arc::new(TempDir::new()?);
|
||||
let db = init_state_db(&home).await?;
|
||||
let memory_root = home.path().join("memories");
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let _thread_a = seed_stage1_output(
|
||||
db.as_ref(),
|
||||
home.path(),
|
||||
now - chrono::Duration::hours(2),
|
||||
"raw memory A",
|
||||
"rollout summary A",
|
||||
"rollout-a",
|
||||
)
|
||||
.await?;
|
||||
|
||||
let rollout_summaries_root = memory_root.join("rollout_summaries");
|
||||
tokio::fs::create_dir_all(&rollout_summaries_root).await?;
|
||||
tokio::fs::write(
|
||||
memory_root.join("raw_memories.md"),
|
||||
"# Raw Memories\n\nraw memory A\n",
|
||||
)
|
||||
.await?;
|
||||
tokio::fs::write(
|
||||
rollout_summaries_root.join("rollout-a.md"),
|
||||
"git_branch: branch-rollout-a\n\nrollout summary A\n",
|
||||
)
|
||||
.await?;
|
||||
reset_git_repository(&memory_root).await?;
|
||||
|
||||
let _thread_b = seed_stage1_output(
|
||||
db.as_ref(),
|
||||
home.path(),
|
||||
now - chrono::Duration::hours(1),
|
||||
"raw memory B",
|
||||
"rollout summary B",
|
||||
"rollout-b",
|
||||
)
|
||||
.await?;
|
||||
|
||||
let phase2 = mount_sse_once(
|
||||
&server,
|
||||
sse(vec![
|
||||
ev_response_created("resp-phase2"),
|
||||
ev_assistant_message("msg-phase2", "phase2 complete"),
|
||||
ev_completed("resp-phase2"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
|
||||
let test = build_test_codex(&server, home.clone()).await?;
|
||||
trigger_memories_startup(&test).await;
|
||||
|
||||
let request = wait_for_single_request(&phase2).await;
|
||||
let prompt = phase2_prompt_text(&request);
|
||||
assert!(
|
||||
prompt.contains("phase2_workspace_diff.md"),
|
||||
"expected workspace diff file in prompt: {prompt}"
|
||||
);
|
||||
|
||||
wait_for_phase2_workspace_reset(&memory_root).await?;
|
||||
let raw_memories = tokio::fs::read_to_string(memory_root.join("raw_memories.md")).await?;
|
||||
assert!(raw_memories.contains("raw memory B"));
|
||||
assert!(!raw_memories.contains("raw memory A"));
|
||||
let rollout_summaries = read_rollout_summary_bodies(&memory_root).await?;
|
||||
assert_eq!(rollout_summaries.len(), 1);
|
||||
assert!(
|
||||
rollout_summaries
|
||||
.iter()
|
||||
.any(|summary| summary.contains("rollout summary B"))
|
||||
);
|
||||
assert!(
|
||||
rollout_summaries
|
||||
.iter()
|
||||
.any(|summary| summary.contains("git_branch: branch-rollout-b"))
|
||||
);
|
||||
assert!(
|
||||
rollout_summaries
|
||||
.iter()
|
||||
.all(|summary| !summary.contains("rollout summary A"))
|
||||
);
|
||||
|
||||
shutdown_test_codex(&test).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn memories_startup_phase2_prunes_old_extension_resources() -> anyhow::Result<()> {
|
||||
let server = start_mock_server().await;
|
||||
let home = Arc::new(TempDir::new()?);
|
||||
let db = init_state_db(&home).await?;
|
||||
let now = chrono::Utc::now();
|
||||
let _thread_id = seed_stage1_output(
|
||||
db.as_ref(),
|
||||
home.path(),
|
||||
now - chrono::Duration::hours(1),
|
||||
"raw memory",
|
||||
"rollout summary",
|
||||
"rollout",
|
||||
)
|
||||
.await?;
|
||||
|
||||
let chronicle_resources = home.path().join("memories/extensions/chronicle/resources");
|
||||
tokio::fs::create_dir_all(&chronicle_resources).await?;
|
||||
tokio::fs::write(
|
||||
home.path()
|
||||
.join("memories/extensions/chronicle/instructions.md"),
|
||||
"instructions",
|
||||
)
|
||||
.await?;
|
||||
let old_file = chronicle_resources.join(format!(
|
||||
"{}-abcd-10min-old.md",
|
||||
(now - chrono::Duration::days(8)).format("%Y-%m-%dT%H-%M-%S")
|
||||
));
|
||||
tokio::fs::write(&old_file, "old resource").await?;
|
||||
let recent_file = chronicle_resources.join(format!(
|
||||
"{}-abcd-10min-recent.md",
|
||||
(now - chrono::Duration::days(6)).format("%Y-%m-%dT%H-%M-%S")
|
||||
));
|
||||
tokio::fs::write(&recent_file, "recent resource").await?;
|
||||
|
||||
let phase2 = mount_sse_once(
|
||||
&server,
|
||||
sse(vec![
|
||||
ev_response_created("resp-phase2"),
|
||||
ev_assistant_message("msg-phase2", "phase2 complete"),
|
||||
ev_completed("resp-phase2"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
|
||||
let test = build_test_codex(&server, home.clone()).await?;
|
||||
trigger_memories_startup(&test).await;
|
||||
|
||||
let request = wait_for_single_request(&phase2).await;
|
||||
let prompt = phase2_prompt_text(&request);
|
||||
assert!(
|
||||
prompt.contains("phase2_workspace_diff.md"),
|
||||
"expected workspace diff file in prompt: {prompt}"
|
||||
);
|
||||
|
||||
wait_for_phase2_workspace_reset(&home.path().join("memories")).await?;
|
||||
wait_for_file_removed(&old_file).await?;
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&old_file).await?,
|
||||
"old extension resource should be pruned"
|
||||
);
|
||||
assert!(
|
||||
tokio::fs::try_exists(&recent_file).await?,
|
||||
"recent extension resource should be retained"
|
||||
);
|
||||
|
||||
shutdown_test_codex(&test).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_input()
|
||||
-> anyhow::Result<()> {
|
||||
let server = start_mock_server().await;
|
||||
let home = Arc::new(TempDir::new()?);
|
||||
let db = init_state_db(&home).await?;
|
||||
db.enqueue_global_consolidation(/*input_watermark*/ 1)
|
||||
.await?;
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let chronicle_resources = home.path().join("memories/extensions/chronicle/resources");
|
||||
tokio::fs::create_dir_all(&chronicle_resources).await?;
|
||||
tokio::fs::write(
|
||||
home.path()
|
||||
.join("memories/extensions/chronicle/instructions.md"),
|
||||
"instructions",
|
||||
)
|
||||
.await?;
|
||||
let old_file = chronicle_resources.join(format!(
|
||||
"{}-abcd-10min-old.md",
|
||||
(now - chrono::Duration::days(8)).format("%Y-%m-%dT%H-%M-%S")
|
||||
));
|
||||
tokio::fs::write(&old_file, "old resource").await?;
|
||||
|
||||
let phase2 = mount_sse_once(
|
||||
&server,
|
||||
sse(vec![
|
||||
ev_response_created("resp-phase2-empty"),
|
||||
ev_assistant_message("msg-phase2-empty", "phase2 complete"),
|
||||
ev_completed("resp-phase2-empty"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
|
||||
let test = build_test_codex(&server, home.clone()).await?;
|
||||
trigger_memories_startup(&test).await;
|
||||
|
||||
let request = wait_for_single_request(&phase2).await;
|
||||
let prompt = phase2_prompt_text(&request);
|
||||
assert!(
|
||||
prompt.contains("phase2_workspace_diff.md"),
|
||||
"expected workspace diff file in prompt: {prompt}"
|
||||
);
|
||||
|
||||
wait_for_file_removed(&old_file).await?;
|
||||
wait_for_phase2_workspace_reset(&home.path().join("memories")).await?;
|
||||
|
||||
shutdown_test_codex(&test).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn memories_startup_phase1_uses_live_thread_service_tier() -> anyhow::Result<()> {
|
||||
let server = start_mock_server().await;
|
||||
let home = Arc::new(TempDir::new()?);
|
||||
let test = build_test_codex(&server, home).await?;
|
||||
assert_eq!(test.config.service_tier, None);
|
||||
|
||||
test.codex
|
||||
.submit(Op::OverrideTurnContext {
|
||||
cwd: None,
|
||||
approval_policy: None,
|
||||
approvals_reviewer: None,
|
||||
sandbox_policy: None,
|
||||
permission_profile: None,
|
||||
windows_sandbox_level: None,
|
||||
model: None,
|
||||
effort: None,
|
||||
summary: None,
|
||||
service_tier: Some(Some(ServiceTier::Fast)),
|
||||
collaboration_mode: None,
|
||||
personality: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
let config_snapshot = wait_for_service_tier(&test, Some(ServiceTier::Fast)).await?;
|
||||
assert_eq!(config_snapshot.service_tier, Some(ServiceTier::Fast));
|
||||
|
||||
let context = crate::runtime::MemoryStartupContext::new(
|
||||
Arc::clone(&test.thread_manager),
|
||||
test.thread_manager.auth_manager(),
|
||||
test.session_configured.session_id,
|
||||
Arc::clone(&test.codex),
|
||||
&test.config,
|
||||
config_snapshot.session_source.clone(),
|
||||
);
|
||||
let request_context = context
|
||||
.stage_one_request_context(
|
||||
&test.config,
|
||||
test.config.model.as_deref().unwrap_or("gpt-5.4-mini"),
|
||||
ReasoningEffort::Low,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(request_context.service_tier, Some(ServiceTier::Fast));
|
||||
|
||||
shutdown_test_codex(&test).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn build_test_codex(
|
||||
server: &wiremock::MockServer,
|
||||
home: Arc<TempDir>,
|
||||
) -> anyhow::Result<TestCodex> {
|
||||
test_codex()
|
||||
.with_home(home)
|
||||
.with_config(|config| {
|
||||
config
|
||||
.features
|
||||
.enable(Feature::Sqlite)
|
||||
.expect("test config should allow feature update");
|
||||
config.memories.max_raw_memories_for_consolidation = 1;
|
||||
})
|
||||
.build(server)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn init_state_db(home: &Arc<TempDir>) -> anyhow::Result<Arc<codex_state::StateRuntime>> {
|
||||
let db =
|
||||
codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".into()).await?;
|
||||
db.mark_backfill_complete(/*last_watermark*/ None).await?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
async fn trigger_memories_startup(test: &TestCodex) {
|
||||
let config_snapshot = test.codex.config_snapshot().await;
|
||||
let mut config = test.config.clone();
|
||||
config
|
||||
.features
|
||||
.enable(Feature::MemoryTool)
|
||||
.expect("test config should allow feature update");
|
||||
start_memories_startup_task(
|
||||
Arc::clone(&test.thread_manager),
|
||||
test.thread_manager.auth_manager(),
|
||||
test.session_configured.session_id,
|
||||
Arc::clone(&test.codex),
|
||||
Arc::new(config),
|
||||
&config_snapshot.session_source,
|
||||
);
|
||||
}
|
||||
|
||||
async fn seed_stage1_output(
|
||||
db: &codex_state::StateRuntime,
|
||||
codex_home: &Path,
|
||||
updated_at: chrono::DateTime<chrono::Utc>,
|
||||
raw_memory: &str,
|
||||
rollout_summary: &str,
|
||||
rollout_slug: &str,
|
||||
) -> anyhow::Result<ThreadId> {
|
||||
let thread_id = ThreadId::new();
|
||||
let mut metadata_builder = codex_state::ThreadMetadataBuilder::new(
|
||||
thread_id,
|
||||
codex_home.join(format!("rollout-{thread_id}.jsonl")),
|
||||
updated_at,
|
||||
SessionSource::Cli,
|
||||
);
|
||||
metadata_builder.cwd = codex_home.join(format!("workspace-{rollout_slug}"));
|
||||
metadata_builder.model_provider = Some("test-provider".to_string());
|
||||
metadata_builder.git_branch = Some(format!("branch-{rollout_slug}"));
|
||||
let metadata = metadata_builder.build("test-provider");
|
||||
db.upsert_thread(&metadata).await?;
|
||||
|
||||
seed_stage1_output_for_existing_thread(
|
||||
db,
|
||||
thread_id,
|
||||
updated_at.timestamp(),
|
||||
raw_memory,
|
||||
rollout_summary,
|
||||
Some(rollout_slug),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(thread_id)
|
||||
}
|
||||
|
||||
async fn wait_for_single_request(mock: &ResponseMock) -> ResponsesRequest {
|
||||
wait_for_request(mock, /*expected_count*/ 1).await.remove(0)
|
||||
}
|
||||
|
||||
async fn wait_for_file_removed(path: &Path) -> anyhow::Result<()> {
|
||||
let deadline = Instant::now() + Duration::from_secs(10);
|
||||
loop {
|
||||
if !tokio::fs::try_exists(path).await? {
|
||||
return Ok(());
|
||||
}
|
||||
assert!(
|
||||
Instant::now() < deadline,
|
||||
"timed out waiting for {} to be removed",
|
||||
path.display()
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_request(mock: &ResponseMock, expected_count: usize) -> Vec<ResponsesRequest> {
|
||||
let deadline = Instant::now() + Duration::from_secs(10);
|
||||
loop {
|
||||
let requests = mock.requests();
|
||||
if requests.len() >= expected_count {
|
||||
return requests;
|
||||
}
|
||||
assert!(
|
||||
Instant::now() < deadline,
|
||||
"timed out waiting for {expected_count} phase2 requests"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_service_tier(
|
||||
test: &TestCodex,
|
||||
expected_service_tier: Option<ServiceTier>,
|
||||
) -> anyhow::Result<codex_core::ThreadConfigSnapshot> {
|
||||
let deadline = Instant::now() + Duration::from_secs(10);
|
||||
loop {
|
||||
let config_snapshot = test.codex.config_snapshot().await;
|
||||
if config_snapshot.service_tier == expected_service_tier {
|
||||
return Ok(config_snapshot);
|
||||
}
|
||||
anyhow::ensure!(
|
||||
Instant::now() < deadline,
|
||||
"timed out waiting for service_tier to become {expected_service_tier:?}, current={:?}",
|
||||
config_snapshot.service_tier
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn phase2_prompt_text(request: &ResponsesRequest) -> String {
|
||||
request
|
||||
.message_input_texts("user")
|
||||
.into_iter()
|
||||
.find(|text| text.contains("Memory workspace diff:"))
|
||||
.expect("phase2 prompt text")
|
||||
}
|
||||
|
||||
async fn wait_for_phase2_workspace_reset(memory_root: &Path) -> anyhow::Result<()> {
|
||||
wait_for_file_removed(&memory_root.join("phase2_workspace_diff.md")).await?;
|
||||
let deadline = Instant::now() + Duration::from_secs(10);
|
||||
loop {
|
||||
if let Ok(diff) = diff_since_latest_init(memory_root).await
|
||||
&& !diff.has_changes()
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
assert!(
|
||||
Instant::now() < deadline,
|
||||
"timed out waiting for clean memory workspace baseline"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed_stage1_output_for_existing_thread(
|
||||
db: &codex_state::StateRuntime,
|
||||
thread_id: ThreadId,
|
||||
updated_at: i64,
|
||||
raw_memory: &str,
|
||||
rollout_summary: &str,
|
||||
rollout_slug: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let owner = ThreadId::new();
|
||||
let claim = db
|
||||
.try_claim_stage1_job(
|
||||
thread_id, owner, updated_at, /*lease_seconds*/ 3_600,
|
||||
/*max_running_jobs*/ 64,
|
||||
)
|
||||
.await?;
|
||||
let ownership_token = match claim {
|
||||
codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
|
||||
other => panic!("unexpected stage-1 claim outcome: {other:?}"),
|
||||
};
|
||||
|
||||
assert!(
|
||||
db.mark_stage1_job_succeeded(
|
||||
thread_id,
|
||||
&ownership_token,
|
||||
updated_at,
|
||||
raw_memory,
|
||||
rollout_summary,
|
||||
rollout_slug,
|
||||
)
|
||||
.await?,
|
||||
"stage-1 success should enqueue global consolidation"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_rollout_summary_bodies(memory_root: &Path) -> anyhow::Result<Vec<String>> {
|
||||
let mut dir = tokio::fs::read_dir(memory_root.join("rollout_summaries")).await?;
|
||||
let mut summaries = Vec::new();
|
||||
while let Some(entry) = dir.next_entry().await? {
|
||||
summaries.push(tokio::fs::read_to_string(entry.path()).await?);
|
||||
}
|
||||
summaries.sort();
|
||||
Ok(summaries)
|
||||
}
|
||||
|
||||
async fn shutdown_test_codex(test: &TestCodex) -> anyhow::Result<()> {
|
||||
test.codex.submit(Op::Shutdown {}).await?;
|
||||
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::ShutdownComplete)).await;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,10 +1,18 @@
|
||||
use super::rollout_summary_file_stem;
|
||||
use crate::ensure_layout;
|
||||
use crate::raw_memories_file;
|
||||
use crate::rebuild_raw_memories_file_from_memories;
|
||||
use crate::rollout_summaries_dir;
|
||||
use crate::sync_rollout_summaries_from_memories;
|
||||
use chrono::TimeZone;
|
||||
use chrono::Utc;
|
||||
use codex_config::types::DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_state::Stage1Output;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::PathBuf;
|
||||
use tempfile::tempdir;
|
||||
|
||||
const FIXED_PREFIX: &str = "2025-02-11T15-35-19-jqmb";
|
||||
|
||||
fn stage1_output_with_slug(thread_id: ThreadId, rollout_slug: Option<&str>) -> Stage1Output {
|
||||
@@ -59,3 +67,83 @@ fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_is_empty() {
|
||||
|
||||
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let root = dir.path().join("memory");
|
||||
ensure_layout(&root).await.expect("ensure layout");
|
||||
|
||||
let keep_id = ThreadId::default().to_string();
|
||||
let drop_id = ThreadId::default().to_string();
|
||||
let keep_path = rollout_summaries_dir(&root).join(format!("{keep_id}.md"));
|
||||
let drop_path = rollout_summaries_dir(&root).join(format!("{drop_id}.md"));
|
||||
tokio::fs::write(&keep_path, "keep")
|
||||
.await
|
||||
.expect("write keep");
|
||||
tokio::fs::write(&drop_path, "drop")
|
||||
.await
|
||||
.expect("write drop");
|
||||
|
||||
let memories = vec![Stage1Output {
|
||||
thread_id: ThreadId::try_from(keep_id.clone()).expect("thread id"),
|
||||
source_updated_at: Utc.timestamp_opt(100, 0).single().expect("timestamp"),
|
||||
raw_memory: "raw memory".to_string(),
|
||||
rollout_summary: "short summary".to_string(),
|
||||
rollout_slug: None,
|
||||
rollout_path: PathBuf::from("/tmp/rollout-100.jsonl"),
|
||||
cwd: PathBuf::from("/tmp/workspace"),
|
||||
git_branch: None,
|
||||
generated_at: Utc.timestamp_opt(101, 0).single().expect("timestamp"),
|
||||
}];
|
||||
|
||||
sync_rollout_summaries_from_memories(
|
||||
&root,
|
||||
&memories,
|
||||
DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION,
|
||||
)
|
||||
.await
|
||||
.expect("sync rollout summaries");
|
||||
rebuild_raw_memories_file_from_memories(
|
||||
&root,
|
||||
&memories,
|
||||
DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION,
|
||||
)
|
||||
.await
|
||||
.expect("rebuild raw memories");
|
||||
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&keep_path)
|
||||
.await
|
||||
.expect("check stale keep path"),
|
||||
"sync should prune stale filename that used thread id only"
|
||||
);
|
||||
assert!(
|
||||
!tokio::fs::try_exists(&drop_path)
|
||||
.await
|
||||
.expect("check stale drop path"),
|
||||
"sync should prune stale filename for dropped thread"
|
||||
);
|
||||
|
||||
let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root))
|
||||
.await
|
||||
.expect("open rollout summaries dir");
|
||||
let mut files = Vec::new();
|
||||
while let Some(entry) = dir.next_entry().await.expect("read dir entry") {
|
||||
files.push(entry.file_name().to_string_lossy().to_string());
|
||||
}
|
||||
files.sort_unstable();
|
||||
assert_eq!(files.len(), 1);
|
||||
let canonical_rollout_summary_file = &files[0];
|
||||
|
||||
let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root))
|
||||
.await
|
||||
.expect("read raw memories");
|
||||
assert!(raw_memories.contains("raw memory"));
|
||||
assert!(raw_memories.contains(&keep_id));
|
||||
assert!(raw_memories.contains("cwd: /tmp/workspace"));
|
||||
assert!(raw_memories.contains("rollout_path: /tmp/rollout-100.jsonl"));
|
||||
assert!(raw_memories.contains(&format!(
|
||||
"rollout_summary_file: {canonical_rollout_summary_file}"
|
||||
)));
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user