Preserve audio across history and tool outputs (#34385)

## Why

Audio attachments could be dropped from tool outputs and reconstructed thread history, while inline base64 payload size could overstate their model context cost.

## What changed

- Preserve remote and local audio attachments in legacy user-message events and thread history.
- Keep audio from function and custom tool outputs when the model supports it and the output fits the truncation budget; otherwise replace it with an omission marker.
- Estimate inline audio cost from its duration, with a serialized-size fallback for malformed or unsupported data URLs, so compaction and truncation use model-visible cost instead of raw base64 size.
- Replace audio in messages and tool outputs with an explanatory placeholder for models without audio input support.

## Testing

Add coverage for history round trips, dynamic tool responses, modality filtering, duration-based context estimates, output-budget truncation, and remote compaction.

GitOrigin-RevId: 5ed4628850a668f8383b7cd73db71b144cc02cd1

codex-rs/core/src/audio_preparation.rs

@@ -3,6 +3,17 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::FunctionCallOutputContentItem;
 use codex_protocol::models::ResponseItem;
+use codex_utils_cache::BlockingLruCache;
+use codex_utils_cache::sha1_digest;
+use codex_utils_output_truncation::approx_token_count;
+use std::io::Cursor;
+use std::num::NonZeroUsize;
+use std::sync::LazyLock;
+use symphonia::core::formats::FormatOptions;
+use symphonia::core::formats::TrackType;
+use symphonia::core::formats::probe::Hint;
+use symphonia::core::io::MediaSourceStream;
+use symphonia::core::meta::MetadataOptions;
 use tracing::warn;
 
 const AUDIO_PROCESSING_ERROR_PLACEHOLDER: &str =
@@ -17,6 +28,15 @@ const UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER: &str =
 /// This matches the Responses API audio input limit.
 const MAX_PROMPT_AUDIO_INPUT_BYTES: usize = 50 * 1024 * 1024;
 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;
+
+static AUDIO_TOKEN_ESTIMATE_CACHE: LazyLock<BlockingLruCache<[u8; 20], usize>> =
+    LazyLock::new(|| {
+        BlockingLruCache::new(
+            NonZeroUsize::new(AUDIO_TOKEN_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN),
+        )
+    });
 
 #[derive(Debug, thiserror::Error)]
 enum AudioPreparationError {
@@ -121,6 +141,65 @@ fn canonical_audio_mime(mime: &str) -> Option<&'static str> {
     }
 }
 
