diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index c4ed2837ec..7adf228b1a 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 9069235c69..0cf6e96e98 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1049,6 +1049,9 @@ "fast_mode": { "type": "boolean" }, + "file_backed_tool_audio": { + "type": "boolean" + }, "gateway_setup_wizard": { "type": "boolean" }, @@ -8205,6 +8208,9 @@ "fast_mode": { "type": "boolean" }, + "file_backed_tool_audio": { + "type": "boolean" + }, "gateway_setup_wizard": { "type": "boolean" }, diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index eaa8f6fa87..e0f2ab5cbc 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -18,7 +18,7 @@ use crate::compact_remote_history::trim_function_call_history_to_fit_context_win use crate::context::world_state::WorldState; use crate::context_manager::ContextManager; use crate::context_manager::rejected_audio; -use crate::context_manager::replace_unresolvable_image_files; +use crate::context_manager::replace_unresolvable_media_files; use crate::responses_metadata::CompactionTurnMetadata; use crate::session::CompactionPrefix; use crate::session::session::Session; @@ -340,7 +340,7 @@ pub(crate) async fn run_background_compaction( ) -> CodexResult { tokio::select! { _ = shutdown_cancellation.cancelled() => return Err(CodexErr::TurnAborted), - () = replace_unresolvable_image_files( + () = replace_unresolvable_media_files( prepared_input.input.iter_mut(), sess.services.image_store.as_ref(), ) => {}, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index ee06ba8edd..42c894f294 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -9,7 +9,7 @@ //! Live usage boundaries advance only when a response supplies usable usage. //! Oversized instructions keep an incomplete excerpt for bounded root review, including //! sources recovered from legacy Guardian checkpoints before their raw history is dropped. -//! Prompt projections replace unavailable file images without changing retained history. +//! Prompt projections replace unavailable file images and tool audio without changing retained history. #[path = "history_rejected_audio.rs"] pub(crate) mod rejected_audio; @@ -25,7 +25,7 @@ use crate::context::world_state::WorldState; use crate::context::world_state::WorldStateSnapshot; use crate::context::world_state::WorldStateUpdate; use crate::context_manager::normalize; -use crate::context_manager::replace_unresolvable_image_files; +use crate::context_manager::replace_unresolvable_media_files; use crate::event_mapping::contextual_dev_fragments; use crate::event_mapping::is_contextual_user_message_content; use crate::event_mapping::parse_turn_item; @@ -680,6 +680,7 @@ impl ContextManager { self.sampling.record(&mut processed); if let ResponseItem::FunctionCallOutput { output, .. } | ResponseItem::CustomToolCallOutput { output, .. } = &mut processed.item + && !metadata.is_some_and(|metadata| metadata.tool_audio_budget_applied) { // The override already includes the tool's serialization allowance. let policy = metadata @@ -716,7 +717,7 @@ impl ContextManager { /// Returns the history prepared for sending to the model. This applies a proper /// normalization and drops un-suited items. Unsupported image and audio content /// is stripped from messages and tool outputs according to `input_modalities`. - /// Unavailable file images are replaced with text placeholders. + /// Unavailable file images and tool audio are replaced with text placeholders. pub(crate) async fn for_prompt( self, input_modalities: &[InputModality], @@ -736,7 +737,7 @@ impl ContextManager { image_store: &dyn AttachmentStore, ) -> Vec { self.normalize_history(input_modalities); - replace_unresolvable_image_files( + replace_unresolvable_media_files( Arc::make_mut(&mut self.items) .iter_mut() .map(|envelope| &mut envelope.item), diff --git a/codex-rs/core/src/context_manager/history_rejected_audio.rs b/codex-rs/core/src/context_manager/history_rejected_audio.rs index 75e7f3d236..229a360f6d 100644 --- a/codex-rs/core/src/context_manager/history_rejected_audio.rs +++ b/codex-rs/core/src/context_manager/history_rejected_audio.rs @@ -58,6 +58,7 @@ pub(crate) fn pending_audio_confirmation( enum FileErrorSource<'a> { Audio(&'a str), InlineAudio, + ToolFileAudio, FileImage, Opaque, } @@ -85,9 +86,12 @@ fn visit_sources<'a>(item: &'a ResponseItem, mut visit: impl FnMut(FileErrorSour | ResponseItem::CustomToolCallOutput { output, .. } => { for part in output.content_items().unwrap_or_default() { match part { - FunctionCallOutputContentItem::InputAudio { .. } => { - visit(FileErrorSource::InlineAudio) - } + FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::File { .. }, + } => visit(FileErrorSource::ToolFileAudio), + FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::Inline { .. }, + } => visit(FileErrorSource::InlineAudio), FunctionCallOutputContentItem::InputImage { image: ImageReference::File { .. }, .. @@ -141,8 +145,8 @@ fn opaque_content(item: &ResponseItem) -> Option<&str> { pub(crate) struct SubmittedFileAudio { file_ids: Vec, candidate: Option, - has_inline_audio: bool, - has_file_image: bool, + has_competing_audio: bool, + has_competing_file: bool, } impl SubmittedFileAudio { @@ -152,14 +156,18 @@ impl SubmittedFileAudio { clear_checkpoints: impl Iterator + Clone, ) -> Option { let mut file_ids = Vec::new(); - let mut has_inline_audio = false; - let mut has_file_image = false; + let mut has_competing_audio = false; + let mut has_competing_file = false; let mut unknown_checkpoint = false; for item in items { visit_sources(item, |source| match source { FileErrorSource::Audio(id) => file_ids.push(id), - FileErrorSource::InlineAudio => has_inline_audio = true, - FileErrorSource::FileImage => has_file_image = true, + FileErrorSource::InlineAudio => has_competing_audio = true, + FileErrorSource::ToolFileAudio => { + has_competing_audio = true; + has_competing_file = true; + } + FileErrorSource::FileImage => has_competing_file = true, FileErrorSource::Opaque => { unknown_checkpoint |= !opaque_content(item).is_some_and(|ciphertext| { clear_checkpoints.clone().any(|known| known == ciphertext) @@ -185,8 +193,8 @@ impl SubmittedFileAudio { Some(Self { file_ids: file_ids.into_iter().map(str::to_owned).collect(), candidate, - has_inline_audio, - has_file_image, + has_competing_audio, + has_competing_file, }) } @@ -195,9 +203,9 @@ impl SubmittedFileAudio { let unique_file = self.file_ids.len() == 1; let attributable = match error { CodexErrorDetails::InvalidAudioInput { per_recording, .. } => { - !self.has_inline_audio && (unique_file || *per_recording) + !self.has_competing_audio && (unique_file || *per_recording) } - _ => unique_file && !self.has_file_image && is_unknown_file_error(error), + _ => unique_file && !self.has_competing_file && is_unknown_file_error(error), }; attributable.then_some(index) } diff --git a/codex-rs/core/src/context_manager/history_rejected_audio_tests.rs b/codex-rs/core/src/context_manager/history_rejected_audio_tests.rs index 6609c63468..22461e697d 100644 --- a/codex-rs/core/src/context_manager/history_rejected_audio_tests.rs +++ b/codex-rs/core/src/context_manager/history_rejected_audio_tests.rs @@ -108,6 +108,38 @@ fn tool_output_audio_makes_an_invalid_audio_error_ambiguous() { } } +#[test] +fn only_file_backed_tool_audio_competes_with_a_user_upload_for_an_unknown_file_error() { + use serde_json::json; + for kind in ["function_call_output", "custom_tool_call_output"] { + for (part, expected) in [ + ( + json!({"type":"input_audio","file_id":"expired_tool_clip"}), + None, + ), + ( + json!({"type":"input_audio","audio_url":"data:audio/wav;base64,AA=="}), + Some("valid_upload"), + ), + ] { + let output = + serde_json::from_value(json!({"type":kind, "call_id":"call", "output":[part]})) + .unwrap(); + let items = [message(vec![file_audio("valid_upload")]), output]; + assert_eq!( + rejected_file_audio_id(&unknown_file_error(), &items).as_deref(), + expected, + "{kind}" + ); + assert_eq!( + rejected_file_audio_id(&audio_error(), &items), + None, + "{kind}" + ); + } + } +} + #[test] fn attributes_audio_error_only_to_one_visible_file_audio() { let error = audio_error(); @@ -145,9 +177,8 @@ fn attributes_audio_error_only_to_one_visible_file_audio() { ); } -#[test] -fn unknown_file_error_counts_the_submitted_projection_not_raw_history() { - let error = CodexErrorDetails::UnexpectedStatus(codex_protocol::error::UnexpectedResponseError { +fn unknown_file_error() -> CodexErrorDetails { + CodexErrorDetails::UnexpectedStatus(codex_protocol::error::UnexpectedResponseError { status: http::StatusCode::NOT_FOUND, body: serde_json::json!({"error": {"message": "Invalid input: Unknown `file_id` in input content."}}).to_string(), url: None, @@ -156,7 +187,12 @@ fn unknown_file_error_counts_the_submitted_projection_not_raw_history() { request_id: None, identity_authorization_error: None, identity_error_code: None, - }); + }) +} + +#[test] +fn unknown_file_error_counts_the_submitted_projection_not_raw_history() { + let error = unknown_file_error(); let unavailable_image_projection = message(vec![ ContentItem::InputText { text: "[Image unavailable]".to_string(), diff --git a/codex-rs/core/src/context_manager/image_availability.rs b/codex-rs/core/src/context_manager/media_availability.rs similarity index 71% rename from codex-rs/core/src/context_manager/image_availability.rs rename to codex-rs/core/src/context_manager/media_availability.rs index c831fd22d1..bb350a234c 100644 --- a/codex-rs/core/src/context_manager/image_availability.rs +++ b/codex-rs/core/src/context_manager/media_availability.rs @@ -1,4 +1,4 @@ -//! Projects unavailable file images into placeholders without changing tool attribution. +//! Projects unavailable images and tool audio into placeholders without changing tool attribution. use std::collections::HashSet; @@ -6,6 +6,7 @@ use codex_attachment_store::AttachmentStore; use codex_context_fragments::AnnotatedContent; use codex_context_fragments::set_annotated_content; use codex_context_fragments::to_annotated_content; +use codex_protocol::models::AudioReference; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; use codex_protocol::models::ContentItemMetadata; @@ -15,12 +16,12 @@ use codex_protocol::models::ResponseItem; use futures::StreamExt; use futures::stream; -const MAX_CONCURRENT_IMAGE_AVAILABILITY_CHECKS: usize = 8; +const MAX_CONCURRENT_FILE_AVAILABILITY_CHECKS: usize = 8; -/// Checks distinct file IDs concurrently and replaces only unavailable image content. -pub(crate) async fn replace_unresolvable_image_files<'a>( +/// Checks distinct file IDs concurrently and replaces only media the store confirms is unavailable. +pub(crate) async fn replace_unresolvable_media_files<'a>( items: impl Iterator, - image_store: &dyn AttachmentStore, + store: &dyn AttachmentStore, ) { let items: Vec<_> = items.collect(); let mut file_ids = HashSet::new(); @@ -44,6 +45,9 @@ pub(crate) async fn replace_unresolvable_image_files<'a>( if let FunctionCallOutputContentItem::InputImage { image: ImageReference::File { file_id }, .. + } + | FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::File { file_id }, } = part { file_ids.insert(file_id.clone()); @@ -56,10 +60,10 @@ pub(crate) async fn replace_unresolvable_image_files<'a>( } let mut checks = stream::iter(file_ids) .map(|file_id| async move { - let resolvable = image_store.is_resolvable(&file_id).await; + let resolvable = store.is_resolvable(&file_id).await; (file_id, resolvable) }) - .buffer_unordered(MAX_CONCURRENT_IMAGE_AVAILABILITY_CHECKS); + .buffer_unordered(MAX_CONCURRENT_FILE_AVAILABILITY_CHECKS); let mut unavailable_file_ids = HashSet::new(); while let Some((file_id, resolvable)) = checks.next().await { if !resolvable { @@ -103,14 +107,22 @@ pub(crate) async fn replace_unresolvable_image_files<'a>( | ResponseItem::CustomToolCallOutput { output, .. } => { if let Some(content) = output.content_items_mut() { for part in content { - if let FunctionCallOutputContentItem::InputImage { - image: ImageReference::File { file_id }, - .. - } = part - && unavailable_file_ids.contains(file_id) - { + let (file_id, placeholder) = match part { + FunctionCallOutputContentItem::InputImage { + image: ImageReference::File { file_id }, + .. + } => (file_id, "[image unavailable]"), + FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::File { file_id }, + } => ( + file_id, + "audio content omitted because its file is unavailable", + ), + _ => continue, + }; + if unavailable_file_ids.contains(file_id) { *part = FunctionCallOutputContentItem::InputText { - text: "[image unavailable]".to_string(), + text: placeholder.into(), }; } } diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 67be01956c..4621b8295b 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -1,5 +1,5 @@ mod history; -mod image_availability; +mod media_availability; mod normalize; mod sampling; pub(crate) mod updates; @@ -11,5 +11,5 @@ 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; +pub(crate) use media_availability::replace_unresolvable_media_files; pub(crate) use sampling::sampled_through; diff --git a/codex-rs/core/src/session/file_audio.rs b/codex-rs/core/src/session/file_audio.rs index a94cca9785..15b46fcfd2 100644 --- a/codex-rs/core/src/session/file_audio.rs +++ b/codex-rs/core/src/session/file_audio.rs @@ -8,13 +8,19 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use codex_attachment_store::AttachmentStore; use codex_attachment_store::FormatSpecificMetadata; use codex_attachment_store::ResolveRequest; +use codex_attachment_store::UploadRequest; +use codex_attachment_store::UploadResult; use codex_history::ResponseItemEnvelope; use codex_protocol::models::AudioReference; -use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ResponseItem; use codex_utils_audio::estimate_file_audio_tokens; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::truncate_measured_tool_output; use futures::StreamExt; use futures::stream; @@ -37,17 +43,206 @@ pub(crate) async fn snapshot(session: &Session) -> Arc { pub(super) fn file_ids(items: &[ResponseItemEnvelope]) -> impl Iterator { items .iter() - .filter_map(|envelope| match &envelope.item { - ResponseItem::Message { content, .. } => Some(content), - _ => None, + .flat_map(|envelope| envelope.item.file_audio_ids()) +} + +/// Only fresh results pass through the store. A replay keeps the representation in the rollout. +pub(super) async fn prepare_new_items( + items: &mut [ResponseItemEnvelope], + store: &dyn AttachmentStore, + thread_id: &str, + ephemeral: bool, + default_token_limit: usize, +) { + for ResponseItemEnvelope { item, metadata } in items { + let (ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. }) = item + else { + codex_utils_audio::prepare_response_items(std::slice::from_mut(item)); + continue; + }; + if !output.contains_audio() { + continue; + } + let metadata = metadata.get_or_insert_default(); + // The producer's output limit can be larger than the history limit (e.g. Code Mode). + // Apply the latter to the whole result while audio still has its inline cost. + let limit = metadata + .history_truncation_token_limit + .unwrap_or(default_token_limit); + let producer_budgeted = std::mem::take(&mut metadata.producer_bounded_tool_audio); + // Validate and measure on a worker. Only clips retained by history need their + // decoded bytes materialized for upload, one at a time. + let Some(parts) = output.content_items_mut() else { + continue; + }; + let parts = std::mem::take(parts); + let preserve_producer_limit = + producer_budgeted && metadata.history_truncation_token_limit.is_none(); + let Ok((mut admitted, repair_limit)) = tokio::task::spawn_blocking(move || { + admit_tool_audio(parts, limit, preserve_producer_limit) }) - .flatten() - .filter_map(|part| match part { - ContentItem::InputAudio { + .await + else { + *output = codex_protocol::models::FunctionCallOutputPayload::from_text( + "audio content omitted because it could not be processed".into(), + ); + continue; + }; + let mut replaced_with_text = false; + for (part, charge) in &mut admitted { + if let FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::Inline { audio_url }, + } = part + { + let url = std::mem::take(audio_url); + *part = store_tool_audio(store, thread_id, ephemeral, url).await; + if matches!(part, FunctionCallOutputContentItem::InputText { .. }) { + *charge = 0; + replaced_with_text = true; + } + } + } + if replaced_with_text { + admitted = + truncate_measured_tool_output(admitted, TruncationPolicy::Tokens(repair_limit)); + } + output.body = codex_protocol::models::FunctionCallOutputBody::ContentItems( + admitted.into_iter().map(|(part, _)| part).collect(), + ); + metadata.tool_audio_budget_applied = true; + } +} + +fn admit_tool_audio( + parts: Vec, + history_limit: usize, + preserve_producer_limit: bool, +) -> (Vec<(FunctionCallOutputContentItem, usize)>, usize) { + let mut measured = Vec::with_capacity(parts.len()); + for mut part in parts { + let charge = match &mut part { + FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::Inline { audio_url }, + } => match codex_utils_audio::prepare_and_measure_inline_audio(audio_url) { + Ok(tokens) => tokens, + Err(error) => { + part = FunctionCallOutputContentItem::InputText { + text: error.placeholder().into(), + }; + 0 + } + }, + FunctionCallOutputContentItem::InputAudio { + audio: AudioReference::File { .. }, + } => estimate_file_audio_tokens(/*duration*/ None), + _ => 0, + }; + measured.push((part, charge)); + } + let admitted = truncate_measured_tool_output(measured, TruncationPolicy::Tokens(history_limit)); + // Admission matches inline history. Only storage can introduce a new placeholder + // after that point. With no explicit producer limit, keep such repairs within + // the size of the result the model would otherwise have received. + let repair_limit = if preserve_producer_limit { + admitted + .iter() + .fold(0_usize, |sum, (part, charge)| { + let cost = match part { + FunctionCallOutputContentItem::InputText { text } => approx_token_count(text), + _ => *charge, + }; + sum.saturating_add(cost) + }) + .min(history_limit) + } else { + history_limit + }; + (admitted, repair_limit) +} + +async fn store_tool_audio( + store: &dyn AttachmentStore, + thread_id: &str, + ephemeral: bool, + inline_url: String, +) -> FunctionCallOutputContentItem { + let (mime, data, inline_url) = match tokio::task::spawn_blocking(move || { + codex_utils_audio::decode_audio(&inline_url).map(|(mime, data)| (mime, data, inline_url)) + }) + .await + { + Ok(Ok(decoded)) => decoded, + Ok(Err(error)) => { + return FunctionCallOutputContentItem::InputText { + text: error.placeholder().into(), + }; + } + Err(_) => { + return FunctionCallOutputContentItem::InputText { + text: "audio content omitted because it could not be processed".into(), + }; + } + }; + // These tags come from Core. MCP servers and Code Mode never supply attachment IDs or policy tags. + let tags = HashMap::from([ + ("usage".to_string(), "tool_audio".to_string()), + ("mime_type".to_string(), mime.to_string()), + ]); + let result = store + .upload(UploadRequest { + thread_id: thread_id.to_owned(), + ephemeral, + file_name: None, + tags, + data, + }) + .await; + tool_audio_after_upload(result, mime, inline_url) +} + +fn tool_audio_after_upload( + result: Result, + mime: &str, + mut inline_url: String, +) -> FunctionCallOutputContentItem { + match result { + Ok(UploadResult::File { file_id }) if !file_id.trim().is_empty() => { + FunctionCallOutputContentItem::InputAudio { audio: AudioReference::File { file_id }, - } => Some(file_id), - _ => None, - }) + } + } + Ok(UploadResult::Inline { bytes }) => FunctionCallOutputContentItem::inline_audio( + codex_utils_audio::audio_data_url(mime, &bytes), + ), + Err(error) if error.kind() == codex_attachment_store::AttachmentStoreErrorKind::Backend => { + tracing::warn!(kind = ?error.kind(), "tool audio upload failed; using inline audio"); + // decode_audio validated the base64 and normalized the MIME alias. + if let Some(separator) = inline_url.find(',') { + inline_url.replace_range(..separator, &format!("data:{mime};base64")); + } + FunctionCallOutputContentItem::inline_audio(inline_url) + } + Err(error) + if error.kind() + == codex_attachment_store::AttachmentStoreErrorKind::InvalidAttachment => + { + FunctionCallOutputContentItem::InputText { + text: "audio content omitted because the attachment store rejected the clip".into(), + } + } + Ok(UploadResult::File { .. }) | Err(_) => { + tracing::warn!("could not store tool audio attachment"); + FunctionCallOutputContentItem::InputText { + text: "audio content omitted because it could not be stored".to_owned(), + } + } + } +} + +pub(super) async fn resolve_pending(session: &Session) { + let ids = std::mem::take(&mut session.state.lock().await.pending_audio_duration_ids); + resolve(session, ids).await; } pub(super) async fn resolve(session: &Session, mut file_ids: BTreeSet) { diff --git a/codex-rs/core/src/session/file_audio_tests.rs b/codex-rs/core/src/session/file_audio_tests.rs index 0e12b5b340..5871dbf048 100644 --- a/codex-rs/core/src/session/file_audio_tests.rs +++ b/codex-rs/core/src/session/file_audio_tests.rs @@ -44,3 +44,123 @@ fn history_sums_measured_and_unknown_file_audio() { ); assert_eq!(history.estimated_tokens_since_last_usage(), 12_300); } + +#[test] +fn transient_upload_failure_retains_inline_audio_but_invalid_attachments_are_omitted() { + use codex_attachment_store::AttachmentStoreError; + use codex_attachment_store::AttachmentStoreErrorKind; + use serde_json::json; + for (kind, expected) in [ + ( + AttachmentStoreErrorKind::Backend, + json!({"type":"input_audio", "audio_url":"data:audio/wav;base64,QQ=="}), + ), + ( + AttachmentStoreErrorKind::InvalidAttachment, + json!({"type":"input_text", "text":"audio content omitted because the attachment store rejected the clip"}), + ), + ] { + let result = Err(AttachmentStoreError::new(kind, "upload failed")); + let part = + tool_audio_after_upload(result, "audio/wav", "data:audio/x-wav;base64,QQ==".into()); + assert_eq!(serde_json::to_value(part).unwrap(), expected); + } +} + +async fn prepared_tool_output( + parts: serde_json::Value, + metadata: codex_history::CodexHarnessMetadata, + limit: usize, +) -> serde_json::Value { + let mut item = ResponseItemEnvelope { + item: serde_json::from_value( + serde_json::json!({"type":"function_call_output","call_id":"clip","output":parts}), + ) + .unwrap(), + metadata: Some(metadata), + }; + prepare_new_items( + std::slice::from_mut(&mut item), + &codex_attachment_store::InlineAttachmentStore, + "thread", + /*ephemeral*/ false, + limit, + ) + .await; + let metadata = item.metadata.unwrap(); + assert!(!metadata.producer_bounded_tool_audio); + assert!(metadata.tool_audio_budget_applied); + assert_eq!( + serde_json::to_value(&metadata).unwrap()["tool_audio_budget_applied"], + true + ); + serde_json::to_value(item.item).unwrap()["output"].take() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn producer_bounded_audio_keeps_the_same_parts_as_inline_history() { + use codex_history::CodexHarnessMetadata; + use serde_json::json; + let clip = json!({"type":"input_audio","audio_url":"data:audio/wav;base64,UklGRiYAAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQIAAAAAAA=="}); + let caption = json!({"type":"input_text","text":"caption ".repeat(100)}); + let policy = TruncationPolicy::Tokens(20); + let limit = codex_utils_output_truncation::with_serialization_allowance(policy).token_budget(); + let malformed = json!({"type":"input_audio","audio_url":format!("data:audio/wav;base64,{}", "!".repeat(8192))}); + for (parts, keep_clip) in [ + (json!([malformed, clip]), true), + (json!([{"type":"input_audio","audio_url":""}, clip]), true), + (json!([clip, caption]), true), + (json!([caption, clip]), false), + ] { + let mut inline = serde_json::from_value( + json!({"type":"function_call_output","call_id":"clip","output":parts}), + ) + .unwrap(); + codex_utils_audio::prepare_response_items(std::slice::from_mut(&mut inline)); + let mut history = ContextManager::new(uuid_v7_item_id_generator_for_test()); + history.record_items([&inline], policy); + let expected = + serde_json::to_value(history.raw_items().last().unwrap()).unwrap()["output"].take(); + assert_eq!(expected.as_array().unwrap().contains(&clip), keep_clip); + let metadata = CodexHarnessMetadata { + producer_bounded_tool_audio: true, + ..Default::default() + }; + assert_eq!(prepared_tool_output(parts, metadata, limit).await, expected); + } +} + +#[tokio::test] +async fn zero_cost_invalid_audio_cannot_expand_into_unbounded_error_text() { + use codex_history::CodexHarnessMetadata; + use serde_json::json; + for producer_budgeted in [false, true] { + let mut parts = vec![json!({"type":"input_text","text":"keep"})]; + parts.extend((0..100).map(|_| json!({"type":"input_audio","audio_url":""}))); + let metadata = CodexHarnessMetadata { + producer_bounded_tool_audio: producer_budgeted, + ..Default::default() + }; + assert_eq!( + prepared_tool_output(parts.into(), metadata, /*limit*/ 1).await, + json!([{"type":"input_text","text":"keep"},{"type":"input_text","text":"[omitted 100 text items ...]"}]) + ); + } +} + +#[tokio::test] +async fn known_producer_limit_keeps_the_replacement_and_following_caption() { + use serde_json::json; + let caption = json!({"type":"input_text","text":"caption"}); + let parts = json!([{"type":"input_audio","audio_url":""}, caption]); + let metadata = codex_history::CodexHarnessMetadata { + producer_bounded_tool_audio: true, + history_truncation_token_limit: Some(100), + ..Default::default() + }; + let error = codex_utils_audio::decode_audio("").unwrap_err(); + assert_eq!( + prepared_tool_output(parts, metadata, /*limit*/ 1).await, + json!([{"type":"input_text","text":error.placeholder()},caption]) + ); +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index a4a4bef0e1..a82aac213c 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3961,9 +3961,25 @@ impl Session { self.services.image_store.as_ref(), )) .await; + let file_backed = turn_context + .config + .features + .enabled(Feature::FileBackedToolAudio); + if file_backed { + Box::pin(file_audio::prepare_new_items( + &mut items, + self.services.image_store.as_ref(), + &self.thread_id.to_string(), + turn_context.config.ephemeral, + with_serialization_allowance(model_info.truncation_policy.into()).token_budget(), + )) + .await; + } for envelope in &mut items { let item = &mut envelope.item; - prepare_audio_response_items(std::slice::from_mut(item)); + if !file_backed { + prepare_audio_response_items(std::slice::from_mut(item)); + } // Most response items get their passthrough turn ID at the durable history boundary. Self::stamp_response_item_for_history(item, &turn_context.sub_id); self.assign_missing_response_item_id(item); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 08e01df6b7..8ebb2e9550 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -678,6 +678,9 @@ pub(crate) async fn run_turn( // Process async hooks only after sampling and its tools have finished. drain_async_hook_results(&sess, &turn_context, /*before_user_prompt*/ false).await; let (has_pending_input, token_status) = async { + // Tool output can add audio without a queued user input. Resolve its + // duration before checking the continuation's compaction threshold. + super::file_audio::resolve_pending(&sess).await; let has_pending_input = sess.input_queue.has_pending_input(&sess.active_turn).await; let token_status = super::context_window::context_window_token_status( @@ -1079,8 +1082,7 @@ pub(crate) async fn run_hooks_and_record_inputs( // 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; + super::file_audio::resolve_pending(sess).await; } Ok(blocked_input && !accepted_user_input) } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index ad1df3f04f..63f0dde03f 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -393,7 +393,7 @@ pub(crate) async fn handle_output_item_done( // The call is already in history; resume needs its output even if we never dispatch. let output = ToolCallRuntime::aborted_response(&call, /*elapsed_secs*/ None); - return Ok(output.into_response()); + return Ok(output.into_response(/*file_backed_tool_audio*/ false)); } runtime.handle_tool_call(call, cancellation_token).await }) diff --git a/codex-rs/core/src/tools/code_mode/output.rs b/codex-rs/core/src/tools/code_mode/output.rs index a85b77ac3e..be3ea1bd59 100644 --- a/codex-rs/core/src/tools/code_mode/output.rs +++ b/codex-rs/core/src/tools/code_mode/output.rs @@ -40,6 +40,10 @@ impl CodeModeToolOutput { } impl ToolOutput for CodeModeToolOutput { + fn output_budget_applied(&self) -> bool { + true + } + fn log_output(&self) -> String { self.output.log_output() } diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 4460563e66..5b2845698d 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -178,6 +178,10 @@ pub struct McpToolOutput { } impl ToolOutput for McpToolOutput { + fn output_budget_applied(&self) -> bool { + true + } + fn log_output(&self) -> String { // Logging has its own budget; do not first apply the model-context budget. let output = self.result.log_output(); diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index ea467edd89..98eeb26405 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -112,6 +112,12 @@ impl ToolCallRuntime { ) -> impl std::future::Future> { let error_call = call.clone(); let source = call.direct_source(); + let file_backed_tool_audio = self + .step_context + .turn + .config + .features + .enabled(codex_features::Feature::FileBackedToolAudio); let recorder = self.session.services.executed_tool_calls.clone(); let item_id_generator = Arc::clone(&self.session.item_id_generator); let recorded_call = recorder.prepare_direct_call(&call, &source, &self.step_context); @@ -135,7 +141,7 @@ impl ToolCallRuntime { { call.set_tool_result_metadata(ToolResultMetadata::new(metadata)); } - result.into_response() + result.into_response(file_backed_tool_audio) } Err(FunctionCallError::Fatal(message)) => return Err(CodexErr::Fatal(message)), Err(other) => { diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 9269b8dd6b..59ff5bbda9 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -35,6 +35,7 @@ use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; use codex_protocol::parse_command::ParsedCommand; use codex_protocol::protocol::EventMsg; use codex_rollout::state_db; @@ -327,7 +328,7 @@ pub(crate) struct AnyToolResult { } impl AnyToolResult { - pub(crate) fn into_response(self) -> ResponseItemEnvelope { + pub(crate) fn into_response(self, file_backed_tool_audio: bool) -> ResponseItemEnvelope { let Self { call_id, payload, @@ -335,12 +336,19 @@ impl AnyToolResult { .. } = self; let history_truncation_token_limit = result.fallback_token_limit_override(); + let item = result.to_response_item(&call_id, &payload).into(); + let producer_bounded_tool_audio = file_backed_tool_audio + && result.output_budget_applied() + && matches!(&item, ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. } if output.contains_audio()); ResponseItemEnvelope { - item: result.to_response_item(&call_id, &payload).into(), - metadata: history_truncation_token_limit.map(|limit| CodexHarnessMetadata { - history_truncation_token_limit: Some(limit), - ..Default::default() - }), + item, + metadata: (producer_bounded_tool_audio || history_truncation_token_limit.is_some()) + .then_some(CodexHarnessMetadata { + producer_bounded_tool_audio, + history_truncation_token_limit, + ..Default::default() + }), } } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 10f445297f..01ce0a6e3a 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -640,7 +640,7 @@ fn post_tool_use_feedback_output_preserves_fallback_token_limit_override( }; assert_eq!( - result.into_response(), + result.into_response(/*file_backed_tool_audio*/ false), ResponseItemEnvelope { item: ResponseItem::from(ResponseInputItem::FunctionCallOutput { call_id: "call-1".to_string(), @@ -654,6 +654,45 @@ fn post_tool_use_feedback_output_preserves_fallback_token_limit_override( ); } +#[test_case::test_case(true; "file-backed audio enabled")] +#[test_case::test_case(false; "file-backed audio disabled")] +fn mcp_audio_producer_limit_is_transient_until_history_admission(file_backed_tool_audio: bool) { + let result = AnyToolResult { + call_id: "clip".into(), + payload: ToolPayload::Function { + arguments: "{}".into(), + }, + result: Box::new(crate::tools::context::McpToolOutput { + result: codex_protocol::mcp::CallToolResult { + content: vec![serde_json::json!({ + "type": "audio", + "mimeType": "audio/wav", + "data": "UklGRiYAAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQIAAAAAAA==", + })], + structured_content: None, + is_error: None, + meta: None, + }, + tool_input: serde_json::json!({}), + result_metadata_capture_allowed: false, + wall_time: Duration::ZERO, + original_image_detail_supported: false, + truncation_policy: TruncationPolicy::Tokens(100), + }), + post_tool_use_payload: None, + }; + + let metadata = result + .into_response(file_backed_tool_audio) + .metadata + .unwrap(); + assert_eq!(metadata.producer_bounded_tool_audio, file_backed_tool_audio); + assert!(!metadata.tool_audio_budget_applied); + let serialized = serde_json::to_value(&metadata).unwrap(); + assert!(serialized.get("producer_bounded_tool_audio").is_none()); + assert!(serialized.get("tool_audio_budget_applied").is_none()); +} + #[test] fn post_tool_use_feedback_output_keeps_code_mode_result_typed() { let result = AnyToolResult { @@ -674,7 +713,7 @@ fn post_tool_use_feedback_output_keeps_code_mode_result_typed() { }; assert_eq!( - result.into_response(), + result.into_response(/*file_backed_tool_audio*/ false), ResponseItemEnvelope::new( ResponseInputItem::FunctionCallOutput { call_id: "call-1".to_string(), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index aa862f9591..bbb1139206 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -692,7 +692,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow ) .await?; - let response = result.into_response().item; + let response = result.into_response(/*file_backed_tool_audio*/ false).item; match response { ResponseItem::FunctionCallOutput { call_id, output, .. diff --git a/codex-rs/core/tests/suite/audio_duration_store.rs b/codex-rs/core/tests/suite/audio_duration_store.rs index caf7750ec1..1678de6e26 100644 --- a/codex-rs/core/tests/suite/audio_duration_store.rs +++ b/codex-rs/core/tests/suite/audio_duration_store.rs @@ -12,22 +12,60 @@ use codex_attachment_store::ResolveFuture; use codex_attachment_store::ResolveRequest; use codex_attachment_store::UploadFuture; use codex_attachment_store::UploadRequest; +use codex_attachment_store::UploadResult; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; #[derive(Default)] pub(super) struct TestAudioDurationStore { pub(super) slow_lookup_started: tokio::sync::Notify, pub(super) measure_short_clips: bool, + pub(super) tool_uploads: Mutex>, + pub(super) tool_file_expired: AtomicBool, + pub(super) tool_duration_reads: AtomicUsize, } impl AttachmentStore for TestAudioDurationStore { fn upload(&self, request: UploadRequest) -> UploadFuture<'_> { + if request + .tags + .get("usage") + .is_some_and(|usage| usage == "tool_audio") + { + self.tool_uploads + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(request); + return Box::pin(async { + Ok(UploadResult::File { + file_id: "file_tool_audio".into(), + }) + }); + } InlineAttachmentStore.upload(request) } + fn is_resolvable<'a>( + &'a self, + file_id: &'a str, + ) -> codex_attachment_store::IsResolvableFuture<'a> { + Box::pin(async move { + file_id != "file_tool_audio" || !self.tool_file_expired.load(Ordering::SeqCst) + }) + } + 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_tool_audio" => { + self.tool_duration_reads.fetch_add(1, Ordering::SeqCst); + Some(FormatSpecificMetadata::Audio(AudioMetadata { + duration_ms: Some(500), + })) + } "file_audio_1" | "file_audio_7" if self.measure_short_clips => { Some(FormatSpecificMetadata::Audio(AudioMetadata { duration_ms: Some(30_000), diff --git a/codex-rs/core/tests/suite/audio_truncation.rs b/codex-rs/core/tests/suite/audio_truncation.rs index 3e24f5939a..c8909c6886 100644 --- a/codex-rs/core/tests/suite/audio_truncation.rs +++ b/codex-rs/core/tests/suite/audio_truncation.rs @@ -1,8 +1,10 @@ +use super::audio_duration_store::TestAudioDurationStore; use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; +use codex_features::Feature; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec; @@ -23,11 +25,11 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; +use std::sync::Arc; const TEST_WAV_SAMPLE_RATE: u32 = 8_000; -const OMITTED_AUDIO_MARKER: &str = "[omitted 1 audio items ...]"; -fn pcm_wav_data_url(sample_count: u32) -> String { +pub(super) fn pcm_wav_data_url(sample_count: u32) -> String { let padding = sample_count % 2; let mut bytes = Vec::new(); bytes.extend_from_slice(b"RIFF"); @@ -49,8 +51,12 @@ fn pcm_wav_data_url(sample_count: u32) -> String { format!("data:audio/wav;base64,{}", BASE64_STANDARD.encode(bytes)) } +#[test_case::test_case(false; "inline")] +#[test_case::test_case(true; "file backed")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted() -> Result<()> { +async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted( + file_backed: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -73,10 +79,21 @@ async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted() -> Result<( ) .await; - let mut builder = test_codex().with_model_info_override("gpt-5.5", |model_info| { - model_info.input_modalities.push(InputModality::Audio); - model_info.truncation_policy = TruncationPolicyConfig::tokens(/*limit*/ 50); - }); + let store = Arc::new(TestAudioDurationStore::default()); + let mut builder = test_codex() + .with_image_store(store.clone()) + .with_config(move |config| { + assert!( + config + .features + .set_enabled(Feature::FileBackedToolAudio, file_backed) + .is_ok() + ); + }) + .with_model_info_override("gpt-5.5", |model_info| { + model_info.input_modalities.push(InputModality::Audio); + model_info.truncation_policy = TruncationPolicyConfig::tokens(/*limit*/ 50); + }); let base_test = builder.build_with_auto_env(&server).await?; let dynamic_tool = DynamicToolSpec::Namespace(DynamicToolNamespaceSpec { name: "codex_app".to_string(), @@ -118,13 +135,24 @@ async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted() -> Result<( else { unreachable!("event guard guarantees DynamicToolCallRequest"); }; + // The four-second clip fits by itself, but the preceding text has consumed 40 + // of the 60 tokens (including the serialization allowance). + let prefix = "note".repeat(40); test.codex .submit(Op::DynamicToolResponse { id: request.call_id, response: DynamicToolResponse { - content_items: vec![DynamicToolCallOutputContentItem::InputAudio { - audio_url: pcm_wav_data_url(/*sample_count*/ 80_000), - }], + content_items: vec![ + DynamicToolCallOutputContentItem::InputText { + text: prefix.clone(), + }, + DynamicToolCallOutputContentItem::InputAudio { + audio_url: pcm_wav_data_url(/*sample_count*/ 32_000), + }, + DynamicToolCallOutputContentItem::InputAudio { + audio_url: pcm_wav_data_url(/*sample_count*/ 80_000), + }, + ], success: true, }, }) @@ -134,6 +162,14 @@ async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted() -> Result<( }) .await; + assert!( + store + .tool_uploads + .lock() + .map_err(|_| anyhow::anyhow!("upload recorder poisoned"))? + .is_empty(), + "over-budget audio must be omitted before upload" + ); let requests = responses_mock.requests(); assert_eq!(requests.len(), 2); let output = requests[1] @@ -145,8 +181,9 @@ async fn dynamic_tool_audio_exceeding_the_output_budget_is_omitted() -> Result<( serde_json::from_value::(output)?, FunctionCallOutputPayload { body: FunctionCallOutputBody::ContentItems(vec![ + FunctionCallOutputContentItem::InputText { text: prefix }, FunctionCallOutputContentItem::InputText { - text: OMITTED_AUDIO_MARKER.to_string(), + text: "[omitted 2 audio items ...]".to_string(), }, ]), success: None, 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 b7e66ed2d2..93b62e9ea9 100644 --- a/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs +++ b/codex-rs/core/tests/suite/scenarios_audio_file_tests.rs @@ -2,8 +2,10 @@ use super::super::audio_duration_store::TestAudioDurationStore; use super::*; +use anyhow::Context as _; use codex_protocol::models::AudioReference; use codex_protocol::openai_models::InputModality; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::turn_input::TurnInputSubmission; use core_test_support::responses; use core_test_support::responses::mount_response_sequence; @@ -82,6 +84,269 @@ async fn uploaded_audio_survives_follow_up_and_resume() -> Result<()> { Ok(()) } +/// The model calls a registered tool directly or via Code Mode. Both persist the store's ID. +#[test_case(ToolMode::Direct, "function_call_output", ThreadHistoryMode::Legacy; "function")] +#[test_case(ToolMode::Direct, "function_call_output", ThreadHistoryMode::Paginated; "paginated function")] +#[test_case(ToolMode::CodeMode, "custom_tool_call_output", ThreadHistoryMode::Legacy; "code_mode")] +#[test_case(ToolMode::CodeMode, "custom_tool_call_output", ThreadHistoryMode::Paginated; "paginated code_mode")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn tool_audio_store_id_survives_rollout_and_next_request( + mode: ToolMode, + output: &str, + history_mode: ThreadHistoryMode, +) -> Result<()> { + use std::sync::atomic::Ordering; + + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let call_id = "recording-call"; + let store = Arc::new(TestAudioDurationStore::default()); + let builder = || { + test_codex() + .with_history_mode(history_mode) + .with_image_store(store.clone()) + .with_config(move |config| { + config.allow_persistent_dynamic_tool_updates = true; + assert!(config.features.enable(Feature::FileBackedToolAudio).is_ok()); + assert!( + config + .features + .set_enabled(Feature::CodeMode, mode == ToolMode::CodeMode) + .is_ok() + ); + }) + .with_model_info_override("gpt-5.5", move |model| { + model.input_modalities.push(InputModality::Audio); + // An unknown file reserves 12,000 tokens; this half-second clip must continue without compaction. + model.auto_compact_token_limit = Some(10_000); + model.tool_mode = Some(mode); + model.use_responses_lite = mode == ToolMode::CodeMode; + model.truncation_policy = + codex_protocol::openai_models::TruncationPolicyConfig::tokens( + /*limit*/ 50, + ); + }) + }; + let test = builder().build_with_auto_env(&server).await?; + test.codex + .replace_dynamic_tools(vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "recording".into(), + description: "Return an audio clip.".into(), + input_schema: json!({"type":"object","properties":{},"additionalProperties":false}), + defer_loading: false, + })]) + .await?; + + let call = if mode == ToolMode::CodeMode { + ev_custom_tool_call( + call_id, + "exec", + r#"audio(await tools.recording({})); text("caption ".repeat(400));"#, + ) + } else { + responses::ev_function_call(call_id, "recording", "{}") + }; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("r1"), call, ev_completed("r1")]), + sse(vec![ + ev_assistant_message("heard", "I received the clip."), + ev_completed("r2"), + ]), + sse(vec![ + ev_assistant_message("expired", "The recording is unavailable."), + ev_completed("r3"), + ]), + sse(vec![ + ev_assistant_message("resumed", "The recording is still unavailable."), + ev_completed("r4"), + ]), + sse(vec![ + ev_assistant_message("restored", "The recording is available again."), + ev_completed("r5"), + ]), + ], + ) + .await; + + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![text( + "Listen to the recording.", + )])) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::DynamicToolCallRequest(call) => Some(call.clone()), + _ => None, + }) + .await; + assert_eq!(request.tool, "recording"); + let wav = super::super::audio_truncation::pcm_wav_data_url(/*sample_count*/ 4_000); + test.codex + .submit(Op::DynamicToolResponse { + id: request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputAudio { + audio_url: wav.clone(), + }], + success: true, + }, + }) + .await?; + let live_parts = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ItemCompleted(event) => match &event.item { + TurnItem::DynamicToolCall(call) if call.tool == "recording" => { + call.content_items.clone() + } + _ => None, + }, + _ => None, + }) + .await; + assert_eq!( + live_parts, + vec![DynamicToolCallOutputContentItem::InputAudio { + audio_url: wav.clone() + }] + ); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + store.tool_duration_reads.load(Ordering::SeqCst), + 1, + "duration resolved during the original tool turn" + ); + store.tool_file_expired.store(true, Ordering::SeqCst); + test.submit_text_turn("Check the recording again.").await?; + let rollout = test.codex.rollout_path().expect("durable thread rollout"); + test.codex.shutdown_and_wait().await?; + // Reopen the durable store in a fresh manager, as a hosted resume would. + let restarted = builder() + .with_home(test.home.clone()) + .build_with_auto_env(&server) + .await?; + let stored = restarted + .thread_store + .load_latest_model_context(codex_thread_store::LoadThreadHistoryParams { + thread_id: test.session_configured.thread_id, + include_archived: true, + }) + .await?; + let resumed = restarted + .thread_manager + .resume_thread_with_history( + restarted.config.clone(), + codex_history::InitialHistory::Resumed(codex_history::ResumedHistory { + history_revision: None, + conversation_id: stored.thread_id, + history: Arc::new(stored.items), + rollout_path: Some(rollout.clone()), + }), + restarted.thread_manager.auth_manager(), + /*parent_trace*/ None, + Default::default(), + ) + .await?; + assert_eq!( + resumed.thread.config_snapshot().await.history_mode, + history_mode + ); + for (expired, prompt) in [ + (true, "Listen after restart."), + (false, "Listen after the file is restored."), + ] { + store.tool_file_expired.store(expired, Ordering::SeqCst); + resumed + .thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![text(prompt)])) + .await?; + wait_for_event(&resumed.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } + let requests = mock.requests(); + assert_eq!(requests.len(), 5); + let expected = json!({"type":"input_audio","file_id":"file_tool_audio"}); + let persisted = std::fs::read_to_string(rollout)?; + let items = persisted + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + let audio_output = + |item: &serde_json::Value| item["type"] == output && item["call_id"] == call_id; + let history = items + .iter() + .map(|line| &line["payload"]) + .find(|item| audio_output(item)) + .expect("persisted tool output"); + let output_parts = history["output"] + .as_array() + .context("persisted tool content")?; + let audio = output_parts + .iter() + .filter(|part| part["type"] == "input_audio") + .collect::>(); + assert_eq!(audio, vec![&expected]); + if mode == ToolMode::CodeMode { + let caption = output_parts + .iter() + .filter_map(|part| part["text"].as_str()) + .find(|text| text.contains("caption")) + .expect("some caption remains after the clip"); + assert!( + caption.len() < 400, + "history must bound the 2,800-byte caption" + ); + } + let mut unavailable = history["output"].clone(); + for part in unavailable + .as_array_mut() + .context("tool content for the expired-file request")? + { + if *part == expected { + *part = json!({"type":"input_text","text":"audio content omitted because its file is unavailable"}); + } + } + for (request, expected_output) in requests[1..].iter().zip([ + &history["output"], + &unavailable, + &unavailable, + &history["output"], + ]) { + let item = request + .input() + .into_iter() + .find(&audio_output) + .expect("tool result after execution and resume"); + assert_eq!(&item["output"], expected_output); + } + { + let uploads = store + .tool_uploads + .lock() + .map_err(|_| anyhow::anyhow!("upload recorder poisoned"))?; + let [upload] = uploads.as_slice() else { + panic!("expected one upload across execution and resume"); + }; + assert_eq!(upload.tags["usage"], "tool_audio"); + assert_eq!(upload.data, codex_utils_audio::decode_audio(&wav)?.1); + } + insta::assert_snapshot!( + format!("file_backed_{output}_trace"), + context_snapshot::format_request_history_snapshot( + "The tool returns file-backed audio. Follow-up and cold resume omit the expired clip, then use the original file again when it is available.", + &requests, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .include_request_settings(), + ) + ); + Ok(()) +} + #[derive(Clone, Copy, PartialEq, Eq)] enum AudioAtCompaction { Available, diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_custom_tool_call_output_trace.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_custom_tool_call_output_trace.snap new file mode 100644 index 0000000000..ef17013129 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_custom_tool_call_output_trace.snap @@ -0,0 +1,136 @@ +--- +source: core/tests/suite/scenarios_audio_file_tests.rs +expression: "context_snapshot::format_request_history_snapshot(\"The tool returns file-backed audio. Follow-up and cold resume omit the expired clip, then use the original file again when it is available.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: The tool returns file-backed audio. Follow-up and cold resume omit the expired clip, then use the original file again when it is available. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-5.5" + parallel_tool_calls: false + prompt_cache_key: "" + reasoning: {"context":"all_turns","effort":"medium"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" +-- request 1 (turn) -- +00:additional_tools/developer (2; hash=3299611bc98d0784): + - namespace/functions; hash=1976557bc4879146 + - custom/exec + - function/wait + - function/exec_command + - function/write_stdin + - function/request_user_input + - custom/apply_patch + - function/view_image + - function/recording + - tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=b573ed69c0b0e6c7 +01:message/developer: + +02:message/developer[2]: + [01] + [02] + Deferred tool namespaces: + - multi_agent_v1: Tools for spawning and managing sub-agents. + +03:message/user: + +04:message/user: + Listen to the recording. +-- request 2 (turn) -- +05:custom_tool_call/exec:audio(await tools.recording({})); text("caption ".repeat(400)); +06:custom_tool_call_output:Script completed + Wall time seconds + Output: + | | caption caption caption caption caption caption caption caption caption caption captio…757 tokens trunca...ption caption caption caption caption caption caption [hash=a35e70d094c6690c] + +## Window 2 (after request 2: input diverged at item 06) +Settings: same as window 1 +-- request 3 (turn) -- +00:additional_tools/developer (2; hash=3299611bc98d0784): + - namespace/functions; hash=1976557bc4879146 + - custom/exec + - function/wait + - function/exec_command + - function/write_stdin + - function/request_user_input + - custom/apply_patch + - function/view_image + - function/recording + - tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=b573ed69c0b0e6c7 +01:message/developer: + +02:message/developer[2]: + [01] + [02] + Deferred tool namespaces: + - multi_agent_v1: Tools for spawning and managing sub-agents. + +03:message/user: + +04:message/user: + Listen to the recording. +05:custom_tool_call/exec:audio(await tools.recording({})); text("caption ".repeat(400)); +06:custom_tool_call_output:Script completed + Wall time seconds + Output: + | audio content omitted because its file is unavailable | caption caption caption caption caption caption caption caption caption caption captio…757 tokens trunca...ption caption caption caption caption caption caption [hash=a35e70d094c6690c] +07:message/assistant: + I received the clip. +08:message/user: + Check the recording again. +-- request 4 (turn) -- +09:message/assistant: + The recording is unavailable. +10:message/user: + +11:message/user: + Listen after restart. + +## Window 3 (after request 4: input diverged at item 06) +Settings: same as window 1 +-- request 5 (turn) -- +00:additional_tools/developer (2; hash=3299611bc98d0784): + - namespace/functions; hash=1976557bc4879146 + - custom/exec + - function/wait + - function/exec_command + - function/write_stdin + - function/request_user_input + - custom/apply_patch + - function/view_image + - function/recording + - tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=b573ed69c0b0e6c7 +01:message/developer: + +02:message/developer[2]: + [01] + [02] + Deferred tool namespaces: + - multi_agent_v1: Tools for spawning and managing sub-agents. + +03:message/user: + +04:message/user: + Listen to the recording. +05:custom_tool_call/exec:audio(await tools.recording({})); text("caption ".repeat(400)); +06:custom_tool_call_output:Script completed + Wall time seconds + Output: + | | caption caption caption caption caption caption caption caption caption caption captio…757 tokens trunca...ption caption caption caption caption caption caption [hash=a35e70d094c6690c] +07:message/assistant: + I received the clip. +08:message/user: + Check the recording again. +09:message/assistant: + The recording is unavailable. +10:message/user: + +11:message/user: + Listen after restart. +12:message/assistant: + The recording is still unavailable. +13:message/user: + Listen after the file is restored. diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_function_call_output_trace.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_function_call_output_trace.snap new file mode 100644 index 0000000000..80f69b422e --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__audio_file__file_backed_function_call_output_trace.snap @@ -0,0 +1,103 @@ +--- +source: core/tests/suite/scenarios_audio_file_tests.rs +expression: "context_snapshot::format_request_history_snapshot(\"The tool returns file-backed audio. Follow-up and cold resume omit the expired clip, then use the original file again when it is available.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: The tool returns file-backed audio. Follow-up and cold resume omit the expired clip, then use the original file again when it is available. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-5.5" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"effort":"medium"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" + tools (8; hash=fb7f1f10f0db22ba): + - function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5 + - function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2 + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba + - function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[path]; hash=1908de9c0be7118a + - function/recording: Return an audio clip.; args=[]; hash=aeaf4746261bcbb1 + - tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=b573ed69c0b0e6c7 + - web_search; hash=f103e327c6870a46 +-- 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: + Listen to the recording. +-- request 2 (turn) -- +04:function_call/recording:{} +05:function_call_output: + +## Window 2 (after request 2: input diverged at item 05) +Settings: same as window 1 +-- 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: + Listen to the recording. +04:function_call/recording:{} +05:function_call_output:audio content omitted because its file is unavailable +06:message/assistant: + I received the clip. +07:message/user: + Check the recording again. +-- request 4 (turn) -- +08:message/assistant: + The recording is unavailable. +09:message/user: + +10:message/user: + Listen after restart. + +## Window 3 (after request 4: input diverged at item 05) +Settings: same as window 1 +-- request 5 (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: + Listen to the recording. +04:function_call/recording:{} +05:function_call_output: +06:message/assistant: + I received the clip. +07:message/user: + Check the recording again. +08:message/assistant: + The recording is unavailable. +09:message/user: + +10:message/user: + Listen after restart. +11:message/assistant: + The recording is still unavailable. +12:message/user: + Listen after the file is restored. diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 669cd73697..bf4742c704 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -402,6 +402,8 @@ pub enum Feature { ImageGenerationAnimation, /// Omit inline image and audio content from app-server item notifications. OmitAppServerNotificationMedia, + /// Allow this host's attachment store to prepare fresh tool audio before persistence. + FileBackedToolAudio, /// Tell the model when a prompt image was resized and include its dimensions. ImageResizeNotice, /// Normalize prompt images with the reviewed image sanitizer. @@ -2036,6 +2038,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::FileBackedToolAudio, + key: "file_backed_tool_audio", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::ImageResizeNotice, key: "image_resize_notice", diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index cbf3934e42..8a2ec57b37 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -126,6 +126,15 @@ pub struct CodexHarnessMetadata { )] pub history_truncation_token_limit: Option, + /// Fresh tool output has already passed a producer limit. Used only to budget + /// storage-error placeholders before the item is recorded. Never persisted. + #[serde(skip)] + pub producer_bounded_tool_audio: bool, + + /// Core applied the history limit while tool audio still had its inline cost. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub tool_audio_budget_applied: bool, + /// Bounded assistant text confirmed by a successful messaging tool result. /// Captured after input hooks; untrusted context, never user authorization. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 97fa5985c0..05273a0cac 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -2532,6 +2532,14 @@ impl Default for FunctionCallOutputBody { } impl FunctionCallOutputPayload { + pub fn contains_audio(&self) -> bool { + self.content_items().is_some_and(|parts| { + parts + .iter() + .any(|part| matches!(part, FunctionCallOutputContentItem::InputAudio { .. })) + }) + } + pub fn from_text(content: String) -> Self { Self { body: FunctionCallOutputBody::Text(content), diff --git a/codex-rs/tools/src/tool_output.rs b/codex-rs/tools/src/tool_output.rs index 01481a1eaf..5f6f80d34d 100644 --- a/codex-rs/tools/src/tool_output.rs +++ b/codex-rs/tools/src/tool_output.rs @@ -29,6 +29,11 @@ pub trait ToolOutput: Send { false } + /// Whether the complete model-visible result was already limited by its producer. + fn output_budget_applied(&self) -> bool { + false + } + /// Overrides history's fallback token limit after tool-specific truncation. /// Include any serialization allowance; history uses this limit unchanged. fn fallback_token_limit_override(&self) -> Option { @@ -95,6 +100,10 @@ where (**self).contains_external_context() } + fn output_budget_applied(&self) -> bool { + (**self).output_budget_applied() + } + fn fallback_token_limit_override(&self) -> Option { (**self).fallback_token_limit_override() } diff --git a/codex-rs/utils/audio/src/audio_preparation_tests.rs b/codex-rs/utils/audio/src/audio_preparation_tests.rs index 07e609ea9f..fc2d6abe53 100644 --- a/codex-rs/utils/audio/src/audio_preparation_tests.rs +++ b/codex-rs/utils/audio/src/audio_preparation_tests.rs @@ -210,3 +210,13 @@ fn file_audio_uses_the_measured_length_or_the_full_file_budget() { assert_eq!(super::estimate_file_audio_tokens(duration), expected); } } + +#[test] +fn audio_data_url_encodes_all_padding_lengths_directly_into_the_url() { + for (bytes, expected) in [(b"a".as_slice(), "YQ=="), (b"ab", "YWI="), (b"abc", "YWJj")] { + assert_eq!( + super::audio_data_url("audio/wav", bytes), + format!("data:audio/wav;base64,{expected}") + ); + } +} diff --git a/codex-rs/utils/audio/src/lib.rs b/codex-rs/utils/audio/src/lib.rs index edd626500a..1a8333d0a2 100644 --- a/codex-rs/utils/audio/src/lib.rs +++ b/codex-rs/utils/audio/src/lib.rs @@ -41,7 +41,7 @@ static AUDIO_TOKEN_ESTIMATE_CACHE: LazyLock &'static str { + pub fn placeholder(&self) -> &'static str { match self { AudioPreparationError::InvalidDataUrl { .. } => AUDIO_PROCESSING_ERROR_PLACEHOLDER, AudioPreparationError::UnsupportedFormat => UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER, @@ -192,7 +192,14 @@ fn estimate_audio_token_count(audio_url: &str) -> usize { /// The same duration-based estimate without accessing the Tokio-dependent cache. /// Suitable for bounded output rendered on a current-thread runtime. fn estimate_audio_token_count_without_cache(audio_url: &str) -> usize { - let Some(duration_seconds) = audio_duration_seconds(audio_url) else { + let Ok((mime, bytes)) = decode_audio(audio_url) else { + return approx_token_count(audio_url); + }; + decoded_audio_token_count(mime, bytes, audio_url) +} + +fn decoded_audio_token_count(mime: &str, bytes: Vec, audio_url: &str) -> usize { + let Some(duration_seconds) = audio_duration_seconds(mime, bytes) else { return approx_token_count(audio_url); }; let token_count = (duration_seconds * f64::from(AUDIO_TOKENS_PER_SECOND)).ceil(); @@ -203,22 +210,7 @@ fn estimate_audio_token_count_without_cache(audio_url: &str) -> usize { } } -fn audio_duration_seconds(audio_url: &str) -> Option { - 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; - } - }; +fn audio_duration_seconds(canonical_mime: &str, bytes: Vec) -> Option { let media_source = MediaSourceStream::new(Box::new(Cursor::new(bytes)), Default::default()); let mut hint = Hint::new(); hint.mime_type(canonical_mime); @@ -247,7 +239,44 @@ fn audio_duration_seconds(audio_url: &str) -> Option { duration_seconds.is_finite().then_some(duration_seconds) } +/// Validate the data URL and measure its decoded duration in the same pass. +pub fn prepare_and_measure_inline_audio( + audio_url: &mut String, +) -> Result { + let (mime, bytes) = decode_audio(audio_url)?; + normalize_validated_header(audio_url, mime); + Ok(decoded_audio_token_count(mime, bytes, audio_url)) +} + fn prepare_audio(audio_url: &mut String) -> Result<(), AudioPreparationError> { + let (mime, _) = decode_audio(audio_url)?; + normalize_validated_header(audio_url, mime); + Ok(()) +} + +// The standard base64 decoder has already validated the payload, including padding. +fn normalize_validated_header(url: &mut String, mime: &str) { + let canonical = format!("data:{mime};base64"); + if let Some((header, _)) = url.split_once(',') + && header != canonical + { + url.replace_range(..header.len(), &canonical); + } +} + +/// Canonical inline representation for a validated payload. +pub fn audio_data_url(mime: &str, data: &[u8]) -> String { + let mut url = + String::with_capacity("data:;base64,".len() + mime.len() + data.len().div_ceil(3) * 4); + url.push_str("data:"); + url.push_str(mime); + url.push_str(";base64,"); + BASE64_STANDARD.encode_string(data, &mut url); + url +} + +/// Parses a size-bounded audio data URL. A file-writing host also verifies its media container. +pub fn decode_audio(audio_url: &str) -> Result<(&'static str, Vec), AudioPreparationError> { if !is_data_url(audio_url) { return Err(AudioPreparationError::InvalidDataUrl { reason: "audio input must be a data URL", @@ -300,9 +329,7 @@ fn prepare_audio(audio_url: &mut String) -> Result<(), AudioPreparationError> { return Err(AudioPreparationError::AudioTooLarge { size: bytes.len() }); } - let encoded = BASE64_STANDARD.encode(bytes); - *audio_url = format!("data:{canonical_mime};base64,{encoded}"); - Ok(()) + Ok((canonical_mime, bytes)) } #[cfg(test)] diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index b0dcbfdd11..e676a9d3c7 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -311,13 +311,47 @@ fn truncate_output_items( } } +/// Apply the mixed-output policy to owned, premeasured parts. Audio keeps its original +/// token charge through truncation and storage. The charge for other parts is zero. +pub fn truncate_measured_tool_output( + items: Vec<(FunctionCallOutputContentItem, usize)>, + policy: TruncationPolicy, +) -> Vec<(FunctionCallOutputContentItem, usize)> { + truncate_measured_output_iter( + items + .into_iter() + .map(|(item, charge)| (Cow::Owned(item), charge)), + policy, + &mut [], + |item, charge| (item, charge), + ) +} + fn truncate_output_iter<'a>( items: impl ExactSizeIterator>, policy: TruncationPolicy, estimate_audio_token_count: impl Fn(&AudioReference) -> usize, boundaries: &mut [OutputBoundary], ) -> Vec { - let mut out: Vec = Vec::with_capacity(items.len()); + let measured = items.map(|item| { + let charge = match item.as_ref() { + FunctionCallOutputContentItem::InputAudio { audio } => { + estimate_audio_token_count(audio) + } + _ => 0, + }; + (item, charge) + }); + truncate_measured_output_iter(measured, policy, boundaries, |item, _| item) +} + +fn truncate_measured_output_iter<'a, T>( + items: impl ExactSizeIterator, usize)>, + policy: TruncationPolicy, + boundaries: &mut [OutputBoundary], + emit: impl Fn(FunctionCallOutputContentItem, usize) -> T, +) -> Vec { + let mut out = Vec::with_capacity(items.len()); let mut remaining_budget = match policy { TruncationPolicy::Bytes(_) => policy.byte_budget(), TruncationPolicy::Tokens(_) => policy.token_budget(), @@ -326,7 +360,7 @@ fn truncate_output_iter<'a>( let mut omitted_audio_items = 0usize; let mut next_boundary = 0; - for (index, item) in items.enumerate() { + for (index, (item, audio_tokens)) in items.enumerate() { let item_out = out.len(); let item_policy = match policy { TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget), @@ -354,7 +388,10 @@ fn truncate_output_iter<'a>( if snippet.is_empty() { omitted_text_items += 1; } else { - out.push(FunctionCallOutputContentItem::InputText { text: snippet }); + out.push(emit( + FunctionCallOutputContentItem::InputText { text: snippet }, + 0, + )); } remaining_budget = 0; false @@ -362,11 +399,10 @@ fn truncate_output_iter<'a>( } FunctionCallOutputContentItem::InputImage { .. } | FunctionCallOutputContentItem::EncryptedContent { .. } => true, - FunctionCallOutputContentItem::InputAudio { audio } => { - let token_cost = estimate_audio_token_count(audio); + FunctionCallOutputContentItem::InputAudio { .. } => { let cost = match policy { - TruncationPolicy::Bytes(_) => approx_bytes_for_tokens(token_cost), - TruncationPolicy::Tokens(_) => token_cost, + TruncationPolicy::Bytes(_) => approx_bytes_for_tokens(audio_tokens), + TruncationPolicy::Tokens(_) => audio_tokens, }; if cost <= remaining_budget { remaining_budget = remaining_budget.saturating_sub(cost); @@ -390,7 +426,7 @@ fn truncate_output_iter<'a>( next_boundary += 1; } if keep { - out.push(item.into_owned()); + out.push(emit(item.into_owned(), audio_tokens)); } } // End boundaries precede generated omission counts: counts are not original @@ -401,14 +437,20 @@ fn truncate_output_iter<'a>( } if omitted_text_items > 0 { - out.push(FunctionCallOutputContentItem::InputText { - text: format!("[omitted {omitted_text_items} text items ...]"), - }); + out.push(emit( + FunctionCallOutputContentItem::InputText { + text: format!("[omitted {omitted_text_items} text items ...]"), + }, + 0, + )); } if omitted_audio_items > 0 { - out.push(FunctionCallOutputContentItem::InputText { - text: format!("[omitted {omitted_audio_items} audio items ...]"), - }); + out.push(emit( + FunctionCallOutputContentItem::InputText { + text: format!("[omitted {omitted_audio_items} audio items ...]"), + }, + 0, + )); } out } diff --git a/codex-rs/utils/output-truncation/src/truncate_tests.rs b/codex-rs/utils/output-truncation/src/truncate_tests.rs index ddb230eb5c..9f8cf0432c 100644 --- a/codex-rs/utils/output-truncation/src/truncate_tests.rs +++ b/codex-rs/utils/output-truncation/src/truncate_tests.rs @@ -8,6 +8,7 @@ use crate::formatted_truncate_text; use crate::mixed_output_at_token_boundaries; use crate::truncate_function_output_items_with_policy; use crate::truncate_mcp_tool_result; +use crate::truncate_measured_tool_output; use crate::truncate_text; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::AudioReference; @@ -577,3 +578,26 @@ fn owned_mixed_output_moves_a_preserved_image_across_text_truncation() { }; assert_eq!(image_url.as_ptr(), url_ptr); } + +#[test] +fn measured_tool_output_keeps_charges_paired_with_retained_clips() { + let clip = FunctionCallOutputContentItem::inline_audio("clip".into()); + let text = FunctionCallOutputContentItem::InputText { + text: "keep".into(), + }; + let omitted = FunctionCallOutputContentItem::InputText { + text: "[omitted 1 audio items ...]".into(), + }; + assert_eq!( + truncate_measured_tool_output( + vec![ + (clip.clone(), 9), + (text.clone(), 0), + (clip.clone(), 2), + (clip.clone(), 1) + ], + TruncationPolicy::Tokens(4) + ), + vec![(text, 0), (clip.clone(), 2), (clip, 1), (omitted, 0)] + ); +}