mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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<str>` 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
This commit is contained in:
1 parent
95aada11c4
commit
91d6f48992
4 files changed
+134
-23
No files matched your search
@@ -224,7 +224,7 @@ fn render_markdown_transcript(cells: &[Arc<dyn HistoryCell>]) -> Result<String,
|
|||||||
let lines = if let Some(user) = cell.as_any().downcast_ref::<UserHistoryCell>() {
|
let lines = if let Some(user) = cell.as_any().downcast_ref::<UserHistoryCell>() {
|
||||||
let (message, _) =
|
let (message, _) =
|
||||||
crate::ide_context::extract_prompt_request_with_offset(&user.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 mut lines = raw_lines_from_source(&message);
|
||||||
let image_count = user.local_image_paths.len() + user.remote_image_urls.len();
|
let image_count = user.local_image_paths.len() + user.remote_image_urls.len();
|
||||||
let image_labels = (0..image_count)
|
let image_labels = (0..image_count)
|
||||||
@@ -273,9 +273,7 @@ fn render_markdown_transcript(cells: &[Arc<dyn HistoryCell>]) -> Result<String,
|
|||||||
markdown.push_str(" ");
|
markdown.push_str(" ");
|
||||||
}
|
}
|
||||||
for span in line.spans {
|
for span in line.spans {
|
||||||
markdown.push_str(&crate::history_cell::sanitize_user_text(
|
markdown.push_str(&crate::history_cell::sanitize_user_text(span.content));
|
||||||
span.content.as_ref(),
|
|
||||||
));
|
|
||||||
}
|
}
|
||||||
markdown.push('\n');
|
markdown.push('\n');
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1131,15 +1131,17 @@ impl ChatComposer {
|
|||||||
/// the next user Enter key, then syncs popup state.
|
/// the next user Enter key, then syncs popup state.
|
||||||
pub fn handle_paste(&mut self, pasted: String) -> bool {
|
pub fn handle_paste(&mut self, pasted: String) -> bool {
|
||||||
let pasted = pasted.replace("\r\n", "\n").replace('\r', "\n");
|
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();
|
let char_count = pasted.chars().count();
|
||||||
if char_count > LARGE_PASTE_CHAR_THRESHOLD {
|
if char_count > LARGE_PASTE_CHAR_THRESHOLD {
|
||||||
let placeholder = self.next_large_paste_placeholder(char_count);
|
let placeholder = self.next_large_paste_placeholder(char_count);
|
||||||
self.draft.textarea.insert_element(&placeholder);
|
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
|
} else if char_count > 1
|
||||||
&& self.image_paste_enabled()
|
&& self.image_paste_enabled()
|
||||||
&& self.handle_paste_image_path(pasted.clone())
|
&& self.handle_paste_image_path(&pasted)
|
||||||
{
|
{
|
||||||
self.draft.textarea.insert_str(" ");
|
self.draft.textarea.insert_str(" ");
|
||||||
} else {
|
} else {
|
||||||
@@ -1150,8 +1152,8 @@ impl ChatComposer {
|
|||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn handle_paste_image_path(&mut self, pasted: String) -> bool {
|
pub fn handle_paste_image_path(&mut self, pasted: &str) -> bool {
|
||||||
let Some(path_buf) = normalize_pasted_path(&pasted) else {
|
let Some(path_buf) = normalize_pasted_path(pasted) else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
use super::markdown_render_cache::MarkdownRenderCache;
|
use super::markdown_render_cache::MarkdownRenderCache;
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use std::borrow::Cow;
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub(crate) struct UserHistoryCell {
|
pub(crate) struct UserHistoryCell {
|
||||||
@@ -13,17 +14,67 @@ pub(crate) struct UserHistoryCell {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Remove CSI sequences and control characters, preserving tabs and newlines.
|
/// Remove CSI sequences and control characters, preserving tabs and newlines.
|
||||||
pub(crate) fn sanitize_user_text(text: &str) -> String {
|
pub(crate) fn sanitize_user_text(text: Cow<'_, str>) -> Cow<'_, str> {
|
||||||
let mut sanitized = String::with_capacity(text.len());
|
fn sanitize_borrowed(text: &str) -> Cow<'_, str> {
|
||||||
let mut chars = text.chars().peekable();
|
let mut remaining = Some(text);
|
||||||
while let Some(ch) = chars.next() {
|
let mut spans = std::iter::from_fn(move || {
|
||||||
if ch == '\x1b' && chars.next_if_eq(&'[').is_some() {
|
let current = remaining.take()?;
|
||||||
let _ = chars.find(|ch| ('@'..='~').contains(ch));
|
let mut escaped = false;
|
||||||
} else if matches!(ch, '\n' | '\t') || !ch.is_control() {
|
let Some((prefix, suffix)) = current.split_once(|ch: char| {
|
||||||
sanitized.push(ch);
|
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.
|
/// Build logical lines for a user message with styled text elements.
|
||||||
@@ -108,8 +159,8 @@ fn trim_trailing_blank_lines(mut lines: Vec<Line<'static>>) -> Vec<Line<'static>
|
|||||||
|
|
||||||
impl HistoryCell for UserHistoryCell {
|
impl HistoryCell for UserHistoryCell {
|
||||||
fn display_lines(&self, width: u16) -> Vec<Line<'static>> {
|
fn display_lines(&self, width: u16) -> Vec<Line<'static>> {
|
||||||
let message = sanitize_user_text(&self.message);
|
let message = sanitize_user_text((&self.message).into());
|
||||||
let text_elements = if message == self.message {
|
let text_elements = if message.as_ref() == self.message {
|
||||||
self.text_elements.as_slice()
|
self.text_elements.as_slice()
|
||||||
} else {
|
} else {
|
||||||
&[]
|
&[]
|
||||||
@@ -154,7 +205,7 @@ impl HistoryCell for UserHistoryCell {
|
|||||||
(!wrapped.is_empty()).then_some(wrapped)
|
(!wrapped.is_empty()).then_some(wrapped)
|
||||||
} else {
|
} else {
|
||||||
let raw_lines = build_user_message_lines_with_elements(
|
let raw_lines = build_user_message_lines_with_elements(
|
||||||
&message,
|
message.as_ref(),
|
||||||
text_elements,
|
text_elements,
|
||||||
style,
|
style,
|
||||||
element_style,
|
element_style,
|
||||||
@@ -198,8 +249,8 @@ impl HistoryCell for UserHistoryCell {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn raw_lines(&self) -> Vec<Line<'static>> {
|
fn raw_lines(&self) -> Vec<Line<'static>> {
|
||||||
let message = sanitize_user_text(&self.message);
|
let message = sanitize_user_text((&self.message).into());
|
||||||
let mut lines = raw_lines_from_source(message.trim_end_matches(['\r', '\n']));
|
let mut lines = raw_lines_from_source(message.as_ref().trim_end_matches(['\r', '\n']));
|
||||||
if !self.remote_image_urls.is_empty() {
|
if !self.remote_image_urls.is_empty() {
|
||||||
if !lines.is_empty() {
|
if !lines.is_empty() {
|
||||||
lines.push(Line::from(""));
|
lines.push(Line::from(""));
|
||||||
|
|||||||
@@ -1,7 +1,67 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
use crate::history_cell::markdown_render_cache::MarkdownRenderCacheKey;
|
use crate::history_cell::markdown_render_cache::MarkdownRenderCacheKey;
|
||||||
|
use assert_matches::assert_matches;
|
||||||
use pretty_assertions::assert_eq;
|
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(
|
fn replace_cached_lines(
|
||||||
cell: &AgentMarkdownCell,
|
cell: &AgentMarkdownCell,
|
||||||
update_key: impl FnOnce(&mut MarkdownRenderCacheKey),
|
update_key: impl FnOnce(&mut MarkdownRenderCacheKey),
|
||||||
|
|||||||
Reference in new issue
Block a user