+pub(crate) fn estimate_audio_token_count(audio_url: &str) -> usize {
+    let key = sha1_digest(audio_url.as_bytes());
+    AUDIO_TOKEN_ESTIMATE_CACHE.get_or_insert_with(key, || {
+        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();
+        if token_count >= usize::MAX as f64 {
+            usize::MAX
+        } else {
+            token_count as usize
+        }
+    })
+}
+
+fn audio_duration_seconds(audio_url: &str) -> Option<f64> {
+    let (metadata, payload) = audio_url.split_once(',')?;
+    let metadata = metadata.get("data:".len()..)?;
+    let mut metadata_parts = metadata.split(';');
+    let canonical_mime = canonical_audio_mime(metadata_parts.next()?)?;
+    if !metadata_parts.any(|part| part.eq_ignore_ascii_case("base64")) {
+        return None;
+    }
+
+    let bytes = match BASE64_STANDARD.decode(payload) {
+        Ok(bytes) => bytes,
+        Err(error) => {
+            tracing::trace!(%error, "failed to decode audio payload for token estimation");
+            return None;
+        }
+    };
+    let media_source = MediaSourceStream::new(Box::new(Cursor::new(bytes)), Default::default());
+    let mut hint = Hint::new();
+    hint.mime_type(canonical_mime);
+    let format = match symphonia::default::get_probe().probe(
+        &hint,
+        media_source,
+        FormatOptions::default(),
+        MetadataOptions::default(),
+    ) {
+        Ok(format) => format,
+        Err(error) => {
+            tracing::trace!(%error, "failed to read audio duration for token estimation");
+            return None;
+        }
+    };
+    let track = format.default_track(TrackType::Audio)?;
+    let timing = track.time_base.zip(track.duration).or_else(|| {
+        format
+            .media_info()
+            .time_base
+            .zip(format.media_info().duration)
+    });
+    let (time_base, duration) = timing?;
+    let duration_seconds =
+        duration.get() as f64 * f64::from(time_base.numer.get()) / f64::from(time_base.denom.get());
+    duration_seconds.is_finite().then_some(duration_seconds)
+}
+
 fn prepare_audio(audio_url: &mut String) -> Result<(), AudioPreparationError> {
     if !is_data_url(audio_url) {
         return Err(AudioPreparationError::InvalidDataUrl {

codex-rs/core/src/context_manager/history.rs

@@ -1,3 +1,4 @@
+use crate::audio_preparation::estimate_audio_token_count;
 use crate::context::ContextualUserFragment;
 use crate::context::world_state::WorldState;
 use crate::context::world_state::WorldStateSnapshot;
@@ -437,7 +438,7 @@ pub(crate) fn truncate_function_output_payload(
             FunctionCallOutputBody::Text(truncate_text(content, policy))
         }
         FunctionCallOutputBody::ContentItems(items) => FunctionCallOutputBody::ContentItems(
-            truncate_function_output_items_with_policy(items, policy),
+            truncate_function_output_items_with_policy(items, policy, estimate_audio_token_count),
         ),
     };
 
@@ -531,14 +532,18 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 {
                 .unwrap_or_default();
             let (image_payload_bytes, image_replacement_bytes) =
                 image_data_url_estimate_adjustment(item);
+            let (audio_payload_bytes, audio_replacement_bytes) =
+                audio_data_url_estimate_adjustment(item);
             let (encrypted_payload_bytes, encrypted_replacement_bytes) =
                 encrypted_function_output_estimate_adjustment(item);
-            // Replace raw base64 payload bytes with a per-image estimate.
+            // Replace raw base64 payload bytes with per-modality estimates.
             // We intentionally preserve the data URL prefix and JSON
             // wrapper bytes already included in `raw`.
             let raw = raw
                 .saturating_sub(image_payload_bytes)
-                .saturating_add(image_replacement_bytes);
+                .saturating_add(image_replacement_bytes)
+                .saturating_sub(audio_payload_bytes)
+                .saturating_add(audio_replacement_bytes);
             raw.saturating_sub(encrypted_payload_bytes)
                 .saturating_add(encrypted_replacement_bytes)
         }
@@ -551,6 +556,16 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 {
 /// We only discount payloads for `data:image/...;base64,...` URLs (case
 /// insensitive markers) and leave everything else at raw serialized size.
 fn parse_base64_image_data_url(url: &str) -> Option<&str> {
+    parse_base64_data_url(url, "image/")
+}
+
+/// Returns the base64 payload for inline audio data URLs that are eligible for
+/// token-estimation discounting.
+fn parse_base64_audio_data_url(url: &str) -> Option<&str> {
+    parse_base64_data_url(url, "audio/")
+}
+
+fn parse_base64_data_url<'a>(url: &'a str, media_type_prefix: &str) -> Option<&'a str> {
     if !url
         .get(.."data:".len())
         .is_some_and(|prefix| prefix.eq_ignore_ascii_case("data:"))
@@ -561,15 +576,15 @@ fn parse_base64_image_data_url(url: &str) -> Option<&str> {
     let metadata = &url[..comma_index];
     let payload = &url[comma_index + 1..];
     // Parse the media type and parameters without decoding. This keeps the
-    // estimator cheap while ensuring we only apply the fixed-cost image
-    // heuristic to image-typed base64 data URLs.
+    // estimator cheap while ensuring we only apply modality heuristics to
+    // appropriately typed base64 data URLs.
     let metadata_without_scheme = &metadata["data:".len()..];
     let mut metadata_parts = metadata_without_scheme.split(';');
     let mime_type = metadata_parts.next().unwrap_or_default();
     let has_base64_marker = metadata_parts.any(|part| part.eq_ignore_ascii_case("base64"));
     if !mime_type
-        .get(.."image/".len())
-        .is_some_and(|prefix| prefix.eq_ignore_ascii_case("image/"))
+        .get(..media_type_prefix.len())
+        .is_some_and(|prefix| prefix.eq_ignore_ascii_case(media_type_prefix))
     {
         return None;
     }
@@ -662,6 +677,50 @@ fn image_data_url_estimate_adjustment(item: &ResponseItem) -> (i64, i64) {
     (payload_bytes, replacement_bytes)
 }
 
+/// Scans one response item for inline base64 audio data URLs and returns:
+/// - total base64 payload bytes to subtract from raw serialized size
+/// - total replacement byte estimate for those audio inputs
+fn audio_data_url_estimate_adjustment(item: &ResponseItem) -> (i64, i64) {
+    let mut payload_bytes = 0i64;
+    let mut replacement_bytes = 0i64;
+
+    let mut accumulate = |audio_url: &str| {
+        if let Some(payload_len) = parse_base64_audio_data_url(audio_url).map(str::len) {
+            payload_bytes =
+                payload_bytes.saturating_add(i64::try_from(payload_len).unwrap_or(i64::MAX));
+            replacement_bytes = replacement_bytes.saturating_add(
+                i64::try_from(approx_bytes_for_tokens(estimate_audio_token_count(
+                    audio_url,
+                )))
+                .unwrap_or(i64::MAX),
+            );
+        }
+    };
+
+    match item {
+        ResponseItem::Message { content, .. } => {
+            for content_item in content {
+                if let ContentItem::InputAudio { audio_url } = content_item {
+                    accumulate(audio_url);
+                }
+            }
+        }
+        ResponseItem::FunctionCallOutput { output, .. }
+        | ResponseItem::CustomToolCallOutput { output, .. } => {
+            if let FunctionCallOutputBody::ContentItems(items) = &output.body {
+                for content_item in items {
+                    if let FunctionCallOutputContentItem::InputAudio { audio_url } = content_item {
+                        accumulate(audio_url);
+                    }
+                }
+            }
+        }
+        _ => {}
+    }
+
+    (payload_bytes, replacement_bytes)
+}
+
 fn encrypted_function_output_estimate_adjustment(item: &ResponseItem) -> (i64, i64) {
     let ResponseItem::FunctionCallOutput { output, .. } = item else {
         return (0, 0);

codex-rs/core/src/context_manager/history_tests.rs

@@ -40,6 +40,31 @@ use regex_lite::Regex;
 
 const EXEC_FORMAT_MAX_BYTES: usize = 10_000;
 const EXEC_FORMAT_MAX_TOKENS: usize = 2_500;
+const TEST_WAV_SAMPLE_RATE: u32 = 8_000;
+
+fn pcm_wav_data_url(sample_count: u32) -> (String, usize) {
+    let padding = sample_count % 2;
+    let mut bytes = Vec::new();
+    bytes.extend_from_slice(b"RIFF");
+    bytes.extend_from_slice(&(36 + sample_count + padding).to_le_bytes());
+    bytes.extend_from_slice(b"WAVEfmt ");
+    bytes.extend_from_slice(&16u32.to_le_bytes());
+    bytes.extend_from_slice(&1u16.to_le_bytes());
+    bytes.extend_from_slice(&1u16.to_le_bytes());
+    bytes.extend_from_slice(&TEST_WAV_SAMPLE_RATE.to_le_bytes());
+    bytes.extend_from_slice(&TEST_WAV_SAMPLE_RATE.to_le_bytes());
+    bytes.extend_from_slice(&1u16.to_le_bytes());
+    bytes.extend_from_slice(&8u16.to_le_bytes());
+    bytes.extend_from_slice(b"data");
+    bytes.extend_from_slice(&sample_count.to_le_bytes());
+    bytes.resize(
+        bytes.len() + sample_count as usize + padding as usize,
+        /*value*/ 0,
+    );
+    let payload = BASE64_STANDARD.encode(bytes);
+    let payload_len = payload.len();
+    (format!("data:audio/wav;base64,{payload}"), payload_len)
+}
 
 fn assistant_msg(text: &str) -> ResponseItem {
     ResponseItem::Message {
@@ -513,6 +538,9 @@ fn for_prompt_strips_media_when_model_does_not_support_it() {
                     image_url: "https://example.com/result.png".to_string(),
                     detail: Some(DEFAULT_IMAGE_DETAIL),
                 },
+                FunctionCallOutputContentItem::InputAudio {
+                    audio_url: "data:audio/mpeg;base64,YXVkaW8=".to_string(),
+                },
             ]),
             internal_chat_message_metadata_passthrough: None,
         },
@@ -537,10 +565,14 @@ fn for_prompt_strips_media_when_model_does_not_support_it() {
                     image_url: "https://example.com/js-repl-result.png".to_string(),
                     detail: Some(DEFAULT_IMAGE_DETAIL),
                 },
+                FunctionCallOutputContentItem::InputAudio {
+                    audio_url: "data:audio/ogg;base64,YXVkaW8=".to_string(),
+                },
             ]),
             internal_chat_message_metadata_passthrough: None,
         },
     ];
+    let fully_supported_items = items.clone();
     let history = create_history_with_items(items);
     let text_only_modalities = vec![InputModality::Text];
     let stripped = history.for_prompt(&text_only_modalities);
@@ -587,6 +619,10 @@ fn for_prompt_strips_media_when_model_does_not_support_it() {
                     text: "image content omitted because you do not support image input"
                         .to_string(),
                 },
+                FunctionCallOutputContentItem::InputText {
+                    text: "audio content omitted because you do not support audio input"
+                        .to_string(),
+                },
             ]),
             internal_chat_message_metadata_passthrough: None,
         },
