From bd3d4d1436bb41b94fd38ba9bdd34d74524e7a9f Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Fri, 25 Sep 2026 14:44:23 +0000 Subject: [PATCH] Preserve user text parts during local compaction (#48115) ## 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 --- codex-rs/core/src/compact.rs | 143 ++++++++++++--------------- codex-rs/core/src/compact_tests.rs | 137 ++++++++++++++++--------- codex-rs/core/tests/suite/compact.rs | 25 +++-- 3 files changed, 172 insertions(+), 133 deletions(-) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 7d7c9c3fff..f9fe78dbaf 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -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 { } #[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, +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, - harness_metadata: Option, + original: &'a ResponseItem, + harness_metadata: Option<&'a CodexHarnessMetadata>, } #[cfg(test)] -pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec { +pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec> { 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 Vec { +) -> Vec> { 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, -) -> Option { +fn compacted_user_message<'a>( + item: &'a ResponseItem, + harness_metadata: Option<&'a CodexHarnessMetadata>, +) -> Option> { 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, - user_messages: &[CompactedUserMessage], + user_messages: &[CompactedUserMessage<'_>], summary_text: &str, ) -> Vec { build_compacted_history_with_limit( @@ -666,11 +656,11 @@ pub(crate) fn build_compacted_history( fn build_compacted_history_with_limit( mut history: Vec, - user_messages: &[CompactedUserMessage], + user_messages: &[CompactedUserMessage<'_>], summary_text: &str, max_tokens: usize, ) -> Vec { - let mut selected_messages: Vec = 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() diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 320b114817..217d6b63fd 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -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 = 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", ); diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 786b8d93d0..0a4494ad62 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -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"], }), );