Budget retained file audio during remote compaction (#59237)

## Why

Remote compaction did not charge retained file audio against its token budget, allowing audio-heavy history to exceed the intended limit.

## What changed

- Charge file audio using `FILE_AUDIO_TOKENS_ESTIMATE`, independently of image budgeting.
- Retain the newest media content that fits, keeping images and audio together with their adjacent labels and preserving content annotations. Stop at oversized media instead of backfilling older content.
- Share media-label detection with user-input event mapping.
- Add download timeouts and a Debian mirror fallback for the musl build's `libcap` source archive.

## Testing

Add unit coverage for audio budgeting in both image-budget modes, atomic media labels, and truncation boundaries. Add integration coverage showing that compaction retains the newest five of seven file-audio items within the retained-message budget, including a mixed-media case.

GitOrigin-RevId: 4b7cd69f53966461fdeec14096712671d480187b
This commit is contained in:
Hung
2026-11-03 21:53:26 +00:00
committed by copyberry
parent f1b0721c88
commit 3780a5f750
11 changed files with 688 additions and 347 deletions
+4 -1
View File
@@ -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}"
+40 -38
View File
@@ -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(&notice.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(&notice.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;
@@ -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<ContentItem>) -> 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<ResponseItem>, max_tokens: usize) -> Vec<ResponseItem> {
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::<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 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 {
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),
);
}
}
@@ -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<ResponseItemEnvelope> {
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::<usize>();
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)
}
@@ -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<usize> {
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<ResponseItemEnvelope> {
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)
}
@@ -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<ContentItem>) -> 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<ResponseItem>, max_tokens: usize) -> Vec<ResponseItem> {
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::<Vec<_>>();
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::<Vec<_>>();
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::<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 image_tokens = parts[1..4]
.iter()
.map(|item| media::content_item_token_count(item, RetainedImageBudget::Enabled))
.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 {
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),
);
}
}
+1
View File
@@ -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;
+9 -21
View File
@@ -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 { .. }
+16
View File
@@ -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<String> {
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<Value> {
self.body_json()["input"]
.as_array()
+134
View File
@@ -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::<Vec<_>>();
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::<Vec<_>>()
);
assert_eq!(
requests[1].message_input_audio_file_ids("user"),
(3..=7)
.map(|number| format!("file_audio_{number}"))
.collect::<Vec<_>>()
);
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(()));
+22
View File
@@ -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,