agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: streamline phase 2 (#11712)

Browse Source
This commit is contained in:
jif-oai
2026-02-13 13:21:11 +00:00
committed by GitHub
Unverified
parent feae389942
commit 36541876f4
21 changed files with 864 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/json/ServerNotification.json
+2 -1
View File
@@ -5997,7 +5997,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json
+4 -2
View File
@@ -9126,7 +9126,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
@@ -14480,7 +14481,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v1/GetConversationSummaryResponse.json
+2 -1
View File
@@ -113,7 +113,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v1/ListConversationsResponse.json
+2 -1
View File
@@ -113,7 +113,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json
+2 -1
View File
@@ -666,7 +666,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json
+2 -1
View File
@@ -472,7 +472,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json
+2 -1
View File
@@ -472,7 +472,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json
+2 -1
View File
@@ -666,7 +666,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json
+2 -1
View File
@@ -472,7 +472,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json
+2 -1
View File
@@ -666,7 +666,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json
+2 -1
View File
@@ -472,7 +472,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json
+2 -1
View File
@@ -472,7 +472,8 @@
{
"enum": [
"review",
"compact"
"compact",
"memory_consolidation"
],
"type": "string"
},
codex-rs/app-server-protocol/schema/typescript/SubAgentSource.ts
+1 -1
View File
@@ -3,4 +3,4 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
import type { ThreadId } from "./ThreadId";
export type SubAgentSource = "review" | "compact" | { "thread_spawn": { parent_thread_id: ThreadId, depth: number, } } | { "other": string };
export type SubAgentSource = "review" | "compact" | { "thread_spawn": { parent_thread_id: ThreadId, depth: number, } } | "memory_consolidation" | { "other": string };
codex-rs/codex-api/src/requests/headers.rs
+3
View File
@@ -17,6 +17,9 @@ pub(crate) fn subagent_header(source: &Option<SessionSource>) -> Option<String>
match sub {
codex_protocol::protocol::SubAgentSource::Review => Some("review".to_string()),
codex_protocol::protocol::SubAgentSource::Compact => Some("compact".to_string()),
codex_protocol::protocol::SubAgentSource::MemoryConsolidation => {
Some("memory_consolidation".to_string())
}
codex_protocol::protocol::SubAgentSource::ThreadSpawn { .. } => {
Some("collab_spawn".to_string())
}
codex-rs/core/src/client.rs
+3
View File
@@ -325,6 +325,9 @@ impl ModelClient {
let subagent = match sub {
crate::protocol::SubAgentSource::Review => "review".to_string(),
crate::protocol::SubAgentSource::Compact => "compact".to_string(),
crate::protocol::SubAgentSource::MemoryConsolidation => {
"memory_consolidation".to_string()
}
crate::protocol::SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(),
crate::protocol::SubAgentSource::Other(label) => label.clone(),
};
codex-rs/core/src/memories/dispatch.rs
-688
View File
@@ -1,688 +0,0 @@
use crate::codex::Session;
use crate::config::Config;
use crate::config::Constrained;
use crate::memories::memory_root;
use crate::memories::metrics;
use crate::memories::phase_two;
use crate::memories::phase2::spawn_phase2_completion_task;
use crate::memories::prompts::build_consolidation_prompt;
use crate::memories::storage::rebuild_raw_memories_file_from_memories;
use crate::memories::storage::sync_rollout_summaries_from_memories;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::user_input::UserInput;
use codex_utils_absolute_path::AbsolutePathBuf;
use std::sync::Arc;
use tracing::debug;
use tracing::info;
use tracing::warn;
//TODO(jif) clean.
fn completion_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)
}
pub(in crate::memories) async fn run_global_memory_consolidation(
session: &Arc<Session>,
config: Arc<Config>,
) -> bool {
let otel_manager = &session.services.otel_manager;
let Some(state_db) = session.services.state_db.as_deref() else {
warn!("state db unavailable; skipping global memory consolidation");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "skipped_state_db_unavailable")],
);
return false;
};
let claim = match state_db
.try_claim_global_phase2_job(session.conversation_id, phase_two::JOB_LEASE_SECONDS)
.await
{
Ok(claim) => claim,
Err(err) => {
warn!("state db try_claim_global_phase2_job failed during memories startup: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_claim")],
);
return false;
}
};
let (ownership_token, claimed_watermark) = match claim {
codex_state::Phase2JobClaimOutcome::Claimed {
ownership_token,
input_watermark,
} => {
otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "claimed")]);
(ownership_token, input_watermark)
}
codex_state::Phase2JobClaimOutcome::SkippedNotDirty => {
debug!("memory phase-2 global lock is up-to-date; skipping consolidation");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "skipped_not_dirty")],
);
return false;
}
codex_state::Phase2JobClaimOutcome::SkippedRunning => {
debug!("memory phase-2 global consolidation already running; skipping");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "skipped_running")],
);
return false;
}
};
let root = memory_root(&config.codex_home);
let consolidation_config = {
let mut consolidation_config = config.as_ref().clone();
consolidation_config.cwd = root.clone();
consolidation_config.permissions.approval_policy =
Constrained::allow_only(AskForApproval::Never);
let mut writable_roots = Vec::new();
match AbsolutePathBuf::from_absolute_path(consolidation_config.codex_home.clone()) {
Ok(codex_home) => writable_roots.push(codex_home),
Err(err) => warn!(
"memory phase-2 consolidation could not add codex_home writable root {}: {err}",
consolidation_config.codex_home.display()
),
}
let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite {
writable_roots,
read_only_access: Default::default(),
network_access: false,
exclude_tmpdir_env_var: false,
exclude_slash_tmp: false,
};
if let Err(err) = consolidation_config
.permissions
.sandbox_policy
.set(consolidation_sandbox_policy)
{
warn!("memory phase-2 consolidation sandbox policy was rejected by constraints: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_sandbox_policy")],
);
let _ = state_db
.mark_global_phase2_job_failed(
&ownership_token,
"consolidation sandbox policy was rejected by constraints",
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
return false;
}
consolidation_config
};
let latest_memories = match state_db
.list_stage1_outputs_for_global(phase_two::MAX_RAW_MEMORIES_FOR_GLOBAL)
.await
{
Ok(memories) => memories,
Err(err) => {
warn!("state db list_stage1_outputs_for_global failed during consolidation: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_load_stage1_outputs")],
);
let _ = state_db
.mark_global_phase2_job_failed(
&ownership_token,
"failed to read stage-1 outputs before global consolidation",
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
return false;
}
};
if !latest_memories.is_empty() {
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_INPUT,
latest_memories.len() as i64,
&[],
);
}
let completion_watermark = completion_watermark(claimed_watermark, &latest_memories);
if let Err(err) = sync_rollout_summaries_from_memories(&root, &latest_memories).await {
warn!("failed syncing local memory artifacts for global consolidation: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_sync_artifacts")],
);
let _ = state_db
.mark_global_phase2_job_failed(
&ownership_token,
"failed syncing local memory artifacts",
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
return false;
}
if let Err(err) = rebuild_raw_memories_file_from_memories(&root, &latest_memories).await {
warn!("failed rebuilding raw memories aggregate for global consolidation: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_rebuild_raw_memories")],
);
let _ = state_db
.mark_global_phase2_job_failed(
&ownership_token,
"failed rebuilding raw memories aggregate",
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
return false;
}
if latest_memories.is_empty() {
debug!("memory phase-2 has no stage-1 outputs; finalized local memory artifacts");
let _ = state_db
.mark_global_phase2_job_succeeded(&ownership_token, completion_watermark)
.await;
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "succeeded_no_input")],
);
return false;
}
let prompt = build_consolidation_prompt(&root);
let input = vec![UserInput::Text {
text: prompt,
text_elements: vec![],
}];
let source = SessionSource::SubAgent(SubAgentSource::Other(
phase_two::MEMORY_CONSOLIDATION_SUBAGENT_LABEL.to_string(),
));
match session
.services
.agent_control
.spawn_agent(consolidation_config, input, Some(source))
.await
{
Ok(consolidation_agent_id) => {
info!(
"memory phase-2 global consolidation agent started: agent_id={consolidation_agent_id}"
);
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "agent_spawned")],
);
spawn_phase2_completion_task(
session.as_ref(),
ownership_token,
completion_watermark,
consolidation_agent_id,
);
true
}
Err(err) => {
warn!("failed to spawn global memory consolidation agent: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_spawn_agent")],
);
let _ = state_db
.mark_global_phase2_job_failed(
&ownership_token,
"failed to spawn consolidation agent",
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
false
}
}
}
#[cfg(test)]
mod tests {
use super::completion_watermark;
use super::run_global_memory_consolidation;
use crate::CodexAuth;
use crate::ThreadManager;
use crate::agent::control::AgentControl;
use crate::codex::Session;
use crate::codex::make_session_and_context;
use crate::config::Config;
use crate::config::test_config;
use crate::memories::memory_root;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_state::Phase2JobClaimOutcome;
use codex_state::Stage1Output;
use codex_state::ThreadMetadataBuilder;
use pretty_assertions::assert_eq;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;
struct DispatchHarness {
_codex_home: TempDir,
config: Arc<Config>,
session: Arc<Session>,
manager: ThreadManager,
state_db: Arc<codex_state::StateRuntime>,
}
impl DispatchHarness {
async fn new() -> Self {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let mut config = test_config();
config.codex_home = codex_home.path().to_path_buf();
config.cwd = config.codex_home.clone();
let config = Arc::new(config);
let state_db = codex_state::StateRuntime::init(
config.codex_home.clone(),
config.model_provider_id.clone(),
None,
)
.await
.expect("initialize state db");
let manager = ThreadManager::with_models_provider_and_home_for_tests(
CodexAuth::from_api_key("dummy"),
config.model_provider.clone(),
config.codex_home.clone(),
);
let (mut session, _turn_context) = make_session_and_context().await;
session.services.state_db = Some(Arc::clone(&state_db));
session.services.agent_control = manager.agent_control();
Self {
_codex_home: codex_home,
config,
session: Arc::new(session),
manager,
state_db,
}
}
async fn seed_stage1_output(&self, source_updated_at: i64) {
let thread_id = ThreadId::new();
let mut metadata_builder = ThreadMetadataBuilder::new(
thread_id,
self.config
.codex_home
.join(format!("rollout-{thread_id}.jsonl")),
Utc::now(),
SessionSource::Cli,
);
metadata_builder.cwd = self.config.cwd.clone();
metadata_builder.model_provider = Some(self.config.model_provider_id.clone());
let metadata = metadata_builder.build(&self.config.model_provider_id);
self.state_db
.upsert_thread(&metadata)
.await
.expect("upsert thread metadata");
let claim = self
.state_db
.try_claim_stage1_job(
thread_id,
self.session.conversation_id,
source_updated_at,
3_600,
64,
)
.await
.expect("claim stage-1 job");
let ownership_token = match claim {
codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
other => panic!("unexpected stage-1 claim outcome: {other:?}"),
};
assert!(
self.state_db
.mark_stage1_job_succeeded(
thread_id,
&ownership_token,
source_updated_at,
"raw memory",
"rollout summary",
)
.await
.expect("mark stage-1 success"),
"stage-1 success should enqueue global consolidation"
);
}
async fn shutdown_threads(&self) {
self.manager
.remove_and_close_all_threads()
.await
.expect("shutdown spawned threads");
}
fn user_input_ops_count(&self) -> usize {
self.manager
.captured_ops()
.into_iter()
.filter(|(_, op)| matches!(op, Op::UserInput { .. }))
.count()
}
}
#[test]
fn completion_watermark_never_regresses_below_claimed_input_watermark() {
let stage1_output = Stage1Output {
thread_id: ThreadId::new(),
source_updated_at: chrono::DateTime::<Utc>::from_timestamp(123, 0)
.expect("valid source_updated_at timestamp"),
raw_memory: "raw memory".to_string(),
rollout_summary: "rollout summary".to_string(),
cwd: PathBuf::from("/tmp/workspace"),
generated_at: chrono::DateTime::<Utc>::from_timestamp(124, 0)
.expect("valid generated_at timestamp"),
};
let completion = completion_watermark(1_000, &[stage1_output]);
assert_eq!(completion, 1_000);
}
#[tokio::test]
async fn dispatch_reclaims_stale_global_lock_and_starts_consolidation() {
let harness = DispatchHarness::new().await;
harness.seed_stage1_output(100).await;
let stale_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 0)
.await
.expect("claim stale global lock");
assert!(
matches!(stale_claim, Phase2JobClaimOutcome::Claimed { .. }),
"stale lock precondition should be claimed"
);
let scheduled =
run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await;
assert!(
scheduled,
"dispatch should reclaim stale lock and spawn one agent"
);
let running_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim while running");
assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning);
let user_input_ops = harness.user_input_ops_count();
assert_eq!(user_input_ops, 1);
let thread_ids = harness.manager.list_thread_ids().await;
assert_eq!(thread_ids.len(), 1);
let subagent = harness
.manager
.get_thread(thread_ids[0])
.await
.expect("get consolidation thread");
let config_snapshot = subagent.config_snapshot().await;
assert_eq!(config_snapshot.approval_policy, AskForApproval::Never);
assert_eq!(config_snapshot.cwd, memory_root(&harness.config.codex_home));
match config_snapshot.sandbox_policy {
SandboxPolicy::WorkspaceWrite { writable_roots, .. } => {
assert!(
writable_roots
.iter()
.any(|root| root.as_path() == harness.config.codex_home.as_path()),
"consolidation subagent should have codex_home as writable root"
);
}
other => panic!("unexpected sandbox policy: {other:?}"),
}
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_schedules_only_one_agent_while_lock_is_running() {
let harness = DispatchHarness::new().await;
harness.seed_stage1_output(200).await;
let first_run =
run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await;
let second_run =
run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await;
assert!(first_run, "first dispatch should schedule consolidation");
assert!(
!second_run,
"second dispatch should skip while the global lock is running"
);
let user_input_ops = harness.user_input_ops_count();
assert_eq!(user_input_ops, 1);
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_with_dirty_job_and_no_stage1_outputs_skips_spawn_and_clears_dirty_flag() {
let harness = DispatchHarness::new().await;
harness
.state_db
.enqueue_global_consolidation(999)
.await
.expect("enqueue global consolidation");
let scheduled =
run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await;
assert!(
!scheduled,
"dispatch should not spawn when no stage-1 outputs are available"
);
assert_eq!(harness.user_input_ops_count(), 0);
let claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after empty dispatch");
assert_eq!(
claim,
Phase2JobClaimOutcome::SkippedNotDirty,
"empty dispatch should finalize global job as up-to-date"
);
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_with_empty_stage1_outputs_rebuilds_local_artifacts() {
let harness = DispatchHarness::new().await;
let root = memory_root(&harness.config.codex_home);
let summaries_dir = rollout_summaries_dir(&root);
tokio::fs::create_dir_all(&summaries_dir)
.await
.expect("create rollout summaries dir");
let stale_summary_path = summaries_dir.join(format!("{}.md", ThreadId::new()));
tokio::fs::write(&stale_summary_path, "stale summary\n")
.await
.expect("write stale rollout summary");
let raw_memories_path = raw_memories_file(&root);
tokio::fs::write(&raw_memories_path, "stale raw memories\n")
.await
.expect("write stale raw memories");
let memory_index_path = root.join("MEMORY.md");
tokio::fs::write(&memory_index_path, "stale memory index\n")
.await
.expect("write stale memory index");
let memory_summary_path = root.join("memory_summary.md");
tokio::fs::write(&memory_summary_path, "stale memory summary\n")
.await
.expect("write stale memory summary");
let stale_skill_file = root.join("skills/demo/SKILL.md");
tokio::fs::create_dir_all(
stale_skill_file
.parent()
.expect("skills subdirectory parent should exist"),
)
.await
.expect("create stale skills dir");
tokio::fs::write(&stale_skill_file, "stale skill\n")
.await
.expect("write stale skill");
harness
.state_db
.enqueue_global_consolidation(999)
.await
.expect("enqueue global consolidation");
let scheduled =
run_global_memory_consolidation(&harness.session, Arc::clone(&harness.config)).await;
assert!(
!scheduled,
"dispatch should skip subagent spawn when no stage-1 outputs are available"
);
assert!(
!tokio::fs::try_exists(&stale_summary_path)
.await
.expect("check stale summary existence"),
"empty consolidation should prune stale rollout summary files"
);
let raw_memories = tokio::fs::read_to_string(&raw_memories_path)
.await
.expect("read rebuilt raw memories");
assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n");
assert!(
!tokio::fs::try_exists(&memory_index_path)
.await
.expect("check memory index existence"),
"empty consolidation should remove stale MEMORY.md"
);
assert!(
!tokio::fs::try_exists(&memory_summary_path)
.await
.expect("check memory summary existence"),
"empty consolidation should remove stale memory_summary.md"
);
assert!(
!tokio::fs::try_exists(&stale_skill_file)
.await
.expect("check stale skill existence"),
"empty consolidation should remove stale skills artifacts"
);
assert!(
!tokio::fs::try_exists(root.join("skills"))
.await
.expect("check skills dir existence"),
"empty consolidation should remove stale skills directory"
);
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_marks_job_for_retry_when_spawn_agent_fails() {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let mut config = test_config();
config.codex_home = codex_home.path().to_path_buf();
config.cwd = config.codex_home.clone();
let config = Arc::new(config);
let state_db = codex_state::StateRuntime::init(
config.codex_home.clone(),
config.model_provider_id.clone(),
None,
)
.await
.expect("initialize state db");
let (mut session, _turn_context) = make_session_and_context().await;
session.services.state_db = Some(Arc::clone(&state_db));
session.services.agent_control = AgentControl::default();
let session = Arc::new(session);
let thread_id = ThreadId::new();
let mut metadata_builder = ThreadMetadataBuilder::new(
thread_id,
config.codex_home.join(format!("rollout-{thread_id}.jsonl")),
Utc::now(),
SessionSource::Cli,
);
metadata_builder.cwd = config.cwd.clone();
metadata_builder.model_provider = Some(config.model_provider_id.clone());
let metadata = metadata_builder.build(&config.model_provider_id);
state_db
.upsert_thread(&metadata)
.await
.expect("upsert thread metadata");
let claim = state_db
.try_claim_stage1_job(thread_id, session.conversation_id, 100, 3_600, 64)
.await
.expect("claim stage-1 job");
let ownership_token = match claim {
codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
other => panic!("unexpected stage-1 claim outcome: {other:?}"),
};
assert!(
state_db
.mark_stage1_job_succeeded(
thread_id,
&ownership_token,
100,
"raw memory",
"rollout summary",
)
.await
.expect("mark stage-1 success"),
"stage-1 success should enqueue global consolidation"
);
let scheduled = run_global_memory_consolidation(&session, Arc::clone(&config)).await;
assert!(
!scheduled,
"dispatch should return false when consolidation subagent cannot be spawned"
);
let retry_claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after spawn failure");
assert_eq!(
retry_claim,
Phase2JobClaimOutcome::SkippedNotDirty,
"spawn failures should leave the job in retry backoff instead of running"
);
}
}
codex-rs/core/src/memories/mod.rs
-3
View File
@@ -4,7 +4,6 @@
//! - Phase 1: select rollouts, extract stage-1 raw memories, persist stage-1 outputs, and enqueue consolidation.
//! - Phase 2: claim a global consolidation lock, materialize consolidation inputs, and dispatch one consolidation agent.
mod dispatch;
mod phase1;
mod phase2;
pub(crate) mod prompts;
@@ -58,8 +57,6 @@ mod phase_one {
/// Phase 2 (aka `Consolidation`).
mod phase_two {
/// Subagent source label used to identify consolidation tasks.
pub(super) const MEMORY_CONSOLIDATION_SUBAGENT_LABEL: &str = "memory_consolidation";
/// Maximum number of recent raw memories retained for global consolidation.
pub(super) const MAX_RAW_MEMORIES_FOR_GLOBAL: usize = 1_024;
/// Lease duration (seconds) for phase-2 consolidation job ownership.
codex-rs/core/src/memories/phase2.rs
+332 -368
View File
@@ -1,136 +1,351 @@
use crate::agent::AgentStatus;
use crate::agent::status::is_final as is_final_agent_status;
use crate::codex::Session;
use crate::config::Config;
use crate::memories::memory_root;
use crate::memories::metrics;
use crate::memories::phase_two;
use crate::memories::prompts::build_consolidation_prompt;
use crate::memories::storage::rebuild_raw_memories_file_from_memories;
use crate::memories::storage::sync_rollout_summaries_from_memories;
use codex_config::Constrained;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::user_input::UserInput;
use codex_state::StateRuntime;
use codex_utils_absolute_path::AbsolutePathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::watch;
use tracing::debug;
use tracing::info;
use tracing::warn;
pub(in crate::memories) fn spawn_phase2_completion_task(
session: &Session,
ownership_token: String,
completion_watermark: i64,
consolidation_agent_id: ThreadId,
) {
let state_db = session.services.state_db.clone();
let agent_control = session.services.agent_control.clone();
let otel_manager = session.services.otel_manager.clone();
tokio::spawn(async move {
let Some(state_db) = state_db else {
return;
};
let status_rx = match agent_control.subscribe_status(consolidation_agent_id).await {
Ok(status_rx) => status_rx,
Err(err) => {
warn!(
"failed to subscribe to global memory consolidation agent {consolidation_agent_id}: {err}"
);
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_subscribe_status")],
);
mark_phase2_failed_with_recovery(
state_db.as_ref(),
&ownership_token,
"failed to subscribe to consolidation agent status",
)
.await;
return;
}
};
let final_status = run_phase2_completion_task(
Arc::clone(&state_db),
ownership_token,
completion_watermark,
consolidation_agent_id,
status_rx,
)
.await;
if matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) {
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "failed_agent_unavailable")],
);
return;
}
if is_phase2_success(&final_status) {
otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "succeeded")],
);
} else {
otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "failed")]);
}
tokio::spawn(async move {
if let Err(err) = agent_control.shutdown_agent(consolidation_agent_id).await {
warn!(
"failed to auto-close global memory consolidation agent {consolidation_agent_id}: {err}"
);
}
});
});
#[derive(Debug, Clone, Default)]
struct Claim {
token: String,
watermark: i64,
}
async fn run_phase2_completion_task(
state_db: Arc<codex_state::StateRuntime>,
ownership_token: String,
completion_watermark: i64,
consolidation_agent_id: ThreadId,
mut status_rx: watch::Receiver<AgentStatus>,
) -> AgentStatus {
let final_status = {
#[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(super) async fn run(session: &Arc<Session>, config: Arc<Config>) {
let Some(db) = session.services.state_db.as_deref() else {
// This should not happen.
return;
};
let root = memory_root(&config.codex_home);
// 1. Claim the job.
let claim = match job::claim(session, db).await {
Ok(claim) => claim,
Err(e) => {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", e)],
);
return;
}
};
// 2. Get the config for the agent
let Some(agent_config) = agent::get_config(config.clone()) else {
// If we can't get the config, we can't consolidate.
tracing::error!("failed to get agent config");
job::failed(session, db, &claim, "failed_sandbox_policy").await;
return;
};
// 3. Query the memories
let raw_memories = match db
.list_stage1_outputs_for_global(phase_two::MAX_RAW_MEMORIES_FOR_GLOBAL)
.await
{
Ok(memories) => memories,
Err(err) => {
tracing::error!("failed to list stage1 outputs from global: {}", err);
job::failed(session, db, &claim, "failed_load_stage1_outputs").await;
return;
}
};
let new_watermark = get_watermark(claim.watermark, &raw_memories);
// 4. Update the file system by syncing the raw memories with the one extracted from DB at
// step 3
// [`rollout_summaries/`]
if let Err(err) = sync_rollout_summaries_from_memories(&root, &raw_memories).await {
tracing::error!("failed syncing local memory artifacts for global consolidation: {err}");
job::failed(session, db, &claim, "failed_sync_artifacts").await;
return;
}
// [`raw_memories.md`]
if let Err(err) = rebuild_raw_memories_file_from_memories(&root, &raw_memories).await {
tracing::error!("failed syncing local memory artifacts for global consolidation: {err}");
job::failed(session, db, &claim, "failed_rebuild_raw_memories").await;
return;
}
if raw_memories.is_empty() {
// We check only after sync of the file system.
job::succeed(session, db, &claim, new_watermark, "succeeded_no_input").await;
return;
}
// 5. Spawn the agent
let prompt = agent::get_prompt(config);
let source = SessionSource::SubAgent(SubAgentSource::MemoryConsolidation);
let thread_id = match session
.services
.agent_control
.spawn_agent(agent_config, prompt, Some(source))
.await
{
Ok(thread_id) => thread_id,
Err(err) => {
tracing::error!("failed to spawn global memory consolidation agent: {err}");
job::failed(session, db, &claim, "failed_spawn_agent").await;
return;
}
};
// 6. Spawn the agent handler.
agent::handle(session, claim, new_watermark, thread_id);
// 7. Metrics and logs.
let counters = Counters {
input: raw_memories.len() as i64,
};
emit_metrics(session, counters);
}
mod job {
use super::*;
pub(super) async fn claim(
session: &Arc<Session>,
db: &StateRuntime,
) -> Result<Claim, &'static str> {
let otel_manager = &session.services.otel_manager;
let claim = db
.try_claim_global_phase2_job(session.conversation_id, phase_two::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,
} => {
otel_manager.counter(metrics::MEMORY_PHASE_TWO_JOBS, 1, &[("status", "claimed")]);
(ownership_token, input_watermark)
}
codex_state::Phase2JobClaimOutcome::SkippedNotDirty => return Err("skipped_not_dirty"),
codex_state::Phase2JobClaimOutcome::SkippedRunning => return Err("skipped_running"),
};
Ok(Claim { token, watermark })
}
pub(super) async fn failed(
session: &Arc<Session>,
db: &StateRuntime,
claim: &Claim,
reason: &'static str,
) {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", reason)],
);
if matches!(
db.mark_global_phase2_job_failed(
&claim.token,
reason,
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await,
Ok(false)
) {
let _ = db
.mark_global_phase2_job_failed_if_unowned(
&claim.token,
reason,
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await;
}
}
pub(super) async fn succeed(
session: &Arc<Session>,
db: &StateRuntime,
claim: &Claim,
completion_watermark: i64,
reason: &'static str,
) {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", reason)],
);
let _ = db
.mark_global_phase2_job_succeeded(&claim.token, completion_watermark)
.await;
}
}
mod agent {
use super::*;
pub(super) fn get_config(config: Arc<Config>) -> Option<Config> {
let root = memory_root(&config.codex_home);
let mut consolidation_config = config.as_ref().clone();
consolidation_config.cwd = root;
// Approval policy
consolidation_config.permissions.approval_policy =
Constrained::allow_only(AskForApproval::Never);
// Sandbox policy
let mut writable_roots = Vec::new();
match AbsolutePathBuf::from_absolute_path(consolidation_config.codex_home.clone()) {
Ok(codex_home) => writable_roots.push(codex_home),
Err(err) => warn!(
"memory phase-2 consolidation could not add codex_home writable root {}: {err}",
consolidation_config.codex_home.display()
),
}
// The consolidation agent only needs local codex_home write access and no network.
let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite {
writable_roots,
read_only_access: Default::default(),
network_access: false,
exclude_tmpdir_env_var: false,
exclude_slash_tmp: false,
};
consolidation_config
.permissions
.sandbox_policy
.set(consolidation_sandbox_policy)
.ok()?;
Some(consolidation_config)
}
pub(super) fn get_prompt(config: Arc<Config>) -> Vec<UserInput> {
let root = memory_root(&config.codex_home);
let prompt = build_consolidation_prompt(&root);
vec![UserInput::Text {
text: prompt,
text_elements: vec![],
}]
}
/// Handle the agent while it is running.
pub(super) fn handle(
session: &Arc<Session>,
claim: Claim,
new_watermark: i64,
thread_id: ThreadId,
) {
let Some(db) = session.services.state_db.clone() else {
return;
};
let session = session.clone();
tokio::spawn(async move {
let agent_control = session.services.agent_control.clone();
// TODO(jif) we might have a very small race here.
let rx = match agent_control.subscribe_status(thread_id).await {
Ok(rx) => rx,
Err(err) => {
tracing::error!("agent_control.subscribe_status failed: {err:?}");
job::failed(&session, &db, &claim, "failed_subscribe_status").await;
return;
}
};
// Loop the agent until we have the final status.
let final_status = loop_agent(
db.clone(),
claim.token.clone(),
new_watermark,
thread_id,
rx,
)
.await;
if matches!(final_status, AgentStatus::Completed(_)) {
job::succeed(&session, &db, &claim, new_watermark, "succeeded").await;
} else {
job::failed(&session, &db, &claim, "failed_agent").await;
}
// Fire and forget close of the agent.
if !matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) {
tokio::spawn(async move {
if let Err(err) = agent_control.shutdown_agent(thread_id).await {
warn!(
"failed to auto-close global memory consolidation agent {thread_id}: {err}"
);
}
});
} else {
tracing::warn!("The agent was already gone");
}
});
}
async fn loop_agent(
db: Arc<StateRuntime>,
token: String,
_new_watermark: i64,
thread_id: ThreadId,
mut rx: watch::Receiver<AgentStatus>,
) -> AgentStatus {
let mut heartbeat_interval =
tokio::time::interval(Duration::from_secs(phase_two::JOB_HEARTBEAT_SECONDS));
heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
let status = status_rx.borrow().clone();
let status = rx.borrow().clone();
if is_final_agent_status(&status) {
break status;
}
tokio::select! {
changed = status_rx.changed() => {
if changed.is_err() {
warn!(
"lost status updates for global memory consolidation agent {consolidation_agent_id}"
update = rx.changed() => {
if update.is_err() {
tracing::warn!(
"lost status updates for global memory consolidation agent {thread_id}"
);
break status;
}
}
_ = heartbeat_interval.tick() => {
match state_db
match db
.heartbeat_global_phase2_job(
&ownership_token,
&token,
phase_two::JOB_LEASE_SECONDS,
)
.await
{
Ok(true) => {}
Ok(false) => {
warn!(
"memory phase-2 heartbeat lost global ownership; finalizing as failure"
);
break AgentStatus::Errored(
"lost global phase-2 ownership during heartbeat".to_string(),
);
}
Err(err) => {
warn!(
"state db heartbeat_global_phase2_job failed during memories startup: {err}"
);
break AgentStatus::Errored(format!(
"phase-2 heartbeat update failed: {err}"
));
@@ -139,281 +354,30 @@ async fn run_phase2_completion_task(
}
}
}
};
}
}
let phase2_success = is_phase2_success(&final_status);
info!(
"memory phase-2 global consolidation complete: agent_id={consolidation_agent_id} success={phase2_success} final_status={final_status:?}"
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) // todo double check the claimed here.
}
fn emit_metrics(session: &Arc<Session>, counters: Counters) {
let otel = session.services.otel_manager.clone();
if counters.input > 0 {
otel.counter(metrics::MEMORY_PHASE_TWO_INPUT, counters.input, &[]);
}
otel.counter(
metrics::MEMORY_PHASE_TWO_JOBS,
1,
&[("status", "agent_spawned")],
);
if phase2_success {
match state_db
.mark_global_phase2_job_succeeded(&ownership_token, completion_watermark)
.await
{
Ok(true) => {}
Ok(false) => {
debug!(
"memory phase-2 success finalization skipped after global ownership changed"
);
}
Err(err) => {
warn!(
"state db mark_global_phase2_job_succeeded failed during memories startup: {err}"
);
}
}
return final_status;
}
let failure_reason = phase2_failure_reason(&final_status);
mark_phase2_failed_with_recovery(state_db.as_ref(), &ownership_token, &failure_reason).await;
warn!(
"memory phase-2 global consolidation agent finished with non-success status: agent_id={consolidation_agent_id} final_status={final_status:?}"
);
final_status
}
async fn mark_phase2_failed_with_recovery(
state_db: &codex_state::StateRuntime,
ownership_token: &str,
failure_reason: &str,
) {
match state_db
.mark_global_phase2_job_failed(
ownership_token,
failure_reason,
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await
{
Ok(true) => {}
Ok(false) => match state_db
.mark_global_phase2_job_failed_if_unowned(
ownership_token,
failure_reason,
phase_two::JOB_RETRY_DELAY_SECONDS,
)
.await
{
Ok(true) => {
debug!(
"memory phase-2 failure finalization applied fallback update for unowned running job"
);
}
Ok(false) => {
debug!(
"memory phase-2 failure finalization skipped after global ownership changed"
);
}
Err(err) => {
warn!(
"state db mark_global_phase2_job_failed_if_unowned failed during memories startup: {err}"
);
}
},
Err(err) => {
warn!("state db mark_global_phase2_job_failed failed during memories startup: {err}");
}
}
}
fn is_phase2_success(final_status: &AgentStatus) -> bool {
matches!(final_status, AgentStatus::Completed(_))
}
fn phase2_failure_reason(final_status: &AgentStatus) -> String {
format!("consolidation agent finished with status {final_status:?}")
}
#[cfg(test)]
mod tests {
use super::is_phase2_success;
use super::phase2_failure_reason;
use super::run_phase2_completion_task;
use crate::agent::AgentStatus;
use codex_protocol::ThreadId;
use codex_state::Phase2JobClaimOutcome;
use pretty_assertions::assert_eq;
use std::sync::Arc;
#[test]
fn phase2_success_only_for_completed_status() {
assert!(is_phase2_success(&AgentStatus::Completed(None)));
assert!(!is_phase2_success(&AgentStatus::Running));
assert!(!is_phase2_success(&AgentStatus::Errored(
"oops".to_string()
)));
}
#[test]
fn phase2_failure_reason_includes_status() {
let status = AgentStatus::Errored("boom".to_string());
let reason = phase2_failure_reason(&status);
assert!(reason.contains("consolidation agent finished with status"));
assert!(reason.contains("boom"));
}
#[tokio::test]
async fn phase2_completion_marks_succeeded_for_completed_status() {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let state_db = Arc::new(
codex_state::StateRuntime::init(
codex_home.path().to_path_buf(),
"test-provider".to_string(),
None,
)
.await
.expect("initialize state runtime"),
);
let owner = ThreadId::new();
state_db
.enqueue_global_consolidation(123)
.await
.expect("enqueue global consolidation");
let claim = state_db
.try_claim_global_phase2_job(owner, 3_600)
.await
.expect("claim global phase-2 job");
let ownership_token = match claim {
Phase2JobClaimOutcome::Claimed {
ownership_token, ..
} => ownership_token,
other => panic!("unexpected phase-2 claim outcome: {other:?}"),
};
let (_status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Completed(None));
run_phase2_completion_task(
Arc::clone(&state_db),
ownership_token.clone(),
123,
ThreadId::new(),
status_rx,
)
.await;
let up_to_date_claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim up-to-date global job");
assert_eq!(up_to_date_claim, Phase2JobClaimOutcome::SkippedNotDirty);
state_db
.enqueue_global_consolidation(124)
.await
.expect("enqueue advanced consolidation watermark");
let rerun_claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim rerun global job");
assert!(
matches!(rerun_claim, Phase2JobClaimOutcome::Claimed { .. }),
"advanced watermark should be claimable after success finalization"
);
}
#[tokio::test]
async fn phase2_completion_marks_failed_when_status_updates_are_lost() {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let state_db = Arc::new(
codex_state::StateRuntime::init(
codex_home.path().to_path_buf(),
"test-provider".to_string(),
None,
)
.await
.expect("initialize state runtime"),
);
state_db
.enqueue_global_consolidation(456)
.await
.expect("enqueue global consolidation");
let claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global phase-2 job");
let ownership_token = match claim {
Phase2JobClaimOutcome::Claimed {
ownership_token, ..
} => ownership_token,
other => panic!("unexpected phase-2 claim outcome: {other:?}"),
};
let (status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Running);
drop(status_tx);
run_phase2_completion_task(
Arc::clone(&state_db),
ownership_token,
456,
ThreadId::new(),
status_rx,
)
.await;
let claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim after failure finalization");
assert_eq!(
claim,
Phase2JobClaimOutcome::SkippedNotDirty,
"failure finalization should leave global job in retry-backoff, not running ownership"
);
}
#[tokio::test]
async fn phase2_completion_heartbeat_loss_does_not_steal_active_other_owner() {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let state_db = Arc::new(
codex_state::StateRuntime::init(
codex_home.path().to_path_buf(),
"test-provider".to_string(),
None,
)
.await
.expect("initialize state runtime"),
);
state_db
.enqueue_global_consolidation(789)
.await
.expect("enqueue global consolidation");
let claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global phase-2 job");
let claimed_token = match claim {
Phase2JobClaimOutcome::Claimed {
ownership_token, ..
} => ownership_token,
other => panic!("unexpected phase-2 claim outcome: {other:?}"),
};
let (_status_tx, status_rx) = tokio::sync::watch::channel(AgentStatus::Running);
run_phase2_completion_task(
Arc::clone(&state_db),
"non-owner-token".to_string(),
789,
ThreadId::new(),
status_rx,
)
.await;
let claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim after heartbeat ownership loss");
assert_eq!(
claim,
Phase2JobClaimOutcome::SkippedRunning,
"heartbeat ownership-loss handling should not steal a live owner lease"
);
assert_eq!(
state_db
.mark_global_phase2_job_succeeded(claimed_token.as_str(), 789)
.await
.expect("mark original owner success"),
true,
"the original owner should still be able to finalize"
);
}
}
codex-rs/core/src/memories/start.rs
+2 -1
View File
@@ -2,6 +2,7 @@ use crate::codex::Session;
use crate::config::Config;
use crate::features::Feature;
use crate::memories::phase1;
use crate::memories::phase2;
use codex_protocol::protocol::SessionSource;
use std::sync::Arc;
use tracing::warn;
@@ -36,6 +37,6 @@ pub(crate) fn start_memories_startup_task(
// Run phase 1.
phase1::run(&session).await;
// Run phase 2.
crate::memories::dispatch::run_global_memory_consolidation(&session, config).await;
phase2::run(&session, config).await;
});
}
codex-rs/core/src/memories/tests.rs
+495
View File
@@ -87,3 +87,498 @@ async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only
assert!(raw_memories.contains(&keep_id));
assert!(raw_memories.contains("cwd: /tmp/workspace"));
}
mod phase2 {
use crate::CodexAuth;
use crate::ThreadManager;
use crate::agent::AgentControl;
use crate::codex::Session;
use crate::codex::make_session_and_context;
use crate::config::Config;
use crate::config::test_config;
use crate::memories::memory_root;
use crate::memories::phase2;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
use chrono::Utc;
use codex_config::Constrained;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_state::Phase2JobClaimOutcome;
use codex_state::Stage1Output;
use codex_state::ThreadMetadataBuilder;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;
fn stage1_output_with_source_updated_at(source_updated_at: i64) -> Stage1Output {
Stage1Output {
thread_id: ThreadId::new(),
source_updated_at: chrono::DateTime::<Utc>::from_timestamp(source_updated_at, 0)
.expect("valid source_updated_at timestamp"),
raw_memory: "raw memory".to_string(),
rollout_summary: "rollout summary".to_string(),
cwd: PathBuf::from("/tmp/workspace"),
generated_at: chrono::DateTime::<Utc>::from_timestamp(source_updated_at + 1, 0)
.expect("valid generated_at timestamp"),
}
}
struct DispatchHarness {
_codex_home: TempDir,
config: Arc<Config>,
session: Arc<Session>,
manager: ThreadManager,
state_db: Arc<codex_state::StateRuntime>,
}
impl DispatchHarness {
async fn new() -> Self {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let mut config = test_config();
config.codex_home = codex_home.path().to_path_buf();
config.cwd = config.codex_home.clone();
let config = Arc::new(config);
let state_db = codex_state::StateRuntime::init(
config.codex_home.clone(),
config.model_provider_id.clone(),
None,
)
.await
.expect("initialize state db");
let manager = ThreadManager::with_models_provider_and_home_for_tests(
CodexAuth::from_api_key("dummy"),
config.model_provider.clone(),
config.codex_home.clone(),
);
let (mut session, _turn_context) = make_session_and_context().await;
session.services.state_db = Some(Arc::clone(&state_db));
session.services.agent_control = manager.agent_control();
Self {
_codex_home: codex_home,
config,
session: Arc::new(session),
manager,
state_db,
}
}
async fn seed_stage1_output(&self, source_updated_at: i64) {
let thread_id = ThreadId::new();
let mut metadata_builder = ThreadMetadataBuilder::new(
thread_id,
self.config
.codex_home
.join(format!("rollout-{thread_id}.jsonl")),
Utc::now(),
SessionSource::Cli,
);
metadata_builder.cwd = self.config.cwd.clone();
metadata_builder.model_provider = Some(self.config.model_provider_id.clone());
let metadata = metadata_builder.build(&self.config.model_provider_id);
self.state_db
.upsert_thread(&metadata)
.await
.expect("upsert thread metadata");
let claim = self
.state_db
.try_claim_stage1_job(
thread_id,
self.session.conversation_id,
source_updated_at,
3_600,
64,
)
.await
.expect("claim stage-1 job");
let ownership_token = match claim {
codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
other => panic!("unexpected stage-1 claim outcome: {other:?}"),
};
assert!(
self.state_db
.mark_stage1_job_succeeded(
thread_id,
&ownership_token,
source_updated_at,
"raw memory",
"rollout summary",
)
.await
.expect("mark stage-1 success"),
"stage-1 success should enqueue global consolidation"
);
}
async fn shutdown_threads(&self) {
self.manager
.remove_and_close_all_threads()
.await
.expect("shutdown spawned threads");
}
fn user_input_ops_count(&self) -> usize {
self.manager
.captured_ops()
.into_iter()
.filter(|(_, op)| matches!(op, Op::UserInput { .. }))
.count()
}
}
#[test]
fn completion_watermark_never_regresses_below_claimed_input_watermark() {
let stage1_output = stage1_output_with_source_updated_at(123);
let completion = phase2::get_watermark(1_000, &[stage1_output]);
pretty_assertions::assert_eq!(completion, 1_000);
}
#[test]
fn completion_watermark_uses_claimed_watermark_when_there_are_no_memories() {
let completion = phase2::get_watermark(777, &[]);
pretty_assertions::assert_eq!(completion, 777);
}
#[test]
fn completion_watermark_uses_latest_memory_timestamp_when_it_is_newer() {
let older = stage1_output_with_source_updated_at(123);
let newer = stage1_output_with_source_updated_at(456);
let completion = phase2::get_watermark(200, &[older, newer]);
pretty_assertions::assert_eq!(completion, 456);
}
#[tokio::test]
async fn dispatch_skips_when_global_job_is_not_dirty() {
let harness = DispatchHarness::new().await;
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
}
#[tokio::test]
async fn dispatch_skips_when_global_job_is_already_running() {
let harness = DispatchHarness::new().await;
harness
.state_db
.enqueue_global_consolidation(123)
.await
.expect("enqueue global consolidation");
let claimed = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim running global lock");
assert!(
matches!(claimed, Phase2JobClaimOutcome::Claimed { .. }),
"precondition should claim the running lock"
);
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
let running_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim while lock is still running");
pretty_assertions::assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning);
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
}
#[tokio::test]
async fn dispatch_reclaims_stale_global_lock_and_starts_consolidation() {
let harness = DispatchHarness::new().await;
harness.seed_stage1_output(100).await;
let stale_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 0)
.await
.expect("claim stale global lock");
assert!(
matches!(stale_claim, Phase2JobClaimOutcome::Claimed { .. }),
"stale lock precondition should be claimed"
);
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
let running_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim while running");
pretty_assertions::assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning);
let user_input_ops = harness.user_input_ops_count();
pretty_assertions::assert_eq!(user_input_ops, 1);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 1);
let subagent = harness
.manager
.get_thread(thread_ids[0])
.await
.expect("get consolidation thread");
let config_snapshot = subagent.config_snapshot().await;
pretty_assertions::assert_eq!(config_snapshot.approval_policy, AskForApproval::Never);
pretty_assertions::assert_eq!(config_snapshot.cwd, memory_root(&harness.config.codex_home));
match config_snapshot.sandbox_policy {
SandboxPolicy::WorkspaceWrite { writable_roots, .. } => {
assert!(
writable_roots
.iter()
.any(|root| root.as_path() == harness.config.codex_home.as_path()),
"consolidation subagent should have codex_home as writable root"
);
}
other => panic!("unexpected sandbox policy: {other:?}"),
}
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_with_empty_stage1_outputs_rebuilds_local_artifacts() {
let harness = DispatchHarness::new().await;
let root = memory_root(&harness.config.codex_home);
let summaries_dir = rollout_summaries_dir(&root);
tokio::fs::create_dir_all(&summaries_dir)
.await
.expect("create rollout summaries dir");
let stale_summary_path = summaries_dir.join(format!("{}.md", ThreadId::new()));
tokio::fs::write(&stale_summary_path, "stale summary\n")
.await
.expect("write stale rollout summary");
let raw_memories_path = raw_memories_file(&root);
tokio::fs::write(&raw_memories_path, "stale raw memories\n")
.await
.expect("write stale raw memories");
let memory_index_path = root.join("MEMORY.md");
tokio::fs::write(&memory_index_path, "stale memory index\n")
.await
.expect("write stale memory index");
let memory_summary_path = root.join("memory_summary.md");
tokio::fs::write(&memory_summary_path, "stale memory summary\n")
.await
.expect("write stale memory summary");
let stale_skill_file = root.join("skills/demo/SKILL.md");
tokio::fs::create_dir_all(
stale_skill_file
.parent()
.expect("skills subdirectory parent should exist"),
)
.await
.expect("create stale skills dir");
tokio::fs::write(&stale_skill_file, "stale skill\n")
.await
.expect("write stale skill");
harness
.state_db
.enqueue_global_consolidation(999)
.await
.expect("enqueue global consolidation");
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
assert!(
!tokio::fs::try_exists(&stale_summary_path)
.await
.expect("check stale summary existence"),
"empty consolidation should prune stale rollout summary files"
);
let raw_memories = tokio::fs::read_to_string(&raw_memories_path)
.await
.expect("read rebuilt raw memories");
pretty_assertions::assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n");
assert!(
!tokio::fs::try_exists(&memory_index_path)
.await
.expect("check memory index existence"),
"empty consolidation should remove stale MEMORY.md"
);
assert!(
!tokio::fs::try_exists(&memory_summary_path)
.await
.expect("check memory summary existence"),
"empty consolidation should remove stale memory_summary.md"
);
assert!(
!tokio::fs::try_exists(&stale_skill_file)
.await
.expect("check stale skill existence"),
"empty consolidation should remove stale skills artifacts"
);
assert!(
!tokio::fs::try_exists(root.join("skills"))
.await
.expect("check skills dir existence"),
"empty consolidation should remove stale skills directory"
);
let next_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after empty consolidation success");
pretty_assertions::assert_eq!(next_claim, Phase2JobClaimOutcome::SkippedNotDirty);
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
harness.shutdown_threads().await;
}
#[tokio::test]
async fn dispatch_marks_job_for_retry_when_sandbox_policy_cannot_be_overridden() {
let harness = DispatchHarness::new().await;
harness
.state_db
.enqueue_global_consolidation(99)
.await
.expect("enqueue global consolidation");
let mut constrained_config = harness.config.as_ref().clone();
constrained_config.permissions.sandbox_policy =
Constrained::allow_only(SandboxPolicy::DangerFullAccess);
phase2::run(&harness.session, Arc::new(constrained_config)).await;
let retry_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after sandbox policy failure");
pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty);
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
}
#[tokio::test]
async fn dispatch_marks_job_for_retry_when_syncing_artifacts_fails() {
let harness = DispatchHarness::new().await;
harness.seed_stage1_output(100).await;
let root = memory_root(&harness.config.codex_home);
tokio::fs::write(&root, "not a directory")
.await
.expect("create file at memory root");
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
let retry_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after sync failure");
pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty);
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
}
#[tokio::test]
async fn dispatch_marks_job_for_retry_when_rebuilding_raw_memories_fails() {
let harness = DispatchHarness::new().await;
harness.seed_stage1_output(100).await;
let root = memory_root(&harness.config.codex_home);
tokio::fs::create_dir_all(raw_memories_file(&root))
.await
.expect("create raw_memories.md as a directory");
phase2::run(&harness.session, Arc::clone(&harness.config)).await;
let retry_claim = harness
.state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after rebuild failure");
pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedNotDirty);
pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0);
let thread_ids = harness.manager.list_thread_ids().await;
pretty_assertions::assert_eq!(thread_ids.len(), 0);
}
#[tokio::test]
async fn dispatch_marks_job_for_retry_when_spawn_agent_fails() {
let codex_home = tempfile::tempdir().expect("create temp codex home");
let mut config = test_config();
config.codex_home = codex_home.path().to_path_buf();
config.cwd = config.codex_home.clone();
let config = Arc::new(config);
let state_db = codex_state::StateRuntime::init(
config.codex_home.clone(),
config.model_provider_id.clone(),
None,
)
.await
.expect("initialize state db");
let (mut session, _turn_context) = make_session_and_context().await;
session.services.state_db = Some(Arc::clone(&state_db));
session.services.agent_control = AgentControl::default();
let session = Arc::new(session);
let thread_id = ThreadId::new();
let mut metadata_builder = ThreadMetadataBuilder::new(
thread_id,
config.codex_home.join(format!("rollout-{thread_id}.jsonl")),
Utc::now(),
SessionSource::Cli,
);
metadata_builder.cwd = config.cwd.clone();
metadata_builder.model_provider = Some(config.model_provider_id.clone());
let metadata = metadata_builder.build(&config.model_provider_id);
state_db
.upsert_thread(&metadata)
.await
.expect("upsert thread metadata");
let claim = state_db
.try_claim_stage1_job(thread_id, session.conversation_id, 100, 3_600, 64)
.await
.expect("claim stage-1 job");
let ownership_token = match claim {
codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token,
other => panic!("unexpected stage-1 claim outcome: {other:?}"),
};
assert!(
state_db
.mark_stage1_job_succeeded(
thread_id,
&ownership_token,
100,
"raw memory",
"rollout summary",
)
.await
.expect("mark stage-1 success"),
"stage-1 success should enqueue global consolidation"
);
phase2::run(&session, Arc::clone(&config)).await;
let retry_claim = state_db
.try_claim_global_phase2_job(ThreadId::new(), 3_600)
.await
.expect("claim global job after spawn failure");
pretty_assertions::assert_eq!(
retry_claim,
Phase2JobClaimOutcome::SkippedNotDirty,
"spawn failures should leave the job in retry backoff instead of running"
);
}
}
codex-rs/protocol/src/protocol.rs
+2
View File
@@ -1825,6 +1825,7 @@ pub enum SubAgentSource {
parent_thread_id: ThreadId,
depth: i32,
},
MemoryConsolidation,
Other(String),
}
@@ -1846,6 +1847,7 @@ impl fmt::Display for SubAgentSource {
match self {
SubAgentSource::Review => f.write_str("review"),
SubAgentSource::Compact => f.write_str("compact"),
SubAgentSource::MemoryConsolidation => f.write_str("memory_consolidation"),
SubAgentSource::ThreadSpawn {
parent_thread_id,
depth,