diff --git a/.github/scripts/install-musl-build-tools.sh b/.github/scripts/install-musl-build-tools.sh index a7d37b7ff9..83b5a81360 100644 --- a/.github/scripts/install-musl-build-tools.sh +++ b/.github/scripts/install-musl-build-tools.sh @@ -36,6 +36,7 @@ libcap_version="2.75" libcap_sha256="de4e7e064c9ba451d5234dd46e897d7c71c96a9ebf9a0c445bc04f4742d83632" libcap_tarball_name="libcap-${libcap_version}.tar.xz" libcap_download_url="https://mirrors.edge.kernel.org/pub/linux/libs/security/linux-privs/libcap2/${libcap_tarball_name}" +libcap_debian_url="https://deb.debian.org/debian/pool/main/libc/libcap2/libcap2_${libcap_version}.orig.tar.xz" # Use the musl toolchain as the Rust linker to avoid Zig injecting its own CRT. if command -v "${arch}-linux-musl-gcc" >/dev/null; then @@ -63,7 +64,9 @@ if [[ ! -f "${libcap_prefix}/lib/libcap.a" ]]; then mkdir -p "${libcap_src_root}" "${libcap_prefix}/lib" "${libcap_prefix}/include/sys" "${libcap_prefix}/include/linux" "${libcap_pkgconfig_dir}" libcap_tarball="${libcap_root}/${libcap_tarball_name}" - curl -fsSL "${libcap_download_url}" -o "${libcap_tarball}" + if ! curl -fsSL --connect-timeout 10 --max-time 60 "${libcap_download_url}" -o "${libcap_tarball}"; then + curl -fsSL --connect-timeout 10 --max-time 60 --retry 2 "${libcap_debian_url}" -o "${libcap_tarball}" + fi echo "${libcap_sha256} ${libcap_tarball}" | sha256sum -c - tar -xJf "${libcap_tarball}" -C "${libcap_src_root}" diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 619af67d14..2abeee12b6 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -47,6 +47,7 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::items::ContextCompactionItem; use codex_protocol::items::TurnItem; use codex_protocol::models::AgentMessageInputContent; +use codex_protocol::models::AudioReference; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; @@ -75,8 +76,8 @@ use attempt::run_remote_compact_v2_attempt; mod installation; use installation::run_compaction; -#[path = "compact_remote_v2_images.rs"] -mod images; +#[path = "compact_remote_v2_media.rs"] +mod media; #[derive(Clone, Copy, PartialEq, Eq)] enum RetainedImageBudget { @@ -640,18 +641,26 @@ fn truncate_retained_messages( } let client_developer = is_client_authored_developer_message(&group.source); - let charge_images = image_budget == RetainedImageBudget::Enabled && !client_developer; - let notice_tokens = group - .attached_notice - .as_ref() - .map_or(0, |notice| message_text_token_count(¬ice.item).max(1)); + let image_budget = if client_developer { + RetainedImageBudget::Disabled + } else { + image_budget + }; + let has_budgeted_image = image_budget == RetainedImageBudget::Enabled + && retained_input_image_count(&group.source.item) > 0; + let has_file_audio = matches!(&group.source.item, ResponseItem::Message { content, .. } + if content.iter().any(|part| matches!(part, ContentItem::InputAudio { + audio: AudioReference::File { .. } + }))); + // Charged media is truncated newest-first as a message boundary. Each + // modality still defines its own grouping and token cost in `media`. + let requires_media_boundary_truncation = has_budgeted_image || has_file_audio; + let notice_tokens = group.attached_notice.as_ref().map_or(0, |notice| { + retained_message_token_count(¬ice.item, RetainedImageBudget::Disabled).max(1) + }); // Client-authored developer messages already charge non-text content via // the serialized estimate. Preserve their text-only boundary correction. - let content_tokens = if charge_images { - message_content_token_count(&group.source.item) - } else { - message_text_token_count(&group.source.item) - }; + let content_tokens = retained_message_token_count(&group.source.item, image_budget); let source_tokens = if client_developer { usize::try_from(estimate_item_token_count(&group.source.item)).unwrap_or(usize::MAX) } else { @@ -674,14 +683,13 @@ fn truncate_retained_messages( } else { available_tokens }; - let image_count = retained_input_image_count(&group.source.item); - if charge_images && image_count > 0 { - // An oversized image can leave no boundary content. Do not backfill + if requires_media_boundary_truncation { + // Oversized media can leave no boundary content. Do not backfill // the remaining budget with older messages in that case. remaining = 0; } - let truncated_item = if charge_images && image_count > 0 { - images::truncate_message_to_token_budget(group.source, content_budget) + let truncated_item = if requires_media_boundary_truncation { + media::truncate_message_to_token_budget(group.source, content_budget, image_budget) } else { truncate_message_text_to_token_budget(group.source, content_budget) }; @@ -695,9 +703,16 @@ fn truncate_retained_messages( let adjusted_budget = content_budget .saturating_sub(item_tokens - available_tokens) .saturating_sub(1); - let Some(adjusted) = + let adjusted = if requires_media_boundary_truncation { + media::truncate_message_to_token_budget( + truncated_item, + adjusted_budget, + image_budget, + ) + } else { truncate_message_text_to_token_budget(truncated_item, adjusted_budget) - else { + }; + let Some(adjusted) = adjusted else { continue; }; if usize::try_from(estimate_item_token_count(&adjusted.item)) @@ -714,7 +729,7 @@ fn truncate_retained_messages( } truncated_reversed.push(truncated_item); remaining = 0; - } else if charge_images && retained_input_image_count(&group.source.item) > 0 { + } else if requires_media_boundary_truncation { remaining = 0; } } @@ -722,28 +737,15 @@ fn truncate_retained_messages( truncated_reversed } -fn message_content_token_count(item: &ResponseItem) -> usize { - let ResponseItem::Message { content, .. } = item else { - return usize::try_from(estimate_item_token_count(item)).unwrap_or(usize::MAX); - }; - - content.iter().map(images::content_item_token_count).sum() -} - -fn message_text_token_count(item: &ResponseItem) -> usize { +fn retained_message_token_count(item: &ResponseItem, image_budget: RetainedImageBudget) -> usize { let ResponseItem::Message { content, .. } = item else { return usize::try_from(estimate_item_token_count(item)).unwrap_or(usize::MAX); }; content .iter() - .map(|item| match item { - ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { - approx_token_count(text) - } - ContentItem::InputImage { .. } | ContentItem::InputAudio { .. } => 0, - }) - .sum() + .map(|part| media::content_item_token_count(part, image_budget)) + .fold(0, usize::saturating_add) } fn truncate_message_text_to_token_budget( @@ -1470,5 +1472,5 @@ mod tests { } #[cfg(test)] -#[path = "compact_remote_v2_image_budget_tests.rs"] -mod image_budget_tests; +#[path = "compact_remote_v2_media_budget_tests.rs"] +mod media_budget_tests; 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 deleted file mode 100644 index cec85187ad..0000000000 --- a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs +++ /dev/null @@ -1,184 +0,0 @@ -use super::*; -use codex_protocol::models::AudioReference; -use codex_protocol::models::ContentItemKind; -use codex_protocol::models::ImageReference; -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; - -fn message(content: Vec) -> ResponseItem { - ResponseItem::Message { - status: None, - encrypted_content: None, - id: None, - role: "user".to_string(), - content, - phase: None, - internal_chat_message_metadata_passthrough: None, - } -} - -fn image() -> ContentItem { - ContentItem::InputImage { - image: ImageReference::Inline { - image_url: "data:image/png;base64,abc".to_string(), - }, - detail: None, - } -} - -fn file_image() -> ContentItem { - ContentItem::InputImage { - image: ImageReference::File { - file_id: "file_123".to_string(), - }, - detail: Some(codex_protocol::models::ImageDetail::Original), - } -} - -fn text(value: &str) -> ContentItem { - ContentItem::InputText { - text: value.to_string(), - } -} - -fn trim(items: Vec, max_tokens: usize) -> Vec { - truncate_retained_messages( - items.into_iter().map(ResponseItemEnvelope::new).collect(), - max_tokens, - RetainedImageBudget::Enabled, - ) - .into_iter() - .map(ResponseItemEnvelope::into_item) - .collect() -} - -#[test] -fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { - let newest = message(vec![text("new")]); - for image in [image(), file_image()] { - let items = vec![ - message(vec![text("old")]), - message(vec![image.clone()]), - newest.clone(), - ]; - let image_tokens = images::content_item_token_count(&image); - for (max_tokens, expected) in [ - ( - image_tokens + 1, - vec![message(vec![image.clone()]), newest.clone()], - ), - (image_tokens, vec![newest.clone()]), - (1, vec![newest.clone()]), - ] { - assert_eq!(trim(items.clone(), max_tokens), expected); - } - } -} - -#[test] -fn later_image_parts_preserve_labels_audio_and_annotations() { - let parts = vec![ - text("earlier text"), - text(&local_image_open_tag_text_with_path( - /*label_number*/ 7, - std::path::Path::new("image.png"), - )), - image(), - text(&image_close_tag_text()), - text("keep"), - ContentItem::InputAudio { - audio: AudioReference::Inline { - 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 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 { - 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]); - } -} - -#[test] -fn image_treatment_preserves_text_only_boundary_behavior() { - let source = ResponseItemEnvelope::new(message(vec![ - text("earlier text"), - text(&"middle".repeat(20)), - text("later text"), - ContentItem::InputAudio { - audio: AudioReference::Inline { - audio_url: "data:audio/wav;base64,abc".to_string(), - }, - }, - ])); - for max_tokens in [0, 1, 4, 16, 100] { - assert_eq!( - truncate_retained_messages( - vec![source.clone()], - max_tokens, - RetainedImageBudget::Enabled, - ), - truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), - ); - } -} - -#[test] -fn image_treatment_preserves_client_developer_boundary_behavior() { - let mut item = message(vec![image(), text(&"a".repeat(1000))]); - let ResponseItem::Message { role, .. } = &mut item else { - unreachable!() - }; - *role = "developer".to_string(); - let source = ResponseItemEnvelope { - item, - metadata: Some(CodexHarnessMetadata { - client_authored: true, - ..Default::default() - }), - }; - for max_tokens in [1871, 1879, 1880, 1950, 2200] { - assert_eq!( - truncate_retained_messages( - vec![source.clone()], - max_tokens, - RetainedImageBudget::Enabled - ), - truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), - ); - } -} diff --git a/codex-rs/core/src/compact_remote_v2_images.rs b/codex-rs/core/src/compact_remote_v2_images.rs deleted file mode 100644 index cd8697813a..0000000000 --- a/codex-rs/core/src/compact_remote_v2_images.rs +++ /dev/null @@ -1,103 +0,0 @@ -use crate::context_manager::estimate_image_reference_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::is_image_close_tag_text; -use codex_protocol::models::is_image_open_tag_text; -use codex_protocol::models::is_local_image_open_tag_text; -use codex_protocol::protocol::TruncationPolicy; -use codex_utils_output_truncation::approx_token_count; -use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; -use codex_utils_output_truncation::truncate_text; - -pub(super) fn content_item_token_count(item: &ContentItem) -> usize { - match item { - ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { - approx_token_count(text) - } - ContentItem::InputImage { image, detail } => usize::try_from( - approx_tokens_from_byte_count_i64(estimate_image_reference_bytes(image, *detail)), - ) - .unwrap_or(usize::MAX), - ContentItem::InputAudio { .. } => 0, - } -} - -/// Retain later parts of an image-containing boundary message, keeping each -/// image and its adjacent harness labels atomic. Text keeps its existing middle -/// truncation policy; audio remains uncharged and is preserved as before. -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 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() { - ContentItem::InputImage { .. } => Some(last), - ContentItem::InputText { text } - if is_image_close_tag_text(text) - && last > 0 - && matches!(content[last - 1].content(), ContentItem::InputImage { .. }) => - { - Some(last - 1) - } - _ => None, - }; - if let Some(image_index) = image_index { - let has_open_tag = image_index > 0 - && matches!( - content[image_index - 1].content(), - 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; - remaining = if fits { remaining - token_count } else { 0 }; - for item in content.drain(start..).rev() { - if fits { - retained.push(item); - } - } - continue; - } - let mut item = content.pop()?; - match item.content_mut() { - ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { - if remaining == 0 { - continue; - } - let token_count = approx_token_count(text); - if token_count <= remaining { - remaining -= token_count; - } else { - *text = truncate_text(text, TruncationPolicy::Tokens(remaining)); - remaining = 0; - } - if !text.is_empty() { - retained.push(item); - } - } - ContentItem::InputAudio { .. } => retained.push(item), - ContentItem::InputImage { .. } => unreachable!("images are handled atomically above"), - } - } - if retained.is_empty() { - return None; - } - retained.reverse(); - set_annotated_content(&mut envelope.item, retained)?; - if let Some(metadata) = &mut envelope.metadata { - metadata.mark_retained_sources_incomplete(); - } - Some(envelope) -} diff --git a/codex-rs/core/src/compact_remote_v2_media.rs b/codex-rs/core/src/compact_remote_v2_media.rs new file mode 100644 index 0000000000..87225451b2 --- /dev/null +++ b/codex-rs/core/src/compact_remote_v2_media.rs @@ -0,0 +1,105 @@ +//! Counts retained media tokens and truncates each media part together with its labels. + +use super::RetainedImageBudget; +use crate::context_manager::FILE_AUDIO_TOKENS_ESTIMATE; +use crate::context_manager::estimate_image_reference_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::AudioReference; +use codex_protocol::models::ContentItem; +use codex_protocol::models::is_closing_media_label; +use codex_protocol::models::is_opening_media_label; +use codex_protocol::protocol::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; +use codex_utils_output_truncation::truncate_text; + +pub(super) fn content_item_token_count( + item: &ContentItem, + image_budget: RetainedImageBudget, +) -> usize { + match item { + ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { + approx_token_count(text) + } + ContentItem::InputImage { .. } if image_budget == RetainedImageBudget::Disabled => 0, + ContentItem::InputImage { image, detail } => usize::try_from( + approx_tokens_from_byte_count_i64(estimate_image_reference_bytes(image, *detail)), + ) + .unwrap_or(usize::MAX), + ContentItem::InputAudio { + audio: AudioReference::File { .. }, + } => FILE_AUDIO_TOKENS_ESTIMATE, + ContentItem::InputAudio { + audio: AudioReference::Inline { .. }, + } => 0, + } +} + +/// Find the trailing image or audio part together with its adjacent harness labels. +fn trailing_media_start(content: &[AnnotatedContent]) -> Option { + let last = content.len().checked_sub(1)?; + let media = match content[last].content() { + ContentItem::InputImage { .. } | ContentItem::InputAudio { .. } => last, + ContentItem::InputText { text } + if last > 0 && is_closing_media_label(text, content[last - 1].content()) => + { + last - 1 + } + _ => return None, + }; + let has_open_tag = media > 0 + && matches!(content[media - 1].content(), ContentItem::InputText { text } if is_opening_media_label(text, content[media].content())); + Some(media - usize::from(has_open_tag)) +} + +/// Keep the newest suffix. A media part and its labels must fit together even +/// when the media itself has zero cost. +pub(super) fn truncate_message_to_token_budget( + mut envelope: ResponseItemEnvelope, + max_tokens: usize, + image_budget: RetainedImageBudget, +) -> Option { + let mut content = to_annotated_content(&mut envelope.item)?; + let mut remaining = max_tokens; + let mut retained = Vec::with_capacity(content.len()); + while remaining > 0 && !content.is_empty() { + if let Some(start) = trailing_media_start(&content) { + let tokens = content[start..] + .iter() + .map(|part| content_item_token_count(part.content(), image_budget)) + .fold(0, usize::saturating_add); + if tokens > remaining { + break; + } + remaining -= tokens; + retained.extend(content.drain(start..).rev()); + continue; + } + let mut item = content.pop()?; + let (ContentItem::InputText { text } | ContentItem::OutputText { text, .. }) = + item.content_mut() + else { + unreachable!("media parts are handled above"); + }; + let tokens = approx_token_count(text); + if tokens > remaining { + *text = truncate_text(text, TruncationPolicy::Tokens(remaining)); + } + remaining = remaining.saturating_sub(tokens); + if !text.is_empty() { + retained.push(item); + } + } + if retained.is_empty() { + return None; + } + retained.reverse(); + set_annotated_content(&mut envelope.item, retained)?; + if let Some(metadata) = &mut envelope.metadata { + metadata.mark_retained_sources_incomplete(); + } + Some(envelope) +} diff --git a/codex-rs/core/src/compact_remote_v2_media_budget_tests.rs b/codex-rs/core/src/compact_remote_v2_media_budget_tests.rs new file mode 100644 index 0000000000..973c162845 --- /dev/null +++ b/codex-rs/core/src/compact_remote_v2_media_budget_tests.rs @@ -0,0 +1,357 @@ +#![cfg(test)] + +use super::*; +use codex_protocol::models::AudioReference; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::ImageReference; +use codex_protocol::models::InternalChatMessageMetadataPassthrough; +use codex_protocol::models::audio_close_tag_text; +use codex_protocol::models::audio_open_tag_text; +use codex_protocol::models::image_close_tag_text; +use codex_protocol::models::local_image_open_tag_text_with_path; +use pretty_assertions::assert_eq; + +fn message(content: Vec) -> ResponseItem { + ResponseItem::Message { + status: None, + encrypted_content: None, + id: None, + role: "user".to_string(), + content, + phase: None, + internal_chat_message_metadata_passthrough: None, + } +} + +fn image() -> ContentItem { + ContentItem::InputImage { + image: ImageReference::Inline { + image_url: "data:image/png;base64,abc".to_string(), + }, + detail: None, + } +} + +fn file_image() -> ContentItem { + ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_string(), + }, + detail: Some(codex_protocol::models::ImageDetail::Original), + } +} + +fn file_audio() -> ContentItem { + ContentItem::InputAudio { + audio: AudioReference::File { + file_id: "file_audio".to_string(), + }, + } +} + +fn text(value: &str) -> ContentItem { + ContentItem::InputText { + text: value.to_string(), + } +} + +fn trim(items: Vec, max_tokens: usize) -> Vec { + truncate_retained_messages( + items.into_iter().map(ResponseItemEnvelope::new).collect(), + max_tokens, + RetainedImageBudget::Enabled, + ) + .into_iter() + .map(ResponseItemEnvelope::into_item) + .collect() +} + +#[test] +fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { + let newest = message(vec![text("new")]); + for image in [image(), file_image()] { + let items = vec![ + message(vec![text("old")]), + message(vec![image.clone()]), + newest.clone(), + ]; + let image_tokens = media::content_item_token_count(&image, RetainedImageBudget::Enabled); + for (max_tokens, expected) in [ + ( + image_tokens + 1, + vec![message(vec![image.clone()]), newest.clone()], + ), + (image_tokens, vec![newest.clone()]), + (1, vec![newest.clone()]), + ] { + assert_eq!(trim(items.clone(), max_tokens), expected); + } + } +} + +#[test] +fn file_audio_is_charged_and_atomic_in_both_image_budget_modes() { + let audio = file_audio(); + let audio_tokens = media::content_item_token_count(&audio, RetainedImageBudget::Enabled); + let newest = message(vec![text("new")]); + let parts = [text("before"), audio, text("after")]; + let kinds = ["before", "audio", "after"]; + let annotated_boundary = |start: usize| { + let mut item = message(parts[start..].to_vec()); + let ResponseItem::Message { + internal_chat_message_metadata_passthrough, + .. + } = &mut item + else { + unreachable!() + }; + *internal_chat_message_metadata_passthrough = + Some(InternalChatMessageMetadataPassthrough { + content_item_kinds: Some( + kinds[start..] + .iter() + .map(|kind| ContentItemKind((*kind).to_string())) + .collect(), + ), + ..Default::default() + }); + item + }; + let boundary = annotated_boundary(0); + let newest_tokens = retained_message_token_count(&newest, RetainedImageBudget::Disabled).max(1); + let after_tokens = + media::content_item_token_count(&text("after"), RetainedImageBudget::Enabled); + for image_budget in [RetainedImageBudget::Disabled, RetainedImageBudget::Enabled] { + for (budget, expected) in [ + ( + retained_message_token_count(&boundary, image_budget) + newest_tokens, + vec![boundary.clone(), newest.clone()], + ), + ( + audio_tokens + after_tokens + newest_tokens, + vec![annotated_boundary(1), newest.clone()], + ), + ( + audio_tokens + after_tokens + newest_tokens - 1, + vec![annotated_boundary(2), newest.clone()], + ), + ] { + let actual = truncate_retained_messages( + vec![ + message(vec![text("older")]), + boundary.clone(), + newest.clone(), + ] + .into_iter() + .map(ResponseItemEnvelope::new) + .collect(), + budget, + image_budget, + ) + .into_iter() + .map(ResponseItemEnvelope::into_item) + .collect::>(); + assert_eq!(actual, expected); + } + } +} + +#[test] +fn file_audio_cutoff_does_not_backfill_older_uncharged_media() { + let inline_audio = ContentItem::InputAudio { + audio: AudioReference::Inline { + audio_url: "data:audio/wav;base64,older".to_string(), + }, + }; + let trailing_text = text("after"); + let boundary = message(vec![ + text(&local_image_open_tag_text_with_path( + /*label_number*/ 7, + std::path::Path::new("older.png"), + )), + image(), + text(&image_close_tag_text()), + text("older recording"), + inline_audio, + text("newer recording"), + file_audio(), + trailing_text.clone(), + ]); + let max_tokens = media::content_item_token_count(&file_audio(), RetainedImageBudget::Enabled) + + media::content_item_token_count(&trailing_text, RetainedImageBudget::Enabled) + - 1; + + for image_budget in [RetainedImageBudget::Disabled, RetainedImageBudget::Enabled] { + let actual = truncate_retained_messages( + vec![ResponseItemEnvelope::new(boundary.clone())], + max_tokens, + image_budget, + ) + .into_iter() + .map(ResponseItemEnvelope::into_item) + .collect::>(); + let [ResponseItem::Message { content, .. }] = actual.as_slice() else { + panic!("expected one truncated user message"); + }; + assert_eq!(content, std::slice::from_ref(&trailing_text)); + } +} + +#[test] +fn older_media_keeps_both_labels_or_is_dropped_after_file_audio_fits() { + let file = file_audio(); + let file_tokens = media::content_item_token_count(&file, RetainedImageBudget::Enabled); + for (open, media, close) in [ + ( + local_image_open_tag_text_with_path( + /*label_number*/ 7, + std::path::Path::new("older.png"), + ), + image(), + image_close_tag_text(), + ), + ( + audio_open_tag_text(), + ContentItem::InputAudio { + audio: AudioReference::Inline { + audio_url: "data:audio/wav;base64,older".to_string(), + }, + }, + audio_close_tag_text(), + ), + ] { + let parts = vec![text(&open), media, text(&close), file.clone()]; + let whole = message(parts.clone()); + let source = message([vec![text("older text")], parts.clone()].concat()); + for image_budget in [RetainedImageBudget::Disabled, RetainedImageBudget::Enabled] { + let whole_tokens = retained_message_token_count(&whole, image_budget); + for (budget, expected) in [ + (file_tokens, vec![file.clone()]), + (whole_tokens - 1, vec![file.clone()]), + (whole_tokens, parts.clone()), + ] { + let actual = truncate_retained_messages( + vec![ResponseItemEnvelope::new(source.clone())], + budget, + image_budget, + ); + let [ + ResponseItemEnvelope { + item: ResponseItem::Message { content, .. }, + .. + }, + ] = actual.as_slice() + else { + panic!("expected one retained user message"); + }; + assert_eq!(content, &expected); + } + } + } +} + +#[test] +fn later_image_parts_preserve_labels_audio_and_annotations() { + let parts = vec![ + text("earlier text"), + text(&local_image_open_tag_text_with_path( + /*label_number*/ 7, + std::path::Path::new("image.png"), + )), + image(), + text(&image_close_tag_text()), + text("keep"), + ContentItem::InputAudio { + audio: AudioReference::Inline { + 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 image_tokens = parts[1..4] + .iter() + .map(|item| media::content_item_token_count(item, RetainedImageBudget::Enabled)) + .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 { + 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]); + } +} + +#[test] +fn image_treatment_preserves_text_only_boundary_behavior() { + let source = ResponseItemEnvelope::new(message(vec![ + text("earlier text"), + text(&"middle".repeat(20)), + text("later text"), + ContentItem::InputAudio { + audio: AudioReference::Inline { + audio_url: "data:audio/wav;base64,abc".to_string(), + }, + }, + ])); + for max_tokens in [0, 1, 4, 16, 100] { + assert_eq!( + truncate_retained_messages( + vec![source.clone()], + max_tokens, + RetainedImageBudget::Enabled, + ), + truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), + ); + } +} + +#[test] +fn image_treatment_preserves_client_developer_boundary_behavior() { + let mut item = message(vec![image(), text(&"a".repeat(1000))]); + let ResponseItem::Message { role, .. } = &mut item else { + unreachable!() + }; + *role = "developer".to_string(); + let source = ResponseItemEnvelope { + item, + metadata: Some(CodexHarnessMetadata { + client_authored: true, + ..Default::default() + }), + }; + for max_tokens in [1871, 1879, 1880, 1950, 2200] { + assert_eq!( + truncate_retained_messages( + vec![source.clone()], + max_tokens, + RetainedImageBudget::Enabled + ), + truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), + ); + } +} diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 3ffc67c061..2f21ff9421 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -4,6 +4,7 @@ mod normalize; pub(crate) mod updates; pub(crate) use history::ContextManager; +pub(crate) use history::FILE_AUDIO_TOKENS_ESTIMATE; pub(crate) use history::HistoryReplacement; pub(crate) use history::estimate_image_reference_bytes; pub(crate) use history::estimate_item_token_count; diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index dfc96bebde..ccc8be7145 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -11,14 +11,8 @@ use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::models::WebSearchAction; -use codex_protocol::models::is_audio_close_tag_text; -use codex_protocol::models::is_audio_open_tag_text; -use codex_protocol::models::is_image_close_tag_text; -use codex_protocol::models::is_image_open_tag_text; -use codex_protocol::models::is_local_audio_close_tag_text; -use codex_protocol::models::is_local_audio_open_tag_text; -use codex_protocol::models::is_local_image_close_tag_text; -use codex_protocol::models::is_local_image_open_tag_text; +use codex_protocol::models::is_closing_media_label; +use codex_protocol::models::is_opening_media_label; use codex_protocol::protocol::APPS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; @@ -103,19 +97,13 @@ pub(crate) fn user_input_parts( .enumerate() .filter(|(idx, content_item)| match content_item { ContentItem::InputText { text } => { - let is_image_label = ((is_local_image_open_tag_text(text) - || is_image_open_tag_text(text)) - && matches!(message.get(idx + 1), Some(ContentItem::InputImage { .. }))) - || (*idx > 0 - && (is_local_image_close_tag_text(text) || is_image_close_tag_text(text)) - && matches!(message.get(idx - 1), Some(ContentItem::InputImage { .. }))); - let is_audio_label = ((is_local_audio_open_tag_text(text) - || is_audio_open_tag_text(text)) - && matches!(message.get(idx + 1), Some(ContentItem::InputAudio { .. }))) - || (*idx > 0 - && (is_local_audio_close_tag_text(text) || is_audio_close_tag_text(text)) - && matches!(message.get(idx - 1), Some(ContentItem::InputAudio { .. }))); - !is_image_label && !is_audio_label + let opens_media = message + .get(idx + 1) + .is_some_and(|part| is_opening_media_label(text, part)); + let closes_media = idx + .checked_sub(1) + .is_some_and(|i| is_closing_media_label(text, &message[i])); + !opens_media && !closes_media } ContentItem::InputImage { .. } | ContentItem::InputAudio { .. } diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index f9e819ddd4..c36062446a 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -317,6 +317,22 @@ impl ResponsesRequest { .collect() } + /// Returns all `input_audio` `file_id` spans from `message` inputs for the provided role. + pub fn message_input_audio_file_ids(&self, role: &str) -> Vec { + self.inputs_of_type("message") + .into_iter() + .filter(|item| item.get("role").and_then(Value::as_str) == Some(role)) + .filter_map(|item| item.get("content").and_then(Value::as_array).cloned()) + .flatten() + .filter(|span| span.get("type").and_then(Value::as_str) == Some("input_audio")) + .filter_map(|span| { + span.get("file_id") + .and_then(Value::as_str) + .map(str::to_owned) + }) + .collect() + } + pub fn input(&self) -> Vec { self.body_json()["input"] .as_array() diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 671fe45d2b..04b846698a 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -19,9 +19,14 @@ use codex_model_provider_info::GCP_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::AgentPath; use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::models::AudioReference; use codex_protocol::models::ContentItem; use codex_protocol::models::ImageReference; use codex_protocol::models::ResponseItem; +use codex_protocol::models::image_close_tag_text; +use codex_protocol::models::image_open_tag_text; +use codex_protocol::models::input_text; +use codex_protocol::openai_models::InputModality; use codex_protocol::protocol::ConversationStartParams; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; @@ -856,6 +861,135 @@ async fn remote_compact_v2_charges_retained_images_to_token_budget( Ok(()) } +#[test_case(false; "audio_only")] +#[test_case(true; "mixed_media_boundary")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_compact_v2_budgets_persisted_file_audio_atomically( + older_image: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config + .features + .set_enabled(Feature::CompactionImageBudget, /*enabled*/ false) + .expect("disable image budgeting for this test"); + }) + .with_model_info_override("gpt-5.5", |model| { + model.input_modalities.push(InputModality::Audio); + }), + ) + .await?; + // Seven maximum-cost audio items exceed the 64k retained-message budget. The + // newest five fit whole at 12,000 tokens each. + let image_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; + let initial_history = (1..=7) + .map(|number| { + let mut content = Vec::new(); + if older_image && number == 3 { + // The file fits, then the older text exhausts the remaining budget. + content.extend([ + input_text(image_open_tag_text()), + ContentItem::InputImage { + image: ImageReference::Inline { + image_url: image_url.to_string(), + }, + detail: None, + }, + input_text(image_close_tag_text()), + input_text("older context ".repeat(4000)), + ]); + } + content.extend([ + input_text(format!("recording {number}")), + ContentItem::InputAudio { + audio: AudioReference::File { + file_id: format!("file_audio_{number}"), + }, + }, + ]); + RolloutItem::ResponseItem( + ResponseItem::Message { + id: None, + role: "user".to_string(), + content, + status: None, + phase: None, + encrypted_content: None, + internal_chat_message_metadata_passthrough: None, + } + .into(), + ) + }) + .collect::>(); + let codex = harness + .test() + .thread_manager + .start_thread(StartThreadOptions { + initial_history: InitialHistory::Forked(initial_history), + ..StartThreadOptions::new(harness.test().config.clone()) + }) + .await? + .thread; + let responses = responses::mount_sse_sequence( + harness.server(), + vec![ + sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "AUDIO_BUDGET_SUMMARY" + }, + }), + responses::ev_completed("compact-audio"), + ]), + sse(vec![ + responses::ev_assistant_message("after-audio", "done"), + responses::ev_completed("after-audio"), + ]), + ], + ) + .await; + + codex.submit(Op::Compact).await?; + wait_for_turn_complete(&codex).await; + codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "continue after audio compaction".to_string(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_turn_complete(&codex).await; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + if older_image { + assert_eq!(requests[0].message_input_image_urls("user").len(), 1); + assert!(requests[1].message_input_image_urls("user").is_empty()); + } + assert_eq!( + requests[0].message_input_audio_file_ids("user"), + (1..=7) + .map(|number| format!("file_audio_{number}")) + .collect::>() + ); + assert_eq!( + requests[1].message_input_audio_file_ids("user"), + (3..=7) + .map(|number| format!("file_audio_{number}")) + .collect::>() + ); + assert_eq!( + requests[1].inputs_of_type("compaction")[0]["encrypted_content"], + "AUDIO_BUDGET_SUMMARY" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_compact_v2_reuses_compaction_trigger_for_followups() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 3e5dfc4ada..cdeac4a598 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -1906,6 +1906,28 @@ pub fn is_audio_close_tag_text(text: &str) -> bool { text == AUDIO_CLOSE_TAG } +/// Whether this text opens a harness label for the adjacent media part. +pub fn is_opening_media_label(text: &str, media: &ContentItem) -> bool { + match media { + ContentItem::InputImage { .. } => { + is_local_image_open_tag_text(text) || is_image_open_tag_text(text) + } + ContentItem::InputAudio { .. } => { + is_local_audio_open_tag_text(text) || is_audio_open_tag_text(text) + } + ContentItem::InputText { .. } | ContentItem::OutputText { .. } => false, + } +} + +/// Whether this text closes a harness label for the adjacent media part. +pub fn is_closing_media_label(text: &str, media: &ContentItem) -> bool { + match media { + ContentItem::InputImage { .. } => is_image_close_tag_text(text), + ContentItem::InputAudio { .. } => is_audio_close_tag_text(text), + ContentItem::InputText { .. } | ContentItem::OutputText { .. } => false, + } +} + fn invalid_image_error_placeholder( path: &std::path::Path, error: impl std::fmt::Display,