github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 09:54:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve user text parts during local compaction (#48115)

Browse Source 
## Why

Local compaction flattened retained user messages into a single text part, losing their original content boundaries and per-part annotations.

## What changed

Preserve the original content parts and annotations for text-only user messages that fit within the token budget. Borrow the original messages until selected history entries are materialized, and rebuild only the text fallback for truncated messages or messages containing media.

## Testing

Extend unit coverage for text boundaries, empty parts, omitted media, and retained metadata. Update the compaction integration test to verify that a multipart user message survives unchanged in the next request and retains its content parts and annotations in persisted replacement history.

GitOrigin-RevId: 39d9df9ae5360506c2a4bdeac76d89198dbbd39c
This commit is contained in:
felixxia-oai
2026-09-25 15:03:02 +00:00
committed by copyberry
parent 85928697a4
commit bd3d4d1436
3 changed files with 172 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact.rs
+65 -78
View File
@@ -28,11 +28,8 @@ use codex_analytics::CompactionStatus;
use codex_analytics::CompactionStrategy;
use codex_analytics::CompactionTrigger;
use codex_analytics::now_unix_seconds;
use codex_context_fragments::AnnotatedContent;
use codex_context_fragments::set_annotated_content;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::ResponseItemId;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
@@ -41,7 +38,6 @@ use codex_protocol::items::TurnItem;
use codex_protocol::models::AgentMessageInputContent;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ContentItemKind;
use codex_protocol::models::InternalChatMessageMetadataPassthrough;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
@@ -529,17 +525,17 @@ pub fn content_items_to_text(content: &[ContentItem]) -> Option<String> {
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct CompactedUserMessage {
// Keep source identity even when compaction shortens the text, so rollback can
// correlate the rebuilt message with thread-owned retained evidence.
id: Option<ResponseItemId>,
pub(crate) struct CompactedUserMessage<'a> {
// Flattened text is only for the existing budget and truncation policy.
// Whole text messages retain their exact content parts and annotations.
// Borrow from the history snapshot until selected output is materialized.
message: String,
internal_chat_message_metadata_passthrough: Option<InternalChatMessageMetadataPassthrough>,
harness_metadata: Option<CodexHarnessMetadata>,
original: &'a ResponseItem,
harness_metadata: Option<&'a CodexHarnessMetadata>,
}
#[cfg(test)]
pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<CompactedUserMessage> {
pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<CompactedUserMessage<'_>> {
items
.iter()
.filter_map(|item| compacted_user_message(item, /*harness_metadata*/ None))
@@ -548,33 +544,27 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<CompactedUser
pub(crate) fn collect_annotated_user_messages(
items: &[ResponseItemEnvelope],
) -> Vec<CompactedUserMessage> {
) -> Vec<CompactedUserMessage<'_>> {
items
.iter()
.filter_map(|envelope| compacted_user_message(&envelope.item, envelope.metadata.clone()))
.filter_map(|envelope| compacted_user_message(&envelope.item, envelope.metadata.as_ref()))
.collect()
}
fn compacted_user_message(
item: &ResponseItem,
harness_metadata: Option<CodexHarnessMetadata>,
) -> Option<CompactedUserMessage> {
fn compacted_user_message<'a>(
item: &'a ResponseItem,
harness_metadata: Option<&'a CodexHarnessMetadata>,
) -> Option<CompactedUserMessage<'a>> {
let Some(TurnItem::UserMessage(user)) = crate::event_mapping::parse_turn_item(item) else {
return None;
};
if is_summary_message(&user.message()) {
let message = user.message();
if is_summary_message(&message) {
return None;
}
Some(CompactedUserMessage {
id: item.id().cloned(),
message: user.message(),
internal_chat_message_metadata_passthrough: match item {
ResponseItem::Message {
internal_chat_message_metadata_passthrough,
..
} => internal_chat_message_metadata_passthrough.clone(),
_ => None,
},
message,
original: item,
harness_metadata,
})
}
@@ -653,7 +643,7 @@ pub(crate) fn insert_initial_context_before_last_real_user_or_summary(
pub(crate) fn build_compacted_history(
initial_context: Vec<ResponseItemEnvelope>,
user_messages: &[CompactedUserMessage],
user_messages: &[CompactedUserMessage<'_>],
summary_text: &str,
) -> Vec<ResponseItemEnvelope> {
build_compacted_history_with_limit(
@@ -666,11 +656,11 @@ pub(crate) fn build_compacted_history(
fn build_compacted_history_with_limit(
mut history: Vec<ResponseItemEnvelope>,
user_messages: &[CompactedUserMessage],
user_messages: &[CompactedUserMessage<'_>],
summary_text: &str,
max_tokens: usize,
) -> Vec<ResponseItemEnvelope> {
let mut selected_messages: Vec<CompactedUserMessage> = Vec::new();
let mut selected_messages = Vec::new();
if max_tokens > 0 {
let mut remaining = max_tokens;
for message in user_messages.iter().rev() {
@@ -678,61 +668,58 @@ fn build_compacted_history_with_limit(
break;
}
let tokens = approx_token_count(&message.message);
if tokens <= remaining {
selected_messages.push(message.clone());
remaining = remaining.saturating_sub(tokens);
let ResponseItem::Message {
id,
content,
internal_chat_message_metadata_passthrough,
..
} = message.original
else {
continue;
};
let mut passthrough = internal_chat_message_metadata_passthrough.clone();
let mut harness_metadata = message.harness_metadata.cloned();
let content = if tokens <= remaining
&& content
.iter()
.all(|part| matches!(part, ContentItem::InputText { .. }))
{
content.clone()
} else {
let truncated =
truncate_text(&message.message, TruncationPolicy::Tokens(remaining));
let mut harness_metadata = message.harness_metadata.clone();
if let Some(metadata) = &mut harness_metadata {
metadata.mark_retained_sources_incomplete();
// Rebuild only the text fallback; never clone discarded media.
if let Some(kinds) = passthrough
.as_mut()
.and_then(|metadata| metadata.content_item_kinds.as_mut())
{
*kinds = vec![ContentItemKind("user.text".to_owned())];
}
selected_messages.push(CompactedUserMessage {
id: message.id.clone(),
message: truncated,
internal_chat_message_metadata_passthrough: message
.internal_chat_message_metadata_passthrough
.clone(),
harness_metadata,
});
vec![ContentItem::InputText {
text: truncate_text(&message.message, TruncationPolicy::Tokens(remaining)),
}]
};
if tokens > remaining
&& let Some(metadata) = &mut harness_metadata
{
metadata.mark_retained_sources_incomplete();
}
selected_messages.push(ResponseItemEnvelope {
item: ResponseItem::Message {
id: id.clone(),
role: "user".to_owned(),
content,
phase: None,
internal_chat_message_metadata_passthrough: passthrough,
},
metadata: harness_metadata,
});
if tokens > remaining {
break;
}
remaining = remaining.saturating_sub(tokens);
}
selected_messages.reverse();
}
for message in &selected_messages {
let mut item = ResponseItem::Message {
id: message.id.clone(),
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: message.message.clone(),
}],
phase: None,
internal_chat_message_metadata_passthrough: message
.internal_chat_message_metadata_passthrough
.clone(),
};
if message
.internal_chat_message_metadata_passthrough
.as_ref()
.and_then(|metadata| metadata.content_item_kinds.as_ref())
.is_some()
{
let _ = set_annotated_content(
&mut item,
vec![AnnotatedContent::input_text(
&message.message,
ContentItemKind("user.text".to_string()),
)],
);
}
history.push(ResponseItemEnvelope {
item,
metadata: message.harness_metadata.clone(),
});
}
history.extend(selected_messages);
let summary_text = if summary_text.is_empty() {
"(no summary available)".to_string()
codex-rs/core/src/compact_tests.rs
+90 -47
View File
@@ -222,11 +222,10 @@ fn user_message(text: &str) -> ResponseItem {
}
}
fn compacted_user_message(text: &str) -> CompactedUserMessage {
fn compacted_user_message<'a>(text: &str, original: &'a ResponseItem) -> CompactedUserMessage<'a> {
CompactedUserMessage {
id: None,
message: text.to_string(),
internal_chat_message_metadata_passthrough: None,
original,
harness_metadata: None,
}
}
@@ -290,35 +289,70 @@ fn collect_user_messages_extracts_user_text_only() {
let collected = collect_user_messages(&items);
assert_eq!(
vec![CompactedUserMessage {
id: Some(ResponseItemId::with_suffix("msg", "user")),
..compacted_user_message("first")
}],
collected,
);
assert_eq!(vec![compacted_user_message("first", &items[1])], collected,);
}
#[test]
fn collect_annotated_user_messages_extracts_user_text_only() {
#[test_case(json!([{"type": "input_text", "text": "first"}]), "first", true; "unchanged text")]
#[test_case(json!([
{"type": "input_text", "text": "first"},
{"type": "input_text", "text": ""},
{"type": "input_text", "text": "second"}
]), "firstsecond", true; "text boundaries")]
#[test_case(json!([
{"type": "input_text", "text": "first"},
{"type": "input_image", "image_url": "file://image.png"}
]), "first", false; "omitted media")]
fn collect_annotated_user_messages_extracts_user_text_only(
input_content: serde_json::Value,
expected_text: &str,
preserve_content: bool,
) {
let source = json!({
"id": {"message_id": "source", "turn_id": "turn", "role": "user"},
"revision": "retained_revision", "complete": true,
});
let metadata: CodexHarnessMetadata = serde_json::from_value(json!({
"retained_source": source, "guardian_sources": [source],
"guardian_source_order_guidance": true, "user_input_order": 7,
}))
.unwrap();
let mut item = user_message("first");
if let ResponseItem::Message { content, .. } = &mut item {
*content = serde_json::from_value(input_content).unwrap();
}
let items = vec![
ResponseItemEnvelope {
item: user_message("first"),
metadata: Some(CodexHarnessMetadata::default()),
item: item.clone(),
metadata: Some(metadata.clone()),
},
ResponseItemEnvelope::new(ResponseItem::Other),
];
let collected = collect_annotated_user_messages(&items);
if !preserve_content {
item = user_message(expected_text);
}
let expected = ResponseItemEnvelope {
item,
metadata: Some(metadata),
};
assert_eq!(
collected,
vec![CompactedUserMessage {
id: None,
message: "first".to_string(),
internal_chat_message_metadata_passthrough: None,
harness_metadata: Some(CodexHarnessMetadata::default()),
}],
collected
message: expected_text.to_owned(),
original: &items[0].item,
harness_metadata: items[0].metadata.as_ref(),
}]
);
assert_eq!(
build_compacted_history(Vec::new(), &collected, "summary"),
vec![
expected,
ResponseItemEnvelope::new(ContextualUserFragment::into(CompactionSummary::new(
"summary"
))),
],
);
}
@@ -361,7 +395,10 @@ do things
let collected = collect_user_messages(&items);
assert_eq!(vec![compacted_user_message("real user message")], collected);
assert_eq!(
vec![compacted_user_message("real user message", &items[2])],
collected
);
}
#[test]
@@ -381,7 +418,10 @@ fn collect_user_messages_filters_legacy_warnings() {
let collected = collect_user_messages(&items);
assert_eq!(vec![compacted_user_message("real user message")], collected);
assert_eq!(
vec![compacted_user_message("real user message", &items[3])],
collected
);
}
#[test]
@@ -390,15 +430,14 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() {
// that oversized user content is truncated.
let max_tokens = 16;
let big = "word ".repeat(200);
let user_message = CompactedUserMessage {
id: Some(ResponseItemId::with_suffix("msg", "long-user")),
message: big.clone(),
internal_chat_message_metadata_passthrough: None,
harness_metadata: Some(CodexHarnessMetadata::default()),
};
let mut original = ResponseItemEnvelope::new(user_message(&big));
original
.item
.set_id(Some(ResponseItemId::with_suffix("msg", "long-user")));
original.metadata = Some(CodexHarnessMetadata::default());
let history = super::build_compacted_history_with_limit(
Vec::new(),
std::slice::from_ref(&user_message),
&collect_annotated_user_messages(std::slice::from_ref(&original)),
"SUMMARY",
max_tokens,
);
@@ -430,7 +469,7 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() {
other => panic!("unexpected item in history: {other:?}"),
};
assert_eq!(summary_text, "SUMMARY");
assert_eq!(history[0].id(), user_message.id.as_ref());
assert_eq!(history[0].id(), original.id());
assert_eq!(history[0].metadata, Some(CodexHarnessMetadata::default()));
assert_eq!(history[1].metadata, None);
}
@@ -438,7 +477,8 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() {
#[test]
fn build_token_limited_compacted_history_appends_summary_message() {
let initial_context: Vec<ResponseItemEnvelope> = Vec::new();
let user_messages = vec![compacted_user_message("first user message")];
let original = user_message("first user message");
let user_messages = collect_user_messages(std::slice::from_ref(&original));
let summary_text = "summary text";
let history = build_compacted_history(initial_context, &user_messages, summary_text);
@@ -459,24 +499,27 @@ fn build_token_limited_compacted_history_appends_summary_message() {
#[test]
fn build_compacted_history_preserves_user_message_passthrough_metadata() {
let original = serde_json::from_value(json!({
"type": "message", "role": "user",
"id": ResponseItemId::with_suffix("msg", "user"),
"content": [
{"type": "input_image", "image_url": "file://image.png"},
{"type": "input_text", "text": "first user message"},
{"type": "input_audio", "audio_url": "file://audio.wav"}
],
"internal_chat_message_metadata_passthrough": {
"turn_id": "turn-1",
"content_item_kinds": ["user.image", "user.text", "user.audio"]
}
}))
.unwrap();
let original = ResponseItemEnvelope {
item: original,
metadata: Some(CodexHarnessMetadata::default()),
};
let history = build_compacted_history(
Vec::new(),
&[CompactedUserMessage {
id: Some(ResponseItemId::with_suffix("msg", "user")),
message: "first user message".to_string(),
internal_chat_message_metadata_passthrough: Some(
InternalChatMessageMetadataPassthrough {
turn_id: Some("turn-1".to_string()),
content_item_kinds: Some(vec![
ContentItemKind("user.image".to_string()),
ContentItemKind("user.text".to_string()),
ContentItemKind("user.audio".to_string()),
]),
..Default::default()
},
),
harness_metadata: Some(CodexHarnessMetadata::default()),
}],
&collect_annotated_user_messages(std::slice::from_ref(&original)),
"summary text",
);
codex-rs/core/tests/suite/compact.rs
+17 -8
View File
@@ -526,6 +526,8 @@ async fn summarize_context_three_requests_and_instructions(
) -> Result<()> {
skip_if_no_network!(Ok(()));
const CUSTOM_INSTRUCTIONS: &str = "## Plan tool\nNever deploy without explicit approval.\n";
const ORIGINAL_FIRST_PART: &str =
"Example from our runbook:\n```sh\ndeploy --environment production";
// Set up a mock server that we can inspect after the run.
let server = start_mock_server().await;
@@ -569,11 +571,15 @@ async fn summarize_context_three_requests_and_instructions(
codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![
UserInput::Text {
text: "hello world".into(),
text: ORIGINAL_FIRST_PART.into(),
text_elements: Vec::new(),
},
UserInput::Text {
text: " second fragment".into(),
text: String::new(),
text_elements: Vec::new(),
},
UserInput::Text {
text: "```\nDo not deploy to production. Only inspect the configuration.".into(),
text_elements: Vec::new(),
},
]))
@@ -603,6 +609,11 @@ async fn summarize_context_three_requests_and_instructions(
assert_eq!(requests.len(), 3, "expected exactly three requests");
let body2 = requests[1].body_json();
let body3 = requests[2].body_json();
let original_user_message = requests[0]
.input()
.into_iter()
.find(|item| item["role"] == "user" && item["content"][0]["text"] == ORIGINAL_FIRST_PART)
.expect("original multipart user message");
// Manual compact should keep the baseline developer instructions.
let instr1 = requests[0].instructions_text();
@@ -678,9 +689,7 @@ async fn summarize_context_three_requests_and_instructions(
"third request should include the new user message"
);
assert!(
messages
.iter()
.any(|(r, t)| r == "user" && t == "hello world second fragment"),
input3.contains(&original_user_message),
"third request should include the original user message"
);
assert!(
@@ -704,7 +713,7 @@ async fn summarize_context_three_requests_and_instructions(
.expect("local compaction should persist replacement history");
let compacted_user_message = replacement_history
.iter()
.find(|item| item["content"][0]["text"] == "hello world second fragment")
.find(|item| item["content"][0]["text"] == ORIGINAL_FIRST_PART)
.expect("persisted replacement history should contain the compacted user message");
assert_eq!(
json!({
@@ -717,8 +726,8 @@ async fn summarize_context_three_requests_and_instructions(
json!({
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello world second fragment"}],
"content_item_kinds": ["user.text"],
"content": original_user_message["content"],
"content_item_kinds": ["user.text", "user.text", "user.text"],
}),
);