Adapt image compaction backport for pre-annotation releases

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::<Vec<_>>();
-    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::<usize>();
     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]);
     }
 }

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<ResponseItemEnvelope> {
-    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::<usize>();
             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)
 }