diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1d27abbfa4..a2ec85317a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2803,6 +2803,8 @@ dependencies = [ "codex-code-mode-protocol", "codex-plugin-host", "codex-protocol", + "codex-utils-audio", + "codex-utils-output-truncation", "deno_core_icudata", "futures", "opentelemetry", diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index b064cfc7f8..f3cd0286f0 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -15,9 +15,17 @@ pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024; /// Stdio connection support for the optional `session/open.outputTruncation` field. /// Each session selects its own setting after the connection hello. pub const BOUNDED_OUTPUT_CAPABILITY: &str = "bounded-output-v1"; -/// Maximum retained bytes between server-truncated observations. +/// Maximum aggregate text bytes retained per server-truncated observation. pub const BOUNDED_OUTPUT_MAX_BYTES: usize = 1024 * 1024; +/// Maximum aggregate inline image bytes retained per server-truncated observation. +pub const BOUNDED_OUTPUT_MAX_IMAGE_BYTES: usize = 8 * 1024 * 1024; + +/// Maximum aggregate encoded audio URL bytes retained per server-truncated observation. +/// Count the whole URL, including the data URL prefix and any base64 payload. +/// This leaves heuristic transport headroom but does not bound the complete response. +pub const BOUNDED_OUTPUT_MAX_AUDIO_BYTES: usize = 62 * 1024 * 1024; + /// Negotiated support for cell execution resource limits on `session/open`. pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits"; diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index e5e72c571c..d70fcd4e18 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -41,13 +41,14 @@ pub enum CodeModeOutputTruncation { /// transmitting output (10,000 tokens by default). Clients must not truncate /// it again or append the script error: status and error metadata remain separate. /// - /// Between observations, the host retains at most - /// [`crate::host::BOUNDED_OUTPUT_MAX_BYTES`] of output payload. It also - /// bounds the storage for retained items separately, without charging it - /// against the payload or observation token budgets. - /// Ordinary text notices report overflow, even for a zero-token budget. - /// Media are indivisible and can be omitted if they exceed the collector. - /// Without overflow, formatting matches Core's client policy. + /// A complete observation keeps up to [`crate::host::BOUNDED_OUTPUT_MAX_BYTES`] + /// of text, [`crate::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES`] of inline image URLs, + /// and [`crate::host::BOUNDED_OUTPUT_MAX_AUDIO_BYTES`] of audio URLs, subject + /// to a separate item limit. While these limits hold, formatting matches + /// Core's client policy. On text overflow the collector keeps a prefix and a + /// rolling tail. Each image or audio clip is admitted if it fits its remaining + /// byte allowance. Eviction at the item limit can free space for later media. + /// Media are indivisible. Bounded loss notices survive even a zero-token budget. Server, } diff --git a/codex-rs/code-mode-runtime/Cargo.toml b/codex-rs/code-mode-runtime/Cargo.toml index 786fad083d..e55086f3c6 100644 --- a/codex-rs/code-mode-runtime/Cargo.toml +++ b/codex-rs/code-mode-runtime/Cargo.toml @@ -17,6 +17,8 @@ base64 = { workspace = true } codex-code-mode-protocol = { workspace = true } codex-plugin-host = { workspace = true } codex-protocol = { workspace = true } +codex-utils-audio = { workspace = true } +codex-utils-output-truncation = { workspace = true } deno_core_icudata = { workspace = true } futures = { workspace = true, features = ["executor", "std"] } opentelemetry = { workspace = true } diff --git a/codex-rs/code-mode-runtime/src/lib.rs b/codex-rs/code-mode-runtime/src/lib.rs index 0a3e29775e..01bfbba5a3 100644 --- a/codex-rs/code-mode-runtime/src/lib.rs +++ b/codex-rs/code-mode-runtime/src/lib.rs @@ -1,4 +1,10 @@ mod cell_actor; +#[allow( + dead_code, + unused_imports, + reason = "Wired into the runtime in the next stack layer" +)] +mod output; mod runtime; mod service; mod session_runtime; diff --git a/codex-rs/code-mode-runtime/src/output/mod.rs b/codex-rs/code-mode-runtime/src/output/mod.rs new file mode 100644 index 0000000000..47a86e9d5d --- /dev/null +++ b/codex-rs/code-mode-runtime/src/output/mod.rs @@ -0,0 +1,124 @@ +//! Bounded retention and per-observation rendering for Code Mode output. +//! +//! Keeps omission accounting inside this module until rendered. No actor or +//! transport state belongs here. + +use codex_code_mode_protocol::CodeModeOutputTruncation; +use codex_code_mode_protocol::FunctionCallOutputContentItem as OutputItem; +use codex_utils_output_truncation::OutputBoundary; + +mod output_collection; +mod output_render; + +pub(crate) use output_collection::BoundedOutput; +pub(crate) use output_render::render_output; + +/// The actor transfers the whole collector at each observation, so a failed +/// synchronous send can restore the original state, including any omissions. +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub(crate) enum OutputCollector { + Client(Vec), + /// Allocate storage only when the observation receives output. + Server(Option>), +} + +impl OutputCollector { + pub(crate) fn new(mode: CodeModeOutputTruncation) -> Self { + match mode { + CodeModeOutputTruncation::Client => Self::Client(Vec::new()), + CodeModeOutputTruncation::Server => Self::Server(None), + } + } + + pub(crate) fn push(&mut self, item: OutputItem) { + match self { + Self::Client(items) => items.push(item), + Self::Server(output) => output.get_or_insert_with(Box::default).push(item), + } + } + + pub(crate) fn take(&mut self) -> Self { + match self { + Self::Client(items) => Self::Client(std::mem::take(items)), + Self::Server(output) => Self::Server(output.take()), + } + } + + pub(crate) fn try_deliver(&mut self, send: impl FnOnce(Self) -> Result<(), Self>) -> bool { + if let Err(undelivered) = send(self.take()) { + *self = undelivered; + return false; + } + true + } + + /// Flatten the retained state only after the observer owns it. + pub(crate) fn into_collected_output(self) -> CollectedOutput { + match self { + Self::Client(items) => items.into(), + Self::Server(Some(output)) => output.into_collected_output(), + Self::Server(None) => CollectedOutput::default(), + } + } +} + +/// Retained content and where the collector had to drop data. Loss is private +/// to the runtime until it is rendered as text after the observation's budget. +#[derive(Clone, Debug, Default, PartialEq)] +pub(crate) struct CollectedOutput { + items: Vec, + gaps: Vec, +} + +impl From> for CollectedOutput { + fn from(items: Vec) -> Self { + Self { + items, + gaps: Vec::new(), + } + } +} + +impl CollectedOutput { + fn push_gap(&mut self, boundary: OutputBoundary, omitted: OutputOmission) { + if let Some(gap) = self.gaps.last_mut() + && gap.boundary == boundary + { + gap.omitted.combine(omitted); + } else { + self.gaps.push(OutputGap { boundary, omitted }); + } + } + + /// Takes the unformatted items. Only for client-side observations, where + /// the caller already owns truncation and the input was not bounded. + pub(crate) fn into_untruncated_items(self) -> Vec { + self.items + } +} + +/// Missing output at a boundary in the retained logical items. A nonzero +/// UTF-8 byte offset marks loss inside a single text item. +#[derive(Clone, Copy, Debug, PartialEq)] +struct OutputGap { + boundary: OutputBoundary, + omitted: OutputOmission, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +struct OutputOmission { + text_bytes: u64, + text_items: u64, + image_items: u64, + audio_items: u64, +} + +impl OutputOmission { + fn combine(&mut self, other: Self) { + self.text_bytes = self.text_bytes.saturating_add(other.text_bytes); + self.text_items = self.text_items.saturating_add(other.text_items); + self.image_items = self.image_items.saturating_add(other.image_items); + self.audio_items = self.audio_items.saturating_add(other.audio_items); + } +} diff --git a/codex-rs/code-mode-runtime/src/output/output_collection.rs b/codex-rs/code-mode-runtime/src/output/output_collection.rs new file mode 100644 index 0000000000..923831b1ae --- /dev/null +++ b/codex-rs/code-mode-runtime/src/output/output_collection.rs @@ -0,0 +1,507 @@ +//! Bounded head/tail retention for raw Code Mode output. Text has a rolling +//! byte allowance. Admit whole images and audio clips in source order against +//! independent byte allowances. Item pressure can free space for later media. + +use std::collections::VecDeque; + +use super::CollectedOutput; +use super::OutputOmission; + +use codex_code_mode_protocol::FunctionCallOutputContentItem as OutputItem; +use codex_utils_output_truncation::OutputBoundary; + +pub(super) use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_AUDIO_BYTES as MAX_AUDIO_BYTES; +pub(super) use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_BYTES as MAX_OUTPUT_BYTES; +pub(super) use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES as MAX_IMAGE_BYTES; + +/// Limit logical items independently of payload bytes. An observation that +/// fits this limit must keep its item boundaries even if our storage layout changes. +pub(crate) const MAX_OUTPUT_ITEMS: usize = 32 * 1024; +const TEXT_BLOCK_BYTES: usize = 4096; + +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +enum RetainedItem { + Content(OutputItem), + Partial(Box), + Gap(OutputOmission), +} + +/// Store a partially evicted suffix in blocks so later writes can release +/// its oldest bytes without copying the rest. Only the first block has an offset. +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +struct PartialText { + blocks: VecDeque, + front_offset: usize, + bytes: usize, + omitted: u64, + continues_prefix: bool, +} + +impl RetainedItem { + fn retain(mut item: OutputItem) -> Self { + let payload = match &mut item { + OutputItem::InputImage { image_url, .. } => image_url, + OutputItem::InputAudio { audio_url } => audio_url, + OutputItem::InputText { text } if text.len() > MAX_OUTPUT_BYTES => { + // Eviction copies only the head and suffix that can survive. + return Self::Content(item); + } + OutputItem::InputText { text } => text, + }; + payload.shrink_to_fit(); + Self::Content(item) + } + + /// Evict bytes after the retained prefix and keep any surviving suffix. + fn text_suffix(text: &str, prefix_len: usize, bytes_to_evict: usize) -> Self { + let tail = &text[prefix_len..]; + let omitted = tail.ceil_char_boundary(bytes_to_evict.min(tail.len())); + let mut suffix = &tail[omitted..]; + if suffix.is_empty() { + return Self::Gap(OutputOmission { + text_bytes: omitted as u64, + text_items: 1, + ..OutputOmission::default() + }); + } + let bytes = suffix.len(); + let mut blocks = VecDeque::new(); + while !suffix.is_empty() { + let end = suffix.floor_char_boundary(suffix.len().min(TEXT_BLOCK_BYTES)); + blocks.push_back(suffix[..end].to_owned()); + suffix = &suffix[end..]; + } + Self::Partial(Box::new(PartialText { + blocks, + front_offset: 0, + bytes, + omitted: omitted as u64, + continues_prefix: prefix_len > 0, + })) + } + + fn content_bytes(&self) -> usize { + match self { + Self::Content(OutputItem::InputText { text }) => text.len(), + Self::Partial(text) => text.bytes, + Self::Content(OutputItem::InputImage { .. } | OutputItem::InputAudio { .. }) + | Self::Gap(_) => 0, + } + } + + fn omission(&self) -> OutputOmission { + let mut loss = OutputOmission::default(); + match self { + Self::Content(OutputItem::InputText { text }) => { + loss.text_bytes = text.len() as u64; + loss.text_items = 1; + } + Self::Content(OutputItem::InputImage { .. }) => loss.image_items = 1, + Self::Content(OutputItem::InputAudio { .. }) => loss.audio_items = 1, + Self::Partial(text) => { + loss.text_bytes = text.omitted.saturating_add(text.bytes as u64); + loss.text_items = 1; + } + Self::Gap(omitted) => loss = *omitted, + } + loss + } + + fn move_into(self, output: &mut CollectedOutput) { + match self { + Self::Content(item) => output.items.push(item), + Self::Gap(loss) => output.push_gap( + OutputBoundary { + item_index: output.items.len(), + text_offset: 0, + }, + loss, + ), + Self::Partial(text) => { + // A split suffix continues the preceding prefix. A suffix whose + // prefix was discarded starts its own logical item. + if !text.continues_prefix { + output.items.push(OutputItem::InputText { + text: String::new(), + }); + } + let item_index = output.items.len() - 1; + let OutputItem::InputText { text: retained } = &mut output.items[item_index] else { + unreachable!("a split suffix follows its text prefix"); + }; + let text_offset = retained.len(); + retained.reserve_exact(text.bytes); + for (i, block) in text.blocks.into_iter().enumerate() { + let start = if i == 0 { text.front_offset } else { 0 }; + retained.push_str(&block[start..]); + } + if text.omitted > 0 { + output.push_gap( + OutputBoundary { + item_index, + text_offset, + }, + OutputOmission { + text_bytes: text.omitted, + ..OutputOmission::default() + }, + ); + } + } + } + } +} + +#[derive(Clone, Debug, Default)] +#[cfg_attr(test, derive(PartialEq))] +struct OpenOutput { + items: VecDeque, + content_bytes: usize, +} + +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub(crate) struct BoundedOutput { + state: Retention, + media: MediaBytes, +} + +#[derive(Clone, Debug, Default)] +#[cfg_attr(test, derive(PartialEq))] +struct MediaBytes { + images: usize, + audio: usize, +} + +impl MediaBytes { + fn admit(&mut self, item: &OutputItem) -> Result<(), OutputOmission> { + let (len, used, limit, omission) = match item { + OutputItem::InputText { .. } => return Ok(()), + OutputItem::InputImage { image_url, .. } => ( + image_url.len(), + &mut self.images, + MAX_IMAGE_BYTES, + OutputOmission { + image_items: 1, + ..OutputOmission::default() + }, + ), + OutputItem::InputAudio { audio_url } => ( + audio_url.len(), + &mut self.audio, + MAX_AUDIO_BYTES, + OutputOmission { + audio_items: 1, + ..OutputOmission::default() + }, + ), + }; + if len > limit - *used { + return Err(omission); + } + *used += len; + Ok(()) + } +} + +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +enum Retention { + /// All individually fitting items are intact, in source order. + Open(OpenOutput), + /// A fixed prefix and a rolling tail with omissions in source order. + HeadTail(HeadTailOutput), +} + +impl Default for BoundedOutput { + fn default() -> Self { + Self { + state: Retention::Open(OpenOutput::default()), + media: MediaBytes::default(), + } + } +} + +#[derive(Clone, Debug)] +#[cfg_attr(test, derive(PartialEq))] +struct HeadTailOutput { + head: Vec, + head_bytes: usize, + tail: VecDeque, + // The first tail item that might still have text bytes. + cursor: usize, + // Adjacent tail gaps occupy one logical item. Repack their physical slots in batches. + adjacent_gaps: usize, + // Text bytes retained in the entire observation, including the head. + content_bytes: usize, + // Gaps at the front of the tail no longer need their own slots. + middle: OutputOmission, +} + +impl HeadTailOutput { + fn from_open(open: OpenOutput) -> Self { + let mut output = Self { + head: Vec::new(), + head_bytes: 0, + tail: open.items, + cursor: 0, + adjacent_gaps: 0, + content_bytes: open.content_bytes, + middle: OutputOmission::default(), + }; + while output.head.len() < MAX_OUTPUT_ITEMS / 2 { + let Some(item) = output.tail.pop_front() else { + break; + }; + let bytes = item.content_bytes(); + let remaining = MAX_OUTPUT_BYTES / 2 - output.head_bytes; + if bytes <= remaining { + output.head_bytes += bytes; + output.head.push(item); + } else if output.content_bytes <= MAX_OUTPUT_BYTES { + // Under item pressure a whole item may borrow unused tail bytes. + output.head_bytes += bytes; + output.head.push(item); + break; + } else { + output.tail.push_front(item); + output.split_front(remaining); + break; + } + } + output + } + + /// Preserve a text prefix and copy only suffix bytes that could still fit. + fn split_front(&mut self, remaining: usize) { + let Some(RetainedItem::Content(OutputItem::InputText { text })) = self.tail.front() else { + return; + }; + let split = text.floor_char_boundary(remaining); + if split == 0 { + return; + } + self.head.push(RetainedItem::Content(OutputItem::InputText { + text: text[..split].to_owned(), + })); + self.head_bytes += split; + let excess = self.content_bytes - MAX_OUTPUT_BYTES; + let suffix = RetainedItem::text_suffix(text, split, excess); + self.content_bytes -= text.len() - split - suffix.content_bytes(); + if matches!(suffix, RetainedItem::Gap(_)) + && matches!(self.tail.get(1), Some(RetainedItem::Gap(_))) + { + self.adjacent_gaps += 1; + } + self.tail[0] = suffix; + } + + fn advance_front(&mut self) -> RetainedItem { + self.cursor = self.cursor.saturating_sub(1); + let Some(item) = self.tail.pop_front() else { + unreachable!("eviction requires a nonempty tail"); + }; + if matches!(item, RetainedItem::Gap(_)) + && matches!(self.tail.front(), Some(RetainedItem::Gap(_))) + { + self.adjacent_gaps -= 1; + } + item + } + + fn trim(&mut self, media: &mut MediaBytes) { + if self.content_bytes > MAX_OUTPUT_BYTES && self.head_bytes > MAX_OUTPUT_BYTES / 2 { + // Only the last prefix item can have borrowed space. Reconsider it + // before the middle gap, whose source position follows that item. + let Some(item) = self.head.pop() else { + unreachable!("borrowed bytes require a prefix item"); + }; + self.head_bytes -= item.content_bytes(); + if self.middle != OutputOmission::default() { + self.tail + .push_front(RetainedItem::Gap(std::mem::take(&mut self.middle))); + } + self.tail.push_front(item); + self.cursor = 0; + let remaining = MAX_OUTPUT_BYTES / 2 - self.head_bytes; + self.split_front(remaining); + } + // Skip retained media and preserve each gap between its source neighbors. + let used = &mut self.content_bytes; + let cursor = &mut self.cursor; + while *used > MAX_OUTPUT_BYTES { + let slot = &mut self.tail[*cursor]; + let old_bytes = slot.content_bytes(); + if old_bytes > 0 { + let excess = *used - MAX_OUTPUT_BYTES; + let remaining = old_bytes.saturating_sub(excess); + match slot { + RetainedItem::Content(OutputItem::InputText { text }) if remaining > 0 => { + *slot = RetainedItem::text_suffix(text, /*prefix_len*/ 0, excess); + } + RetainedItem::Partial(text) if remaining > 0 => { + // Later writes may evict more of this same suffix. + while text.bytes > remaining { + let front = &text.blocks[0][text.front_offset..]; + let removed = + front.ceil_char_boundary((text.bytes - remaining).min(front.len())); + text.bytes -= removed; + text.omitted = text.omitted.saturating_add(removed as u64); + if removed == front.len() { + text.blocks.pop_front(); + text.front_offset = 0; + } else { + text.front_offset += removed; + } + } + if text.bytes == 0 { + *slot = RetainedItem::Gap(slot.omission()); + } + } + _ => *slot = RetainedItem::Gap(slot.omission()), + } + let new_bytes = slot.content_bytes(); + *used -= old_bytes - new_bytes; + if new_bytes > 0 { + break; + } + self.adjacent_gaps += usize::from( + *cursor > 0 && matches!(self.tail[*cursor - 1], RetainedItem::Gap(_)), + ) + usize::from(matches!( + self.tail.get(*cursor + 1), + Some(RetainedItem::Gap(_)) + )); + } + *cursor += 1; + } + // Each sweep removes at least this many slots. Tiny writes cannot + // repeatedly scan the tail, even when a live image brackets many gaps. + if self.adjacent_gaps >= MAX_OUTPUT_ITEMS.div_ceil(8) { + let slots = self.tail.make_contiguous(); + for i in 1..slots.len() { + let (prior, rest) = slots.split_at_mut(i); + if let (RetainedItem::Gap(left), RetainedItem::Gap(right)) = + (&mut prior[i - 1], &mut rest[0]) + { + right.combine(std::mem::take(left)); + } + } + self.tail.retain(|item| !matches!(item, RetainedItem::Gap(loss) if *loss == OutputOmission::default())); + self.cursor = 0; + self.adjacent_gaps = 0; + } + let split_suffix = if matches!(self.tail.front(), Some(RetainedItem::Partial(text)) if text.continues_prefix) + { + // The suffix shares its head item's slot, and no middle gap exists yet. + if self.head.len() + self.tail.len() - self.adjacent_gaps <= MAX_OUTPUT_ITEMS + 1 { + return; + } + Some(self.advance_front()) + } else { + None + }; + // A middle gap costs a slot only once it exists. The first eviction + // creates that gap, so recompute the logical item count on each pass. + while self.head.len() + self.tail.len() - self.adjacent_gaps + + usize::from(self.middle != OutputOmission::default()) + > MAX_OUTPUT_ITEMS + { + let item = self.advance_front(); + self.content_bytes -= item.content_bytes(); + match &item { + RetainedItem::Content(OutputItem::InputImage { image_url, .. }) => { + media.images -= image_url.len(); + } + RetainedItem::Content(OutputItem::InputAudio { audio_url }) => { + media.audio -= audio_url.len(); + } + _ => {} + } + self.middle.combine(item.omission()); + } + while matches!(self.tail.front(), Some(RetainedItem::Gap(_))) { + let item = self.advance_front(); + self.middle.combine(item.omission()); + } + if let Some(suffix) = split_suffix { + self.tail + .push_front(RetainedItem::Gap(std::mem::take(&mut self.middle))); + self.tail.push_front(suffix); + self.cursor = 0; + } + } +} + +impl BoundedOutput { + pub(crate) fn push(&mut self, item: OutputItem) { + if let Err(loss) = self.media.admit(&item) { + self.push_retained(RetainedItem::Gap(loss)); + // Adding the loss marker can evict earlier media from a full observation. + // Replace only this item's marker if that freed enough space. + if let Retention::HeadTail(split) = &mut self.state + && let Some(slot) = split.tail.back_mut() + && matches!(slot, RetainedItem::Gap(omitted) if *omitted == loss) + && self.media.admit(&item).is_ok() + { + *slot = RetainedItem::retain(item); + } + return; + } + self.push_retained(RetainedItem::retain(item)); + } + + fn push_retained(&mut self, item: RetainedItem) { + let bytes = item.content_bytes(); + let (tail, retained) = match &mut self.state { + Retention::Open(open) => (&mut open.items, &mut open.content_bytes), + Retention::HeadTail(split) => (&mut split.tail, &mut split.content_bytes), + }; + if let RetainedItem::Gap(loss) = &item + && let Some(RetainedItem::Gap(last)) = tail.back_mut() + { + last.combine(*loss); + return; + } + tail.push_back(item); + *retained += bytes; + if let Retention::Open(open) = &mut self.state { + if open.content_bytes <= MAX_OUTPUT_BYTES && open.items.len() <= MAX_OUTPUT_ITEMS { + return; + } + self.state = Retention::HeadTail(HeadTailOutput::from_open(std::mem::take(open))); + } + if let Retention::HeadTail(split) = &mut self.state { + split.trim(&mut self.media); + } + } + + fn collect(items: impl Iterator) -> CollectedOutput { + let mut output = CollectedOutput::default(); + for item in items { + item.move_into(&mut output); + } + output + } + + pub(crate) fn into_collected_output(self) -> CollectedOutput { + match self.state { + Retention::Open(open) => Self::collect(open.items.into_iter()), + Retention::HeadTail(split) => Self::collect( + split + .head + .into_iter() + .chain( + (split.middle != OutputOmission::default()) + .then_some(RetainedItem::Gap(split.middle)), + ) + .chain(split.tail), + ), + } + } +} + +#[cfg(test)] +#[path = "output_collection_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-runtime/src/output/output_collection_tests.rs b/codex-rs/code-mode-runtime/src/output/output_collection_tests.rs new file mode 100644 index 0000000000..f82d6ed18c --- /dev/null +++ b/codex-rs/code-mode-runtime/src/output/output_collection_tests.rs @@ -0,0 +1,301 @@ +//! Regression coverage for bounded output retention. + +#![cfg(test)] + +use super::*; +use crate::output::OutputCollector; +use codex_code_mode_protocol::CodeModeOutputTruncation; +use codex_code_mode_protocol::ImageDetail; +use pretty_assertions::assert_eq; + +fn text(value: &str) -> OutputItem { + OutputItem::InputText { + text: value.to_owned(), + } +} + +fn image(image_url: String) -> OutputItem { + OutputItem::InputImage { + image_url, + detail: Some(ImageDetail::Original), + } +} + +fn audio(audio_url: String) -> OutputItem { + OutputItem::InputAudio { audio_url } +} + +fn collect(items: impl IntoIterator) -> BoundedOutput { + let mut output = BoundedOutput::default(); + for item in items { + output.push(item); + } + output +} + +fn retained_capacity(collector: &BoundedOutput) -> usize { + let capacity = |item: &RetainedItem| match item { + RetainedItem::Content(OutputItem::InputText { text }) => text.capacity(), + RetainedItem::Content(OutputItem::InputImage { image_url, .. }) => image_url.capacity(), + RetainedItem::Content(OutputItem::InputAudio { audio_url }) => audio_url.capacity(), + RetainedItem::Gap(_) => 0, + RetainedItem::Partial(text) => text.blocks.iter().map(String::capacity).sum(), + }; + match &collector.state { + Retention::Open(open) => open.items.iter().map(capacity).sum(), + Retention::HeadTail(split) => split.head.iter().chain(&split.tail).map(capacity).sum(), + } +} + +#[test] +fn all_32k_empty_writes_fit_before_the_item_limit_creates_a_gap() { + let mut output = collect((0..32_768).map(|_| text(""))); + assert!(output.clone().into_collected_output() == vec![text(""); 32_768].into()); + + output.push(text("last")); + let snapshot = output.into_collected_output(); + assert_eq!(snapshot.items.len(), 32_767); + assert_eq!(snapshot.items.last(), Some(&text("last"))); + let [gap] = snapshot.gaps.as_slice() else { + panic!("expected one item-limit gap") + }; + assert_eq!(gap.boundary.item_index, 16_384); + assert_eq!(gap.omitted.text_items, 2); +} + +#[test] +fn item_pressure_keeps_a_split_head_items_suffix() { + let half = "x".repeat(MAX_OUTPUT_BYTES / 2); + let retained = text(&half.repeat(2)); + let mut output = collect([text(&format!("{half}y{half}"))]); + for _ in 0..MAX_OUTPUT_ITEMS - 1 { + output.push(text("")); + } + assert!(output.clone().into_collected_output().items.first() == Some(&retained)); + + output.push(text("")); + let snapshot = output.into_collected_output(); + assert!(snapshot.items.first() == Some(&retained)); + assert_eq!(snapshot.items.len(), MAX_OUTPUT_ITEMS - 1); + assert_eq!(snapshot.gaps.len(), 2); + assert_eq!(snapshot.gaps[1].boundary.item_index, 1); + assert_eq!(snapshot.gaps[1].omitted.text_items, 2); +} + +#[test] +fn fitting_items_release_spare_capacity() { + let items = [text("first"), image("image".into()), audio("audio".into())]; + let output = collect(items.iter().cloned().map(|mut item| { + let payload = match &mut item { + OutputItem::InputText { text } => text, + OutputItem::InputImage { image_url, .. } => image_url, + OutputItem::InputAudio { audio_url } => audio_url, + }; + payload.reserve(2 * MAX_OUTPUT_BYTES); + item + })); + assert!(retained_capacity(&output) <= MAX_OUTPUT_BYTES); + assert_eq!(output.into_collected_output(), items.to_vec().into()); +} + +#[test] +fn rolling_unicode_suffix_survives_tiny_writes_without_unbounded_backing() { + let half = MAX_OUTPUT_BYTES / 2; + let original = format!("{}{}z", "a".repeat(half - 1), "😀".repeat(half)); + let mut output = collect([text(&original)]); + let mut last_block = None; + for _ in 0..10_001 { + output.push(text("z")); + let Retention::HeadTail(split) = &output.state else { + panic!("expected truncated output") + }; + let Some(RetainedItem::Partial(front)) = split.tail.front() else { + panic!("expected a partial suffix") + }; + let ptr = front.blocks.back().unwrap().as_ptr(); + assert_eq!(*last_block.get_or_insert(ptr), ptr); + assert!(retained_capacity(&output) < MAX_OUTPUT_BYTES + TEXT_BLOCK_BYTES); + } + let head = original.floor_char_boundary(half); + let tail = original.ceil_char_boundary(original.len() - (MAX_OUTPUT_BYTES - head - 10_001)); + let snapshot = output.into_collected_output(); + assert_eq!(snapshot.items.len(), 10_002); + assert!(snapshot.items[0] == text(&[&original[..head], &original[tail..]].concat())); + let [gap] = snapshot.gaps.as_slice() else { + panic!("expected one split") + }; + assert_eq!( + gap.boundary, + OutputBoundary { + item_index: 0, + text_offset: head + } + ); + assert_eq!(gap.omitted.text_bytes, (tail - head) as u64); +} + +#[test] +fn byte_pressure_reclaims_space_borrowed_by_the_first_item() { + let half = MAX_OUTPUT_BYTES / 2; + let suffix = text(&"z".repeat(half)); + let first = text(&"a".repeat(half + 8)); + let mut output = collect([first.clone()]); + for _ in 0..MAX_OUTPUT_ITEMS { + output.push(text("")); + } + assert!(output.clone().into_collected_output().items[0] == first); + + output.push(suffix.clone()); + let snapshot = output.into_collected_output(); + assert_eq!(snapshot.items.last(), Some(&suffix)); + let [gap] = snapshot.gaps.as_slice() else { + panic!("expected one middle gap") + }; + assert!(snapshot.items[0] == text(&"a".repeat(half))); + assert_eq!(gap.boundary.item_index, 1); + assert_eq!(gap.omitted.text_bytes, 8); +} + +#[test] +fn image_admission_uses_remaining_bytes_and_reclaims_them_on_eviction() { + let half = MAX_OUTPUT_ITEMS / 2; + let mut output = collect((0..half).map(|_| text(""))); + let first = image("i".repeat(MAX_IMAGE_BYTES - 2)); + let last = image("zz".into()); + for item in [first.clone(), image("xxx".into()), last.clone()] { + output.push(item); + } + let snapshot = output.clone().into_collected_output(); + assert!(snapshot.items.ends_with(&[first, last])); + let [gap] = snapshot.gaps.as_slice() else { + panic!("expected one rejected image") + }; + assert_eq!(gap.boundary.item_index, half + 1); + assert_eq!(gap.omitted.image_items, 1); + + for _ in 0..MAX_OUTPUT_ITEMS { + output.push(text("")); + } + let full = image("j".repeat(MAX_IMAGE_BYTES)); + output.push(full.clone()); + assert!(output.into_collected_output().items.last() == Some(&full)); +} + +#[test] +fn image_admission_reuses_space_freed_by_the_same_write() { + let half = MAX_OUTPUT_ITEMS / 2; + let mut output = collect((0..half).map(|_| text(""))); + output.push(image("i".repeat(MAX_IMAGE_BYTES))); + for _ in 0..half - 1 { + output.push(text("")); + } + let replacement = image("replacement".into()); + output.push(replacement.clone()); + let snapshot = output.into_collected_output(); + assert_eq!(snapshot.items.last(), Some(&replacement)); + assert_eq!(snapshot.gaps.len(), 1); + assert_eq!(snapshot.gaps[0].omitted.image_items, 1); +} + +#[test] +fn one_extra_text_byte_keeps_a_large_image_in_either_order() { + let half = MAX_OUTPUT_BYTES / 2; + let original = format!("{}b{}", "a".repeat(half), "c".repeat(half - 1)); + let retained = text(&[&original[..half], &original[half + 1..]].concat()); + let large_image = image("i".repeat(MAX_IMAGE_BYTES / 2 + 1)); + for image_index in [0, 1] { + let mut items = vec![text(&original)]; + items.insert(image_index, large_image.clone()); + items.push(text("z")); + let snapshot = collect(items).into_collected_output(); + let mut expected = vec![retained.clone(), text("z")]; + expected.insert(image_index, large_image.clone()); + assert!(snapshot.items == expected); + let [gap] = snapshot.gaps.as_slice() else { + panic!("expected one omitted byte") + }; + assert_eq!( + gap.boundary, + OutputBoundary { + item_index: 1 - image_index, + text_offset: MAX_OUTPUT_BYTES / 2 + } + ); + assert_eq!(gap.omitted.text_bytes, 1); + } +} + +#[test] +fn output_item_types_have_independent_byte_allowances() { + let items = vec![ + text(&"t".repeat(MAX_OUTPUT_BYTES)), + image("i".repeat(MAX_IMAGE_BYTES)), + audio("a".into()), + ]; + assert_eq!(collect(items.clone()).into_collected_output(), items.into()); +} + +#[test] +fn failed_delivery_keeps_the_observation_until_success() { + // An intact image can move through a failed send and a retry without copying. + let url = "i".repeat(MAX_IMAGE_BYTES); + let ptr = url.as_ptr(); + let mut output = OutputCollector::new(CodeModeOutputTruncation::Server); + output.push(image(url)); + assert!(!output.try_deliver(Err)); + assert!(output.try_deliver(|snapshot| { + let snapshot = snapshot.into_collected_output(); + let [OutputItem::InputImage { image_url, .. }] = snapshot.items.as_slice() else { + panic!("expected the original image") + }; + assert_eq!(image_url.as_ptr(), ptr); + Ok(()) + })); + + // The image allowance has reset. A failed send must also preserve a loss notice. + let full = image("j".repeat(MAX_IMAGE_BYTES)); + output.push(full.clone()); + output.push(text(&"x".repeat(2 * MAX_OUTPUT_BYTES))); + let mut expected = output.clone().into_collected_output(); + assert!(expected.items.first() == Some(&full)); + assert_eq!(expected.gaps.len(), 1); + assert!(!output.try_deliver(Err)); + output.push(text("")); + expected.items.push(text("")); + assert!(output.try_deliver(|snapshot| { + assert!(snapshot.into_collected_output() == expected); + Ok(()) + })); + assert_eq!(output.into_collected_output(), CollectedOutput::default()); +} + +#[test] +fn repeated_eviction_keeps_recent_writes_and_gaps_on_each_side_of_an_image() { + let half = "p".repeat(MAX_OUTPUT_BYTES / 2); + let image = image("i".repeat(2 * MAX_OUTPUT_BYTES)); + let mut output = collect([text(&half), text("b"), image.clone(), text(&half)]); + let chunk = "x".repeat(64); + let writes = 2 * MAX_OUTPUT_ITEMS; + for _ in 0..writes { + output.push(text(&chunk)); + } + let snapshot = output.into_collected_output(); + let recent = MAX_OUTPUT_BYTES / 2 / chunk.len(); + assert_eq!(snapshot.items.len(), 2 + recent); + assert!(snapshot.items.starts_with(&[text(&half), image])); + assert!( + snapshot.items[2..] + .iter() + .all(|item| matches!(item, OutputItem::InputText { text } if text == &chunk)) + ); + let [before, after] = snapshot.gaps.as_slice() else { + panic!("expected gaps around the image") + }; + assert_eq!( + (before.boundary.item_index, after.boundary.item_index), + (1, 2) + ); + assert_eq!(before.omitted.text_bytes, 1); + assert_eq!(after.omitted.text_bytes, (writes * chunk.len()) as u64); + assert_eq!(after.omitted.text_items, (1 + writes - recent) as u64); +} diff --git a/codex-rs/code-mode-runtime/src/output/output_render.rs b/codex-rs/code-mode-runtime/src/output/output_render.rs new file mode 100644 index 0000000000..cfff2c10d3 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/output/output_render.rs @@ -0,0 +1,198 @@ +//! Applies the outer Code Mode response budget at an observation boundary. + +use codex_code_mode_protocol::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; +use codex_code_mode_protocol::FunctionCallOutputContentItem as OutputItem; +use codex_code_mode_protocol::ImageDetail; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::ImageDetail as ProtocolImageDetail; +use codex_protocol::models::ImageReference; +use codex_utils_audio::estimate_audio_token_count; +use codex_utils_output_truncation::OutputBoundary; +use codex_utils_output_truncation::formatted_output_at_token_boundaries; +use codex_utils_output_truncation::mixed_output_at_token_boundaries; + +use crate::output::CollectedOutput; +use crate::output::OutputGap; +use crate::output::OutputOmission; + +// Even with four full-width u64 counts per notice, this adds less than 1 KiB. +const MAX_POSITIONED_NOTICES: usize = 4; + +/// Renders the output since the preceding observation, using that request's +/// own budget. The terminal error participates in the same budget as text +/// emitted by the script; the response still carries its error status separately. +/// +/// Image fidelity is preserved here. Model-specific image detail selection +/// still belongs to the caller and does not affect the truncation budget. +/// +/// This is the same output Core would render as long as the session collector +/// has retained the whole observation. Collector omissions stay outside this +/// budget: Core's warning describes the retained text alone. The shared +/// formatter moves collector boundaries along with their source content. +/// If losses are too scattered for individual notices, report their totals once. +pub(crate) fn render_output( + output: CollectedOutput, + max_output_tokens: Option, +) -> Vec { + let CollectedOutput { items, gaps } = output; + let mut boundaries: Vec<_> = gaps.iter().map(|gap| gap.boundary).collect(); + let mut content = Vec::with_capacity(items.len()); + for item in items { + match item { + OutputItem::InputText { text } => { + content.push(FunctionCallOutputContentItem::InputText { text }); + } + OutputItem::InputImage { image_url, detail } => { + content.push(FunctionCallOutputContentItem::InputImage { + image: ImageReference::Inline { image_url }, + detail: detail.map(|detail| match detail { + ImageDetail::Auto => ProtocolImageDetail::Auto, + ImageDetail::Low => ProtocolImageDetail::Low, + ImageDetail::High => ProtocolImageDetail::High, + ImageDetail::Original => ProtocolImageDetail::Original, + }), + }); + } + OutputItem::InputAudio { audio_url } => { + content.push(FunctionCallOutputContentItem::InputAudio { audio_url }); + } + } + } + let max_tokens = max_output_tokens.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL); + // Core budgets mixed output in order. Keep that policy even if the + // collector removed every image or audio item in the observation. + let formatted = if gaps + .iter() + .any(|gap| gap.omitted.image_items > 0 || gap.omitted.audio_items > 0) + { + mixed_output_at_token_boundaries( + content, + max_tokens, + estimate_audio_token_count, + &mut boundaries, + ) + } else { + formatted_output_at_token_boundaries( + content, + max_tokens, + estimate_audio_token_count, + &mut boundaries, + ) + }; + insert_loss_notices(formatted, boundaries, gaps) +} + +fn insert_loss_notices( + mut formatted: Vec, + boundaries: Vec, + gaps: Vec, +) -> Vec { + if gaps.is_empty() { + return formatted.into_iter().map(into_code_mode_output).collect(); + } + // One extra slot detects when an observation-wide summary is needed. + let mut notices = [((0, 0), OutputOmission::default()); MAX_POSITIONED_NOTICES + 1]; + let mut notice_count = 0; + let mut total = OutputOmission::default(); + let mut summary_position = (formatted.len(), 0); + for (gap, boundary) in gaps.into_iter().zip(boundaries) { + total.combine(gap.omitted); + let OutputBoundary { + item_index, + text_offset, + } = boundary; + // The end of a text item and the start of its successor are one position. + let position = if text_offset != 0 + && let Some(FunctionCallOutputContentItem::InputText { text }) = + formatted.get(item_index) + && text_offset == text.len() + { + (item_index + 1, 0) + } else { + (item_index, text_offset) + }; + // The collector can split one original text item. Even an observation-wide + // summary must keep its two fragments from reading as contiguous text. + if gap.boundary.text_offset != 0 { + summary_position = position; + } + if notice_count == notices.len() { + continue; + } + if let Some((last, omitted)) = notices[..notice_count].last_mut() + && *last == position + { + omitted.combine(gap.omitted); + } else { + notices[notice_count] = (position, gap.omitted); + notice_count += 1; + } + } + let scope = if notice_count == notices.len() { + notices[0] = (summary_position, total); + notice_count = 1; + "across this observation, " + } else { + "" + }; + // Drain from the end so earlier positions stay valid as text is split. + let mut reversed = Vec::with_capacity(formatted.len() + 2 * notice_count); + for &((index, offset), omission) in notices[..notice_count].iter().rev() { + let after = index + usize::from(offset != 0); + reversed.extend(formatted.drain(after..).rev().map(into_code_mode_output)); + if offset != 0 { + let FunctionCallOutputContentItem::InputText { text } = &mut formatted[index] else { + unreachable!("only text has an inner byte position"); + }; + if offset < text.len() { + reversed.push(OutputItem::InputText { + text: text.split_off(offset), + }); + } + } + reversed.push(OutputItem::InputText { + text: format!( + "[Code Mode session output limit reached: {scope}omitted {} text bytes; discarded {} text segments, {} images, {} audio items]", + omission.text_bytes, + omission.text_items, + omission.image_items, + omission.audio_items + ), + }); + } + reversed.extend(formatted.into_iter().rev().map(into_code_mode_output)); + reversed.reverse(); + reversed +} + +fn into_code_mode_output(item: FunctionCallOutputContentItem) -> OutputItem { + match item { + FunctionCallOutputContentItem::InputText { text } => OutputItem::InputText { text }, + FunctionCallOutputContentItem::InputImage { + image: ImageReference::Inline { image_url }, + detail, + } => OutputItem::InputImage { + image_url, + detail: detail.map(|detail| match detail { + ProtocolImageDetail::Auto => ImageDetail::Auto, + ProtocolImageDetail::Low => ImageDetail::Low, + ProtocolImageDetail::High => ImageDetail::High, + ProtocolImageDetail::Original => ImageDetail::Original, + }), + }, + FunctionCallOutputContentItem::InputAudio { audio_url } => { + OutputItem::InputAudio { audio_url } + } + FunctionCallOutputContentItem::EncryptedContent { .. } + | FunctionCallOutputContentItem::InputImage { + image: ImageReference::File { .. }, + .. + } => { + unreachable!("output truncation must not introduce new content types") + } + } +} + +#[cfg(test)] +#[path = "output_render_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-runtime/src/output/output_render_tests.rs b/codex-rs/code-mode-runtime/src/output/output_render_tests.rs new file mode 100644 index 0000000000..fe8b6a78de --- /dev/null +++ b/codex-rs/code-mode-runtime/src/output/output_render_tests.rs @@ -0,0 +1,285 @@ +#![cfg(test)] + +use super::render_output; +use crate::output::CollectedOutput; +use crate::output::OutputGap; +use crate::output::OutputOmission; +use crate::output::output_collection::BoundedOutput; +use crate::output::output_collection::MAX_AUDIO_BYTES; +use crate::output::output_collection::MAX_OUTPUT_BYTES; +use codex_code_mode_protocol::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; +use codex_code_mode_protocol::FunctionCallOutputContentItem as Item; +use codex_code_mode_protocol::ImageDetail; +use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES; +use codex_utils_output_truncation::OutputBoundary; +use pretty_assertions::assert_eq; + +fn text(value: &str) -> Item { + Item::InputText { + text: value.to_owned(), + } +} + +#[test] +fn missing_budget_uses_the_protocol_default() { + let limit = DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; + for (extra, count) in [(0, 2), (1, 1)] { + let items = vec![text("a"), text(&"b".repeat(4 * limit - 2 + extra))]; + let rendered = render_output(items.clone().into(), /*max_output_tokens*/ None); + assert_eq!(rendered.len(), count); + assert_eq!(rendered, render_output(items.into(), Some(limit))); + } +} + +#[test] +fn typed_media_round_trips_with_the_audio_estimate() { + let mut images: Vec<_> = [ + None, + Some(ImageDetail::Auto), + Some(ImageDetail::Low), + Some(ImageDetail::High), + Some(ImageDetail::Original), + ] + .map(|detail| Item::InputImage { + image_url: "img:one".to_owned(), + detail, + }) + .into(); + // One 16-bit PCM sample at 24 kHz costs one token. + let wav = + "data:audio/wav;base64,UklGRiYAAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQIAAAAAAA=="; + let mut items = images.clone(); + items.push(Item::InputAudio { + audio_url: wav.to_owned(), + }); + assert_eq!(render_output(items.clone().into(), Some(1)), items); + images.push(text("[omitted 1 audio items ...]")); + assert_eq!(render_output(items.into(), Some(0)), images); +} + +#[test] +fn prior_image_or_audio_loss_forces_mixed_formatting_without_surviving_media() { + for (images, audio) in [(1, 0), (0, 1)] { + let collected = CollectedOutput { + items: vec![text("1234"), text("5678")], + gaps: vec![OutputGap { + boundary: OutputBoundary { + item_index: 1, + text_offset: 0, + }, + omitted: OutputOmission { + text_bytes: 11, + text_items: 2, + image_items: images, + audio_items: audio, + }, + }], + }; + let notice = text(&format!( + "[Code Mode session output limit reached: omitted 11 text bytes; discarded 2 text segments, {images} images, {audio} audio items]" + )); + assert_eq!( + render_output(collected, Some(1)), + [text("1234"), notice, text("[omitted 1 text items ...]")] + ); + } +} + +#[test] +fn alternating_lost_audio_and_empty_text_yield_one_notice_at_zero_budget() { + // Start at the renderer boundary so the regression doesn't allocate 16 GiB of rejected audio. + let count = 16_384; + let collected = CollectedOutput { + items: vec![text(""); count], + gaps: (0..count) + .map(|item_index| OutputGap { + boundary: OutputBoundary { + item_index, + text_offset: 0, + }, + omitted: OutputOmission { + audio_items: 1, + ..OutputOmission::default() + }, + }) + .collect(), + }; + assert_eq!( + render_output(collected, Some(0)), + [text( + "[Code Mode session output limit reached: omitted 0 text bytes; discarded 0 text segments, 0 images, 16384 audio items]" + )] + ); +} + +#[test] +fn scattered_loss_is_summarized_once_without_joining_text_across_a_gap() { + let items: Vec<_> = (0..128) + .map(|n| Item::InputImage { + image_url: format!("img:{n}"), + detail: None, + }) + .collect(); + let collected = CollectedOutput { + gaps: (0..items.len()) + .map(|item_index| OutputGap { + boundary: OutputBoundary { + item_index, + text_offset: 0, + }, + omitted: OutputOmission { + text_bytes: 5, + text_items: 1, + image_items: 2, + audio_items: 3, + }, + }) + .collect(), + items: items.clone(), + }; + let mut expected = items.clone(); + expected.push(text( + "[Code Mode session output limit reached: across this observation, omitted 640 text bytes; discarded 128 text segments, 256 images, 384 audio items]", + )); + assert_eq!(render_output(collected.clone(), Some(0)), expected); + + let mut collected = collected; + collected.gaps.push(OutputGap { + boundary: OutputBoundary { + item_index: collected.items.len(), + text_offset: "left".len(), + }, + omitted: OutputOmission { + text_bytes: 9, + ..OutputOmission::default() + }, + }); + collected.items.push(text("leftright")); + let mut expected = items; + expected.extend([ + text("left"), + text( + "[Code Mode session output limit reached: across this observation, omitted 649 text bytes; discarded 128 text segments, 256 images, 384 audio items]", + ), + text("right"), + ]); + assert_eq!(render_output(collected, Some(16)), expected); +} + +#[test] +fn text_only_split_keeps_the_loss_notice_inside_the_rendered_text() { + let mut out = BoundedOutput::default(); + let before = "🦀".repeat(MAX_OUTPUT_BYTES / 8); + let after = "🎻".repeat(MAX_OUTPUT_BYTES / 8); + out.push(text(&format!("{before}lost{after}"))); + let rendered = render_output(out.into_collected_output(), Some(2)); + let [Item::InputText { text: prefix }, notice, suffix] = rendered.as_slice() else { + panic!("expected a split") + }; + assert!( + prefix.starts_with("Warning: truncated output") + && prefix.contains('🦀') + && !prefix.contains('🎻') + ); + assert_eq!( + *notice, + text( + "[Code Mode session output limit reached: omitted 4 text bytes; discarded 0 text segments, 0 images, 0 audio items]" + ) + ); + assert_eq!(*suffix, text("🎻")); +} + +#[test] +fn loss_notices_remain_in_source_order_across_budgets_and_unrelated_text() { + let before = "🦀".repeat(MAX_OUTPUT_BYTES / 8); + let after = "🎻".repeat(MAX_OUTPUT_BYTES / 8); + let original = format!("{before}{}{after}", "x".repeat(10_000)); + let image = Item::InputImage { + image_url: "data:image/png;base64,AA==".to_owned(), + detail: Some(ImageDetail::Original), + }; + let audio_notice = text( + "[Code Mode session output limit reached: omitted 0 text bytes; discarded 0 text segments, 0 images, 1 audio items]", + ); + let image_notice = text( + "[Code Mode session output limit reached: omitted 0 text bytes; discarded 0 text segments, 1 images, 0 audio items]", + ); + // Making room for the extra byte evicts one four-byte character from the suffix. + for (neighbor, lost) in [(None, 10000), (Some(""), 10000), (Some("x"), 10004)] { + let split_notice = text(&format!( + "[Code Mode session output limit reached: omitted {lost} text bytes; discarded 0 text segments, 0 images, 0 audio items]" + )); + let mut out = BoundedOutput::default(); + out.push(Item::InputAudio { + audio_url: "a".repeat(MAX_AUDIO_BYTES + 1), + }); + out.push(text(&original)); + out.push(Item::InputImage { + image_url: "i".repeat(BOUNDED_OUTPUT_MAX_IMAGE_BYTES + 1), + detail: None, + }); + out.push(image.clone()); + if let Some(neighbor) = neighbor { + out.push(text(neighbor)); + } + let collected = out.into_collected_output(); + let positions: Vec<_> = collected + .gaps + .iter() + .map(|gap| (gap.boundary.item_index, gap.boundary.text_offset)) + .collect(); + assert_eq!(positions, [(0, 0), (0, before.len()), (1, 0)]); + for budget in [0, 1, 2, MAX_OUTPUT_BYTES] { + let rendered = render_output(collected.clone(), Some(budget)); + if budget == 0 { + assert_eq!( + &rendered[..2], + &[ + text(&format!( + "[Code Mode session output limit reached: omitted {lost} text bytes; discarded 0 text segments, 1 images, 1 audio items]" + )), + image.clone() + ] + ); + continue; + } + if budget == 1 { + // Neither four-byte character fits its two-byte half. The split after + // Core's marker and the loss before the image are now the same position. + assert_eq!(&rendered[0], &audio_notice); + assert_eq!( + &rendered[2..4], + &[ + text(&format!( + "[Code Mode session output limit reached: omitted {lost} text bytes; discarded 0 text segments, 1 images, 0 audio items]" + )), + image.clone() + ] + ); + continue; + } + let [ + first, + Item::InputText { text: left }, + between, + Item::InputText { text: right }, + last, + preserved, + tail @ .., + ] = rendered.as_slice() + else { + panic!( + "expected three notices around split text, neighbor={neighbor:?}, budget={budget}" + ) + }; + assert_eq!( + (first, between, last, preserved), + (&audio_notice, &split_notice, &image_notice, &image) + ); + assert!(left.starts_with('🦀') && !left.contains('🎻')); + assert!(right.ends_with('🎻') && !right.contains('🦀')); + assert_eq!(tail.len(), usize::from(neighbor == Some("x"))); + } + } +} diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index 7c49956547..c0a9fbdc58 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -139,6 +139,7 @@ pub fn truncate_function_output_payload( /// are zeroed. Positions preceding an earlier boundary move to that position, /// preserving caller order. On return, boundaries refer to the rendered items; /// an end boundary precedes any generated omission notices. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct OutputBoundary { pub item_index: usize, pub text_offset: usize,