core/core-agent-ide
8
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

nit: memory truncation (#11479)

Browse source 
Use existing truncation for memories
This commit is contained in:
jif-oai 2026-02-11 21:11:57 +00:00 • committed by GitHub
parent 444324175e
commit de6f2ef746
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 58 additions and 194 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
codex-rs/core/src/codex.rs
View file

@ -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)]

23
codex-rs/core/src/memories/mod.rs
View file

@ -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<String>,
}
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.

68
codex-rs/core/src/memories/prompts.rs
View file

@ -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<String> {
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<String> {
@ -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));
}
}

99
codex-rs/core/src/memories/stage_one.rs
View file

@ -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<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"));
@ -23,6 +18,24 @@ 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!({
@ -118,24 +131,8 @@ fn normalize_stage_one_output(mut output: StageOneOutput) -> Result<StageOneOutp
));
}
output.raw_memory = normalize_raw_memory_structure(&redact_secrets(&output.raw_memory));
output.rollout_summary = redact_secrets(&compact_whitespace(&output.rollout_summary));
if output.raw_memory.len() > 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 {

5
codex-rs/core/src/memories/startup/extract.rs
View file

@ -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,

3
codex-rs/core/src/memories/storage.rs
View file

@ -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

50
codex-rs/core/src/memories/text.rs
View file

@ -1,50 +0,0 @@
pub(super) fn compact_whitespace(input: &str) -> String {
input.split_whitespace().collect::<Vec<_>>().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..]
}