From 6924ce636b2186948f2332ac9460b5871b50aa2f Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 8 Sep 2026 10:53:55 +0000 Subject: [PATCH] Prioritize human evidence in memory v2 extraction (#43799) ## Why Memory extraction has a limited rollout budget. Select evidence by source so human input and assistant final responses take priority over commentary, context, and tool output. ## What changed - Use tiered selection for memory v2, preferring newer evidence within each tier and rendering selected items in chronological order with omission markers. - Keep answered `request_user_input` questions alongside human replies, label evidence sources, replace message media with placeholders, and bound tool output and individual rows. - Add typed memory context fragments with bounded messages and Unicode-safe extraction chunking. ## Testing Add tests for memory v2 startup using tiered input and for extraction chunks preserving Unicode evidence within message bounds. GitOrigin-RevId: 1b2bba4e3264b4fda271e429ad3b7d5f69caed59 --- codex-rs/core/src/context/memory.rs | 63 ++++ codex-rs/core/src/context/memory_tests.rs | 21 ++ codex-rs/core/src/context/mod.rs | 2 + codex-rs/memories/write/src/lib.rs | 1 + codex-rs/memories/write/src/phase1.rs | 106 +----- codex-rs/memories/write/src/prompts.rs | 18 +- codex-rs/memories/write/src/rollout_input.rs | 305 ++++++++++++++++++ .../memories/write/src/rollout_input_tests.rs | 38 +++ codex-rs/memories/write/src/startup_tests.rs | 15 + 9 files changed, 465 insertions(+), 104 deletions(-) create mode 100644 codex-rs/core/src/context/memory.rs create mode 100644 codex-rs/core/src/context/memory_tests.rs create mode 100644 codex-rs/memories/write/src/rollout_input.rs create mode 100644 codex-rs/memories/write/src/rollout_input_tests.rs diff --git a/codex-rs/core/src/context/memory.rs b/codex-rs/core/src/context/memory.rs new file mode 100644 index 0000000000..3b1aefc1d6 --- /dev/null +++ b/codex-rs/core/src/context/memory.rs @@ -0,0 +1,63 @@ +//! Typed, bounded v2 memory context. Extraction chunks preserve selected evidence in order. + +use super::ContextualUserFragment; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::ResponseItem; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; + +/// Memory-owned model context, capped below 10k tokens even for byte-heavy text. +pub enum MemoryContextFragment { + ReadInstructions(String), + ExtractionEvidence(String), +} + +impl MemoryContextFragment { + /// Splits already-budgeted extraction input without dropping evidence. + pub fn extraction_messages(mut text: &str) -> Vec { + let mut messages = Vec::new(); + while !text.is_empty() { + let end = text.floor_char_boundary(text.len().min(8_900)); + messages.push(ContextualUserFragment::into(Self::ExtractionEvidence( + text[..end].to_string(), + ))); + text = &text[end..]; + } + messages + } +} + +impl ContextualUserFragment for MemoryContextFragment { + fn role(&self) -> &'static str { + match self { + Self::ReadInstructions(_) => "developer", + Self::ExtractionEvidence(_) => "user", + } + } + fn content_kind(&self) -> ContentItemKind { + ContentItemKind( + match self { + Self::ReadInstructions(_) => "memories.instructions", + Self::ExtractionEvidence(_) => "memories.extraction_evidence", + } + .to_string(), + ) + } + fn requires_separate_message(&self) -> bool { + true + } + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + fn body(&self) -> String { + let (Self::ReadInstructions(text) | Self::ExtractionEvidence(text)) = self; + truncate_text(text, TruncationPolicy::Bytes(8_900)) + } +} + +#[cfg(test)] +#[path = "memory_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/context/memory_tests.rs b/codex-rs/core/src/context/memory_tests.rs new file mode 100644 index 0000000000..b8b687533b --- /dev/null +++ b/codex-rs/core/src/context/memory_tests.rs @@ -0,0 +1,21 @@ +use super::*; +use codex_protocol::models::ContentItem; +use pretty_assertions::assert_eq; + +#[test] +fn extraction_chunks_preserve_unicode_evidence_with_bounded_messages() { + let evidence = "User correction: 🐈\n".repeat(2_000); + let mut reconstructed = String::new(); + for message in MemoryContextFragment::extraction_messages(&evidence) { + let ResponseItem::Message { role, content, .. } = message else { + panic!("message") + }; + assert_eq!(role, "user"); + let [ContentItem::InputText { text }] = content.as_slice() else { + panic!("text") + }; + assert!(text.len() < 9_000); + reconstructed.push_str(text); + } + assert_eq!(reconstructed, evidence); +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 23b5a97e9f..832cbce3b3 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -24,6 +24,7 @@ mod internal_model_context; mod legacy_apply_patch_exec_command_warning; mod legacy_model_mismatch_warning; mod legacy_unified_exec_process_limit_warning; +mod memory; mod model_switch_instructions; mod multi_agent_mode_instructions; mod multi_agent_role_instructions; @@ -87,6 +88,7 @@ pub use internal_model_context::InvalidInternalContextSource; pub(crate) use legacy_apply_patch_exec_command_warning::LegacyApplyPatchExecCommandWarning; pub(crate) use legacy_model_mismatch_warning::LegacyModelMismatchWarning; pub(crate) use legacy_unified_exec_process_limit_warning::LegacyUnifiedExecProcessLimitWarning; +pub use memory::MemoryContextFragment; pub(crate) use model_switch_instructions::ModelSwitchInstructions; pub(crate) use multi_agent_mode_instructions::MultiAgentModeInstructions; pub(crate) use multi_agent_role_instructions::MultiAgentRoleInstructions; diff --git a/codex-rs/memories/write/src/lib.rs b/codex-rs/memories/write/src/lib.rs index c4e4da6463..616753603f 100644 --- a/codex-rs/memories/write/src/lib.rs +++ b/codex-rs/memories/write/src/lib.rs @@ -11,6 +11,7 @@ mod metrics; mod phase1; mod phase2; mod prompts; +mod rollout_input; mod runtime; mod start; mod storage; diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 9f7d7ff1ea..a66e36d542 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -3,12 +3,14 @@ use crate::metrics::MEMORY_PHASE_ONE_E2E_MS; use crate::metrics::MEMORY_PHASE_ONE_JOBS; use crate::metrics::MEMORY_PHASE_ONE_OUTPUT; use crate::metrics::MEMORY_PHASE_ONE_TOKEN_USAGE; +use crate::rollout_input::sanitize_response_item_for_memories; use crate::runtime::MemoryStartupContext; use crate::runtime::StageOneRequestContext; use codex_config::types::MemoriesConfig; use codex_core::Prompt; use codex_core::RolloutRecorder; use codex_core::config::Config; +use codex_protocol::MemoryVersion; use codex_protocol::ResponseItemId; use codex_protocol::error::CodexErr; use codex_protocol::models::BaseInstructions; @@ -17,7 +19,6 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; use codex_rollout::INTERACTIVE_SESSION_SOURCES; use codex_rollout::RolloutItem; -use codex_rollout::should_persist_response_item_for_memories; use codex_secrets::redact_secrets; use futures::StreamExt; use serde::Deserialize; @@ -287,7 +288,13 @@ mod job { stage_one_context: &StageOneRequestContext, ) -> anyhow::Result<(StageOneOutput, Option)> { let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; - let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?; + let rollout_contents = match config.memories.version { + MemoryVersion::V1 => serialize_filtered_rollout_response_items(&rollout_items)?, + MemoryVersion::V2 => crate::rollout_input::serialize_tiered_input( + &rollout_items, + crate::prompts::rollout_token_limit(&stage_one_context.model_info), + )?, + }; let mut prompt = Prompt::default(); prompt.input = vec![ResponseItem::Message { @@ -426,105 +433,10 @@ mod job { Ok(redact_secrets(serialized)) } - fn sanitize_response_item_for_memories(item: &ResponseItem) -> Option { - let ResponseItem::Message { - id, - role, - content, - phase, - internal_chat_message_metadata_passthrough: metadata, - } = item - else { - return should_persist_response_item_for_memories(item).then(|| item.clone()); - }; - - if role == "developer" { - return None; - } - - if role != "user" { - return Some(item.clone()); - } - - let content = content - .iter() - .filter(|content_item| !is_memory_excluded_contextual_user_fragment(content_item)) - .cloned() - .collect::>(); - if content.is_empty() { - return None; - } - - Some(ResponseItem::Message { - id: id.clone(), - role: role.clone(), - content, - phase: phase.clone(), - internal_chat_message_metadata_passthrough: metadata.clone(), - }) - } - - fn is_memory_excluded_contextual_user_fragment(content_item: &ContentItem) -> bool { - let ContentItem::InputText { text } = content_item else { - return false; - }; - - matches_marked_fragment(text, "# AGENTS.md instructions", "") - || matches_marked_fragment(text, "", "") - } - - fn matches_marked_fragment(text: &str, start_marker: &str, end_marker: &str) -> bool { - let trimmed = text.trim_start(); - let starts_with_marker = trimmed - .get(..start_marker.len()) - .is_some_and(|candidate| candidate.eq_ignore_ascii_case(start_marker)); - let trimmed = trimmed.trim_end(); - let ends_with_marker = trimmed - .get(trimmed.len().saturating_sub(end_marker.len())..) - .is_some_and(|candidate| candidate.eq_ignore_ascii_case(end_marker)); - starts_with_marker && ends_with_marker - } - #[cfg(test)] mod tests { use super::*; - #[test] - fn classifies_memory_excluded_fragments() { - let cases = [ - ( - "# AGENTS.md instructions for /tmp\n\n\nbody\n", - true, - ), - ( - "# AGENTS.md instructions\n\n\nbody\n", - true, - ), - ( - "\ndemo\nskills/demo/SKILL.md\nbody\n", - true, - ), - ( - "\n/tmp\n", - false, - ), - ( - "{\"agent_id\":\"a\",\"status\":\"completed\"}", - false, - ), - ]; - - for (text, expected) in cases { - assert_eq!( - is_memory_excluded_contextual_user_fragment(&ContentItem::InputText { - text: text.to_string(), - }), - expected, - "{text}", - ); - } - } - #[test] fn output_schema_requires_rollout_slug_and_keeps_it_nullable() { let schema = output_schema(); diff --git a/codex-rs/memories/write/src/prompts.rs b/codex-rs/memories/write/src/prompts.rs index e0019f610f..9b536c4b87 100644 --- a/codex-rs/memories/write/src/prompts.rs +++ b/codex-rs/memories/write/src/prompts.rs @@ -95,6 +95,16 @@ fn render_memory_extensions_block(template: &Template, memory_extensions_root: & }) } +pub(crate) fn rollout_token_limit(model_info: &ModelInfo) -> usize { + model_info + .resolved_context_window() + .and_then(|limit| (limit > 0).then_some(limit)) + .map(|limit| limit.saturating_mul(model_info.effective_context_window_percent) / 100) + .map(|limit| (limit.saturating_mul(crate::stage_one::CONTEXT_WINDOW_PERCENT) / 100).max(1)) + .and_then(|limit| usize::try_from(limit).ok()) + .unwrap_or(crate::stage_one::DEFAULT_ROLLOUT_TOKEN_LIMIT) +} + /// Builds the stage-1 user message containing rollout metadata and content. /// /// Large rollout payloads are truncated to 70% of the active model's effective @@ -105,13 +115,7 @@ pub fn build_stage_one_input_message( rollout_cwd: &Path, rollout_contents: &str, ) -> anyhow::Result { - let rollout_token_limit = model_info - .resolved_context_window() - .and_then(|limit| (limit > 0).then_some(limit)) - .map(|limit| limit.saturating_mul(model_info.effective_context_window_percent) / 100) - .map(|limit| (limit.saturating_mul(crate::stage_one::CONTEXT_WINDOW_PERCENT) / 100).max(1)) - .and_then(|limit| usize::try_from(limit).ok()) - .unwrap_or(crate::stage_one::DEFAULT_ROLLOUT_TOKEN_LIMIT); + let rollout_token_limit = rollout_token_limit(model_info); let truncated_rollout_contents = truncate_text( rollout_contents, TruncationPolicy::Tokens(rollout_token_limit), diff --git a/codex-rs/memories/write/src/rollout_input.rs b/codex-rs/memories/write/src/rollout_input.rs new file mode 100644 index 0000000000..1e99325cfd --- /dev/null +++ b/codex-rs/memories/write/src/rollout_input.rs @@ -0,0 +1,305 @@ +//! Selects v2 extraction evidence by provenance, then restores chronological order. +//! Human replies retain their questions; message media is replaced with placeholders. +//! Budgets use the same token estimate as the rest of the memory pipeline. + +use codex_protocol::ToolName; +use codex_protocol::models::ContentItem; +use codex_protocol::models::MessagePhase; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::request_user_input::RequestUserInputResponse; +use codex_rollout::RolloutItem; +use codex_rollout::should_persist_response_item_for_memories; +use codex_secrets::redact_secrets; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_bytes_for_tokens; +use codex_utils_output_truncation::truncate_text; +use std::collections::HashMap; + +const OMITTED: &str = "[... response items omitted ...]\n"; +const TRUNCATION_RESERVE_BYTES: usize = 96; +const TOOL_OUTPUT_TOKENS: usize = 2_000; +const MAX_ROW_BYTES: usize = 10_000; + +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum Tier { + Human, + Final, + OtherAgent, + Commentary, + Context, + Tool, +} + +/// Select newest evidence first within each tier, but render in source order. +pub(crate) fn serialize_tiered_input( + items: &[RolloutItem], + token_limit: usize, +) -> anyhow::Result { + let mut rows = Vec::new(); + let mut user_input_calls = HashMap::new(); + for item in items { + let item = match item { + RolloutItem::ResponseItem(item) => sanitize_response_item_for_memories(&item.item), + RolloutItem::InterAgentCommunication(message) => Some(message.to_model_input_item()), + RolloutItem::SessionMeta(_) + | RolloutItem::InterAgentCommunicationMetadata { .. } + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::RealtimeItem(_) + | RolloutItem::WorldState(_) + | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) + | RolloutItem::RetainedContext(_) + | RolloutItem::EventMsg(_) => None, + }; + let Some(mut item) = item else { continue }; + if let ResponseItem::Message { content, .. } = &mut item { + for part in content { + match part { + ContentItem::InputImage { .. } => { + *part = ContentItem::InputText { + text: "[image omitted]".into(), + }; + } + ContentItem::InputAudio { .. } => { + *part = ContentItem::InputText { + text: "[audio omitted]".into(), + }; + } + ContentItem::InputText { .. } | ContentItem::OutputText { .. } => {} + } + } + } + let question = match &item { + ResponseItem::FunctionCall { + name, + namespace, + call_id, + arguments, + .. + } if name == "request_user_input" + && ToolName::new(namespace.clone(), name).is_default_namespace() => + { + user_input_calls.insert(call_id.clone(), arguments.clone()); + None + } + ResponseItem::FunctionCallOutput { + call_id: Some(call_id), + output, + .. + } if user_input_calls.contains_key(call_id) + && output + .text_content() + .and_then(|text| serde_json::from_str::(text).ok()) + .is_some_and(|reply| { + reply + .answers + .values() + .any(|answer| answer.answers.iter().any(|text| !text.trim().is_empty())) + }) => + { + user_input_calls.remove(call_id) + } + _ => None, + }; + let tier = match &item { + ResponseItem::FunctionCallOutput { .. } if question.is_some() => Tier::Human, + ResponseItem::AgentMessage { .. } => Tier::OtherAgent, + ResponseItem::Message { + role, + content, + phase, + internal_chat_message_metadata_passthrough, + .. + } => { + let agent_content = internal_chat_message_metadata_passthrough + .as_ref() + .and_then(|metadata| metadata.content_item_kinds.as_ref()) + .is_some_and(|kinds| { + kinds.iter().any(|kind| kind.0.starts_with("multi_agent.")) + }) + || InterAgentCommunication::is_message_content(content) + || content.iter().any(|part| match part { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + let text = text.trim_start(); + text.starts_with("") + || (text.starts_with("Message Type:") + && text + .lines() + .nth(1) + .is_some_and(|line| line.starts_with("Task name:")) + && text + .lines() + .nth(2) + .is_some_and(|line| line.starts_with("Sender:")) + && text.lines().nth(3) == Some("Payload:")) + } + _ => false, + }); + if agent_content { + Tier::OtherAgent + } else if role == "user" { + let contextual = internal_chat_message_metadata_passthrough.as_ref() + .and_then(|metadata| metadata.content_item_kinds.as_ref()) + .is_some_and(|kinds| !kinds.is_empty() && kinds.iter().all(|kind| !kind.0.starts_with("user."))) + || content.iter().any(|part| matches!(part, ContentItem::InputText { text } + if matches_marked_fragment(text, "", ""))); + if contextual { + Tier::Context + } else { + Tier::Human + } + } else if matches!(phase, Some(MessagePhase::Commentary)) { + Tier::Commentary + } else { + Tier::Final + } + } + _ => Tier::Tool, + }; + let label = match tier { + Tier::Human => "human user", + Tier::Final => "assistant final", + Tier::Commentary => "assistant commentary", + Tier::OtherAgent => "other agent", + Tier::Context => "harness context", + Tier::Tool => "tool", + }; + let text = serde_json::to_string(&item)?; + let mut text = redact_secrets(match question { + Some(question) => format!("Assistant question: {question}\nHuman reply: {text}"), + None => text, + }); + if tier == Tier::Tool + && matches!( + item, + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::ToolSearchOutput { .. } + ) + { + text = truncate_text( + &text, + TruncationPolicy::Bytes( + approx_bytes_for_tokens(TOOL_OUTPUT_TOKENS) + .saturating_sub(TRUNCATION_RESERVE_BYTES), + ), + ); + } + text = truncate_text( + &text, + TruncationPolicy::Bytes(MAX_ROW_BYTES - TRUNCATION_RESERVE_BYTES), + ); + rows.push((tier, format!("[{label}]\n{text}\n"))); + } + let mut remaining = approx_bytes_for_tokens(token_limit).saturating_sub(OMITTED.len()); + let mut selected = vec![None; rows.len()]; + for tier in [ + Tier::Human, + Tier::Final, + Tier::OtherAgent, + Tier::Commentary, + Tier::Context, + Tier::Tool, + ] { + for (index, (row_tier, text)) in rows.iter().enumerate().rev() { + if *row_tier != tier || remaining <= OMITTED.len() + TRUNCATION_RESERVE_BYTES { + continue; + } + let budget = remaining - OMITTED.len(); + let text = if text.len() <= budget { + text.clone() + } else { + truncate_text( + text, + TruncationPolicy::Bytes(budget - TRUNCATION_RESERVE_BYTES), + ) + }; + remaining = remaining.saturating_sub(text.len() + OMITTED.len()); + selected[index] = Some(text); + } + } + let mut rendered = String::new(); + let mut gap = false; + for row in selected { + match row { + Some(text) => { + rendered.push_str(&text); + gap = false; + } + None if !gap => { + rendered.push_str(OMITTED); + gap = true; + } + None => {} + } + } + if rendered.len() > approx_bytes_for_tokens(token_limit) { + return Ok(String::new()); + } + Ok(rendered) +} + +pub(crate) fn sanitize_response_item_for_memories(item: &ResponseItem) -> Option { + let ResponseItem::Message { + id, + role, + content, + phase, + internal_chat_message_metadata_passthrough: metadata, + } = item + else { + return should_persist_response_item_for_memories(item).then(|| item.clone()); + }; + + if role == "developer" { + return None; + } + + if role != "user" { + return Some(item.clone()); + } + + let content = content + .iter() + .filter(|content_item| !is_memory_excluded_contextual_user_fragment(content_item)) + .cloned() + .collect::>(); + if content.is_empty() { + return None; + } + + Some(ResponseItem::Message { + id: id.clone(), + role: role.clone(), + content, + phase: phase.clone(), + internal_chat_message_metadata_passthrough: metadata.clone(), + }) +} + +fn is_memory_excluded_contextual_user_fragment(content_item: &ContentItem) -> bool { + let ContentItem::InputText { text } = content_item else { + return false; + }; + + matches_marked_fragment(text, "# AGENTS.md instructions", "") + || matches_marked_fragment(text, "", "") +} + +fn matches_marked_fragment(text: &str, start_marker: &str, end_marker: &str) -> bool { + let trimmed = text.trim_start(); + let starts_with_marker = trimmed + .get(..start_marker.len()) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(start_marker)); + let trimmed = trimmed.trim_end(); + let ends_with_marker = trimmed + .get(trimmed.len().saturating_sub(end_marker.len())..) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(end_marker)); + starts_with_marker && ends_with_marker +} + +#[cfg(test)] +#[path = "rollout_input_tests.rs"] +mod tests; diff --git a/codex-rs/memories/write/src/rollout_input_tests.rs b/codex-rs/memories/write/src/rollout_input_tests.rs new file mode 100644 index 0000000000..c2f4baafd8 --- /dev/null +++ b/codex-rs/memories/write/src/rollout_input_tests.rs @@ -0,0 +1,38 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn classifies_memory_excluded_fragments() { + let cases = [ + ( + "# AGENTS.md instructions for /tmp\n\n\nbody\n", + true, + ), + ( + "# AGENTS.md instructions\n\n\nbody\n", + true, + ), + ( + "\ndemo\nskills/demo/SKILL.md\nbody\n", + true, + ), + ( + "\n/tmp\n", + false, + ), + ( + "{\"agent_id\":\"a\",\"status\":\"completed\"}", + false, + ), + ]; + + for (text, expected) in cases { + assert_eq!( + is_memory_excluded_contextual_user_fragment(&ContentItem::InputText { + text: text.to_string(), + }), + expected, + "{text}", + ); + } +} diff --git a/codex-rs/memories/write/src/startup_tests.rs b/codex-rs/memories/write/src/startup_tests.rs index 4771d42029..87190d1d68 100644 --- a/codex-rs/memories/write/src/startup_tests.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -1248,3 +1248,18 @@ async fn shutdown_test_codex(test: &TestCodex) -> anyhow::Result<()> { wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::ShutdownComplete)).await; Ok(()) } + +#[tokio::test] +async fn memories_startup_phase1_v2_uses_tiered_input() -> anyhow::Result<()> { + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let mut memories = startup_test_memories_config(); + memories.version = codex_protocol::MemoryVersion::V2; + let request = run_memory_phase_one_model_request_test(&server, home, memories).await?; + assert!( + request.body_json()["input"] + .to_string() + .contains("[human user]") + ); + Ok(()) +}