From 91d6f48992ad8db636b3ca52a3a36c2fb6d75537 Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Wed, 12 Aug 2026 16:23:55 +0000 Subject: [PATCH] Avoid allocations when sanitizing TUI user text (#38214) ## Why Sanitizing user messages and pasted text always allocated a new string, even when the input contained no terminal control characters. ## What changed - Return a `Cow` from `sanitize_user_text` so clean borrowed input remains borrowed. - Reuse owned input buffers when sanitization leaves the text intact or only removes content at the edges. - Update transcript rendering, history cells, and paste handling to consume the borrowed-or-owned result without unnecessary cloning. ## Testing Add sanitizer coverage for clean text, control sequences, Unicode control characters, incomplete escape sequences, buffer reuse, and multi-fragment output. GitOrigin-RevId: 9e39cbc0d83c5791941b37a7eb4bb918c1a05f68 --- codex-rs/tui/src/app/transcript_export.rs | 6 +- codex-rs/tui/src/bottom_pane/chat_composer.rs | 12 +-- codex-rs/tui/src/history_cell/messages.rs | 81 +++++++++++++++---- .../tui/src/history_cell/messages_tests.rs | 60 ++++++++++++++ 4 files changed, 135 insertions(+), 24 deletions(-) diff --git a/codex-rs/tui/src/app/transcript_export.rs b/codex-rs/tui/src/app/transcript_export.rs index 81cf6e73bd..ca04127efa 100644 --- a/codex-rs/tui/src/app/transcript_export.rs +++ b/codex-rs/tui/src/app/transcript_export.rs @@ -224,7 +224,7 @@ fn render_markdown_transcript(cells: &[Arc]) -> Result() { let (message, _) = crate::ide_context::extract_prompt_request_with_offset(&user.message); - let message = crate::history_cell::sanitize_user_text(message); + let message = crate::history_cell::sanitize_user_text(message.into()); let mut lines = raw_lines_from_source(&message); let image_count = user.local_image_paths.len() + user.remote_image_urls.len(); let image_labels = (0..image_count) @@ -273,9 +273,7 @@ fn render_markdown_transcript(cells: &[Arc]) -> Result bool { let pasted = pasted.replace("\r\n", "\n").replace('\r', "\n"); - let pasted = sanitize_user_text(&pasted); + let pasted = sanitize_user_text(pasted.into()); let char_count = pasted.chars().count(); if char_count > LARGE_PASTE_CHAR_THRESHOLD { let placeholder = self.next_large_paste_placeholder(char_count); self.draft.textarea.insert_element(&placeholder); - self.draft.pending_pastes.push((placeholder, pasted)); + self.draft + .pending_pastes + .push((placeholder, pasted.into_owned())); } else if char_count > 1 && self.image_paste_enabled() - && self.handle_paste_image_path(pasted.clone()) + && self.handle_paste_image_path(&pasted) { self.draft.textarea.insert_str(" "); } else { @@ -1150,8 +1152,8 @@ impl ChatComposer { true } - pub fn handle_paste_image_path(&mut self, pasted: String) -> bool { - let Some(path_buf) = normalize_pasted_path(&pasted) else { + pub fn handle_paste_image_path(&mut self, pasted: &str) -> bool { + let Some(path_buf) = normalize_pasted_path(pasted) else { return false; }; diff --git a/codex-rs/tui/src/history_cell/messages.rs b/codex-rs/tui/src/history_cell/messages.rs index 7cc099d981..28015dfee8 100644 --- a/codex-rs/tui/src/history_cell/messages.rs +++ b/codex-rs/tui/src/history_cell/messages.rs @@ -2,6 +2,7 @@ use super::markdown_render_cache::MarkdownRenderCache; use super::*; +use std::borrow::Cow; #[derive(Debug)] pub(crate) struct UserHistoryCell { @@ -13,17 +14,67 @@ pub(crate) struct UserHistoryCell { } /// Remove CSI sequences and control characters, preserving tabs and newlines. -pub(crate) fn sanitize_user_text(text: &str) -> String { - let mut sanitized = String::with_capacity(text.len()); - let mut chars = text.chars().peekable(); - while let Some(ch) = chars.next() { - if ch == '\x1b' && chars.next_if_eq(&'[').is_some() { - let _ = chars.find(|ch| ('@'..='~').contains(ch)); - } else if matches!(ch, '\n' | '\t') || !ch.is_control() { - sanitized.push(ch); - } +pub(crate) fn sanitize_user_text(text: Cow<'_, str>) -> Cow<'_, str> { + fn sanitize_borrowed(text: &str) -> Cow<'_, str> { + let mut remaining = Some(text); + let mut spans = std::iter::from_fn(move || { + let current = remaining.take()?; + let mut escaped = false; + let Some((prefix, suffix)) = current.split_once(|ch: char| { + escaped = ch == '\x1b'; + escaped || ch.is_control() && !matches!(ch, '\n' | '\t') + }) else { + return Some(current); + }; + + remaining = if escaped && let Some(sequence) = suffix.strip_prefix('[') { + sequence + .split_once(|ch: char| ('@'..='~').contains(&ch)) + .map(|(_, tail)| tail) + } else { + Some(suffix) + }; + Some(prefix) + }) + .filter(|span| !span.is_empty()); + + let first = spans.next().unwrap_or_default(); + let Some(second) = spans.next() else { + return Cow::Borrowed(first); + }; + + Cow::Owned([first, second].into_iter().chain(spans).fold( + String::with_capacity(text.len()), + |mut acc, span| { + acc.push_str(span); + acc + }, + )) + } + + match text { + Cow::Borrowed(text) => sanitize_borrowed(text), + Cow::Owned(mut text) => match sanitize_borrowed(&text) { + Cow::Owned(sanitized) => Cow::Owned(sanitized), + Cow::Borrowed(retained) => { + if retained.is_empty() { + text.clear(); + } else if retained.len() != text.len() { + // Cannot underflow because retained is a subslice of text. + // I'd normally assert that here but this crate denies + // clippy::expect_used. + let start = retained.as_ptr().addr() - text.as_ptr().addr(); + let end = start + retained.len(); + + // Truncate before draining because truncate is constant + // time while drain is linear in the size of the receiver. + text.truncate(end); + drop(text.drain(..start)); + } + Cow::Owned(text) + } + }, } - sanitized } /// Build logical lines for a user message with styled text elements. @@ -108,8 +159,8 @@ fn trim_trailing_blank_lines(mut lines: Vec>) -> Vec impl HistoryCell for UserHistoryCell { fn display_lines(&self, width: u16) -> Vec> { - let message = sanitize_user_text(&self.message); - let text_elements = if message == self.message { + let message = sanitize_user_text((&self.message).into()); + let text_elements = if message.as_ref() == self.message { self.text_elements.as_slice() } else { &[] @@ -154,7 +205,7 @@ impl HistoryCell for UserHistoryCell { (!wrapped.is_empty()).then_some(wrapped) } else { let raw_lines = build_user_message_lines_with_elements( - &message, + message.as_ref(), text_elements, style, element_style, @@ -198,8 +249,8 @@ impl HistoryCell for UserHistoryCell { } fn raw_lines(&self) -> Vec> { - let message = sanitize_user_text(&self.message); - let mut lines = raw_lines_from_source(message.trim_end_matches(['\r', '\n'])); + let message = sanitize_user_text((&self.message).into()); + let mut lines = raw_lines_from_source(message.as_ref().trim_end_matches(['\r', '\n'])); if !self.remote_image_urls.is_empty() { if !lines.is_empty() { lines.push(Line::from("")); diff --git a/codex-rs/tui/src/history_cell/messages_tests.rs b/codex-rs/tui/src/history_cell/messages_tests.rs index 900d05d4ee..c031a21147 100644 --- a/codex-rs/tui/src/history_cell/messages_tests.rs +++ b/codex-rs/tui/src/history_cell/messages_tests.rs @@ -1,7 +1,67 @@ use super::*; use crate::history_cell::markdown_render_cache::MarkdownRenderCacheKey; +use assert_matches::assert_matches; use pretty_assertions::assert_eq; +#[test] +fn sanitizer_borrows_clean_text_and_removes_control_sequences() { + for (text, expected) in [ + ("clean\ttext\n", "clean\ttext\n"), + ("\x07before", "before"), + ("before\x07", "before"), + ("\x1b[31mbefore", "before"), + ("before\x1b[31m", "before"), + ("before\x1b[31", "before"), + ("\x07[31m", "[31m"), + ("\x07", ""), + ] { + assert_matches!( + sanitize_user_text(text.into()), + Cow::Borrowed(sanitized) => assert_eq!(sanitized, expected) + ); + } + assert_matches!( + sanitize_user_text("before\x1b[31mafter\x07".into()), + Cow::Owned(sanitized) => assert_eq!(sanitized, "beforeafter") + ); + assert_eq!(sanitize_user_text("é\u{85}中".into()), "é中"); + assert_eq!(sanitize_user_text("before\x1bafter".into()), "beforeafter"); +} + +#[test] +fn sanitizer_preserves_owned_buffer_for_clean_and_edge_trimmed_text() { + for (text, expected) in [ + ("clean\ttext\n", "clean\ttext\n"), + ("\x07before", "before"), + ("before\x07", "before"), + ("\x07before\x07", "before"), + ("\x1b[31mbefore", "before"), + ("before\x1b[31m", "before"), + ("\x07", ""), + ] { + let owned = text.to_string(); + let original_pointer = owned.as_ptr(); + let original_capacity = owned.capacity(); + + assert_matches!(sanitize_user_text(owned.into()), Cow::Owned(sanitized) => { + assert_eq!(sanitized, expected); + assert_eq!(sanitized.as_ptr(), original_pointer); + assert_eq!(sanitized.capacity(), original_capacity); + }) + } +} + +#[test] +fn sanitizer_preallocates_owned_multi_fragment_text() { + let text = "before\x1b[31mafter\x07".to_string(); + let original_length = text.len(); + + assert_matches!(sanitize_user_text(text.into()), Cow::Owned(sanitized) => { + assert_eq!(sanitized, "beforeafter"); + assert!(sanitized.capacity() >= original_length, "{} >= {}", sanitized.capacity(), original_length); + }) +} + fn replace_cached_lines( cell: &AgentMarkdownCell, update_key: impl FnOnce(&mut MarkdownRenderCacheKey),