Avoid rerendering streamed code fences (#39061)

## Why

The streaming TUI rerendered the entire mutable Markdown block whenever a
newline-bearing delta arrived. For long fenced code blocks, this repeatedly
parsed and syntax-highlighted the growing fence.

## What changed

- Detect open, top-level, language-tagged code fences and append newly
  completed lines directly to the existing render.
- Preserve Syntect parser state across appended lines so multiline syntax and
  the active theme remain consistent.
- Fall back to the canonical Markdown renderer for ambiguous fence syntax,
  theme changes, incomplete lines, and highlighting limit transitions.

## Testing

Added coverage that compares streamed output with full renders across syntax
styles, Unicode, links, blank lines, fence variants, theme changes, and
highlighting limits. A long-fence test also verifies that the rendered prefix
is retained as new lines arrive.

GitOrigin-RevId: 036db2b7bc93607083b8001f8b4fa3cbb58c713e

codex-rs/tui/src/render/highlight.rs

@@ -47,6 +47,11 @@ use syntect::parsing::SyntaxSet;
 use syntect::util::LinesWithEndings;
 use two_face::theme::EmbeddedThemeName;
 
+#[path = "highlight_streaming.rs"]
+mod streaming;
+
+pub(crate) use streaming::StreamingCodeHighlighter;
+
 // -- Global singletons -------------------------------------------------------
 
 static SYNTAX_SET: OnceLock<SyntaxSet> = OnceLock::new();
@@ -630,25 +635,30 @@ fn highlight_to_line_spans_with_theme(
 
     for line in LinesWithEndings::from(code) {
         let ranges = h.highlight_line(line, syntax_set()).ok()?;
-        let mut spans: Vec<Span<'static>> = Vec::new();
-        for (style, text) in ranges {
-            // Strip trailing line endings (LF and CR) since we handle line
-            // breaks ourselves.  CRLF inputs would otherwise leave a stray \r.
-            let text = text.trim_end_matches(['\n', '\r']);
-            if text.is_empty() {
-                continue;
-            }
-            spans.push(Span::styled(text.to_string(), convert_style(style)));
-        }
-        if spans.is_empty() {
-            spans.push(Span::raw(String::new()));
-        }
-        lines.push(spans);
+        lines.push(highlighted_line_spans(ranges));
     }
 
     Some(lines)
 }
 
+/// Convert one highlighted line into the spans shared by full and streaming renders.
+///
+/// Source line endings are omitted, while empty lines retain their canonical empty span.
+fn highlighted_line_spans(ranges: Vec<(SyntectStyle, &str)>) -> Vec<Span<'static>> {
+    let mut spans = Vec::new();
+    for (style, text) in ranges {
+        // Line breaks are represented by the surrounding Line, not its spans.
+        let text = text.trim_end_matches(['\n', '\r']);
+        if !text.is_empty() {
+            spans.push(Span::styled(text.to_string(), convert_style(style)));
+        }
+    }
+    if spans.is_empty() {
+        spans.push(Span::raw(String::new()));
+    }
+    spans
+}
+
 /// Parse `code` using syntect for `lang` and return per-line styled spans.
 /// Each inner Vec represents one source line.  Returns None when the language
 /// is not recognized or the input exceeds safety limits.

codex-rs/tui/src/render/highlight_streaming.rs

