From 1685f2d62132b04fb385bdb70678f3967900e83d Mon Sep 17 00:00:00 2001 From: Hung Date: Wed, 4 Nov 2026 00:06:35 +0000 Subject: [PATCH] Use measured durations to budget uploaded audio (#59268) ## Why Uploaded audio always reserved 12,000 tokens per file, even for short clips. This could discard recordings during compaction that would fit within the retained-message budget. ## What changed - Add optional audio duration metadata to attachments and cache measured durations per session, resolving metadata for new input and replayed history. - Use measured durations for history estimates and compaction budgets at 10 tokens per second, rounded up and capped at 12,000 tokens. Keep the full budget for missing or unusable metadata. - Bound metadata lookups to 250 milliseconds per batch with up to eight concurrent requests, and record accepted input before lookup so interruption preserves the batch. ## Testing Add coverage for duration rounding and fallback estimates, mixed measured and unknown audio, retention through compaction and cold resume, and interruption during metadata lookup. GitOrigin-RevId: 26fe320c7024944d9d4c371d5fdcda84442b5cc5 --- codex-rs/attachment-store/src/lib.rs | 10 + codex-rs/core/src/compact_remote_history.rs | 25 +- codex-rs/core/src/compact_remote_v2.rs | 79 +++++-- .../src/compact_remote_v2_installation.rs | 1 + codex-rs/core/src/compact_remote_v2_media.rs | 10 +- .../compact_remote_v2_media_budget_tests.rs | 81 +++++-- codex-rs/core/src/context_manager/history.rs | 52 +++-- .../core/src/context_manager/history_tests.rs | 4 + codex-rs/core/src/context_manager/mod.rs | 2 +- codex-rs/core/src/session/file_audio.rs | 100 ++++++++ codex-rs/core/src/session/file_audio_tests.rs | 45 ++++ codex-rs/core/src/session/mod.rs | 10 +- codex-rs/core/src/session/turn.rs | 6 + codex-rs/core/src/state/session.rs | 4 + .../core/tests/suite/audio_duration_store.rs | 59 +++++ codex-rs/core/tests/suite/compact_remote.rs | 80 ++++--- codex-rs/core/tests/suite/mod.rs | 1 + .../tests/suite/scenarios_audio_file_tests.rs | 213 +++++++++++++++--- ...ured_audio_compaction_and_cold_resume.snap | 57 +++++ .../audio/src/audio_preparation_tests.rs | 16 ++ codex-rs/utils/audio/src/lib.rs | 17 +- 21 files changed, 755 insertions(+), 117 deletions(-) create mode 100644 codex-rs/core/src/session/file_audio.rs create mode 100644 codex-rs/core/src/session/file_audio_tests.rs create mode 100644 codex-rs/core/tests/suite/audio_duration_store.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__measured_audio_compaction_and_cold_resume.snap diff --git a/codex-rs/attachment-store/src/lib.rs b/codex-rs/attachment-store/src/lib.rs index 2c0fcb3753..9bc70388e2 100644 --- a/codex-rs/attachment-store/src/lib.rs +++ b/codex-rs/attachment-store/src/lib.rs @@ -135,6 +135,8 @@ impl fmt::Debug for AttachmentMetadata { pub enum FormatSpecificMetadata { /// Metadata for an image attachment. Image(ImageMetadata), + /// Metadata for an audio attachment. + Audio(AudioMetadata), } /// Metadata specific to an image attachment. @@ -148,6 +150,14 @@ pub struct ImageMetadata { pub height: Option, } +/// Metadata specific to an audio attachment. +#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +pub struct AudioMetadata { + /// Measured audio duration in milliseconds, when known by the store. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub duration_ms: Option, +} + /// Representation callers should retain after uploading an attachment. #[derive(Clone, Deserialize, Eq, PartialEq, Serialize)] pub enum UploadResult { diff --git a/codex-rs/core/src/compact_remote_history.rs b/codex-rs/core/src/compact_remote_history.rs index 82fc5e3f0e..3a6f1500f3 100644 --- a/codex-rs/core/src/compact_remote_history.rs +++ b/codex-rs/core/src/compact_remote_history.rs @@ -4,7 +4,8 @@ use crate::compact_retained_tool_outputs::is_retained_tool_output; use crate::context::ContextualUserFragment; use crate::context::ImageResizeNotice; use crate::context_manager::ContextManager; -use crate::context_manager::estimate_item_token_count; +use crate::context_manager::estimate_item_token_count_with_audio; +use crate::session::FileAudioDurations; use crate::session::turn_context::TurnContext; use codex_features::Feature; use codex_history::ResponseItemEnvelope; @@ -29,10 +30,16 @@ impl> HistoryItemGroup { std::iter::once(self.source).chain(self.attached_notice) } - pub(crate) fn estimated_token_count(&self) -> i128 { - let source_tokens = i128::from(estimate_item_token_count(self.source.borrow())); + pub(crate) fn estimated_token_count(&self, audio_durations: &FileAudioDurations) -> i128 { + let source_tokens = i128::from(estimate_item_token_count_with_audio( + self.source.borrow(), + audio_durations, + )); let notice_tokens = self.attached_notice.as_ref().map_or(0, |notice| { - i128::from(estimate_item_token_count(notice.borrow())) + i128::from(estimate_item_token_count_with_audio( + notice.borrow(), + audio_durations, + )) }); source_tokens.saturating_add(notice_tokens) } @@ -80,9 +87,10 @@ pub(crate) fn trim_function_call_history_to_fit_context_window( // saturation in the normal history estimator. let base_tokens = i128::try_from(approx_token_count(&base_instructions.text)).unwrap_or(i128::MAX); + let audio_durations = &history.file_audio_durations; let original_items = history.annotated_items(); let mut estimated_tokens = history_item_groups(original_items.iter().map(|item| &item.item)) - .map(|group| group.estimated_token_count()) + .map(|group| group.estimated_token_count(audio_durations)) .fold(base_tokens, i128::saturating_add); let initial_estimated_tokens = i64::try_from(estimated_tokens).unwrap_or(i64::MAX); let mut rewritten_items = Vec::new(); @@ -116,8 +124,11 @@ pub(crate) fn trim_function_call_history_to_fit_context_window( break; }; estimated_tokens = estimated_tokens - .saturating_sub(group.estimated_token_count()) - .saturating_add(i128::from(estimate_item_token_count(&rewritten_item.item))); + .saturating_sub(group.estimated_token_count(audio_durations)) + .saturating_add(i128::from(estimate_item_token_count_with_audio( + &rewritten_item.item, + audio_durations, + ))); consumed_items += group_item_count; rewritten_items.push(rewritten_item); } diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 2abeee12b6..bbdc304984 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::sync::OnceLock; use crate::compact_retained_tool_outputs::is_retained_tool_output; +use crate::session::FileAudioDurations; use crate::Prompt; use crate::ResponseStream; @@ -18,6 +19,7 @@ use crate::compact_remote_history::HistoryItemGroup; use crate::compact_remote_history::history_item_groups; use crate::context::world_state::WorldState; use crate::context_manager::estimate_item_token_count; +use crate::context_manager::estimate_item_token_count_with_audio; use crate::hook_runtime::PostCompactHookOutcome; use crate::hook_runtime::PreCompactHookOutcome; use crate::hook_runtime::run_post_compact_hooks; @@ -493,6 +495,7 @@ fn build_v2_compacted_history( compaction_output: ResponseItem, features: &Features, image_budget: RetainedImageBudget, + audio_durations: &FileAudioDurations, ) -> (Vec, usize) { debug_assert_eq!(prompt_input.len(), prompt_input_metadata.len()); let prompt_input = prompt_input @@ -508,8 +511,12 @@ fn build_v2_compacted_history( .filter(|group| is_retained_for_remote_compaction_v2(&group.source, features)) .flat_map(HistoryItemGroup::into_items) .collect::>(); - let mut retained = - truncate_retained_messages(retained, RETAINED_MESSAGE_TOKEN_BUDGET, image_budget); + let mut retained = truncate_retained_messages( + retained, + RETAINED_MESSAGE_TOKEN_BUDGET, + image_budget, + audio_durations, + ); let retained_image_count = retained .iter() .map(|envelope| retained_input_image_count(&envelope.item)) @@ -620,14 +627,21 @@ fn retained_input_image_count(item: &ResponseItem) -> usize { pub(crate) fn truncate_retained_messages_for_remote_compaction( items: Vec, max_tokens: usize, + audio_durations: &FileAudioDurations, ) -> Vec { - truncate_retained_messages(items, max_tokens, RetainedImageBudget::Disabled) + truncate_retained_messages( + items, + max_tokens, + RetainedImageBudget::Disabled, + audio_durations, + ) } fn truncate_retained_messages( items: Vec, max_tokens: usize, image_budget: RetainedImageBudget, + audio_durations: &FileAudioDurations, ) -> Vec { let mut remaining = max_tokens; let mut truncated_reversed = Vec::with_capacity(items.len()); @@ -656,13 +670,23 @@ fn truncate_retained_messages( // 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) + retained_message_token_count( + ¬ice.item, + RetainedImageBudget::Disabled, + audio_durations, + ) + .max(1) }); // Client-authored developer messages already charge non-text content via // the serialized estimate. Preserve their text-only boundary correction. - let content_tokens = retained_message_token_count(&group.source.item, image_budget); + let content_tokens = + retained_message_token_count(&group.source.item, image_budget, audio_durations); let source_tokens = if client_developer { - usize::try_from(estimate_item_token_count(&group.source.item)).unwrap_or(usize::MAX) + usize::try_from(estimate_item_token_count_with_audio( + &group.source.item, + audio_durations, + )) + .unwrap_or(usize::MAX) } else { content_tokens.max(1) }; @@ -689,7 +713,12 @@ fn truncate_retained_messages( remaining = 0; } let truncated_item = if requires_media_boundary_truncation { - media::truncate_message_to_token_budget(group.source, content_budget, image_budget) + media::truncate_message_to_token_budget( + group.source, + content_budget, + image_budget, + audio_durations, + ) } else { truncate_message_text_to_token_budget(group.source, content_budget) }; @@ -697,8 +726,11 @@ fn truncate_retained_messages( continue; }; if client_developer { - let item_tokens = usize::try_from(estimate_item_token_count(&truncated_item.item)) - .unwrap_or(usize::MAX); + let item_tokens = usize::try_from(estimate_item_token_count_with_audio( + &truncated_item.item, + audio_durations, + )) + .unwrap_or(usize::MAX); if item_tokens > available_tokens { let adjusted_budget = content_budget .saturating_sub(item_tokens - available_tokens) @@ -708,6 +740,7 @@ fn truncate_retained_messages( truncated_item, adjusted_budget, image_budget, + audio_durations, ) } else { truncate_message_text_to_token_budget(truncated_item, adjusted_budget) @@ -715,8 +748,11 @@ fn truncate_retained_messages( let Some(adjusted) = adjusted else { continue; }; - if usize::try_from(estimate_item_token_count(&adjusted.item)) - .unwrap_or(usize::MAX) + if usize::try_from(estimate_item_token_count_with_audio( + &adjusted.item, + audio_durations, + )) + .unwrap_or(usize::MAX) > available_tokens { continue; @@ -737,14 +773,19 @@ fn truncate_retained_messages( truncated_reversed } -fn retained_message_token_count(item: &ResponseItem, image_budget: RetainedImageBudget) -> usize { +fn retained_message_token_count( + item: &ResponseItem, + image_budget: RetainedImageBudget, + audio_durations: &FileAudioDurations, +) -> usize { let ResponseItem::Message { content, .. } = item else { - return usize::try_from(estimate_item_token_count(item)).unwrap_or(usize::MAX); + return usize::try_from(estimate_item_token_count_with_audio(item, audio_durations)) + .unwrap_or(usize::MAX); }; content .iter() - .map(|part| media::content_item_token_count(part, image_budget)) + .map(|part| media::content_item_token_count(part, image_budget, audio_durations)) .fold(0, usize::saturating_add) } @@ -867,6 +908,7 @@ mod tests { output, &Features::with_defaults(), RetainedImageBudget::Disabled, + &FileAudioDurations::default(), ) } @@ -885,6 +927,7 @@ mod tests { raw(truncate_retained_messages_for_remote_compaction( annotated(items), max_tokens, + &FileAudioDurations::default(), )) } @@ -992,6 +1035,7 @@ mod tests { output.clone(), &features, RetainedImageBudget::Disabled, + &FileAudioDurations::default(), ); let mut expected = vec![ ResponseItemEnvelope { @@ -1030,8 +1074,11 @@ mod tests { metadata: Some(CodexHarnessMetadata::default()), }; - let truncated = - truncate_retained_messages_for_remote_compaction(vec![item], /*max_tokens*/ 4); + let truncated = truncate_retained_messages_for_remote_compaction( + vec![item], + /*max_tokens*/ 4, + &FileAudioDurations::default(), + ); assert_eq!(truncated.len(), 1); assert_eq!(truncated[0].metadata, Some(CodexHarnessMetadata::default())); diff --git a/codex-rs/core/src/compact_remote_v2_installation.rs b/codex-rs/core/src/compact_remote_v2_installation.rs index 0b03d89681..f14c10a089 100644 --- a/codex-rs/core/src/compact_remote_v2_installation.rs +++ b/codex-rs/core/src/compact_remote_v2_installation.rs @@ -201,6 +201,7 @@ async fn build_and_install_compaction( } else { RetainedImageBudget::Disabled }, + crate::session::file_audio::snapshot(sess).await.as_ref(), ); let reviewer_compaction_hash = reviewer_compaction_hash(sess, turn_context).await; let installed = if let Some((prefix, model_info, cutoff_world_state)) = prefix { diff --git a/codex-rs/core/src/compact_remote_v2_media.rs b/codex-rs/core/src/compact_remote_v2_media.rs index 87225451b2..49027e3a6d 100644 --- a/codex-rs/core/src/compact_remote_v2_media.rs +++ b/codex-rs/core/src/compact_remote_v2_media.rs @@ -1,8 +1,8 @@ //! 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 crate::session::FileAudioDurations; use codex_context_fragments::AnnotatedContent; use codex_context_fragments::set_annotated_content; use codex_context_fragments::to_annotated_content; @@ -19,6 +19,7 @@ use codex_utils_output_truncation::truncate_text; pub(super) fn content_item_token_count( item: &ContentItem, image_budget: RetainedImageBudget, + audio_durations: &FileAudioDurations, ) -> usize { match item { ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { @@ -30,8 +31,8 @@ pub(super) fn content_item_token_count( ) .unwrap_or(usize::MAX), ContentItem::InputAudio { - audio: AudioReference::File { .. }, - } => FILE_AUDIO_TOKENS_ESTIMATE, + audio: AudioReference::File { file_id }, + } => audio_durations.tokens(file_id), ContentItem::InputAudio { audio: AudioReference::Inline { .. }, } => 0, @@ -61,6 +62,7 @@ pub(super) fn truncate_message_to_token_budget( mut envelope: ResponseItemEnvelope, max_tokens: usize, image_budget: RetainedImageBudget, + audio_durations: &FileAudioDurations, ) -> Option { let mut content = to_annotated_content(&mut envelope.item)?; let mut remaining = max_tokens; @@ -69,7 +71,7 @@ pub(super) fn truncate_message_to_token_budget( if let Some(start) = trailing_media_start(&content) { let tokens = content[start..] .iter() - .map(|part| content_item_token_count(part.content(), image_budget)) + .map(|part| content_item_token_count(part.content(), image_budget, audio_durations)) .fold(0, usize::saturating_add); if tokens > remaining { break; 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 index 973c162845..ef92fb2ddc 100644 --- a/codex-rs/core/src/compact_remote_v2_media_budget_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_media_budget_tests.rs @@ -60,6 +60,7 @@ fn trim(items: Vec, max_tokens: usize) -> Vec { items.into_iter().map(ResponseItemEnvelope::new).collect(), max_tokens, RetainedImageBudget::Enabled, + &FileAudioDurations::default(), ) .into_iter() .map(ResponseItemEnvelope::into_item) @@ -75,7 +76,11 @@ fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { message(vec![image.clone()]), newest.clone(), ]; - let image_tokens = media::content_item_token_count(&image, RetainedImageBudget::Enabled); + let image_tokens = media::content_item_token_count( + &image, + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ); for (max_tokens, expected) in [ ( image_tokens + 1, @@ -92,7 +97,11 @@ fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { #[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 audio_tokens = media::content_item_token_count( + &audio, + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ); let newest = message(vec![text("new")]); let parts = [text("before"), audio, text("after")]; let kinds = ["before", "audio", "after"]; @@ -118,13 +127,25 @@ fn file_audio_is_charged_and_atomic_in_both_image_budget_modes() { 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); + let newest_tokens = retained_message_token_count( + &newest, + RetainedImageBudget::Disabled, + &FileAudioDurations::default(), + ) + .max(1); + let after_tokens = media::content_item_token_count( + &text("after"), + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ); for image_budget in [RetainedImageBudget::Disabled, RetainedImageBudget::Enabled] { for (budget, expected) in [ ( - retained_message_token_count(&boundary, image_budget) + newest_tokens, + retained_message_token_count( + &boundary, + image_budget, + &FileAudioDurations::default(), + ) + newest_tokens, vec![boundary.clone(), newest.clone()], ), ( @@ -147,6 +168,7 @@ fn file_audio_is_charged_and_atomic_in_both_image_budget_modes() { .collect(), budget, image_budget, + &FileAudioDurations::default(), ) .into_iter() .map(ResponseItemEnvelope::into_item) @@ -177,15 +199,22 @@ fn file_audio_cutoff_does_not_backfill_older_uncharged_media() { 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; + let max_tokens = media::content_item_token_count( + &file_audio(), + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ) + media::content_item_token_count( + &trailing_text, + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ) - 1; for image_budget in [RetainedImageBudget::Disabled, RetainedImageBudget::Enabled] { let actual = truncate_retained_messages( vec![ResponseItemEnvelope::new(boundary.clone())], max_tokens, image_budget, + &FileAudioDurations::default(), ) .into_iter() .map(ResponseItemEnvelope::into_item) @@ -200,7 +229,11 @@ fn file_audio_cutoff_does_not_backfill_older_uncharged_media() { #[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); + let file_tokens = media::content_item_token_count( + &file, + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ); for (open, media, close) in [ ( local_image_open_tag_text_with_path( @@ -224,7 +257,8 @@ fn older_media_keeps_both_labels_or_is_dropped_after_file_audio_fits() { 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); + let whole_tokens = + retained_message_token_count(&whole, image_budget, &FileAudioDurations::default()); for (budget, expected) in [ (file_tokens, vec![file.clone()]), (whole_tokens - 1, vec![file.clone()]), @@ -234,6 +268,7 @@ fn older_media_keeps_both_labels_or_is_dropped_after_file_audio_fits() { vec![ResponseItemEnvelope::new(source.clone())], budget, image_budget, + &FileAudioDurations::default(), ); let [ ResponseItemEnvelope { @@ -285,7 +320,13 @@ fn later_image_parts_preserve_labels_audio_and_annotations() { }); let image_tokens = parts[1..4] .iter() - .map(|item| media::content_item_token_count(item, RetainedImageBudget::Enabled)) + .map(|item| { + media::content_item_token_count( + item, + RetainedImageBudget::Enabled, + &FileAudioDurations::default(), + ) + }) .sum::(); for (max_tokens, start) in [(image_tokens, 4), (image_tokens + 1, 1)] { let mut expected = source.clone(); @@ -324,8 +365,13 @@ fn image_treatment_preserves_text_only_boundary_behavior() { vec![source.clone()], max_tokens, RetainedImageBudget::Enabled, + &FileAudioDurations::default() + ), + truncate_retained_messages_for_remote_compaction( + vec![source.clone()], + max_tokens, + &FileAudioDurations::default() ), - truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), ); } } @@ -349,9 +395,14 @@ fn image_treatment_preserves_client_developer_boundary_behavior() { truncate_retained_messages( vec![source.clone()], max_tokens, - RetainedImageBudget::Enabled + RetainedImageBudget::Enabled, + &FileAudioDurations::default() + ), + truncate_retained_messages_for_remote_compaction( + vec![source.clone()], + max_tokens, + &FileAudioDurations::default() ), - truncate_retained_messages_for_remote_compaction(vec![source.clone()], max_tokens), ); } } diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 936fbadb72..7824aa19a6 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -31,6 +31,7 @@ use crate::event_mapping::is_contextual_user_message_content; use crate::event_mapping::parse_turn_item; use crate::guardian::GUARDIAN_MAX_ROOT_MESSAGE_TOKENS; use crate::guardian::guardian_truncate_text; +use crate::session::FileAudioDurations; use crate::session::turn_context::TurnContext; use crate::utils::json::serialized_json_bytes; use base64::Engine; @@ -91,11 +92,6 @@ use codex_features::Feature; static GUARDIAN_REVIEW_CONTEXT_REVISION: AtomicU64 = AtomicU64::new(/*v*/ 1); -// Responses validates uploaded audio against its 20-minute duration limit. -// File durations are unavailable locally, so budget that maximum at the existing -// inline estimate of 10 audio tokens per second. -pub(crate) const FILE_AUDIO_TOKENS_ESTIMATE: usize = 20 * 60 * 10; - fn next_guardian_review_context_revision() -> u64 { GUARDIAN_REVIEW_CONTEXT_REVISION.fetch_add(/*val*/ 1, Ordering::Relaxed) } @@ -113,6 +109,7 @@ pub(crate) struct ContextManager { /// File IDs rejected by Responses. Retained history stays append-only; only /// request projections omit these references. rejected_audio_file_ids: Arc>, + pub(crate) file_audio_durations: Arc, /// Reviewer policy travels with the history snapshot, independently of capture. guardian_review_mode: GuardianContextMode, guardian_include_additional_context: bool, @@ -277,6 +274,7 @@ impl ContextManager { review_history: None, retained_context: Arc::default(), rejected_audio_file_ids: Arc::default(), + file_audio_durations: Arc::default(), guardian_review_mode: GuardianContextMode::ThreadOwned, guardian_include_additional_context: Feature::GuardianIncludeAdditionalContext .default_enabled(), @@ -795,7 +793,9 @@ impl ContextManager { let items_tokens = self .items .iter() - .map(|envelope| estimate_item_token_count(&envelope.item)) + .map(|envelope| { + estimate_item_token_count_with_audio(&envelope.item, &self.file_audio_durations) + }) .fold(0i64, i64::saturating_add); Some(base_tokens.saturating_add(items_tokens)) @@ -1068,7 +1068,9 @@ impl ContextManager { } ) }) - .map(|envelope| estimate_item_token_count(&envelope.item)) + .map(|envelope| { + estimate_item_token_count_with_audio(&envelope.item, &self.file_audio_durations) + }) .fold(0i64, i64::saturating_add) } @@ -1103,7 +1105,7 @@ impl ContextManager { pub(crate) fn estimated_tokens_since_last_usage(&self) -> i64 { self.items_since_last_usage() - .map(estimate_item_token_count) + .map(|item| estimate_item_token_count_with_audio(item, &self.file_audio_durations)) .fold(0i64, i64::saturating_add) } @@ -1236,8 +1238,17 @@ fn estimate_encrypted_function_output_length(encoded_len: usize) -> usize { /// Counts content directly, excluding transport IDs, metadata, and outer JSON escaping. /// Original-detail file images use the maximum patch count. pub(crate) fn estimate_item_token_count(item: &ResponseItem) -> i64 { - let model_visible_bytes = estimate_response_item_model_visible_bytes(item); - approx_tokens_from_byte_count_i64(model_visible_bytes) + estimate_item_token_count_with_audio(item, &FileAudioDurations::default()) +} + +pub(crate) fn estimate_item_token_count_with_audio( + item: &ResponseItem, + audio_durations: &FileAudioDurations, +) -> i64 { + approx_tokens_from_byte_count_i64(estimate_response_item_model_visible_bytes( + item, + audio_durations, + )) } /// Approximate model-visible byte cost for one image input. @@ -1263,7 +1274,10 @@ static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock = Laz ) }); -fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { +fn estimate_response_item_model_visible_bytes( + item: &ResponseItem, + audio_durations: &FileAudioDurations, +) -> i64 { match item { ResponseItem::Message { content, .. } => content .iter() @@ -1274,7 +1288,9 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { ContentItem::InputImage { image, detail } => { estimate_image_reference_bytes(image, *detail) } - ContentItem::InputAudio { audio } => estimate_audio_reference_bytes(audio), + ContentItem::InputAudio { audio } => { + estimate_audio_reference_bytes(audio, audio_durations) + } }) .fold(0i64, i64::saturating_add), ResponseItem::AgentMessage { @@ -1487,13 +1503,15 @@ fn estimate_audio_bytes(audio_url: &str) -> i64 { .unwrap_or(i64::MAX) } -/// Audio estimate for callers that only have the reference. Uploaded file durations are not -/// available locally, so file references use the maximum duration accepted by Responses. -fn estimate_audio_reference_bytes(audio: &AudioReference) -> i64 { +fn estimate_audio_reference_bytes( + audio: &AudioReference, + audio_durations: &FileAudioDurations, +) -> i64 { match audio { AudioReference::Inline { audio_url } => estimate_audio_bytes(audio_url), - AudioReference::File { .. } => { - i64::try_from(approx_bytes_for_tokens(FILE_AUDIO_TOKENS_ESTIMATE)).unwrap_or(i64::MAX) + AudioReference::File { file_id } => { + i64::try_from(approx_bytes_for_tokens(audio_durations.tokens(file_id))) + .unwrap_or(i64::MAX) } } } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index e92a0ff02f..36b7a37795 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -150,6 +150,10 @@ fn raw_items(history: &ContextManager) -> Vec { history.raw_items().cloned().collect() } +fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { + super::estimate_response_item_model_visible_bytes(item, &FileAudioDurations::default()) +} + #[test] fn delivered_assistant_context_invalidates_reviews_without_changing_authorization() { let revisions = |history: &ContextManager| { diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 2f21ff9421..9149702d9e 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -4,10 +4,10 @@ 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; +pub(crate) use history::estimate_item_token_count_with_audio; pub(crate) use history::is_user_turn_boundary; pub(crate) use history::rejected_audio; pub(crate) use image_availability::replace_unresolvable_image_files; diff --git a/codex-rs/core/src/session/file_audio.rs b/codex-rs/core/src/session/file_audio.rs new file mode 100644 index 0000000000..a94cca9785 --- /dev/null +++ b/codex-rs/core/src/session/file_audio.rs @@ -0,0 +1,100 @@ +//! Caches trusted file-audio durations for one session. Metadata lookup is best effort +//! and time-bounded. Accepted input is recorded before lookup, and unmeasured +//! files reserve the full audio token budget. + +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use codex_attachment_store::FormatSpecificMetadata; +use codex_attachment_store::ResolveRequest; +use codex_history::ResponseItemEnvelope; +use codex_protocol::models::AudioReference; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_utils_audio::estimate_file_audio_tokens; +use futures::StreamExt; +use futures::stream; + +use crate::session::Session; + +/// Thread-local metadata. It is never serialized into a rollout or a provider request. +#[derive(Clone, Debug, Default)] +pub(crate) struct FileAudioDurations(BTreeMap); + +impl FileAudioDurations { + pub(crate) fn tokens(&self, file_id: &str) -> usize { + estimate_file_audio_tokens(self.0.get(file_id).copied()) + } +} + +pub(crate) async fn snapshot(session: &Session) -> Arc { + Arc::clone(&session.state.lock().await.history.file_audio_durations) +} + +pub(super) fn file_ids(items: &[ResponseItemEnvelope]) -> impl Iterator { + items + .iter() + .filter_map(|envelope| match &envelope.item { + ResponseItem::Message { content, .. } => Some(content), + _ => None, + }) + .flatten() + .filter_map(|part| match part { + ContentItem::InputAudio { + audio: AudioReference::File { file_id }, + } => Some(file_id), + _ => None, + }) +} + +pub(super) async fn resolve(session: &Session, mut file_ids: BTreeSet) { + if file_ids.is_empty() { + return; + } + { + let state = session.state.lock().await; + file_ids.retain(|id| !state.history.file_audio_durations.0.contains_key(id)); + } + let mut lookups = stream::iter(file_ids) + .map(|id| async move { + let result = session + .services + .image_store + .resolve(ResolveRequest { + file_id: &id, + tags: HashMap::new(), + download_url_ttl: None, + }) + .await; + (id, result) + }) + .buffer_unordered(8); + let mut resolved = BTreeMap::new(); + // Metadata is optional. Bound the entire batch, keep completed lookups on timeout, + // and allow future attachments to retry a missing value. + let _ = tokio::time::timeout(Duration::from_millis(250), async { + while let Some((id, result)) = lookups.next().await { + if let Ok(metadata) = result + && let Some(FormatSpecificMetadata::Audio(audio)) = metadata.format_specific + && let Some(duration_ms) = audio.duration_ms + && duration_ms > 0 + { + resolved.insert(id, Duration::from_millis(duration_ms)); + } + } + }) + .await; + if !resolved.is_empty() { + let mut state = session.state.lock().await; + Arc::make_mut(&mut state.history.file_audio_durations) + .0 + .extend(resolved); + } +} + +#[cfg(test)] +#[path = "file_audio_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/session/file_audio_tests.rs b/codex-rs/core/src/session/file_audio_tests.rs new file mode 100644 index 0000000000..a3cff4108b --- /dev/null +++ b/codex-rs/core/src/session/file_audio_tests.rs @@ -0,0 +1,45 @@ +//! Checks history estimates with and without trusted file-audio durations. +#![cfg(test)] + +use super::*; +use crate::context_manager::ContextManager; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::user_input::UserInput; +use codex_utils_output_truncation::TruncationPolicy; +use pretty_assertions::assert_eq; + +#[test] +fn history_sums_measured_and_unknown_file_audio() { + let mut history = ContextManager::new(); + let item = ResponseInputItem::from( + ["short", "unknown"] + .map(|file_id| UserInput::Audio { + audio: AudioReference::File { + file_id: file_id.into(), + }, + }) + .to_vec(), + ); + history.record_items( + [&ResponseItem::from(item)], + TruncationPolicy::Tokens(10_000), + ); + let base = BaseInstructions { + text: String::new(), + provenance: None, + }; + assert_eq!( + history.estimate_token_count_with_base_instructions(&base), + Some(24_000) + ); + history.file_audio_durations = Arc::new(FileAudioDurations(BTreeMap::from([( + "short".into(), + Duration::from_secs(30), + )]))); + assert_eq!( + history.estimate_token_count_with_base_instructions(&base), + Some(12_300) + ); + assert_eq!(history.estimated_tokens_since_last_usage(), 12_300); +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 064d8e975b..b620d57db5 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -245,6 +245,9 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::exec_output::StreamOutput; pub(crate) mod async_compaction; +pub(crate) mod file_audio; +pub(crate) use file_audio::FileAudioDurations; + mod code_mode_warning; mod config_refresh; pub(crate) mod context_window; @@ -1956,7 +1959,8 @@ impl Session { .zip(metadata) .map(|(item, metadata)| ResponseItemEnvelope { item, metadata }) .collect(); - // Ordinary replay needs no reviewer or network refresh. Only an opaque + file_audio::resolve(self, file_audio::file_ids(&history).cloned().collect()).await; + // Ordinary replay needs no reviewer refresh. Only an opaque // checkpoint can make reconstruction depend on reviewer compatibility. let reviewer_compaction_hash = if codex_history::CompactionCheckpoint::latest(&history) .is_some_and(codex_history::CompactionCheckpoint::is_usable) @@ -4053,6 +4057,9 @@ impl Session { }); } } + state + .pending_audio_duration_ids + .extend(file_audio::file_ids(&items).cloned()); state .history .record_annotated_items(&mut items, model_info.truncation_policy.into()); @@ -5283,6 +5290,7 @@ impl Session { .cloned() .collect(), crate::compact_remote_v2::RETAINED_MESSAGE_TOKEN_BUDGET, + &history.file_audio_durations, ); let window = { let mut state = self.state.lock().await; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 769917524f..c32b4472d9 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1024,6 +1024,12 @@ pub(crate) async fn run_hooks_and_record_inputs( .await; } } + // The batch has left the input queue. Persist every accepted item before the + // optional lookup can suspend this task long enough for an interrupt to abort it. + if !input.is_empty() { + let ids = std::mem::take(&mut sess.state.lock().await.pending_audio_duration_ids); + super::file_audio::resolve(sess, ids).await; + } Ok(blocked_input && !accepted_user_input) } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 94c8874687..f35bb2a188 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -5,6 +5,7 @@ use codex_protocol::models::BaseInstructionsProvenance; #[cfg(test)] use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; +use std::collections::BTreeSet; use std::collections::HashSet; use std::collections::VecDeque; use std::sync::Arc; @@ -77,6 +78,8 @@ pub(crate) struct SessionState { /// Persisted origin of the session base instructions, when known. pub(crate) base_instructions_provenance: Option, pub(crate) history: ContextManager, + /// Newly recorded files awaiting best-effort lookup after the input batch is durable. + pub(crate) pending_audio_duration_ids: BTreeSet, /// Cancels work bound to discarded history or a superseded Guardian evidence policy. pub(crate) history_reset: CancellationToken, pub(crate) latest_rate_limits: Option, @@ -132,6 +135,7 @@ impl SessionState { session_configuration, base_instructions_provenance: None, history, + pending_audio_duration_ids: BTreeSet::new(), history_reset: CancellationToken::new(), latest_rate_limits: None, latest_token_usage_record: None, diff --git a/codex-rs/core/tests/suite/audio_duration_store.rs b/codex-rs/core/tests/suite/audio_duration_store.rs new file mode 100644 index 0000000000..caf7750ec1 --- /dev/null +++ b/codex-rs/core/tests/suite/audio_duration_store.rs @@ -0,0 +1,59 @@ +//! Supplies measured clips and unavailable metadata for Core audio budgeting scenarios. + +use codex_attachment_store::AttachmentMetadata; +use codex_attachment_store::AttachmentStore; +use codex_attachment_store::AttachmentStoreError; +use codex_attachment_store::AttachmentStoreErrorKind; +use codex_attachment_store::AudioMetadata; +use codex_attachment_store::FormatSpecificMetadata; +use codex_attachment_store::ImageMetadata; +use codex_attachment_store::InlineAttachmentStore; +use codex_attachment_store::ResolveFuture; +use codex_attachment_store::ResolveRequest; +use codex_attachment_store::UploadFuture; +use codex_attachment_store::UploadRequest; + +#[derive(Default)] +pub(super) struct TestAudioDurationStore { + pub(super) slow_lookup_started: tokio::sync::Notify, + pub(super) measure_short_clips: bool, +} + +impl AttachmentStore for TestAudioDurationStore { + fn upload(&self, request: UploadRequest) -> UploadFuture<'_> { + InlineAttachmentStore.upload(request) + } + + fn resolve<'a>(&'a self, request: ResolveRequest<'a>) -> ResolveFuture<'a> { + Box::pin(async move { + assert!(request.download_url_ttl.is_none()); + let format_specific = match request.file_id { + "file_audio_1" | "file_audio_7" if self.measure_short_clips => { + Some(FormatSpecificMetadata::Audio(AudioMetadata { + duration_ms: Some(30_000), + })) + } + "file_audio_2" => Some(FormatSpecificMetadata::Image(ImageMetadata::default())), + "file_audio_3" => { + return Err(AttachmentStoreError::new( + AttachmentStoreErrorKind::Backend, + "metadata unavailable", + )); + } + "file_audio_4" => Some(FormatSpecificMetadata::Audio(AudioMetadata { + duration_ms: Some(0), + })), + "file_audio_5" => { + self.slow_lookup_started.notify_one(); + return std::future::pending().await; + } + "file_audio_6" => Some(FormatSpecificMetadata::Audio(AudioMetadata::default())), + _ => None, + }; + Ok(AttachmentMetadata { + format_specific, + ..AttachmentMetadata::default() + }) + }) + } +} diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 04b846698a..faa68b07f0 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -1,5 +1,6 @@ //! Exercises streamed remote compaction, including the retained checkpoint and continuing history. +use super::audio_duration_store::TestAudioDurationStore; use anyhow::Context; use anyhow::Result; use base64::Engine; @@ -861,32 +862,41 @@ async fn remote_compact_v2_charges_retained_images_to_token_budget( Ok(()) } -#[test_case(false; "audio_only")] -#[test_case(true; "mixed_media_boundary")] +#[test_case(None, 3, false; "audio_only")] +#[test_case(None, 3, true; "mixed_media_boundary")] +#[test_case(Some(false), 3, false; "metadata_unavailable")] +#[test_case(Some(true), 1, false; "with_partial_metadata")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_compact_v2_budgets_persisted_file_audio_atomically( + measure_short_clips: Option, + first_retained: u32, 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 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); + }); + let builder = if let Some(measure_short_clips) = measure_short_clips { + builder.with_image_store(Arc::new(TestAudioDurationStore { + measure_short_clips, + ..Default::default() + })) + } else { + builder + }; + let harness = TestCodexHarness::with_builder(builder).await?; + // A 64k budget fits five unknown files. Measuring the other two clips at + // 30 seconds each leaves room to retain all seven. let image_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; - let initial_history = (1..=7) + let initial_history = (1..=6) .map(|number| { let mut content = Vec::new(); if older_image && number == 3 { @@ -937,6 +947,10 @@ async fn remote_compact_v2_budgets_persisted_file_audio_atomically( let responses = responses::mount_sse_sequence( harness.server(), vec![ + sse(vec![ + responses::ev_assistant_message("first", "heard it"), + responses::ev_completed("first"), + ]), sse(vec![ json!({ "type": "response.output_item.done", @@ -955,6 +969,14 @@ async fn remote_compact_v2_budgets_persisted_file_audio_atomically( ) .await; + codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Audio { + audio: AudioReference::File { + file_id: "file_audio_7".into(), + }, + }])) + .await?; + wait_for_turn_complete(&codex).await; codex.submit(Op::Compact).await?; wait_for_turn_complete(&codex).await; codex @@ -966,10 +988,10 @@ async fn remote_compact_v2_budgets_persisted_file_audio_atomically( wait_for_turn_complete(&codex).await; let requests = responses.requests(); - assert_eq!(requests.len(), 2); + assert_eq!(requests.len(), 3); 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!(requests[2].message_input_image_urls("user").is_empty()); } assert_eq!( requests[0].message_input_audio_file_ids("user"), @@ -977,14 +999,20 @@ async fn remote_compact_v2_budgets_persisted_file_audio_atomically( .map(|number| format!("file_audio_{number}")) .collect::>() ); + let retained: Vec<_> = requests[2] + .input() + .into_iter() + .flat_map(|item| item["content"].as_array().cloned().unwrap_or_default()) + .filter(|part| part["type"] == "input_audio") + .collect(); assert_eq!( - requests[1].message_input_audio_file_ids("user"), - (3..=7) - .map(|number| format!("file_audio_{number}")) + retained, + (first_retained..=7) + .map(|number| json!({"type":"input_audio", "file_id": format!("file_audio_{number}")})) .collect::>() ); assert_eq!( - requests[1].inputs_of_type("compaction")[0]["encrypted_content"], + requests[2].inputs_of_type("compaction")[0]["encrypted_content"], "AUDIO_BUDGET_SUMMARY" ); Ok(()) diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 70d5f144a1..7d10cbb900 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -53,6 +53,7 @@ mod apply_patch_serialization; mod apply_patch_system_aliases; #[cfg(not(target_os = "windows"))] mod approvals; +mod audio_duration_store; mod audio_truncation; mod auth_recovery_policy; mod auto_review; diff --git a/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs b/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs index 00b4a4efd7..b7e66ed2d2 100644 --- a/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs +++ b/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs @@ -1,8 +1,10 @@ //! Uploaded audio references remain intact across model requests and user-message history. +use super::super::audio_duration_store::TestAudioDurationStore; use super::*; use codex_protocol::models::AudioReference; use codex_protocol::openai_models::InputModality; +use codex_protocol::turn_input::TurnInputSubmission; use core_test_support::responses; use core_test_support::responses::mount_response_sequence; use pretty_assertions::assert_eq; @@ -83,11 +85,13 @@ async fn uploaded_audio_survives_follow_up_and_resume() -> Result<()> { #[derive(Clone, Copy, PartialEq, Eq)] enum AudioAtCompaction { Available, + Measured, MissingAtStore, MissingAtResponses, } #[test_case(AudioAtCompaction::Available, &[true, true, true]; "available")] +#[test_case(AudioAtCompaction::Measured, &[true, true, true]; "measured duration")] #[test_case(AudioAtCompaction::MissingAtStore, &[true, false, false]; "expired before compaction")] #[test_case(AudioAtCompaction::MissingAtResponses, &[true, true, false, false]; "expired during compaction")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -118,24 +122,36 @@ async fn uploaded_audio_survives_compaction_and_cold_resume( ]); let responses = mount_response_sequence(&server, replies).await; let store = Arc::new(RecordingFileAttachmentStore::default()); + let measured = availability == AudioAtCompaction::Measured; let builder = || { - test_codex() - .with_image_store(store.clone()) + let builder = test_codex() .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_model_info_override("gpt-5.5", |model| { model.input_modalities.push(InputModality::Audio); - }) + }); + if measured { + builder.with_image_store(Arc::new(TestAudioDurationStore { + measure_short_clips: true, + ..Default::default() + })) + } else { + builder.with_image_store(store.clone()) + } }; let test = builder().build_with_auto_env(&server).await?; + let ids: Vec = if measured { + (1..=7).map(|n| format!("file_audio_{n}")).collect() + } else { + vec!["file_audio_compacted".into()] + }; + let mut input = vec![text("Retain this recording through compaction.")]; + input.extend(ids.iter().map(|id| UserInput::Audio { + audio: AudioReference::File { + file_id: id.clone(), + }, + })); test.codex - .start_or_steer_turn(TurnInputRequest::user_input(vec![ - text("Retain this recording through compaction."), - UserInput::Audio { - audio: AudioReference::File { - file_id: "file_audio_compacted".to_string(), - }, - }, - ])) + .start_or_steer_turn(TurnInputRequest::user_input(input)) .await?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) @@ -159,7 +175,10 @@ async fn uploaded_audio_survives_compaction_and_cold_resume( let resumed = builder().restart_with_auto_env(&server, &test).await?; let mut input = vec![text("Use the retained recording after cold resume.")]; - if availability != AudioAtCompaction::Available { + if matches!( + availability, + AudioAtCompaction::MissingAtStore | AudioAtCompaction::MissingAtResponses + ) { // Reintroduce the ID to exercise the checkpoint's rejection record. input.push(UserInput::Audio { audio: AudioReference::File { @@ -178,15 +197,24 @@ async fn uploaded_audio_survives_compaction_and_cold_resume( let requests = responses.requests(); assert_eq!(requests.len(), expected_audio.len()); + let expected: Vec<_> = ids + .iter() + .map(|id| json!({"type":"input_audio","file_id":id})) + .collect(); for (request, expect_audio) in requests.iter().zip(expected_audio) { - let audio = request + let audio: Vec<_> = request .input() .into_iter() .flat_map(|item| item["content"].as_array().cloned().unwrap_or_default()) - .find(|part| part["type"] == "input_audio"); + .filter(|part| part["type"] == "input_audio") + .collect(); assert_eq!( - audio, - expect_audio.then(|| json!({"type": "input_audio", "file_id": "file_audio_compacted"})), + audio.as_slice(), + if *expect_audio { + expected.as_slice() + } else { + &[] + } ); assert!( request @@ -202,16 +230,25 @@ async fn uploaded_audio_survives_compaction_and_cold_resume( .inputs_of_type("compaction")[0]["encrypted_content"], "UPLOADED_AUDIO_COMPACTION_CHECKPOINT" ); - if availability == AudioAtCompaction::Available { - insta::assert_snapshot!( + let (snapshot, description) = match availability { + AudioAtCompaction::Available => ( "uploaded_audio_compaction_and_cold_resume", - context_snapshot::format_request_history_snapshot( - "The uploaded audio file ID survives remote compaction and a cold resume in canonical message content.", - &requests, - &ContextSnapshotOptions::default().rewrite_known_segments(), - ) - ); - } + "The uploaded audio file ID survives remote compaction and a cold resume in canonical message content.", + ), + AudioAtCompaction::Measured => ( + "measured_audio_compaction_and_cold_resume", + "Two measured clips and five unmeasured files fit in the retained budget. All seven audio file IDs survive compaction and cold resume.", + ), + AudioAtCompaction::MissingAtStore | AudioAtCompaction::MissingAtResponses => return Ok(()), + }; + insta::assert_snapshot!( + snapshot, + context_snapshot::format_request_history_snapshot( + description, + &requests, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); Ok(()) } @@ -650,7 +687,12 @@ async fn unknown_file_after_two_recordings_were_accepted( .await; assert!(done.error.is_none(), "{done:?}"); } - assert!(store.resolves.lock().expect("probes").is_empty()); + // Each new recording gets one duration lookup. Recovery probes start after the 404. + { + let resolves = store.resolves.lock().expect("store lookups"); + assert_eq!(resolves.get("file_a"), Some(&1)); + assert_eq!(resolves.get("file_b"), Some(&1)); + } store .resolve_errors .lock() @@ -667,6 +709,14 @@ async fn unknown_file_after_two_recordings_were_accepted( _ => None, }) .await; + { + let resolves = store + .resolves + .lock() + .expect("store lookups after rejection"); + assert_eq!(resolves.get("file_a"), Some(&2)); + assert_eq!(resolves.get("file_b"), Some(&2)); + } assert_eq!( error.rejected_audio_file_ids, if confirmed { @@ -703,8 +753,6 @@ async fn unknown_file_after_two_recordings_were_accepted( assert!(input.contains("file_a")); assert_eq!(input.contains("file_b"), !confirmed); } - let probes = store.resolves.lock().expect("probes"); - assert!(probes.contains_key("file_a") && probes.contains_key("file_b")); Ok(()) } @@ -1347,3 +1395,112 @@ async fn unsupported_uploaded_audio_fails_before_inference_without_poisoning_fol ); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn interrupt_during_audio_lookup_preserves_steered_batch_on_follow_up_and_resume() +-> Result<()> { + skip_if_no_network!(Ok(())); + let store = Arc::new(TestAudioDurationStore::default()); + let (release, gate) = tokio::sync::oneshot::channel(); + let (streaming, _) = start_streaming_sse_server(vec![ + vec![StreamingSseChunk { + gate: Some(gate), + body: sse(vec![ev_completed("active")]), + }], + vec![StreamingSseChunk { + gate: None, + body: sse(vec![ev_completed("follow_up")]), + }], + vec![StreamingSseChunk { + gate: None, + body: sse(vec![ev_completed("resumed")]), + }], + ]) + .await; + let config_server = start_mock_server().await; + let builder = || { + let base_url = format!("{}/v1", streaming.uri()); + test_codex() + .with_image_store(store.clone()) + .with_model_info_override("gpt-5.5", |model| { + model.input_modalities.push(InputModality::Audio) + }) + .with_config(move |config| { + config.model_provider.base_url = Some(base_url); + config.model_provider.supports_websockets = false; + config + .features + .disable(Feature::InstantInterrupt) + .expect("queue steering"); + config + .features + .disable(Feature::EnableRequestCompression) + .expect("inspect request bodies"); + }) + }; + let test = builder().build_with_auto_env(&config_server).await?; + let TurnInputSubmission::Started { turn_id, .. } = test + .codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![text("Start working.")])) + .await? + else { + anyhow::bail!("expected an active turn"); + }; + let timeout = Duration::from_secs(10); + tokio::time::timeout(timeout, streaming.wait_for_request_count(/*count*/ 1)).await?; + test.codex + .steer_turn( + TurnInputRequest::user_input(vec![ + text("Keep this recording."), + UserInput::Audio { + audio: AudioReference::File { + file_id: "file_audio_5".into(), + }, + }, + ]), + turn_id.clone(), + ) + .await?; + test.codex + .steer_turn( + TurnInputRequest::user_input(vec![text("Keep the rest of the batch.")]), + turn_id, + ) + .await?; + release.send(()).expect("release active response"); + tokio::time::timeout(timeout, store.slow_lookup_started.notified()).await?; + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + + test.submit_text_turn("Follow up.").await?; + let resumed = builder() + .restart_with_auto_env(&config_server, &test) + .await?; + resumed.submit_text_turn("Continue after restart.").await?; + let requests = streaming.requests().await; + assert_eq!(requests.len(), 3); + for body in &requests[1..] { + let request: serde_json::Value = serde_json::from_slice(body)?; + let parts = request["input"] + .as_array() + .expect("request input") + .iter() + .filter(|item| item["role"] == "user") + .flat_map(|item| item["content"].as_array().into_iter().flatten()) + .collect::>(); + for expected in [ + json!({"type":"input_text", "text":"Keep this recording."}), + json!({"type":"input_audio", "file_id":"file_audio_5"}), + json!({"type":"input_text", "text":"Keep the rest of the batch."}), + ] { + assert!( + parts.contains(&&expected), + "missing {expected} on follow-up or resume" + ); + } + } + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__measured_audio_compaction_and_cold_resume.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__measured_audio_compaction_and_cold_resume.snap new file mode 100644 index 0000000000..9d8cbd2b02 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__measured_audio_compaction_and_cold_resume.snap @@ -0,0 +1,57 @@ +--- +source: core/tests/suite/scenarios_audio_file_tests.rs +assertion_line: 170 +expression: "context_snapshot::format_request_history_snapshot(if measured\n{\n \"Two measured clips and five unmeasured files fit in the retained budget. All seven audio file IDs survive compaction and cold resume.\"\n} else\n{\n \"The uploaded audio file ID survives remote compaction and a cold resume in canonical message content.\"\n}, &requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Two measured clips and five unmeasured files fit in the retained budget. All seven audio file IDs survive compaction and cold resume. + +## Window 1 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer[2]: + [01] + [02] + Deferred tool namespaces: + - multi_agent_v1: Tools for spawning and managing sub-agents. + +02:message/user: + +03:message/user[8]: + [01] Retain this recording through compaction. + [02] + [03] + [04] + [05] + [06] + [07] + [08] +-- request 2 (compaction) -- +04:compaction_trigger + +## Window 2 (after request 2: input diverged at item 04) +-- request 3 (turn) -- +00:message/developer: + +01:message/developer[2]: + [01] + [02] + Deferred tool namespaces: + - multi_agent_v1: Tools for spawning and managing sub-agents. + +02:message/user: + +03:message/user[8]: + [01] Retain this recording through compaction. + [02] + [03] + [04] + [05] + [06] + [07] + [08] +04:compaction:encrypted=true; chars=36; hash=09ef79034851b5bf +05:message/user: + +06:message/user: + Use the retained recording after cold resume. diff --git a/codex-rs/utils/audio/src/audio_preparation_tests.rs b/codex-rs/utils/audio/src/audio_preparation_tests.rs index c277338a26..cbc0d5b4c7 100644 --- a/codex-rs/utils/audio/src/audio_preparation_tests.rs +++ b/codex-rs/utils/audio/src/audio_preparation_tests.rs @@ -194,3 +194,19 @@ fn preparation_errors_map_to_expected_placeholders() { assert_eq!(error.placeholder(), expected); } } + +#[test] +fn file_audio_uses_the_measured_length_or_the_full_file_budget() { + use std::time::Duration; + for (duration, expected) in [ + (None, 12_000), + (Some(Duration::ZERO), 12_000), + (Some(Duration::from_nanos(1)), 1), + (Some(Duration::from_secs(2)), 20), + (Some(Duration::from_millis(2_001)), 21), + (Some(Duration::from_secs(20 * 60)), 12_000), + (Some(Duration::MAX), 12_000), + ] { + assert_eq!(super::estimate_file_audio_tokens(duration), expected); + } +} diff --git a/codex-rs/utils/audio/src/lib.rs b/codex-rs/utils/audio/src/lib.rs index 791fefe610..fa58098260 100644 --- a/codex-rs/utils/audio/src/lib.rs +++ b/codex-rs/utils/audio/src/lib.rs @@ -31,7 +31,7 @@ const UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER: &str = const MAX_PROMPT_AUDIO_BASE64_BYTES: usize = MAX_PROMPT_AUDIO_INPUT_BYTES.div_ceil(3) * 4; const AUDIO_TOKEN_ESTIMATE_CACHE_SIZE: usize = 32; -const AUDIO_TOKENS_PER_SECOND: f64 = 10.0; +const AUDIO_TOKENS_PER_SECOND: u32 = 10; static AUDIO_TOKEN_ESTIMATE_CACHE: LazyLock>>> = LazyLock::new(|| { @@ -147,6 +147,19 @@ fn canonical_audio_mime(mime: &str) -> Option<&'static str> { } } +/// Budget one file reference at the inline audio rate. Missing or unusable metadata +/// reserves the full 20 minutes accepted by Responses. Admission belongs to Responses. +pub fn estimate_file_audio_tokens(duration: Option) -> usize { + const MAX_DURATION: std::time::Duration = std::time::Duration::from_secs(20 * 60); + let duration = duration + .filter(|value| !value.is_zero()) + .unwrap_or(MAX_DURATION) + .min(MAX_DURATION); + let tokens = + (duration.as_nanos() * u128::from(AUDIO_TOKENS_PER_SECOND)).div_ceil(1_000_000_000); + usize::try_from(tokens).unwrap_or(usize::MAX) +} + /// Estimates audio tokens from decoded duration, falling back to the data URL size. pub fn estimate_audio_token_count(audio_url: &str) -> usize { let key = sha1_digest(audio_url.as_bytes()); @@ -160,7 +173,7 @@ pub fn estimate_audio_token_count_without_cache(audio_url: &str) -> usize { let Some(duration_seconds) = audio_duration_seconds(audio_url) else { return approx_token_count(audio_url); }; - let token_count = (duration_seconds * AUDIO_TOKENS_PER_SECOND).ceil(); + let token_count = (duration_seconds * f64::from(AUDIO_TOKENS_PER_SECOND)).ceil(); if token_count >= usize::MAX as f64 { usize::MAX } else {