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"], }), );