@@ -611,11 +647,23 @@ fn for_prompt_strips_media_when_model_does_not_support_it() {
                     text: "image content omitted because you do not support image input"
                         .to_string(),
                 },
+                FunctionCallOutputContentItem::InputText {
+                    text: "audio content omitted because you do not support audio input"
+                        .to_string(),
+                },
             ]),
             internal_chat_message_metadata_passthrough: None,
         },
     ];
     assert_eq!(stripped, expected);
+    assert_eq!(
+        create_history_with_items(fully_supported_items.clone()).for_prompt(&[
+            InputModality::Text,
+            InputModality::Image,
+            InputModality::Audio,
+        ]),
+        fully_supported_items
+    );
 
     // With image support, images are preserved
     let modalities = default_input_modalities();
@@ -1949,6 +1997,121 @@ fn image_data_url_payload_does_not_dominate_custom_tool_call_output_estimate() {
     assert!(estimated < raw_len);
 }
 
+#[test]
+fn audio_data_url_payload_does_not_dominate_message_estimate() {
+    let (audio_url, payload_len) = pcm_wav_data_url(/*sample_count*/ 801);
+    let item = ResponseItem::Message {
+        id: None,
+        role: "user".to_string(),
+        content: vec![ContentItem::InputAudio { audio_url }],
+        phase: None,
+        internal_chat_message_metadata_passthrough: None,
+    };
+
+    let raw_len = serde_json::to_string(&item).unwrap().len() as i64;
+    let estimated = estimate_response_item_model_visible_bytes(&item);
+    let expected = raw_len - payload_len as i64 + approx_bytes_for_tokens(/*tokens*/ 2) as i64;
+
+    assert_eq!(estimated, expected);
+    assert!(estimated < raw_len);
+}
+
+#[test]
+fn audio_data_url_payload_does_not_dominate_function_call_output_estimate() {
+    let (audio_url, payload_len) = pcm_wav_data_url(/*sample_count*/ 800);
+    let item = ResponseItem::FunctionCallOutput {
+        id: None,
+        call_id: "call-audio".to_string(),
+        output: FunctionCallOutputPayload::from_content_items(vec![
+            FunctionCallOutputContentItem::InputAudio { audio_url },
+        ]),
+        internal_chat_message_metadata_passthrough: None,
+    };
+
+    let raw_len = serde_json::to_string(&item).unwrap().len() as i64;
+    let estimated = estimate_response_item_model_visible_bytes(&item);
+    let expected = raw_len - payload_len as i64 + approx_bytes_for_tokens(/*tokens*/ 1) as i64;
+
+    assert_eq!(estimated, expected);
+    assert!(estimated < raw_len);
+}
+
+#[test]
+fn audio_data_url_payload_does_not_dominate_custom_tool_call_output_estimate() {
+    let (audio_url, payload_len) = pcm_wav_data_url(/*sample_count*/ 80_000);
+    let item = ResponseItem::CustomToolCallOutput {
+        id: None,
+        call_id: "call-custom-audio".to_string(),
+        name: None,
+        output: FunctionCallOutputPayload::from_content_items(vec![
+            FunctionCallOutputContentItem::InputAudio { audio_url },
+        ]),
+        internal_chat_message_metadata_passthrough: None,
+    };
+
+    let raw_len = serde_json::to_string(&item).unwrap().len() as i64;
+    let estimated = estimate_response_item_model_visible_bytes(&item);
+    let expected = raw_len - payload_len as i64 + approx_bytes_for_tokens(/*tokens*/ 100) as i64;
+
+    assert_eq!(estimated, expected);
+    assert!(estimated < raw_len);
+}
+
+#[test]
+fn malformed_audio_data_url_falls_back_to_whole_url_size_cost() {
+    let payload = "A".repeat(/*n*/ 100_000);
+    let audio_url = format!("data:audio/wav;base64,{payload}");
+    let fallback_bytes = approx_bytes_for_tokens(approx_token_count(&audio_url)) as i64;
+    let item = ResponseItem::Message {
+        id: None,
+        role: "user".to_string(),
+        content: vec![ContentItem::InputAudio { audio_url }],
+        phase: None,
+        internal_chat_message_metadata_passthrough: None,
+    };
+
+    let raw_len = serde_json::to_string(&item).unwrap().len() as i64;
+    let estimated = estimate_response_item_model_visible_bytes(&item);
+
+    assert_eq!(estimated, raw_len - payload.len() as i64 + fallback_bytes);
+}
+
+#[test]
+fn record_items_omits_audio_that_exceeds_the_output_budget() {
+    let (audio_url, _) = pcm_wav_data_url(/*sample_count*/ 80_000);
+    let item = ResponseItem::FunctionCallOutput {
+        id: None,
+        call_id: "call-audio".to_string(),
+        output: FunctionCallOutputPayload {
+            body: FunctionCallOutputBody::ContentItems(vec![
+                FunctionCallOutputContentItem::InputAudio { audio_url },
+            ]),
+            success: Some(true),
+        },
+        internal_chat_message_metadata_passthrough: None,
+    };
+    let mut history = ContextManager::new();
+
+    history.record_items([&item], TruncationPolicy::Tokens(50));
+
+    assert_eq!(
+        history.items,
+        vec![ResponseItem::FunctionCallOutput {
+            id: None,
+            call_id: "call-audio".to_string(),
+            output: FunctionCallOutputPayload {
+                body: FunctionCallOutputBody::ContentItems(vec![
+                    FunctionCallOutputContentItem::InputText {
+                        text: "[omitted 1 audio items ...]".to_string(),
+                    },
+                ]),
+                success: Some(true),
+            },
+            internal_chat_message_metadata_passthrough: None,
+        }]
+    );
+}
+
 #[test]
 fn non_base64_image_urls_are_unchanged() {
     let message_item = ResponseItem::Message {

codex-rs/core/src/context_manager/normalize.rs

@@ -369,7 +369,7 @@ pub(crate) fn strip_images_when_unsupported(
     }
 }
 
-/// Strip audio content from messages when the model does not support audio.
+/// Strip audio content from messages and tool outputs when the model does not support audio.
 /// When `input_modalities` contains `InputModality::Audio`, no stripping is performed.
 pub(crate) fn strip_audio_when_unsupported(
     input_modalities: &[InputModality],
@@ -380,14 +380,32 @@ pub(crate) fn strip_audio_when_unsupported(
     }
 
     for item in items.iter_mut() {
-        if let ResponseItem::Message { content, .. } = item {
-            for content_item in content.iter_mut() {
-                if matches!(content_item, ContentItem::InputAudio { .. }) {
-                    *content_item = ContentItem::InputText {
-                        text: AUDIO_CONTENT_OMITTED_PLACEHOLDER.to_string(),
-                    };
+        match item {
+            ResponseItem::Message { content, .. } => {
+                for content_item in content.iter_mut() {
+                    if matches!(content_item, ContentItem::InputAudio { .. }) {
+                        *content_item = ContentItem::InputText {
+                            text: AUDIO_CONTENT_OMITTED_PLACEHOLDER.to_string(),
+                        };
+                    }
                 }
             }
+            ResponseItem::FunctionCallOutput { output, .. }
+            | ResponseItem::CustomToolCallOutput { output, .. } => {
+                if let Some(content_items) = output.content_items_mut() {
+                    for content_item in content_items.iter_mut() {
+                        if matches!(
+                            content_item,
+                            FunctionCallOutputContentItem::InputAudio { .. }
+                        ) {
+                            *content_item = FunctionCallOutputContentItem::InputText {
+                                text: AUDIO_CONTENT_OMITTED_PLACEHOLDER.to_string(),
+                            };
+                        }
+                    }
+                }
+            }
+            _ => {}
         }
     }
 }

codex-rs/core/src/tools/code_mode/mod.rs

@@ -21,6 +21,7 @@ use serde_json::Value as JsonValue;
 use tokio::sync::OnceCell;
 use tokio_util::sync::CancellationToken;
 
+use crate::audio_preparation::estimate_audio_token_count;
 use crate::function_tool::FunctionCallError;
 use crate::original_image_detail::can_request_original_image_detail;
 use crate::original_image_detail::sanitize_original_image_detail as sanitize_image_detail_items;
@@ -269,7 +270,7 @@ fn truncate_code_mode_result(
         return truncated_items;
     }
 
-    truncate_function_output_items_with_policy(&items, policy)
+    truncate_function_output_items_with_policy(&items, policy, estimate_audio_token_count)
 }
 
 async fn call_nested_tool(
@@ -427,6 +428,20 @@ mod tests {
         );
     }
 
+    #[test]
+    fn over_budget_audio_output_is_omitted() {
+        let items = vec![FunctionCallOutputContentItem::InputAudio {
+            audio_url: format!("data:audio/wav;base64,{}", "A".repeat(100)),
+        }];
+
+        assert_eq!(
+            truncate_code_mode_result(items, Some(5)),
+            vec![FunctionCallOutputContentItem::InputText {
+                text: "[omitted 1 audio items ...]".to_string(),
+            }]
+        );
+    }
+
     #[tokio::test]
     async fn missing_process_host_falls_back_to_in_process_session() {
         let service = CodeModeService::new(Arc::new(