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
This commit is contained in:
jif
2026-09-08 11:12:01 +00:00
committed by copyberry
parent 5371951292
commit 6924ce636b
9 changed files with 465 additions and 104 deletions
+63
View File
@@ -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<ResponseItem> {
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;
+21
View File
@@ -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);
}
+2
View File
@@ -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;
+1
View File
@@ -11,6 +11,7 @@ mod metrics;
mod phase1;
mod phase2;
mod prompts;
mod rollout_input;
mod runtime;
mod start;
mod storage;
+9 -97
View File
@@ -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<TokenUsage>)> {
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<ResponseItem> {
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::<Vec<_>>();
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", "</INSTRUCTIONS>")
|| matches_marked_fragment(text, "<skill>", "</skill>")
}
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<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"# AGENTS.md instructions\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
true,
),
(
"<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>",
false,
),
(
"<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>",
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();
+11 -7
View File
@@ -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<String> {
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),
@@ -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<String> {
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::<RequestUserInputResponse>(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("<subagent_notification>")
|| (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, "<environment_context>", "</environment_context>")));
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<ResponseItem> {
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::<Vec<_>>();
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", "</INSTRUCTIONS>")
|| matches_marked_fragment(text, "<skill>", "</skill>")
}
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;
@@ -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<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"# AGENTS.md instructions\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
true,
),
(
"<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>",
false,
),
(
"<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>",
false,
),
];
for (text, expected) in cases {
assert_eq!(
is_memory_excluded_contextual_user_fragment(&ContentItem::InputText {
text: text.to_string(),
}),
expected,
"{text}",
);
}
}
@@ -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(())
}