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: drop and clean from phase 1 (#11605)

Browse Source 
This PR is mostly cleaning and simplifying phase 1 of memories
This commit is contained in:
jif-oai
2026-02-12 17:23:00 +00:00
committed by GitHub
Unverified
parent ba6f7a9e15
commit f741fad5c0
15 changed files with 544 additions and 767 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1 -1
View File
@@ -1692,6 +1692,7 @@ dependencies = [
"codex-utils-home-dir",
"codex-utils-pty",
"codex-utils-readiness",
"codex-utils-sanitizer",
"codex-utils-string",
"codex-windows-sandbox",
"core-foundation 0.9.4",
@@ -1720,7 +1721,6 @@ dependencies = [
"predicates",
"pretty_assertions",
"rand 0.9.2",
"regex",
"regex-lite",
"reqwest",
"rmcp",
codex-rs/Cargo.toml
-1
View File
@@ -340,7 +340,6 @@ ignored = [
"icu_provider",
"openssl-sys",
"codex-utils-readiness",
"codex-utils-sanitizer",
"codex-secrets",
]
codex-rs/core/Cargo.toml
+1 -1
View File
@@ -49,6 +49,7 @@ codex-utils-absolute-path = { workspace = true }
codex-utils-home-dir = { workspace = true }
codex-utils-pty = { workspace = true }
codex-utils-readiness = { workspace = true }
codex-utils-sanitizer = { workspace = true }
codex-utils-string = { workspace = true }
codex-windows-sandbox = { package = "codex-windows-sandbox", path = "../windows-sandbox-rs" }
dirs = { workspace = true }
@@ -67,7 +68,6 @@ notify = { workspace = true }
once_cell = { workspace = true }
os_info = { workspace = true }
rand = { workspace = true }
regex = { workspace = true }
regex-lite = { workspace = true }
reqwest = { workspace = true, features = ["json", "stream"] }
rmcp = { workspace = true, default-features = false, features = [
codex-rs/core/src/memories/startup/dispatch.rs → codex-rs/core/src/memories/dispatch.rs
+5 -3
View File
@@ -4,8 +4,8 @@ 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::startup::phase2::spawn_phase2_completion_task;
use crate::memories::storage::rebuild_raw_memories_file_from_memories;
use crate::memories::storage::sync_rollout_summaries_from_memories;
use codex_protocol::protocol::AskForApproval;
@@ -19,6 +19,8 @@ use tracing::debug;
use tracing::info;
use tracing::warn;
//TODO(jif) clean.
fn completion_watermark(
claimed_watermark: i64,
latest_memories: &[codex_state::Stage1Output],
@@ -31,7 +33,7 @@ fn completion_watermark(
.max(claimed_watermark)
}
pub(super) async fn run_global_memory_consolidation(
pub(in crate::memories) async fn run_global_memory_consolidation(
session: &Arc<Session>,
config: Arc<Config>,
) -> bool {
@@ -261,7 +263,6 @@ pub(super) async fn run_global_memory_consolidation(
#[cfg(test)]
mod tests {
use super::completion_watermark;
use super::memory_root;
use super::run_global_memory_consolidation;
use crate::CodexAuth;
use crate::ThreadManager;
@@ -270,6 +271,7 @@ mod tests {
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;
codex-rs/core/src/memories/mod.rs
+17 -11
View File
@@ -4,16 +4,20 @@
//! - 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;
mod stage_one;
mod startup;
mod start;
mod storage;
#[cfg(test)]
mod tests;
use std::path::Path;
use std::path::PathBuf;
/// Starts the memory startup pipeline for eligible root sessions.
/// This is the single entrypoint that `codex` uses to trigger memory startup.
///
/// This is the entry point to read and understand this module.
pub(crate) use start::start_memories_startup_task;
mod artifacts {
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
@@ -22,10 +26,12 @@ mod artifacts {
/// Phase 1 (startup extraction).
mod phase_one {
/// Prompt used for phase 1.
pub(super) const PROMPT: &str = include_str!("../../templates/memories/stage_one_system.md");
/// Maximum number of rollout candidates processed per startup pass.
pub(super) const MAX_ROLLOUTS_PER_STARTUP: usize = 64;
/// Concurrency cap for startup memory extraction and consolidation scheduling.
pub(super) const CONCURRENCY_LIMIT: usize = MAX_ROLLOUTS_PER_STARTUP;
pub(super) const CONCURRENCY_LIMIT: usize = 64;
/// Fallback stage-1 rollout truncation limit (tokens) when model metadata
/// does not include a valid context window.
pub(super) const DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT: usize = 150_000;
@@ -46,6 +52,8 @@ mod phase_one {
pub(super) const JOB_LEASE_SECONDS: i64 = 3_600;
/// Backoff delay (seconds) before retrying a failed stage-1 extraction job.
pub(super) const JOB_RETRY_DELAY_SECONDS: i64 = 3_600;
/// Maximum number of threads to scan.
pub(super) const THREAD_SCAN_LIMIT: usize = 5_000;
}
/// Phase 2 (aka `Consolidation`).
@@ -74,6 +82,9 @@ mod metrics {
pub(super) const MEMORY_PHASE_TWO_INPUT: &str = "codex.memory.phase2.input";
}
use std::path::Path;
use std::path::PathBuf;
pub fn memory_root(codex_home: &Path) -> PathBuf {
codex_home.join("memories")
}
@@ -89,8 +100,3 @@ fn raw_memories_file(root: &Path) -> PathBuf {
async fn ensure_layout(root: &Path) -> std::io::Result<()> {
tokio::fs::create_dir_all(rollout_summaries_dir(root)).await
}
/// Starts the memory startup pipeline for eligible root sessions.
///
/// This is the single entrypoint that `codex` uses to trigger memory startup.
pub(crate) use startup::start_memories_startup_task;
codex-rs/core/src/memories/phase1.rs
+466
View File
@@ -0,0 +1,466 @@
use crate::Prompt;
use crate::RolloutRecorder;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::error::CodexErr;
use crate::memories::metrics;
use crate::memories::phase_one;
use crate::memories::prompts::build_stage_one_input_message;
use crate::rollout::INTERACTIVE_SESSION_SOURCES;
use crate::rollout::policy::should_persist_response_item_for_memories;
use codex_api::ResponseEvent;
use codex_otel::OtelManager;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::RolloutItem;
use codex_utils_sanitizer::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;
#[derive(Clone, Debug)]
pub(in crate::memories) struct Phase1RequestContext {
pub(in crate::memories) model_info: ModelInfo,
pub(in crate::memories) otel_manager: OtelManager,
pub(in crate::memories) reasoning_effort: Option<ReasoningEffortConfig>,
pub(in crate::memories) reasoning_summary: ReasoningSummaryConfig,
pub(in crate::memories) turn_metadata_header: Option<String>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum PhaseOneJobOutcome {
SucceededWithOutput,
SucceededNoOutput,
Failed,
}
struct PhaseOneOutcomeCounts {
claimed: usize,
succeeded_with_output: usize,
succeeded_no_output: usize,
failed: usize,
}
/// 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 accepted from stage-1 output for forward compatibility.
/// This is currently ignored by downstream storage and naming, which remain thread-id based.
#[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(in crate::memories) async fn run(session: &Arc<Session>) {
// 1. Claim startup job.
let Some(claimed_candidates) = claim_startup_jobs(session).await else {
return;
};
if claimed_candidates.is_empty() {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "skipped_no_candidates")],
);
return;
}
// 2. Build request.
let stage_one_context = build_request_context(session).await;
// 3. Run the parallel sampling.
let outcomes = run_jobs(session, claimed_candidates, stage_one_context).await;
// 4. Metrics and logs.
let counts = count_outcomes(outcomes);
emit_metrics(session, &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
);
}
/// 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" },
"raw_memory": { "type": "string" }
},
"required": ["rollout_summary", "rollout_slug", "raw_memory"],
"additionalProperties": false
})
}
impl Phase1RequestContext {
pub(in crate::memories) fn from_turn_context(
turn_context: &TurnContext,
turn_metadata_header: Option<String>,
) -> Self {
Self {
model_info: turn_context.model_info.clone(),
otel_manager: turn_context.otel_manager.clone(),
reasoning_effort: turn_context.reasoning_effort,
reasoning_summary: turn_context.reasoning_summary,
turn_metadata_header,
}
}
}
async fn claim_startup_jobs(session: &Arc<Session>) -> Option<Vec<codex_state::Stage1JobClaim>> {
let Some(state_db) = session.services.state_db.as_deref() 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(
session.conversation_id,
codex_state::Stage1StartupClaimParams {
scan_limit: phase_one::THREAD_SCAN_LIMIT,
max_claimed: phase_one::MAX_ROLLOUTS_PER_STARTUP,
max_age_days: phase_one::MAX_ROLLOUT_AGE_DAYS,
min_rollout_idle_hours: phase_one::MIN_ROLLOUT_IDLE_HOURS,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: phase_one::JOB_LEASE_SECONDS,
},
)
.await
{
Ok(claims) => Some(claims),
Err(err) => {
warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}");
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "failed_claim")],
);
None
}
}
}
async fn build_request_context(session: &Arc<Session>) -> Phase1RequestContext {
let turn_context = session.new_default_turn().await;
Phase1RequestContext::from_turn_context(
turn_context.as_ref(),
turn_context.resolve_turn_metadata_header().await,
)
}
async fn run_jobs(
session: &Arc<Session>,
claimed_candidates: Vec<codex_state::Stage1JobClaim>,
stage_one_context: Phase1RequestContext,
) -> Vec<PhaseOneJobOutcome> {
futures::stream::iter(claimed_candidates.into_iter())
.map(|claim| {
let session = Arc::clone(session);
let stage_one_context = stage_one_context.clone();
async move { job::run(session.as_ref(), claim, &stage_one_context).await }
})
.buffer_unordered(phase_one::CONCURRENCY_LIMIT)
.collect::<Vec<_>>()
.await
}
mod job {
use super::*;
pub(in crate::memories) async fn run(
session: &Session,
claim: codex_state::Stage1JobClaim,
stage_one_context: &Phase1RequestContext,
) -> PhaseOneJobOutcome {
let thread = claim.thread;
let stage_one_output = match sample(
session,
&thread.rollout_path,
&thread.cwd,
stage_one_context,
)
.await
{
Ok(output) => output,
Err(reason) => {
result::failed(
session,
thread.id,
&claim.ownership_token,
&reason.to_string(),
)
.await;
return PhaseOneJobOutcome::Failed;
}
};
if stage_one_output.raw_memory.is_empty() || stage_one_output.rollout_summary.is_empty() {
return result::no_output(session, thread.id, &claim.ownership_token).await;
}
result::success(
session,
thread.id,
&claim.ownership_token,
thread.updated_at.timestamp(),
&stage_one_output.raw_memory,
&stage_one_output.rollout_summary,
)
.await
}
/// Extract the rollout and perform the actual sampling.
async fn sample(
session: &Session,
rollout_path: &Path,
rollout_cwd: &Path,
stage_one_context: &Phase1RequestContext,
) -> anyhow::Result<StageOneOutput> {
let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?;
let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?;
let prompt = 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,
)?,
}],
end_turn: None,
phase: None,
}],
tools: Vec::new(),
parallel_tool_calls: false,
base_instructions: BaseInstructions {
text: phase_one::PROMPT.to_string(),
},
personality: None,
output_schema: Some(output_schema()),
};
let mut client_session = session.services.model_client.new_session();
let mut stream = client_session
.stream(
&prompt,
&stage_one_context.model_info,
&stage_one_context.otel_manager,
stage_one_context.reasoning_effort,
stage_one_context.reasoning_summary,
stage_one_context.turn_metadata_header.as_deref(),
)
.await?;
// TODO(jif) we should have a shared helper somewhere for this.
// Unwrap the stream.
let mut result = String::new();
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 ResponseItem::Message { content, .. } = item
&& let Some(text) = crate::compact::content_items_to_text(&content)
{
result.push_str(&text);
}
}
ResponseEvent::Completed { .. } => break,
_ => {}
}
}
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);
Ok(output)
}
mod result {
use super::*;
pub(in crate::memories) async fn failed(
session: &Session,
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) = session.services.state_db.as_deref() {
let _ = state_db
.mark_stage1_job_failed(
thread_id,
ownership_token,
reason,
phase_one::JOB_RETRY_DELAY_SECONDS,
)
.await;
}
}
pub(in crate::memories) async fn no_output(
session: &Session,
thread_id: codex_protocol::ThreadId,
ownership_token: &str,
) -> PhaseOneJobOutcome {
let Some(state_db) = session.services.state_db.as_deref() else {
return PhaseOneJobOutcome::Failed;
};
if state_db
.mark_stage1_job_succeeded_no_output(thread_id, ownership_token)
.await
.unwrap_or(false)
{
PhaseOneJobOutcome::SucceededNoOutput
} else {
PhaseOneJobOutcome::Failed
}
}
pub(in crate::memories) async fn success(
session: &Session,
thread_id: codex_protocol::ThreadId,
ownership_token: &str,
source_updated_at: i64,
raw_memory: &str,
rollout_summary: &str,
) -> PhaseOneJobOutcome {
let Some(state_db) = session.services.state_db.as_deref() else {
return PhaseOneJobOutcome::Failed;
};
if state_db
.mark_stage1_job_succeeded(
thread_id,
ownership_token,
source_updated_at,
raw_memory,
rollout_summary,
)
.await
.unwrap_or(false)
{
PhaseOneJobOutcome::SucceededWithOutput
} else {
PhaseOneJobOutcome::Failed
}
}
}
/// Serializes filtered stage-1 memory items for prompt inclusion.
fn serialize_filtered_rollout_response_items(
items: &[RolloutItem],
) -> crate::error::Result<String> {
let filtered = items
.iter()
.filter_map(|item| {
if let RolloutItem::ResponseItem(item) = item
&& should_persist_response_item_for_memories(item)
{
Some(item.clone())
} else {
None
}
})
.collect::<Vec<_>>();
serde_json::to_string(&filtered).map_err(|err| {
CodexErr::InvalidRequest(format!("failed to serialize rollout memory: {err}"))
})
}
}
fn count_outcomes(outcomes: Vec<PhaseOneJobOutcome>) -> PhaseOneOutcomeCounts {
let succeeded_with_output = outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::SucceededWithOutput))
.count();
let succeeded_no_output = outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::SucceededNoOutput))
.count();
let failed = outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::Failed))
.count();
PhaseOneOutcomeCounts {
claimed: outcomes.len(),
succeeded_with_output,
succeeded_no_output,
failed,
}
}
fn emit_metrics(session: &Session, counts: &PhaseOneOutcomeCounts) {
if counts.claimed > 0 {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
counts.claimed as i64,
&[("status", "claimed")],
);
}
if counts.succeeded_with_output > 0 {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
counts.succeeded_with_output as i64,
&[("status", "succeeded")],
);
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_OUTPUT,
counts.succeeded_with_output as i64,
&[],
);
}
if counts.succeeded_no_output > 0 {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
counts.succeeded_no_output as i64,
&[("status", "succeeded_no_output")],
);
}
if counts.failed > 0 {
session.services.otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
counts.failed as i64,
&[("status", "failed")],
);
}
}
codex-rs/core/src/memories/startup/phase2.rs → codex-rs/core/src/memories/phase2.rs
+1 -1
View File
@@ -11,7 +11,7 @@ use tracing::debug;
use tracing::info;
use tracing::warn;
pub(super) fn spawn_phase2_completion_task(
pub(in crate::memories) fn spawn_phase2_completion_task(
session: &Session,
ownership_token: String,
completion_watermark: i64,
codex-rs/core/src/memories/prompts.rs
+3 -1
View File
@@ -30,7 +30,6 @@ struct MemoryToolDeveloperInstructionsTemplate<'a> {
}
/// Builds the consolidation subagent prompt for a specific memory root.
///
pub(super) fn build_consolidation_prompt(memory_root: &Path) -> String {
let memory_root = memory_root.display().to_string();
let template = ConsolidationPromptTemplate {
@@ -74,6 +73,9 @@ pub(super) fn build_stage_one_input_message(
.render()?)
}
/// Build prompt used for read path. This prompt must be added to the developer instructions. In
/// case of large memory files, the `memory_summary.md` is truncated at
/// [phase_one::MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT].
pub(crate) async fn build_memory_tool_developer_instructions(codex_home: &Path) -> Option<String> {
let base_path = memory_root(codex_home);
let memory_summary_path = base_path.join("memory_summary.md");
codex-rs/core/src/memories/stage_one.rs
-210
View File
@@ -1,210 +0,0 @@
use crate::error::CodexErr;
use crate::error::Result;
use once_cell::sync::Lazy;
use regex::Regex;
use serde::Deserialize;
use serde_json::Value;
use serde_json::json;
/// System prompt for stage-1 raw memory extraction.
pub(super) const RAW_MEMORY_PROMPT: &str =
include_str!("../../templates/memories/stage_one_system.md");
static OPENAI_KEY_REGEX: Lazy<Regex> = Lazy::new(|| compile_regex(r"sk-[A-Za-z0-9]{20,}"));
static AWS_ACCESS_KEY_ID_REGEX: Lazy<Regex> = Lazy::new(|| compile_regex(r"\bAKIA[0-9A-Z]{16}\b"));
static BEARER_TOKEN_REGEX: Lazy<Regex> =
Lazy::new(|| compile_regex(r"(?i)\bBearer\s+[A-Za-z0-9._\-]{16,}\b"));
static SECRET_ASSIGNMENT_REGEX: Lazy<Regex> = Lazy::new(|| {
compile_regex(r#"(?i)\b(api[_-]?key|token|secret|password)\b(\s*[:=]\s*)(["']?)[^\s"']{8,}"#)
});
/// Parsed stage-1 model output payload.
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub(super) 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 accepted from stage-1 output for forward compatibility.
///
/// This is currently ignored by downstream storage and naming, which remain
/// thread-id based.
#[serde(default, rename = "rollout_slug")]
pub(crate) _rollout_slug: Option<String>,
}
/// JSON schema used to constrain stage-1 model output.
pub(super) fn stage_one_output_schema() -> Value {
json!({
"type": "object",
"properties": {
"rollout_summary": { "type": "string" },
"rollout_slug": { "type": "string" },
"raw_memory": { "type": "string" }
},
"required": ["rollout_summary", "rollout_slug", "raw_memory"],
"additionalProperties": false
})
}
/// Parses and normalizes stage-1 model output into a typed payload.
///
/// Accepts plain JSON objects, fenced JSON, and object snippets embedded in
/// extra text, then enforces redaction and size limits.
pub(super) fn parse_stage_one_output(raw: &str) -> Result<StageOneOutput> {
let parsed = parse_json_object_loose(raw)?;
let output: StageOneOutput = serde_json::from_value(parsed).map_err(|err| {
CodexErr::InvalidRequest(format!("invalid stage-1 memory output JSON payload: {err}"))
})?;
normalize_stage_one_output(output)
}
fn parse_json_object_loose(raw: &str) -> Result<Value> {
let raw = raw.trim();
if let Ok(value) = serde_json::from_str::<Value>(raw)
&& value.is_object()
{
return Ok(value);
}
if let Some(fenced) = raw
.strip_prefix("```json")
.and_then(|s| s.strip_suffix("```"))
.map(str::trim)
&& let Ok(value) = serde_json::from_str::<Value>(fenced)
&& value.is_object()
{
return Ok(value);
}
if let Some(fenced) = raw
.strip_prefix("```")
.and_then(|s| s.strip_suffix("```"))
.map(str::trim)
&& let Ok(value) = serde_json::from_str::<Value>(fenced)
&& value.is_object()
{
return Ok(value);
}
if let (Some(start), Some(end)) = (raw.find('{'), raw.rfind('}'))
&& start < end
{
let snippet = &raw[start..=end];
if let Ok(value) = serde_json::from_str::<Value>(snippet)
&& value.is_object()
{
return Ok(value);
}
}
Err(CodexErr::InvalidRequest(
"unable to parse stage-1 memory JSON output".to_string(),
))
}
fn normalize_stage_one_output(mut output: StageOneOutput) -> Result<StageOneOutput> {
output.raw_memory = output.raw_memory.trim().to_string();
output.rollout_summary = output.rollout_summary.trim().to_string();
output._rollout_slug = output
._rollout_slug
.map(|slug| slug.trim().to_string())
.filter(|slug| !slug.is_empty());
if output.raw_memory.is_empty() && output.rollout_summary.is_empty() {
// Empty pair is a deliberate "no meaningful signal" sentinel.
return Ok(output);
}
if output.raw_memory.is_empty() {
return Err(CodexErr::InvalidRequest(
"stage-1 memory output missing raw_memory".to_string(),
));
}
if output.rollout_summary.is_empty() {
return Err(CodexErr::InvalidRequest(
"stage-1 memory output missing rollout_summary".to_string(),
));
}
output.raw_memory = redact_secrets(&output.raw_memory);
output.rollout_summary = redact_secrets(&output.rollout_summary);
Ok(output)
}
fn redact_secrets(input: &str) -> String {
let redacted = OPENAI_KEY_REGEX.replace_all(input, "[REDACTED_SECRET]");
let redacted = AWS_ACCESS_KEY_ID_REGEX.replace_all(&redacted, "[REDACTED_SECRET]");
let redacted = BEARER_TOKEN_REGEX.replace_all(&redacted, "Bearer [REDACTED_SECRET]");
SECRET_ASSIGNMENT_REGEX
.replace_all(&redacted, "$1$2$3[REDACTED_SECRET]")
.to_string()
}
fn compile_regex(pattern: &str) -> Regex {
match Regex::new(pattern) {
Ok(regex) => regex,
// Panic is ok thanks to `load_regex` test.
Err(err) => panic!("invalid regex pattern `{pattern}`: {err}"),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn load_regex() {
// The goal of this test is just to compile all the regex to prevent the panic
let _ = redact_secrets("secret");
}
#[test]
fn normalize_stage_one_output_redacts_summary() {
let output = StageOneOutput {
raw_memory: "Token: sk-abcdefghijklmnopqrstuvwxyz123456\nBearer abcdefghijklmnopqrstuvwxyz012345".to_string(),
rollout_summary: "password = mysecret123456\n\nsmall".to_string(),
_rollout_slug: None,
};
let normalized = normalize_stage_one_output(output).expect("normalized");
assert!(normalized.raw_memory.contains("[REDACTED_SECRET]"));
assert!(!normalized.rollout_summary.contains("mysecret123456"));
assert_eq!(
normalized.rollout_summary,
"password = [REDACTED_SECRET]\n\nsmall"
);
}
#[test]
fn normalize_stage_one_output_allows_empty_pair_for_skip() {
let output = StageOneOutput {
raw_memory: String::new(),
rollout_summary: String::new(),
_rollout_slug: None,
};
let normalized = normalize_stage_one_output(output).expect("normalized");
assert_eq!(normalized.raw_memory, "");
assert_eq!(normalized.rollout_summary, "");
}
#[test]
fn normalize_stage_one_output_rejects_partial_empty_values() {
let output = StageOneOutput {
raw_memory: String::new(),
rollout_summary: "summary".to_string(),
_rollout_slug: None,
};
let err = normalize_stage_one_output(output).expect_err("should reject");
assert_eq!(err.to_string(), "stage-1 memory output missing raw_memory");
}
}
codex-rs/core/src/memories/start.rs
+41
View File
@@ -0,0 +1,41 @@
use crate::codex::Session;
use crate::config::Config;
use crate::features::Feature;
use crate::memories::phase1;
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(crate) fn start_memories_startup_task(
session: &Arc<Session>,
config: Arc<Config>,
source: &SessionSource,
) {
if config.ephemeral
|| !config.features.enabled(Feature::MemoryTool)
|| matches!(source, SessionSource::SubAgent(_))
{
return;
}
if session.services.state_db.is_none() {
warn!("state db unavailable for memories startup pipeline; skipping");
return;
}
let weak_session = Arc::downgrade(session);
tokio::spawn(async move {
let Some(session) = weak_session.upgrade() else {
return;
};
// Run phase 1.
phase1::run(&session).await;
// Run phase 2.
crate::memories::dispatch::run_global_memory_consolidation(&session, config).await;
});
}
codex-rs/core/src/memories/startup/extract.rs
-199
View File
@@ -1,199 +0,0 @@
use crate::client_common::Prompt;
use crate::client_common::ResponseEvent;
use crate::client_common::ResponseStream;
use crate::codex::Session;
use crate::error::CodexErr;
use crate::error::Result as CodexResult;
use crate::rollout::RolloutRecorder;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use futures::StreamExt;
use tracing::warn;
use crate::memories::prompts::build_stage_one_input_message;
use crate::memories::stage_one::RAW_MEMORY_PROMPT;
use crate::memories::stage_one::StageOneOutput;
use crate::memories::stage_one::parse_stage_one_output;
use crate::memories::stage_one::stage_one_output_schema;
use crate::memories::startup::StageOneRequestContext;
use crate::rollout::policy::should_persist_response_item_for_memories;
use codex_protocol::protocol::RolloutItem;
use std::path::Path;
pub(super) async fn extract_stage_one_output(
session: &Session,
rollout_path: &Path,
rollout_cwd: &Path,
stage_one_context: &StageOneRequestContext,
) -> Result<StageOneOutput, &'static str> {
let (rollout_items, _thread_id, parse_errors) =
match RolloutRecorder::load_rollout_items(rollout_path).await {
Ok(result) => result,
Err(err) => {
warn!(
"failed to load rollout {} for memories: {err}",
rollout_path.display()
);
return Err("failed to load rollout");
}
};
if parse_errors > 0 {
warn!(
"rollout {} had {parse_errors} parse errors while preparing stage-1 memory input",
rollout_path.display()
);
}
let rollout_contents = match serialize_filtered_rollout_response_items(&rollout_items) {
Ok(contents) => contents,
Err(err) => {
warn!(
"failed to prepare filtered rollout payload {} for memories: {err}",
rollout_path.display()
);
return Err("failed to serialize filtered rollout");
}
};
let prompt = 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,
)
.map_err(|_e| "error while building the prompt")?,
}],
end_turn: None,
phase: None,
}],
tools: Vec::new(),
parallel_tool_calls: false,
base_instructions: BaseInstructions {
text: RAW_MEMORY_PROMPT.to_string(),
},
personality: None,
output_schema: Some(stage_one_output_schema()),
};
let mut client_session = session.services.model_client.new_session();
let mut stream = match client_session
.stream(
&prompt,
&stage_one_context.model_info,
&stage_one_context.otel_manager,
stage_one_context.reasoning_effort,
stage_one_context.reasoning_summary,
stage_one_context.turn_metadata_header.as_deref(),
)
.await
{
Ok(stream) => stream,
Err(err) => {
warn!(
"stage-1 memory request failed for rollout {}: {err}",
rollout_path.display()
);
return Err("stage-1 memory request failed");
}
};
let output_text = match collect_response_text_until_completed(&mut stream).await {
Ok(text) => text,
Err(err) => {
warn!(
"failed while waiting for stage-1 memory response for rollout {}: {err}",
rollout_path.display()
);
return Err("stage-1 memory response stream failed");
}
};
match parse_stage_one_output(&output_text) {
Ok(output) => Ok(output),
Err(err) => {
warn!(
"invalid stage-1 memory payload for rollout {}: {err}",
rollout_path.display()
);
Err("invalid stage-1 memory payload")
}
}
}
async fn collect_response_text_until_completed(stream: &mut ResponseStream) -> CodexResult<String> {
let mut output_text = String::new();
loop {
let Some(event) = stream.next().await else {
return Err(CodexErr::Stream(
"stream closed before response.completed".to_string(),
None,
));
};
match event? {
ResponseEvent::OutputTextDelta(delta) => output_text.push_str(&delta),
ResponseEvent::OutputItemDone(item) => {
if output_text.is_empty()
&& let ResponseItem::Message { content, .. } = item
&& let Some(text) = crate::compact::content_items_to_text(&content)
{
output_text.push_str(&text);
}
}
ResponseEvent::Completed { .. } => return Ok(output_text),
_ => {}
}
}
}
/// Serializes filtered stage-1 memory items for prompt inclusion.
fn serialize_filtered_rollout_response_items(
items: &[RolloutItem],
) -> crate::error::Result<String> {
let filtered = items
.iter()
.filter_map(|item| {
if let RolloutItem::ResponseItem(item) = item
&& should_persist_response_item_for_memories(item)
{
Some(item.clone())
} else {
None
}
})
.collect::<Vec<_>>();
serde_json::to_string(&filtered).map_err(|err| {
CodexErr::InvalidRequest(format!("failed to serialize rollout memory: {err}"))
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn serialize_filtered_rollout_response_items_keeps_response_items_only() {
let input = vec![RolloutItem::ResponseItem(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "user input".to_string(),
}],
end_turn: None,
phase: None,
})];
let serialized = serialize_filtered_rollout_response_items(&input).expect("serialize");
let parsed: Vec<ResponseItem> = serde_json::from_str(&serialized).expect("deserialize");
pretty_assertions::assert_eq!(parsed.len(), 1);
assert!(matches!(parsed[0], ResponseItem::Message { .. }));
}
}
codex-rs/core/src/memories/startup/mod.rs
-306
View File
@@ -1,306 +0,0 @@
mod dispatch;
mod extract;
mod phase2;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::config::Config;
use crate::error::Result as CodexResult;
use crate::features::Feature;
use crate::memories::metrics;
use crate::memories::phase_one;
use crate::rollout::INTERACTIVE_SESSION_SOURCES;
use codex_otel::OtelManager;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::SessionSource;
use futures::StreamExt;
use std::sync::Arc;
use tracing::info;
use tracing::warn;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum PhaseOneJobOutcome {
SucceededWithOutput,
SucceededNoOutput,
Failed,
}
pub(super) const PHASE_ONE_THREAD_SCAN_LIMIT: usize = 5_000;
#[derive(Clone)]
struct StageOneRequestContext {
model_info: ModelInfo,
otel_manager: OtelManager,
reasoning_effort: Option<ReasoningEffortConfig>,
reasoning_summary: ReasoningSummaryConfig,
turn_metadata_header: Option<String>,
}
impl StageOneRequestContext {
fn from_turn_context(turn_context: &TurnContext, turn_metadata_header: Option<String>) -> Self {
Self {
model_info: turn_context.model_info.clone(),
otel_manager: turn_context.otel_manager.clone(),
reasoning_effort: turn_context.reasoning_effort,
reasoning_summary: turn_context.reasoning_summary,
turn_metadata_header,
}
}
}
/// 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(crate) fn start_memories_startup_task(
session: &Arc<Session>,
config: Arc<Config>,
source: &SessionSource,
) {
if config.ephemeral
|| !config.features.enabled(Feature::MemoryTool)
|| matches!(source, SessionSource::SubAgent(_))
{
return;
}
let weak_session = Arc::downgrade(session);
tokio::spawn(async move {
let Some(session) = weak_session.upgrade() else {
return;
};
if let Err(err) = run_memories_startup_pipeline(&session, config).await {
warn!("memories startup pipeline failed: {err}");
}
});
}
/// Runs the startup memory pipeline.
///
/// Phase 1 selects rollout candidates, performs stage-1 extraction requests in
/// parallel, persists stage-1 outputs, and enqueues consolidation work.
///
/// Phase 2 claims a global consolidation lock and spawns one consolidation agent.
pub(super) async fn run_memories_startup_pipeline(
session: &Arc<Session>,
config: Arc<Config>,
) -> CodexResult<()> {
let otel_manager = &session.services.otel_manager;
let Some(state_db) = session.services.state_db.as_deref() else {
warn!("state db unavailable for memories startup pipeline; skipping");
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "skipped_state_db_unavailable")],
);
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "skipped_state_db_unavailable")],
);
return Ok(());
};
let allowed_sources = INTERACTIVE_SESSION_SOURCES
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>();
let claimed_candidates = match state_db
.claim_stage1_jobs_for_startup(
session.conversation_id,
codex_state::Stage1StartupClaimParams {
scan_limit: PHASE_ONE_THREAD_SCAN_LIMIT,
max_claimed: phase_one::MAX_ROLLOUTS_PER_STARTUP,
max_age_days: phase_one::MAX_ROLLOUT_AGE_DAYS,
min_rollout_idle_hours: phase_one::MIN_ROLLOUT_IDLE_HOURS,
allowed_sources: allowed_sources.as_slice(),
lease_seconds: phase_one::JOB_LEASE_SECONDS,
},
)
.await
{
Ok(claims) => claims,
Err(err) => {
warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}");
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "failed_claim")],
);
Vec::new()
}
};
let claimed_count = claimed_candidates.len();
if claimed_count == 0 {
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
1,
&[("status", "skipped_no_candidates")],
);
}
let mut phase_one_outcomes = Vec::new();
if claimed_count > 0 {
let turn_context = session.new_default_turn().await;
let stage_one_context = StageOneRequestContext::from_turn_context(
turn_context.as_ref(),
turn_context.resolve_turn_metadata_header().await,
);
phase_one_outcomes = futures::stream::iter(claimed_candidates.into_iter())
.map(|claim| {
let session = Arc::clone(session);
let stage_one_context = stage_one_context.clone();
async move {
let thread = claim.thread;
let stage_one_output = match extract::extract_stage_one_output(
session.as_ref(),
&thread.rollout_path,
&thread.cwd,
&stage_one_context,
)
.await
{
Ok(output) => output,
Err(reason) => {
if let Some(state_db) = session.services.state_db.as_deref() {
let _ = state_db
.mark_stage1_job_failed(
thread.id,
&claim.ownership_token,
reason,
phase_one::JOB_RETRY_DELAY_SECONDS,
)
.await;
}
return PhaseOneJobOutcome::Failed;
}
};
let Some(state_db) = session.services.state_db.as_deref() else {
return PhaseOneJobOutcome::Failed;
};
if stage_one_output.raw_memory.is_empty()
&& stage_one_output.rollout_summary.is_empty()
{
return if state_db
.mark_stage1_job_succeeded_no_output(thread.id, &claim.ownership_token)
.await
.unwrap_or(false)
{
PhaseOneJobOutcome::SucceededNoOutput
} else {
PhaseOneJobOutcome::Failed
};
}
if state_db
.mark_stage1_job_succeeded(
thread.id,
&claim.ownership_token,
thread.updated_at.timestamp(),
&stage_one_output.raw_memory,
&stage_one_output.rollout_summary,
)
.await
.unwrap_or(false)
{
PhaseOneJobOutcome::SucceededWithOutput
} else {
PhaseOneJobOutcome::Failed
}
}
})
.buffer_unordered(phase_one::CONCURRENCY_LIMIT)
.collect::<Vec<PhaseOneJobOutcome>>()
.await;
}
let succeeded_with_output_count = phase_one_outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::SucceededWithOutput))
.count();
let succeeded_no_output_count = phase_one_outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::SucceededNoOutput))
.count();
let failed_count = phase_one_outcomes
.iter()
.filter(|outcome| matches!(outcome, PhaseOneJobOutcome::Failed))
.count();
let succeeded_count = succeeded_with_output_count + succeeded_no_output_count;
if claimed_count > 0 {
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
claimed_count as i64,
&[("status", "claimed")],
);
}
if succeeded_with_output_count > 0 {
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
succeeded_with_output_count as i64,
&[("status", "succeeded")],
);
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_OUTPUT,
succeeded_with_output_count as i64,
&[],
);
}
if succeeded_no_output_count > 0 {
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
succeeded_no_output_count as i64,
&[("status", "succeeded_no_output")],
);
}
if failed_count > 0 {
otel_manager.counter(
metrics::MEMORY_PHASE_ONE_JOBS,
failed_count as i64,
&[("status", "failed")],
);
}
info!(
"memory stage-1 extraction complete: {} job(s) claimed, {} succeeded ({} with output, {} no output), {} failed",
claimed_count,
succeeded_count,
succeeded_with_output_count,
succeeded_no_output_count,
failed_count
);
let consolidation_job_count =
usize::from(dispatch::run_global_memory_consolidation(session, config).await);
info!(
"memory consolidation dispatch complete: {} job(s) scheduled",
consolidation_job_count
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::run_memories_startup_pipeline;
use crate::codex::make_session_and_context;
use crate::config::test_config;
use std::sync::Arc;
#[tokio::test]
async fn startup_pipeline_is_noop_when_state_db_is_unavailable() {
let (session, _turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let config = Arc::new(test_config());
run_memories_startup_pipeline(&session, config)
.await
.expect("startup pipeline should skip cleanly without state db");
}
}
codex-rs/core/src/memories/storage.rs
+2
View File
@@ -9,6 +9,8 @@ use crate::memories::phase_two;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
//TODO(jif) clean.
/// Rebuild `raw_memories.md` from DB-backed stage-1 outputs.
pub(super) async fn rebuild_raw_memories_file_from_memories(
root: &Path,
codex-rs/core/src/memories/tests.rs
+1 -33
View File
@@ -1,4 +1,3 @@
use super::stage_one::parse_stage_one_output;
use super::storage::rebuild_raw_memories_file_from_memories;
use super::storage::sync_rollout_summaries_from_memories;
use crate::memories::ensure_layout;
@@ -20,40 +19,9 @@ fn memory_root_uses_shared_global_path() {
assert_eq!(memory_root(&codex_home), codex_home.join("memories"));
}
#[test]
fn parse_stage_one_output_accepts_fenced_json() {
let raw = "```json\n{\"raw_memory\":\"abc\",\"rollout_summary\":\"short\"}\n```";
let parsed = parse_stage_one_output(raw).expect("parsed");
assert!(parsed.raw_memory.contains("abc"));
assert_eq!(parsed.rollout_summary, "short");
}
#[test]
fn parse_stage_one_output_rejects_legacy_keys() {
let raw = r#"{"rawMemory":"abc","summary":"short"}"#;
assert!(parse_stage_one_output(raw).is_err());
}
#[test]
fn parse_stage_one_output_accepts_empty_pair_for_skip() {
let raw = r#"{"raw_memory":"","rollout_summary":""}"#;
let parsed = parse_stage_one_output(raw).expect("parsed");
assert_eq!(parsed.raw_memory, "");
assert_eq!(parsed.rollout_summary, "");
}
#[test]
fn parse_stage_one_output_accepts_optional_rollout_slug() {
let raw = r#"{"raw_memory":"abc","rollout_summary":"short","rollout_slug":"my-slug"}"#;
let parsed = parse_stage_one_output(raw).expect("parsed");
assert!(parsed.raw_memory.contains("abc"));
assert_eq!(parsed.rollout_summary, "short");
assert_eq!(parsed._rollout_slug, Some("my-slug".to_string()));
}
#[test]
fn stage_one_output_schema_requires_all_declared_properties() {
let schema = super::stage_one::stage_one_output_schema();
let schema = crate::memories::phase1::output_schema();
let properties = schema
.get("properties")
.and_then(Value::as_object)
codex-rs/utils/sanitizer/BUILD.bazel
+6
View File
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "sanitizer",
crate_name = "codex_utils_sanitizer",
)