From de6f2ef746a12b680b207f7abebc2da7cb9bb042 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 11 Feb 2026 21:11:57 +0000 Subject: [PATCH] nit: memory truncation (#11479) Use existing truncation for memories --- codex-rs/core/src/codex.rs | 4 +- codex-rs/core/src/memories/mod.rs | 23 +---- codex-rs/core/src/memories/prompts.rs | 68 ++++--------- codex-rs/core/src/memories/stage_one.rs | 99 ++++++------------- codex-rs/core/src/memories/startup/extract.rs | 5 +- codex-rs/core/src/memories/storage.rs | 3 +- codex-rs/core/src/memories/text.rs | 50 ---------- 7 files changed, 58 insertions(+), 194 deletions(-) delete mode 100644 codex-rs/core/src/memories/text.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4a6c8cff1..21daadb74 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2308,8 +2308,7 @@ impl Session { } // Add developer instructions for memories. if let Some(memory_prompt) = - memories::build_memory_tool_developer_instructions(&turn_context.config.codex_home) - .await + build_memory_tool_developer_instructions(&turn_context.config.codex_home).await && turn_context.features.enabled(Feature::MemoryTool) { items.push(DeveloperInstructions::new(memory_prompt).into()); @@ -5140,6 +5139,7 @@ pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) - }) } +use crate::memories::prompts::build_memory_tool_developer_instructions; #[cfg(test)] pub(crate) use tests::make_session_and_context; #[cfg(test)] diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs index 80834a510..9755a75e9 100644 --- a/codex-rs/core/src/memories/mod.rs +++ b/codex-rs/core/src/memories/mod.rs @@ -4,16 +4,14 @@ //! - 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 prompts; +pub(crate) mod prompts; mod stage_one; mod startup; mod storage; -mod text; #[cfg(test)] mod tests; -use serde::Deserialize; use std::path::Path; use std::path::PathBuf; @@ -42,24 +40,6 @@ const PHASE_TWO_JOB_RETRY_DELAY_SECONDS: i64 = 3_600; /// Heartbeat interval (seconds) for phase-2 running jobs. const PHASE_TWO_JOB_HEARTBEAT_SECONDS: u64 = 30; -/// Parsed stage-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")] - raw_memory: String, - /// Compact summary line used for routing and indexing. - #[serde(rename = "rollout_summary")] - 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")] - _rollout_slug: Option, -} - pub fn memory_root(codex_home: &Path) -> PathBuf { codex_home.join("memories") } @@ -76,7 +56,6 @@ async fn ensure_layout(root: &Path) -> std::io::Result<()> { tokio::fs::create_dir_all(rollout_summaries_dir(root)).await } -pub(crate) use prompts::build_memory_tool_developer_instructions; /// Starts the memory startup pipeline for eligible root sessions. /// /// This is the single entrypoint that `codex` uses to trigger memory startup. diff --git a/codex-rs/core/src/memories/prompts.rs b/codex-rs/core/src/memories/prompts.rs index c098905ab..eaf5de3cd 100644 --- a/codex-rs/core/src/memories/prompts.rs +++ b/codex-rs/core/src/memories/prompts.rs @@ -1,14 +1,11 @@ -use super::text::prefix_at_char_boundary; -use super::text::suffix_at_char_boundary; use crate::memories::memory_root; +use crate::truncate::TruncationPolicy; +use crate::truncate::truncate_text; use askama::Template; use std::path::Path; use tokio::fs; use tracing::warn; -// TODO(jif) use proper truncation -const MAX_ROLLOUT_BYTES_FOR_PROMPT: usize = 100_000; - #[derive(Template)] #[template(path = "memories/consolidation.md", escape = "none")] struct ConsolidationPromptTemplate<'a> { @@ -51,29 +48,18 @@ pub(super) fn build_stage_one_input_message( rollout_path: &Path, rollout_cwd: &Path, rollout_contents: &str, -) -> String { - let (rollout_contents, truncated) = truncate_rollout_for_prompt(rollout_contents); - if truncated { - warn!( - "truncated rollout {} for stage-1 memory prompt to {} bytes", - rollout_path.display(), - MAX_ROLLOUT_BYTES_FOR_PROMPT - ); - } +) -> anyhow::Result { + let truncated_rollout_contents = + truncate_text(rollout_contents, TruncationPolicy::Tokens(150_000)); let rollout_path = rollout_path.display().to_string(); let rollout_cwd = rollout_cwd.display().to_string(); - let template = StageOneInputTemplate { + Ok(StageOneInputTemplate { rollout_path: &rollout_path, rollout_cwd: &rollout_cwd, - rollout_contents: &rollout_contents, - }; - template.render().unwrap_or_else(|err| { - warn!("failed to render memories stage-one input template: {err}"); - format!( - "Analyze this rollout and produce JSON with `raw_memory`, `rollout_summary`, and optional `rollout_slug`.\n\nrollout_context:\n- rollout_path: {rollout_path}\n- rollout_cwd: {rollout_cwd}\n\nrendered conversation:\n{rollout_contents}" - ) - }) + rollout_contents: &truncated_rollout_contents, + } + .render()?) } pub(crate) async fn build_memory_tool_developer_instructions(codex_home: &Path) -> Option { @@ -95,36 +81,24 @@ pub(crate) async fn build_memory_tool_developer_instructions(codex_home: &Path) template.render().ok() } -fn truncate_rollout_for_prompt(input: &str) -> (String, bool) { - if input.len() <= MAX_ROLLOUT_BYTES_FOR_PROMPT { - return (input.to_string(), false); - } - - let marker = "\n\n[... ROLLOUT TRUNCATED FOR MEMORY EXTRACTION ...]\n\n"; - let marker_len = marker.len(); - let budget_without_marker = MAX_ROLLOUT_BYTES_FOR_PROMPT.saturating_sub(marker_len); - let head_budget = budget_without_marker / 3; - let tail_budget = budget_without_marker.saturating_sub(head_budget); - let head = prefix_at_char_boundary(input, head_budget); - let tail = suffix_at_char_boundary(input, tail_budget); - let truncated = format!("{head}{marker}{tail}"); - - (truncated, true) -} - #[cfg(test)] mod tests { use super::*; #[test] - fn truncate_rollout_for_prompt_keeps_head_and_tail() { + fn build_stage_one_input_message_truncates_rollout_with_standard_policy() { let input = format!("{}{}{}", "a".repeat(700_000), "middle", "z".repeat(700_000)); - let (truncated, was_truncated) = truncate_rollout_for_prompt(&input); + let expected_truncated = truncate_text(&input, TruncationPolicy::Tokens(150_000)); + let message = build_stage_one_input_message( + Path::new("/tmp/rollout.jsonl"), + Path::new("/tmp"), + &input, + ) + .unwrap(); - assert!(was_truncated); - assert!(truncated.contains("[... ROLLOUT TRUNCATED FOR MEMORY EXTRACTION ...]")); - assert!(truncated.starts_with('a')); - assert!(truncated.ends_with('z')); - assert!(truncated.len() <= MAX_ROLLOUT_BYTES_FOR_PROMPT + 32); + assert!(expected_truncated.contains("tokens truncated")); + assert!(expected_truncated.starts_with('a')); + assert!(expected_truncated.ends_with('z')); + assert!(message.contains(&expected_truncated)); } } diff --git a/codex-rs/core/src/memories/stage_one.rs b/codex-rs/core/src/memories/stage_one.rs index 378a5569d..95aaba27e 100644 --- a/codex-rs/core/src/memories/stage_one.rs +++ b/codex-rs/core/src/memories/stage_one.rs @@ -2,18 +2,13 @@ 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; -use super::StageOneOutput; -use super::text::compact_whitespace; -use super::text::truncate_text_for_storage; - /// System prompt for stage-1 raw memory extraction. pub(super) const RAW_MEMORY_PROMPT: &str = include_str!("../../templates/memories/stage_one_system.md"); -const MAX_STAGE_ONE_RAW_MEMORY_CHARS: usize = 300_000; -const MAX_STAGE_ONE_SUMMARY_CHARS: usize = 1_200; static OPENAI_KEY_REGEX: Lazy = Lazy::new(|| compile_regex(r"sk-[A-Za-z0-9]{20,}")); static AWS_ACCESS_KEY_ID_REGEX: Lazy = Lazy::new(|| compile_regex(r"\bAKIA[0-9A-Z]{16}\b")); @@ -23,6 +18,24 @@ static SECRET_ASSIGNMENT_REGEX: Lazy = 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, +} + /// JSON schema used to constrain stage-1 model output. pub(super) fn stage_one_output_schema() -> Value { json!({ @@ -118,24 +131,8 @@ fn normalize_stage_one_output(mut output: StageOneOutput) -> Result MAX_STAGE_ONE_RAW_MEMORY_CHARS { - output.raw_memory = truncate_text_for_storage( - &output.raw_memory, - MAX_STAGE_ONE_RAW_MEMORY_CHARS, - "\n\n[... RAW MEMORY TRUNCATED ...]\n\n", - ); - } - - if output.rollout_summary.len() > MAX_STAGE_ONE_SUMMARY_CHARS { - output.rollout_summary = truncate_text_for_storage( - &output.rollout_summary, - MAX_STAGE_ONE_SUMMARY_CHARS, - " [...summary truncated...]", - ); - } + output.raw_memory = redact_secrets(&output.raw_memory); + output.rollout_summary = redact_secrets(&output.rollout_summary); Ok(output) } @@ -150,42 +147,10 @@ fn redact_secrets(input: &str) -> String { .to_string() } -fn normalize_raw_memory_structure(input: &str) -> String { - if has_raw_memory_structure(input) { - return input.to_string(); - } - - format!( - "# Raw Memory\n\ -Memory context: extracted from rollout (normalized fallback structure).\n\ -User preferences: none observed\n\n\ -## Task: Extracted Memory\n\ -Outcome: uncertain\n\ -Key steps:\n\ -- Review raw notes captured below.\n\ -Things that did not work / things that can be improved:\n\ -- Not clearly captured in structured form.\n\ -Reusable knowledge:\n\ -- Re-validate critical claims against the current rollout.\n\ -Pointers and references (annotate why each item matters):\n\ -- Raw memory notes included below.\n\n\ -### Raw memory notes\n\ -{input}\n" - ) -} - -fn has_raw_memory_structure(input: &str) -> bool { - let trimmed = input.trim(); - trimmed.starts_with('#') - && (trimmed.contains("Memory context:") || trimmed.contains("Trace context:")) - && trimmed.contains("User preferences:") - && trimmed.contains("## Task:") - && trimmed.contains("Outcome:") -} - 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}"), } } @@ -195,7 +160,13 @@ mod tests { use super::*; #[test] - fn normalize_stage_one_output_redacts_and_compacts_summary() { + 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(), @@ -208,20 +179,10 @@ mod tests { assert!(!normalized.rollout_summary.contains("mysecret123456")); assert_eq!( normalized.rollout_summary, - "password = [REDACTED_SECRET] small" + "password = [REDACTED_SECRET]\n\nsmall" ); } - #[test] - fn normalize_raw_memory_structure_wraps_unstructured_content() { - let normalized = normalize_raw_memory_structure("loose notes only"); - assert!(normalized.starts_with("# Raw Memory")); - assert!(normalized.contains("Memory context:")); - assert!(normalized.contains("## Task:")); - assert!(normalized.contains("Outcome: uncertain")); - assert!(normalized.contains("loose notes only")); - } - #[test] fn normalize_stage_one_output_allows_empty_pair_for_skip() { let output = StageOneOutput { diff --git a/codex-rs/core/src/memories/startup/extract.rs b/codex-rs/core/src/memories/startup/extract.rs index 91a0d2e65..b3833b8a2 100644 --- a/codex-rs/core/src/memories/startup/extract.rs +++ b/codex-rs/core/src/memories/startup/extract.rs @@ -12,9 +12,9 @@ use futures::StreamExt; use tracing::warn; use super::StageOneRequestContext; -use crate::memories::StageOneOutput; 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::rollout::policy::should_persist_response_item; @@ -61,7 +61,8 @@ pub(super) async fn extract_stage_one_output( id: None, role: "user".to_string(), content: vec![ContentItem::InputText { - text: build_stage_one_input_message(rollout_path, rollout_cwd, &rollout_contents), + text: build_stage_one_input_message(rollout_path, rollout_cwd, &rollout_contents) + .map_err(|_e| "error while building the prompt")?, }], end_turn: None, phase: None, diff --git a/codex-rs/core/src/memories/storage.rs b/codex-rs/core/src/memories/storage.rs index 5b67e229d..a366623b9 100644 --- a/codex-rs/core/src/memories/storage.rs +++ b/codex-rs/core/src/memories/storage.rs @@ -8,7 +8,6 @@ use super::MAX_RAW_MEMORIES_FOR_GLOBAL; use super::ensure_layout; use super::raw_memories_file; use super::rollout_summaries_dir; -use super::text::compact_whitespace; /// Rebuild `raw_memories.md` from DB-backed stage-1 outputs. pub(super) async fn rebuild_raw_memories_file_from_memories( @@ -139,7 +138,7 @@ async fn write_rollout_summary_for_thread( .map_err(|err| std::io::Error::other(format!("format rollout summary: {err}")))?; writeln!(body) .map_err(|err| std::io::Error::other(format!("format rollout summary: {err}")))?; - body.push_str(&compact_whitespace(&memory.rollout_summary)); + body.push_str(&memory.rollout_summary); body.push('\n'); tokio::fs::write(path, body).await diff --git a/codex-rs/core/src/memories/text.rs b/codex-rs/core/src/memories/text.rs deleted file mode 100644 index 213804ad3..000000000 --- a/codex-rs/core/src/memories/text.rs +++ /dev/null @@ -1,50 +0,0 @@ -pub(super) fn compact_whitespace(input: &str) -> String { - input.split_whitespace().collect::>().join(" ") -} - -pub(super) fn truncate_text_for_storage(input: &str, max_bytes: usize, marker: &str) -> String { - if input.len() <= max_bytes { - return input.to_string(); - } - - let budget_without_marker = max_bytes.saturating_sub(marker.len()); - let head_budget = budget_without_marker / 2; - let tail_budget = budget_without_marker.saturating_sub(head_budget); - let head = prefix_at_char_boundary(input, head_budget); - let tail = suffix_at_char_boundary(input, tail_budget); - - format!("{head}{marker}{tail}") -} - -pub(super) fn prefix_at_char_boundary(input: &str, max_bytes: usize) -> &str { - if max_bytes >= input.len() { - return input; - } - - let mut end = 0; - for (idx, _) in input.char_indices() { - if idx > max_bytes { - break; - } - end = idx; - } - - &input[..end] -} - -pub(super) fn suffix_at_char_boundary(input: &str, max_bytes: usize) -> &str { - if max_bytes >= input.len() { - return input; - } - - let start_limit = input.len().saturating_sub(max_bytes); - let mut start = input.len(); - for (idx, _) in input.char_indices().rev() { - if idx < start_limit { - break; - } - start = idx; - } - - &input[start..] -}