diff --git a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs index 30f44b952e..9569571cc7 100644 --- a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs @@ -1,6 +1,4 @@ use super::*; -use codex_protocol::models::ContentItemKind; -use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::image_close_tag_text; use codex_protocol::models::local_image_open_tag_text_with_path; use pretty_assertions::assert_eq; @@ -61,7 +59,7 @@ fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { } #[test] -fn later_image_parts_preserve_labels_audio_and_annotations() { +fn later_image_parts_preserve_labels_and_audio() { let parts = vec![ text("earlier text"), text(&local_image_open_tag_text_with_path( @@ -75,41 +73,17 @@ fn later_image_parts_preserve_labels_audio_and_annotations() { audio_url: "data:audio/wav;base64,abc".to_string(), }, ]; - let kinds = (0..parts.len()) - .map(|i| ContentItemKind(format!("part.{i}"))) - .collect::>(); - let mut source = message(parts.clone()); - let ResponseItem::Message { - internal_chat_message_metadata_passthrough, - .. - } = &mut source - else { - unreachable!() - }; - *internal_chat_message_metadata_passthrough = Some(InternalChatMessageMetadataPassthrough { - turn_id: Some("turn-1".to_string()), - content_item_kinds: Some(kinds.clone()), - ..Default::default() - }); + let source = message(parts.clone()); let image_tokens = parts[1..4] .iter() .map(images::content_item_token_count) .sum::(); for (max_tokens, start) in [(image_tokens, 4), (image_tokens + 1, 1)] { let mut expected = source.clone(); - let ResponseItem::Message { - content, - internal_chat_message_metadata_passthrough, - .. - } = &mut expected - else { + let ResponseItem::Message { content, .. } = &mut expected else { unreachable!() }; *content = parts[start..].to_vec(); - internal_chat_message_metadata_passthrough - .as_mut() - .unwrap() - .content_item_kinds = Some(kinds[start..].to_vec()); assert_eq!(trim(vec![source.clone()], max_tokens), vec![expected]); } } diff --git a/codex-rs/core/src/compact_remote_v2_images.rs b/codex-rs/core/src/compact_remote_v2_images.rs index d2fecb03a2..63cf9309c0 100644 --- a/codex-rs/core/src/compact_remote_v2_images.rs +++ b/codex-rs/core/src/compact_remote_v2_images.rs @@ -1,9 +1,7 @@ use crate::context_manager::estimate_image_bytes; -use codex_context_fragments::AnnotatedContent; -use codex_context_fragments::set_annotated_content; -use codex_context_fragments::to_annotated_content; use codex_history::ResponseItemEnvelope; use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::models::is_image_close_tag_text; use codex_protocol::models::is_image_open_tag_text; use codex_protocol::models::is_local_image_open_tag_text; @@ -32,17 +30,20 @@ pub(super) fn truncate_message_to_token_budget( mut envelope: ResponseItemEnvelope, max_tokens: usize, ) -> Option { - let mut content = to_annotated_content(&mut envelope.item)?; + let ResponseItem::Message { content, .. } = &mut envelope.item else { + return None; + }; + let mut content = std::mem::take(content); let mut remaining = max_tokens; let mut retained = Vec::with_capacity(content.len()); while !content.is_empty() { let last = content.len() - 1; - let image_index = match content[last].content() { + let image_index = match &content[last] { ContentItem::InputImage { .. } => Some(last), ContentItem::InputText { text } if is_image_close_tag_text(text) && last > 0 - && matches!(content[last - 1].content(), ContentItem::InputImage { .. }) => + && matches!(&content[last - 1], ContentItem::InputImage { .. }) => { Some(last - 1) } @@ -51,14 +52,13 @@ pub(super) fn truncate_message_to_token_budget( if let Some(image_index) = image_index { let has_open_tag = image_index > 0 && matches!( - content[image_index - 1].content(), + &content[image_index - 1], ContentItem::InputText { text } if is_local_image_open_tag_text(text) || is_image_open_tag_text(text) ); let start = image_index - usize::from(has_open_tag); let token_count = content[start..] .iter() - .map(AnnotatedContent::content) .map(content_item_token_count) .sum::(); let fits = token_count <= remaining; @@ -71,7 +71,7 @@ pub(super) fn truncate_message_to_token_budget( continue; } let mut item = content.pop()?; - match item.content_mut() { + match &mut item { ContentItem::InputText { text } | ContentItem::OutputText { text } => { if remaining == 0 { continue; @@ -95,6 +95,9 @@ pub(super) fn truncate_message_to_token_budget( return None; } retained.reverse(); - set_annotated_content(&mut envelope.item, retained)?; + let ResponseItem::Message { content, .. } = &mut envelope.item else { + unreachable!("message content was already extracted above"); + }; + *content = retained; Some(envelope) }