@@ -0,0 +1,117 @@
+//! Append-only syntax highlighting for complete lines in an open code fence.
+
+use super::MAX_HIGHLIGHT_LINE_BYTES;
+use super::exceeds_highlight_limits;
+use super::find_syntax;
+use super::highlighted_line_spans;
+use super::syntax_set;
+use super::syntax_theme_revision;
+use super::theme_lock;
+use ratatui::text::Line;
+use syntect::easy::HighlightLines;
+use syntect::highlighting::HighlightState;
+use syntect::parsing::ParseState;
+use syntect::util::LinesWithEndings;
+
+/// Retains only Syntect's parser state, subject to the ordinary highlighting limits.
+///
+/// An absent state means the whole fence already uses the permanent plain-text fallback because
+/// its language is unknown or a size limit was exceeded. The caller owns the rendered prefix.
+pub(crate) struct StreamingCodeHighlighter {
+    state: Option<HighlightedCode>,
+}
+
+struct HighlightedCode {
+    bytes: usize,
+    lines: usize,
+    theme_revision: u64,
+    syntax: (HighlightState, ParseState),
+}
+
+impl StreamingCodeHighlighter {
+    /// Reconstruct highlighting state after the canonical renderer emitted `code`.
+    ///
+    /// Unknown languages and oversized input retain a permanent plain-text fallback. A changed
+    /// theme or failed replay returns `None`, requiring the caller to render the whole fence again.
+    pub(crate) fn new(code: &str, lang: &str, theme_revision: u64) -> Option<Self> {
+        let theme = match theme_lock().read() {
+            Ok(theme) => theme,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        if theme_revision != syntax_theme_revision() {
+            return None;
+        }
+        let lines = code.lines().count();
+        let syntax = find_syntax(lang).filter(|_| {
+            !exceeds_highlight_limits(code.len(), lines)
+                && !code
+                    .lines()
+                    .any(|line| line.len() > MAX_HIGHLIGHT_LINE_BYTES)
+        });
+        let Some(syntax) = syntax else {
+            return Some(Self { state: None });
+        };
+        let mut highlighter = HighlightLines::new(syntax, &theme);
+        for line in LinesWithEndings::from(code) {
+            highlighter.highlight_line(line, syntax_set()).ok()?;
+        }
+        Some(Self {
+            state: Some(HighlightedCode {
+                bytes: code.len(),
+                lines,
+                theme_revision,
+                syntax: highlighter.state(),
+            }),
+        })
+    }
+
+    /// Consume the previous state and highlight just the newly committed complete lines.
+    ///
+    /// Every appended line must be newline-terminated to preserve Syntect's parser state.
+    /// Returning `None` discards the state: the caller must render the whole fence again, which
+    /// also removes old colors when the aggregate input crosses a highlighting limit.
+    pub(crate) fn append(mut self, appended: &str) -> Option<(Self, Vec<Line<'static>>)> {
+        if !appended.ends_with('\n') {
+            return None;
+        }
+        let Some(mut state) = self.state.take() else {
+            let lines = appended
+                .lines()
+                .map(|line| Line::from(line.to_string()))
+                .collect();
+            return Some((self, lines));
+        };
+        let bytes = state.bytes.checked_add(appended.len())?;
+        let lines = state.lines.checked_add(appended.lines().count())?;
+        if exceeds_highlight_limits(bytes, lines)
+            || appended
+                .lines()
+                .any(|line| line.len() > MAX_HIGHLIGHT_LINE_BYTES)
+        {
+            return None;
+        }
+        let theme = match theme_lock().read() {
+            Ok(theme) => theme,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        if state.theme_revision != syntax_theme_revision() {
+            return None;
+        }
+        let (highlight_state, parse_state) = state.syntax;
+        let mut highlighter = HighlightLines::from_state(&theme, highlight_state, parse_state);
+        let mut rendered = Vec::with_capacity(lines - state.lines);
+        for line in LinesWithEndings::from(appended) {
+            let ranges = highlighter.highlight_line(line, syntax_set()).ok()?;
+            rendered.push(Line::from(highlighted_line_spans(ranges)));
+        }
+        state.bytes = bytes;
+        state.lines = lines;
+        state.syntax = highlighter.state();
+        self.state = Some(state);
+        Some((self, rendered))
+    }
+}
+
+#[cfg(test)]
+#[path = "highlight_streaming_tests.rs"]
+mod tests;

codex-rs/tui/src/render/highlight_streaming_tests.rs

@@ -0,0 +1,75 @@
+use super::super::MAX_HIGHLIGHT_BYTES;
+use super::super::MAX_HIGHLIGHT_LINE_BYTES;
+use super::super::MAX_HIGHLIGHT_LINES;
+use super::super::highlight_code_to_lines;
+use super::super::syntax_theme_revision;
+use super::StreamingCodeHighlighter;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn appended_lines_preserve_multiline_syntax_state() {
+    let first = "fn main() {\n    /* comment\n";
+    let appended = "       continues */\n    println!(\"界 e\u{301}\");\n}\n";
+    let code = format!("{first}{appended}");
+    let highlighter =
+        StreamingCodeHighlighter::new(first, "rust", syntax_theme_revision()).unwrap();
+    let (_, lines) = highlighter.append(appended).unwrap();
+
+    assert_eq!(
+        lines,
+        highlight_code_to_lines(&code, "rust")
+            .into_iter()
+            .skip(first.lines().count())
+            .collect::<Vec<_>>(),
+    );
+}
+
+#[test]
+fn empty_and_unknown_language_fences_append() {
+    for lang in ["rust", "unknown-language"] {
+        let highlighter = StreamingCodeHighlighter::new("", lang, syntax_theme_revision()).unwrap();
+        let (_, lines) = highlighter.append("fn main() {}\n\n").unwrap();
+        assert_eq!(lines, highlight_code_to_lines("fn main() {}\n\n", lang));
+    }
+}
+
+#[test]
+fn incomplete_lines_and_theme_changes_discard_state() {
+    let highlighter =
+        StreamingCodeHighlighter::new("fn first() {}\n", "rust", syntax_theme_revision()).unwrap();
+    assert!(highlighter.append("fn partial()").is_none());
+
+    let mut highlighter =
+        StreamingCodeHighlighter::new("fn first() {}\n", "rust", syntax_theme_revision()).unwrap();
+    let state = highlighter.state.as_mut().unwrap();
+    state.theme_revision = state.theme_revision.wrapping_sub(1);
+    assert!(highlighter.append("fn second() {}\n").is_none());
+}
+
+#[test]
+fn crossing_each_highlighting_limit_requires_a_canonical_render() {
+    let cases = [
+        (
+            String::new(),
+            format!("{}\n", "x".repeat(MAX_HIGHLIGHT_LINE_BYTES + 1)),
+        ),
+        ("x\n".repeat(MAX_HIGHLIGHT_LINES), "x\n".to_string()),
+        (
+            format!("{}\n", "x".repeat(1023)).repeat(MAX_HIGHLIGHT_BYTES / 1024),
+            "x\n".to_string(),
+        ),
+    ];
+    for (initial, appended) in cases {
+        let highlighter =
+            StreamingCodeHighlighter::new(&initial, "rust", syntax_theme_revision()).unwrap();
+        assert!(highlighter.state.is_some());
+        assert!(highlighter.append(&appended).is_none());
+
+        let complete = format!("{initial}{appended}");
+        let highlighter =
+            StreamingCodeHighlighter::new(&complete, "rust", syntax_theme_revision()).unwrap();
+        assert!(highlighter.state.is_none());
+        let (_, lines) = highlighter.append("next\n").unwrap();
+        assert_eq!(lines, highlight_code_to_lines("next\n", "unknown-language"));
+    }
+}