diff --git a/codex-rs/tui/src/markdown.rs b/codex-rs/tui/src/markdown.rs index 8880a23194..0f9b7345ba 100644 --- a/codex-rs/tui/src/markdown.rs +++ b/codex-rs/tui/src/markdown.rs @@ -134,6 +134,11 @@ pub(crate) fn render_streaming_markdown_agent_with_links_and_cwd( // Fence unwrapping removes opening/closing lines. A normalized tail that is still a raw // suffix necessarily begins after those removed lines, so its boundary can safely be // mapped back to the raw source; otherwise leave the transformed block mutable. + rendered.pending_math_start = rendered.pending_math_start.map(|boundary| { + markdown_source + .strip_suffix(&normalized[boundary..]) + .map_or(0, str::len) + }); rendered.last_top_level_block_start = rendered .last_top_level_block_start .and_then(|boundary| markdown_source.strip_suffix(&normalized[boundary..])) diff --git a/codex-rs/tui/src/markdown_render.rs b/codex-rs/tui/src/markdown_render.rs index a80abc192d..91ea4370b1 100644 --- a/codex-rs/tui/src/markdown_render.rs +++ b/codex-rs/tui/src/markdown_render.rs @@ -339,7 +339,7 @@ pub(crate) fn render_markdown_lines_with_width_cwd_and_hidden_link_destinations( let mut options = Options::empty(); options.insert(Options::ENABLE_STRIKETHROUGH); options.insert(Options::ENABLE_TABLES); - let math = math::MathMarkdown::new(input, options); + let math = math::MathMarkdown::new(input, options, width); let parser = DecodedTextMerge::new( math.events(Parser::new_ext(&math.markdown, options).into_offset_iter()), ); diff --git a/codex-rs/tui/src/markdown_render/math.rs b/codex-rs/tui/src/markdown_render/math.rs index 741d88cc37..3ac9a1d7df 100644 --- a/codex-rs/tui/src/markdown_render/math.rs +++ b/codex-rs/tui/src/markdown_render/math.rs @@ -1,7 +1,8 @@ //! Recognize math before Markdown consumes TeX escapes, retaining exact source offsets. //! //! Only the rendering copy is masked. Unsupported expressions are restored verbatim; code, -//! links, HTML, and display equations stay under the ordinary Markdown renderer. +//! links, and HTML stay under the ordinary Markdown renderer. Standalone display tracking outlives +//! the conversion budget; rejected prose-prefixed openers have bounded pairing lookahead. use itertools::Either; use pulldown_cmark::Event; @@ -17,14 +18,16 @@ const MAX_MATH_BYTES: usize = 4096; pub(super) struct MathMarkdown<'a> { pub(super) markdown: Cow<'a, str>, + pub(super) pending_start: Option, pub(super) display_ranges: Vec>, replacements: Vec<(Range, String)>, } impl<'a> MathMarkdown<'a> { - pub(super) fn new(input: &'a str, options: Options) -> Self { + pub(super) fn new(input: &'a str, options: Options, width: Option) -> Self { let mut result = Self { markdown: Cow::Borrowed(input), + pending_start: None, display_ranges: Vec::new(), replacements: Vec::new(), }; @@ -37,7 +40,11 @@ impl<'a> MathMarkdown<'a> { .iter() .map(|(_, def)| def.span.clone()) .collect(); + let mut containers = Vec::new(); protected.extend(parser.into_offset_iter().filter_map(|(event, range)| { + if matches!(event, Event::Start(Tag::List(_) | Tag::BlockQuote)) { + containers.push(range.clone()); + } matches!( event, Event::Code(_) @@ -48,16 +55,22 @@ impl<'a> MathMarkdown<'a> { .then_some(range) })); protected.sort_unstable_by_key(|range| range.start); - let mut protected = protected.into_iter().peekable(); + let mut protected = protected.iter().peekable(); + containers.sort_unstable_by_key(|range| range.start); + let mut containers = containers.into_iter().peekable(); let mut offset = 0; let mut scanned = 0; let mut line_start = 0; + let mut line_has_text = false; while offset < input.len() { if let Some(index) = input[scanned..offset].rfind('\n') { line_start = scanned + index + 1; + line_has_text = false; } + line_has_text |= !input[scanned.max(line_start)..offset].trim().is_empty(); scanned = offset; while protected.next_if(|range| range.end <= offset).is_some() {} + while containers.next_if(|range| range.end <= offset).is_some() {} if let Some(range) = protected.peek() && range.contains(&offset) { @@ -96,30 +109,102 @@ impl<'a> MathMarkdown<'a> { .map(|(i, _)| i) .find(|i| *i >= MAX_MATH_BYTES) .unwrap_or(body.len()); - // Display boundaries outlive the conversion budget; a distant closer still owns its opener. - let search = if display { body } else { &body[..limit] }; - let end = search.match_indices(close).find_map(|(index, _)| { + let rejected_display = display && line_has_text; + // Only standalone displays can retain an arbitrarily distant closer. + let search = if display && !rejected_display { + body + } else { + &body[..limit] + }; + let mut multiline_close = false; + let mut closing_protected = protected.clone(); + let mut rejected_close = false; + let end = search.match_indices(&close[..1]).find_map(|(index, _)| { let end = offset + index; - (!escaped(input, end)).then_some(end) - }); - if display { - if open == "$$" && end.is_none() && !input[line_start..start].trim().is_empty() { - continue; + if !search[index..].starts_with(close) || escaped(input, end) { + return None; + } + if display { + while closing_protected + .next_if(|range| range.end <= end) + .is_some() + {} + if closing_protected + .peek() + .is_some_and(|range| range.start < end + close.len()) + { + return None; + } + if !rejected_display + && input[end + close.len()..] + .chars() + .take_while(|ch| *ch != '\n') + .any(|ch| !ch.is_whitespace()) + { + if multiline_close || input[offset..end].contains('\n') { + multiline_close = true; + return None; + } + rejected_close = true; + } + } + Some(end) + }); + if display && (!rejected_display || end.is_some() || body.len() < MAX_MATH_BYTES) { + result + .display_ranges + .push(start..end.map_or(input.len(), |end| end + close.len())); + } + // Retain rejected pairing only within the lookahead window, so shell PID dollars + // cannot keep the entire streamed response mutable in the rendering cache. + if rejected_display || rejected_close { + if let Some(end) = end + && !input[offset..end].trim().is_empty() + { + // A later line's opening $$ must remain available to its own equation. + if rejected_display + && open == "$$" + && input[offset..end] + .rsplit_once('\n') + .is_some_and(|(_, prefix)| prefix.trim().is_empty()) + && input[end + close.len()..] + .chars() + .take_while(|ch| *ch != '\n') + .any(|ch| !ch.is_whitespace()) + { + continue; + } + offset = end + close.len(); + } else if end.is_none() && open == "\\[" && body.len() < MAX_MATH_BYTES { + break; } - // Exclude display equations from inline recognition, including unfinished ones. - offset = end.map_or(input.len(), |end| end + close.len()); - result.display_ranges.push(start..offset); continue; } let Some(end) = end else { + if display { + if body.len() < MAX_MATH_BYTES { + result.pending_start.get_or_insert(line_start); + } + let span = start..input.len(); + result + .markdown + .to_mut() + .replace_range(span.clone(), &"$".repeat(span.len())); + result.replacements.push((span, input[start..].to_owned())); + break; + } continue; }; let span = start..end + close.len(); + let formula = &input[offset..end]; + // Every matched display owns its closer, including rejected expressions. + if display { + offset = span.end; + } if protected.peek().is_some_and(|range| range.start < span.end) { continue; } - let formula = &input[offset..end]; - if formula.contains('\n') { + if !display && formula.contains('\n') { continue; } if open == "$" { @@ -136,8 +221,27 @@ impl<'a> MathMarkdown<'a> { continue; } } - let rendered = - render::render(formula).unwrap_or_else(|| input[span.clone()].to_owned()); + let rendered = if formula.len() < MAX_MATH_BYTES { + render::render(formula, display) + } else { + None + }; + let rendered = rendered + .filter(|text| { + // Nested Markdown prefixes have their own width; keep spatial layouts at the top level. + // Never wrap a spatial layout into misleading pieces. + !text.contains('\n') + || !containers + .peek() + .is_some_and(|range| range.contains(&start)) + && width.is_none_or(|width| { + text.lines().all(|line| { + crate::width::display_width(line) + <= width.saturating_sub(/*rhs*/ 4) + }) + }) + }) + .unwrap_or_else(|| input[span.clone()].to_owned()); // Dollars are ordinary text in the Markdown parser and cannot form an HTML tag. result .markdown @@ -204,3 +308,7 @@ fn escaped(input: &str, offset: usize) -> bool { #[cfg(test)] #[path = "math_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "math_display_tests.rs"] +mod display_tests; diff --git a/codex-rs/tui/src/markdown_render/math/render.rs b/codex-rs/tui/src/markdown_render/math/render.rs index 4b1ba3f29e..6dc1659149 100644 --- a/codex-rs/tui/src/markdown_render/math/render.rs +++ b/codex-rs/tui/src/markdown_render/math/render.rs @@ -2,18 +2,77 @@ use crate::width::display_width; -pub(super) fn render(source: &str) -> Option { +const MAX_ROWS: usize = 16; +const MAX_COLUMNS: usize = 256; + +struct Layout { + rows: Vec, + baseline: usize, +} + +impl Layout { + fn text(text: impl Into) -> Self { + Self { + rows: vec![text.into()], + baseline: 0, + } + } + + fn width(&self) -> usize { + self.rows + .iter() + .map(|row| display_width(row)) + .max() + .unwrap_or_default() + } + + fn join(self, right: Self) -> Option { + let baseline = self.baseline.max(right.baseline); + let height = (baseline + self.rows.len() - self.baseline) + .max(baseline + right.rows.len() - right.baseline); + // Separate neighboring fraction bars so their numerators cannot become one number. + let width = self.width() + usize::from(self.rows.len() > 1 && right.rows.len() > 1); + if height > MAX_ROWS || width + right.width() > MAX_COLUMNS { + return None; + } + let mut rows = vec![String::new(); height]; + for (index, row) in rows.iter_mut().enumerate() { + if let Some(left) = index + .checked_sub(baseline - self.baseline) + .and_then(|i| self.rows.get(i)) + { + row.push_str(left); + } + row.push_str(&" ".repeat(width - display_width(row))); + if let Some(right) = index + .checked_sub(baseline - right.baseline) + .and_then(|i| right.rows.get(i)) + { + row.push_str(right); + } + } + Some(Self { rows, baseline }) + } + + fn single(&self) -> Option<&str> { + (self.rows.len() == 1).then(|| self.rows[0].as_str()) + } +} + +pub(super) fn render(source: &str, display: bool) -> Option { let mut parser = MathParser { remaining: source.trim(), depth: 0, + display, }; let result = parser.sequence(/*group*/ false)?; - (!result.trim().is_empty()).then_some(result) + (!result.rows.iter().all(|row| row.trim().is_empty())).then(|| result.rows.join("\n")) } struct MathParser<'a> { remaining: &'a str, depth: usize, + display: bool, } impl MathParser<'_> { @@ -23,12 +82,12 @@ impl MathParser<'_> { Some(ch) } - fn sequence(&mut self, group: bool) -> Option { + fn sequence(&mut self, group: bool) -> Option { if self.depth >= 32 { return None; } self.depth += 1; - let mut result = String::new(); + let mut result = Layout::text(""); let mut scripts = 0; let mut has_base = false; while let Some(ch) = self.remaining.chars().next() { @@ -52,18 +111,15 @@ impl MathParser<'_> { let atom = self.atom()?; if !ch.is_whitespace() && ch != '^' && ch != '_' { // Flattening a compound base would change the scope of a following script. - has_base = atom.chars().count() == 1; - } - result.push_str(&atom); - if display_width(&result) > 256 { - return None; + has_base = atom.single().is_some_and(|text| text.chars().count() == 1); } + result = result.join(atom)?; } self.depth -= 1; (!group).then_some(result) } - fn argument(&mut self) -> Option { + fn argument(&mut self) -> Option { self.remaining = self.remaining.trim_start(); if self.remaining.starts_with(['^', '_']) { return None; @@ -71,7 +127,7 @@ impl MathParser<'_> { self.atom() } - fn atom(&mut self) -> Option { + fn atom(&mut self) -> Option { if self.depth >= 32 { return None; } @@ -81,7 +137,7 @@ impl MathParser<'_> { result } - fn atom_inner(&mut self) -> Option { + fn atom_inner(&mut self) -> Option { let ch = self.take()?; match ch { '{' => self.sequence(/*group*/ true), @@ -100,23 +156,23 @@ impl MathParser<'_> { ) }; let mut output = String::new(); - for value in arg.chars() { + for value in arg.single()?.chars() { let index = plain.chars().position(|ch| ch == value)?; output.push(alphabet.chars().nth(index)?); } - Some(output) + Some(Layout::text(output)) } '}' | '$' | '%' | '#' | '&' | '`' => None, ch if ch.is_whitespace() => { self.remaining = self.remaining.trim_start(); - Some(String::from(" ")) + Some(Layout::text(" ")) } ch if ch.is_control() => None, - ch => Some(ch.to_string()), + ch => Some(Layout::text(ch.to_string())), } } - fn command(&mut self) -> Option { + fn command(&mut self) -> Option { let length = self .remaining .bytes() @@ -124,11 +180,11 @@ impl MathParser<'_> { .count(); if length == 0 { return match self.take()? { - ',' | ';' | ':' | ' ' => Some(String::from(" ")), - '!' => Some(String::new()), - '{' => Some(String::from("{")), - '}' => Some(String::from("}")), - '|' => Some(String::from("‖")), + ',' | ';' | ':' | ' ' => Some(Layout::text(" ")), + '!' => Some(Layout::text("")), + '{' => Some(Layout::text("{")), + '}' => Some(Layout::text("}")), + '|' => Some(Layout::text("‖")), _ => None, }; } @@ -138,18 +194,47 @@ impl MathParser<'_> { "frac" | "dfrac" | "tfrac" => { let numerator = self.argument()?; let denominator = self.argument()?; - Some(format!("(({numerator})/({denominator}))")) + if !self.display { + return Some(Layout::text(format!( + "(({})/({}))", + numerator.single()?, + denominator.single()? + ))); + } + // Nested bars need a richer layout to preserve fraction hierarchy. + let numerator = numerator.single()?; + let denominator = denominator.single()?; + let width = display_width(numerator) + .max(display_width(denominator)) + .max(/*other*/ 1); + if width > MAX_COLUMNS { + return None; + } + Some(Layout { + rows: vec![ + format!( + "{}{numerator}", + " ".repeat((width - display_width(numerator)) / 2) + ), + "─".repeat(width), + format!( + "{}{denominator}", + " ".repeat((width - display_width(denominator)) / 2) + ), + ], + baseline: 1, + }) } "sqrt" => { if self.remaining.trim_start().starts_with('[') { return None; } let radicand = self.argument()?; - Some(format!("√({radicand})")) + Some(Layout::text(format!("√({})", radicand.single()?))) } "mathbb" => { let arg = self.argument()?; - let text = match arg.as_str() { + let text = match arg.single()? { "R" => "ℝ", "C" => "ℂ", "N" => "ℕ", @@ -158,7 +243,7 @@ impl MathParser<'_> { "P" => "ℙ", _ => return None, }; - Some(String::from(text)) + Some(Layout::text(text)) } "mathrm" | "mathbf" | "mathit" => self.argument(), "text" | "operatorname" => { @@ -171,21 +256,21 @@ impl MathParser<'_> { return None; } self.remaining = &self.remaining[end + 1..]; - Some(String::from(text)) + Some(Layout::text(text)) } "left" | "right" => { self.remaining = self.remaining.trim_start(); match self.take()? { - '.' => Some(String::new()), - ch @ ('(' | ')' | '[' | ']' | '|') => Some(ch.to_string()), + '.' => Some(Layout::text("")), + ch @ ('(' | ')' | '[' | ']' | '|') => Some(Layout::text(ch.to_string())), _ => None, } } - "quad" | "qquad" => Some(String::from(" ")), + "quad" | "qquad" => Some(Layout::text(" ")), "sin" | "cos" | "tan" | "log" | "ln" | "exp" | "lim" | "max" | "min" => { - Some(String::from(name)) + Some(Layout::text(name)) } - _ => symbol(name).map(String::from), + _ => symbol(name).map(Layout::text), } } } diff --git a/codex-rs/tui/src/markdown_render/math_display_tests.rs b/codex-rs/tui/src/markdown_render/math_display_tests.rs new file mode 100644 index 0000000000..b0c2288661 --- /dev/null +++ b/codex-rs/tui/src/markdown_render/math_display_tests.rs @@ -0,0 +1,157 @@ +use super::MathMarkdown; +use super::render::render; +use crate::markdown_render::render_markdown_text_with_width; +use pretty_assertions::assert_eq; +use pulldown_cmark::Options; + +fn plain(source: &str, width: usize) -> String { + render_markdown_text_with_width(source, Some(width)) + .lines + .iter() + .map(|line| { + line.spans + .iter() + .map(|span| span.content.as_ref()) + .collect::() + }) + .collect::>() + .join("\n") +} + +#[test] +fn unicode_math_snapshot() { + let source = r"The spectrum is $\lambda_1 \leq \lambda_2$ on $\mathbb{R}^n$. + +$$ +x=\frac{-b\pm\sqrt{b^2-4ac}}{2a} +$$ + +$$\frac{1}{2}\frac{3}{4}$$ + +\[ +\int_0^1 x^2\,dx=\frac{1}{3} +\] + +$$ +x +$$ trailing +$$ + +Inline \(\alpha^2 + \beta_{10}\), root $\sqrt{x^2+y^2}$, +and fraction $\frac{a}{b}$. + +Unsupported: $\begin{matrix}a&b\end{matrix}$. + +Code: `$\alpha$`; money: $5.00 and $10.00; shell: $HOME. + +```latex +\frac{a}{b} +``` + +Shell examples: $HOME and echo $$. + +After rejected equations: $\alpha$. + +$$\beta$$"; + insta::assert_snapshot!(plain(source, /*width*/ 80)); +} + +#[test] +fn unicode_math_narrow_layout_stays_meaningful() { + insta::assert_snapshot!(plain( + "$$\\frac{a+b}{c+d}$$\n\nWords $\\alpha^2$ more words.", + /*width*/ 12 + )); + for width in [80, 14] { + for prefix in [ + "- a\n - b\n - c\n\n ", + "- a\n - b\n - c\n", + "> > > Quote\n", + ] { + let formula = r"$$\frac{1234567890}{x}$$"; + assert_eq!( + plain(&format!("{prefix}{formula}"), width), + plain(&format!("{prefix}`{formula}`"), width) + ); + } + } + // A fraction too wide to retain its geometry stays source, using ordinary text wrapping. + assert_eq!( + plain("$$\\frac{abcdefghij}{k}$$", /*width*/ 12), + plain("`$$\\frac{abcdefghij}{k}$$`", /*width*/ 12) + ); +} + +#[test] +fn unicode_math_pending_display_tracks_original_offset() { + assert_eq!( + MathMarkdown::new("Prose\n\n$$\nx^2\n\n", Options::empty(), Some(80)).pending_start, + Some(7) + ); + assert_eq!( + MathMarkdown::new("```\n$$\n", Options::empty(), Some(80)).pending_start, + None + ); + assert_eq!( + MathMarkdown::new("echo $$\n", Options::empty(), Some(80)).pending_start, + None + ); +} + +#[test] +fn unicode_math_nested_fractions_and_oversized_pending_stay_source() { + for source in [r"\frac{\frac{a}{b}}{c}", r"\frac{a}{\frac{b}{c}}"] { + assert_eq!(render(source, /*display*/ true), None); + } + let source = format!("$$\n{}", "x".repeat(/*n*/ 5000)); + assert_eq!( + MathMarkdown::new(&source, Options::empty(), Some(80)).pending_start, + None + ); +} + +#[test] +fn unicode_math_oversized_display_stays_literal() { + for (open, close) in [("$$", "$$"), ("\\[", "\\]")] { + for ending in ["", close] { + let source = format!("{open}\n{}{ending}", "# x\n- y\n".repeat(/*n*/ 600)); + assert_eq!(plain(&source, /*width*/ 80), source.trim_end()); + } + } +} + +#[test] +fn unicode_math_pending_source_and_pid_boundaries() { + for (open, close) in [("$$", "$$"), ("\\[", "\\]")] { + assert_eq!( + plain(&format!("{open}\nx^2\n\n+y^2\n{close}"), /*width*/ 80), + "x² +y²" + ); + let source = format!("{open}\nx\n{close} trailing\n{close}\n\nAfter $\\alpha$."); + assert_eq!( + plain(&source, /*width*/ 80), + format!("{open}\nx\n{close} trailing\n{close}\n\nAfter α.") + ); + } + for source in ["$$\n# x\n", "$$\n- x\n", "\\[\n# x\n", "\\[\n- x\n"] { + assert_eq!(plain(source, /*width*/ 80), source.trim_end()); + } + let source = format!("{}After $\\alpha$.", "prose \\[\n".repeat(/*n*/ 1000)); + let math = MathMarkdown::new(&source, Options::empty(), Some(80)); + assert_eq!(math.display_ranges.len(), 1); + assert!(math.replacements.is_empty()); + assert_eq!( + plain( + "echo $$\n\n$$\nx^2\n$$\n\nAfter $\\alpha$.", + /*width*/ 80 + ), + "echo $$\n\nx²\n\nAfter α." + ); + assert_eq!( + plain( + "$$\nprice=\\$$$\n\nAfter $\\alpha$.\n\n$$\\beta$$", + /*width*/ 80 + ), + "$$\nprice=\\$$$\n\nAfter α.\n\nβ" + ); +} diff --git a/codex-rs/tui/src/markdown_render/math_tests.rs b/codex-rs/tui/src/markdown_render/math_tests.rs index 9affeacc56..dba122e882 100644 --- a/codex-rs/tui/src/markdown_render/math_tests.rs +++ b/codex-rs/tui/src/markdown_render/math_tests.rs @@ -59,10 +59,6 @@ fn unicode_math_preserves_markdown_contexts() { ] { assert_eq!(plain(source, /*width*/ 80), expected); } - // This stage leaves display equations and TeX fences to ordinary Markdown rendering. - for source in [r"$$x^2$$", "$$\nx^2\n$$", r"\[x^2\]", "```latex\nx^2\n```"] { - assert!(!plain(source, /*width*/ 80).contains('²')); - } } #[test] @@ -84,7 +80,7 @@ fn unicode_math_bounds_and_unsupported_input() { r"\text{\alpha}", r"\begin{matrix}a&b\end{matrix}", ] { - assert_eq!(render(source), None, "{source}"); + assert_eq!(render(source, /*display*/ false), None, "{source}"); assert_eq!( plain(&format!("\\({source}\\)"), /*width*/ 80), format!("\\({source}\\)") @@ -95,7 +91,7 @@ fn unicode_math_bounds_and_unsupported_input() { "x".repeat(/*n*/ 300), format!("{}x", "\\sqrt".repeat(/*n*/ 100)), ] { - assert_eq!(render(&source), None); + assert_eq!(render(&source, /*display*/ false), None); } } diff --git a/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_narrow_layout_stays_meaningful.snap b/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_narrow_layout_stays_meaningful.snap new file mode 100644 index 0000000000..3fc310bf09 --- /dev/null +++ b/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_narrow_layout_stays_meaningful.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/markdown_render/math_display_tests.rs +assertion_line: 46 +expression: "plain(\"$$\\\\frac{a+b}{c+d}$$\\n\\nWords $\\\\alpha^2$ more words.\", 12)" +--- +a+b +─── +c+d + +Words α² +more words. diff --git a/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_snapshot.snap b/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_snapshot.snap new file mode 100644 index 0000000000..bc09976c91 --- /dev/null +++ b/codex-rs/tui/src/markdown_render/snapshots/codex_tui__markdown_render__math__display_tests__unicode_math_snapshot.snap @@ -0,0 +1,37 @@ +--- +source: tui/src/markdown_render/math_display_tests.rs +expression: "plain(source, 80)" +--- +The spectrum is λ₁ ≤ λ₂ on ℝⁿ. + + -b±√(b²-4ac) +x=──────────── + 2a + +1 3 +─ ─ +2 4 + + 1 +∫₀¹ x² dx=─ + 3 + +$$ +x +$$ trailing +$$ + +Inline α² + β₁₀, root √(x²+y²), +and fraction ((a)/(b)). + +Unsupported: $\begin{matrix}a&b\end{matrix}$. + +Code: $\alpha$; money: $5.00 and $10.00; shell: $HOME. + +\frac{a}{b} + +Shell examples: $HOME and echo $$. + +After rejected equations: α. + +β diff --git a/codex-rs/tui/src/markdown_render/streaming.rs b/codex-rs/tui/src/markdown_render/streaming.rs index 75fefdae7a..001fced6be 100644 --- a/codex-rs/tui/src/markdown_render/streaming.rs +++ b/codex-rs/tui/src/markdown_render/streaming.rs @@ -19,6 +19,8 @@ use std::path::Path; pub(crate) struct StreamingMarkdownRender { /// Styled output produced by the same parser pass that collected the metadata below. pub(crate) lines: Vec, + /// Source line containing an unfinished display equation, which must not enter scrollback. + pub(crate) pending_math_start: Option, /// Byte offset of the final top-level block when at least one earlier block exists. pub(crate) last_top_level_block_start: Option, /// Whether a reference definition can retroactively change another block's rendering. @@ -41,7 +43,7 @@ pub(crate) fn render_streaming_markdown_lines_with_width_and_cwd( options.insert(Options::ENABLE_STRIKETHROUGH); options.insert(Options::ENABLE_TABLES); let citations = FileCitations::new(input, options); - let math = MathMarkdown::new(&citations.markdown, options); + let math = MathMarkdown::new(&citations.markdown, options, width); let parser = Parser::new_ext(&math.markdown, options); let has_reference_link_definition = parser.reference_definitions().iter().next().is_some(); let parser = TopLevelBlockTracker { @@ -55,14 +57,15 @@ pub(crate) fn render_streaming_markdown_lines_with_width_and_cwd( writer.run(); StreamingMarkdownRender { lines: writer.text, + pending_math_start: math.pending_start, last_top_level_block_start: (writer.iter.block_count > 1) .then_some(writer.iter.last_start) - // A cached suffix must not mistake a display closer for a new opener. - .filter(|boundary| { + .filter(|start| math.pending_start.is_none_or(|pending| *start <= pending)) + .filter(|start| { !math .display_ranges .iter() - .any(|range| range.start < *boundary && *boundary < range.end) + .any(|range| range.start < *start && *start < range.end) }), has_reference_link_definition, first_top_level_block_is_html: writer.iter.first_is_html, diff --git a/codex-rs/tui/src/streaming/controller.rs b/codex-rs/tui/src/streaming/controller.rs index d89095442b..1e5ab2cacc 100644 --- a/codex-rs/tui/src/streaming/controller.rs +++ b/codex-rs/tui/src/streaming/controller.rs @@ -52,6 +52,7 @@ use std::time::Duration; use std::time::Instant; use super::StreamState; +use super::prose_preview::PreviewMode; use super::prose_preview::ProsePreview; use super::render::StreamingRender; use super::render::render_source; @@ -168,12 +169,21 @@ impl StreamCore { } fn refresh_preview(&mut self) -> bool { - if self.holdback_scanner.allows_prose_preview() { + let pending = self.state.collector.pending_source(); + let prose_preview = self.holdback_scanner.allows_prose_preview(); + let math = self.render_mode == HistoryRenderMode::Rich + && (self.render.pending_math_start.is_some() + || prose_preview && (pending.starts_with("$$") || pending.starts_with("\\["))); + if math || prose_preview { self.preview.update( - self.state.collector.pending_source(), + pending, self.width, &self.cwd, - self.render_mode, + if math { + PreviewMode::Math + } else { + PreviewMode::Prose(self.render_mode) + }, self.inline_visualization_context.as_ref(), ) } else { @@ -278,9 +288,9 @@ impl StreamCore { let had_live_tail = self.has_tail(); self.width = width; self.state.collector.set_width(width); - self.refresh_preview(); let source = self.state.collector.committed_source(); if source.is_empty() { + self.refresh_preview(); return; } @@ -291,6 +301,7 @@ impl StreamCore { self.render_mode, self.inline_visualization_context.as_ref(), ); + self.refresh_preview(); self.emitted_stable_len = self.emitted_stable_len.min(self.render.lines.len()); if had_pending_queue && self.emitted_stable_len == self.render.lines.len() @@ -331,9 +342,9 @@ impl StreamCore { let had_pending_queue = self.state.queued_len() > 0; let had_live_tail = self.has_tail(); self.render_mode = render_mode; - self.refresh_preview(); let source = self.state.collector.committed_source(); if source.is_empty() { + self.refresh_preview(); return; } @@ -344,6 +355,7 @@ impl StreamCore { self.render_mode, self.inline_visualization_context.as_ref(), ); + self.refresh_preview(); self.emitted_stable_len = self.emitted_stable_len.min(self.render.lines.len()); if had_pending_queue && self.emitted_stable_len == self.render.lines.len() @@ -419,7 +431,7 @@ impl StreamCore { /// column widths. For `PendingHeader`, only content from the speculative /// header line onward is kept mutable so earlier prose can continue /// streaming. When no table is detected, everything flows directly to - /// stable. This is the core decision point for the holdback mechanism. + /// stable. Unclosed display math also stays mutable until its closing delimiter arrives. fn active_tail_budget_lines(&mut self) -> usize { if self.render_mode == HistoryRenderMode::Raw { return 0; @@ -439,7 +451,11 @@ impl StreamCore { elapsed_us = scan_start.elapsed().as_micros(), "table holdback decision", ); - tail_budget + let math_budget = self + .render + .pending_math_start + .map_or(0, |start| self.tail_budget_from_source_start(start)); + tail_budget.max(math_budget) } /// Convert a raw-source boundary into the number of rendered tail lines. @@ -787,6 +803,10 @@ impl PlanStreamController { } } +#[cfg(test)] +#[path = "math_tests.rs"] +mod math_tests; + #[cfg(test)] #[path = "controller_preview_tests.rs"] mod preview_tests; diff --git a/codex-rs/tui/src/streaming/math_tests.rs b/codex-rs/tui/src/streaming/math_tests.rs new file mode 100644 index 0000000000..3226278f12 --- /dev/null +++ b/codex-rs/tui/src/streaming/math_tests.rs @@ -0,0 +1,195 @@ +use super::StreamCore; +use super::render_source; +use crate::history_cell::HistoryRenderMode; +use crate::markdown::render_streaming_markdown_agent_with_links_and_cwd; +use pretty_assertions::assert_eq; + +#[test] +fn unicode_math_stream_holds_display_until_closed_and_preserves_source() { + let cwd = std::env::temp_dir(); + for mode in [HistoryRenderMode::Rich, HistoryRenderMode::Raw] { + for (delimiters, prefix) in [(("$$", "$$"), "Before.\n\n"), (("\\[", "\\]"), "Before.\n")] { + let mut stream = StreamCore::new( + Some(60), + &cwd, + mode, + /*inline_visualization_context*/ None, + ); + let mut emitted = Vec::new(); + let mut source = String::new(); + for chunk in [ + prefix, + delimiters.0, + "\n", + "\\frac{a+b}{c}\n", + "\n", + delimiters.1, + "\n\n", + "After $\\alpha^2$.\n", + ] { + source.push_str(chunk); + stream.push_delta(chunk); + emitted.extend(stream.tick_batch(usize::MAX)); + if mode == HistoryRenderMode::Rich && source.ends_with("\\frac{a+b}{c}\n") { + assert_eq!( + emitted, + render_source( + prefix, + Some(60), + &cwd, + mode, + /*inline_visualization_context*/ None + ) + ); + } + } + let (remaining, raw) = stream.finalize_remaining(); + emitted.extend(remaining); + assert_eq!(raw, source); + assert_eq!( + emitted, + render_source( + &source, + Some(60), + &cwd, + mode, + /*inline_visualization_context*/ None + ) + ); + } + } +} + +#[test] +fn unicode_math_raw_preview_preserves_line_on_newline() { + let cwd = std::env::temp_dir(); + for open in ["$$", "\\["] { + let mut stream = StreamCore::new( + Some(12), + &cwd, + HistoryRenderMode::Raw, + /*inline_visualization_context*/ None, + ); + let source = format!("{open}\\frac{{abcdefghijk}}{{lmnop}}"); + stream.push_delta(&source); + let preview = stream.current_tail_lines(); + assert_eq!(preview.len(), 1); + assert_eq!(preview[0].line.to_string(), source); + if open == "$$" { + insta::assert_snapshot!("unicode_math_raw_preview", preview[0].line.to_string()); + } + stream.push_delta("\n"); + let (lines, raw) = stream.finalize_remaining(); + assert_eq!(lines, preview); + assert_eq!(raw, format!("{source}\n")); + } +} + +#[test] +fn unicode_math_unfinished_display_is_visible_and_reflows() { + let cwd = std::env::temp_dir(); + for (open, close) in [("$$", "$$"), (r"\[", r"\]")] { + let mut stream = StreamCore::new( + Some(80), + &cwd, + HistoryRenderMode::Rich, + /*inline_visualization_context*/ None, + ); + for (chunk, expected) in [(&open[..1], &open[..1]), (&open[1..], open)] { + stream.push_delta(chunk); + assert_eq!(stream.current_tail_lines()[0].line.to_string(), expected); + } + stream.push_delta("\n \\frac{|a+b+c+d+e+f|}{g+h}"); + for width in [80, 20, 80] { + stream.set_width(Some(width)); + let lines = stream.current_tail_lines(); + assert!( + lines + .iter() + .any(|line| line.line.to_string().contains("frac")) + ); + assert!(lines.iter().all(|line| line.line.width() <= width)); + if open == "$$" && width == 20 { + insta::assert_snapshot!( + "unicode_math_live_source_narrow", + lines + .iter() + .map(|line| line.line.to_string()) + .collect::>() + .join("\n") + ); + } + } + stream.push_delta(&format!("\n{close}\n")); + let (lines, source) = stream.finalize_remaining(); + assert_eq!( + lines, + render_source( + &source, + Some(80), + &cwd, + HistoryRenderMode::Rich, + /*inline_visualization_context*/ None + ) + ); + } +} + +#[test] +fn unicode_math_rejected_display_keeps_its_closer_and_following_text() { + let cwd = std::env::temp_dir(); + for (open, close) in [("$$", "$$"), ("\\[", "\\]")] { + for (prefix, body) in [ + ("", "x".repeat(/*n*/ 5000)), + ("", "`x`".to_owned()), + ("", format!("`{close}\nx`")), + ("", "[x](https://example.com)".to_owned()), + ("Equation: ", "x^2".to_owned()), + ] { + let mut render = super::super::render::StreamingRender::new(); + let mut source = String::new(); + for chunk in [ + format!("{prefix}{open}\n\n"), + format!("{body}\n\n"), + format!("{close}\n\n"), + "After $\\alpha$.\n\n".to_owned(), + "$$\\beta$$\n".to_owned(), + ] { + source.push_str(&chunk); + render.append( + &source, + &chunk, + Some(40), + &cwd, + HistoryRenderMode::Rich, + /*inline_visualization_context*/ None, + ); + let expected = render_streaming_markdown_agent_with_links_and_cwd( + &source, + Some(40), + Some(&cwd), + ); + assert_eq!( + (&render.lines, render.pending_math_start), + (&expected.lines, expected.pending_math_start) + ); + if source.len() > 4096 || !prefix.is_empty() { + assert_eq!(render.pending_math_start, None); + } + } + assert!( + render + .lines + .iter() + .any(|line| line.line.to_string().contains("After α.")) + ); + assert!( + render + .lines + .iter() + .any(|line| line.line.to_string().contains('β')) + ); + assert_eq!(render.pending_math_start, None); + } + } +} diff --git a/codex-rs/tui/src/streaming/prose_preview.rs b/codex-rs/tui/src/streaming/prose_preview.rs index 65609345b9..94a848a330 100644 --- a/codex-rs/tui/src/streaming/prose_preview.rs +++ b/codex-rs/tui/src/streaming/prose_preview.rs @@ -1,5 +1,6 @@ //! Bounded, disposable previews of unterminated prose. Preview lines never enter scrollback; -//! newline commitment and finalization render the original source independently. +//! newline commitment and finalization render the original source independently. Math previews +//! show wrapped source until closure, without interpreting TeX punctuation as Markdown. use super::render::render_source; use crate::history_cell::HistoryRenderMode; @@ -10,6 +11,11 @@ use std::path::Path; const MAX_PREVIEW_BYTES: usize = 8192; +pub(super) enum PreviewMode { + Prose(HistoryRenderMode), + Math, +} + /// Tracks one append-only incomplete line; reset when a newline is committed. #[derive(Default)] pub(super) struct ProsePreview { @@ -25,7 +31,7 @@ impl ProsePreview { source: &str, width: Option, cwd: &Path, - render_mode: HistoryRenderMode, + mode: PreviewMode, inline_visualization_context: Option<&InlineVisualizationContext>, ) -> bool { // Only scan newly arrived bytes, including on very long single-line responses. @@ -33,24 +39,31 @@ impl ProsePreview { self.scanned_len = source.len(); // Indented and quoted lines may belong to nested code blocks. Keep their // existing newline holdback instead of guessing at the missing block context. - if !(self.has_pipe - || source.starts_with([' ', '\t', '>']) - || source.starts_with("```") - || source.starts_with("~~~") - || matches!(source, "`" | "``" | "~" | "~~")) + if matches!(mode, PreviewMode::Math) + || !(self.has_pipe + || source.starts_with([' ', '\t', '>']) + || source.starts_with("```") + || source.starts_with("~~~") + || matches!(source, "`" | "``" | "~" | "~~")) { self.safe_len = source.len(); } // Retain the last safe text when tokens reveal structure, but still reflow it. let source = &source[..self.safe_len]; let start = source.ceil_char_boundary(source.len().saturating_sub(MAX_PREVIEW_BYTES)); - let mut lines = render_source( - &source[start..], - width, - cwd, - render_mode, - inline_visualization_context, - ); + let mut lines = match mode { + PreviewMode::Prose(render_mode) => render_source( + &source[start..], + width, + cwd, + render_mode, + inline_visualization_context, + ), + PreviewMode::Math => textwrap::wrap(&source[start..], width.unwrap_or(usize::MAX)) + .into_iter() + .map(|line| HyperlinkLine::new(Line::from(line.into_owned()))) + .collect(), + }; if start > 0 { lines.insert(0, HyperlinkLine::new(Line::from("…"))); } diff --git a/codex-rs/tui/src/streaming/render.rs b/codex-rs/tui/src/streaming/render.rs index 1996332fff..ff082ba856 100644 --- a/codex-rs/tui/src/streaming/render.rs +++ b/codex-rs/tui/src/streaming/render.rs @@ -22,6 +22,7 @@ use std::path::Path; /// re-rendered as committed source arrives. pub(super) struct StreamingRender { pub(super) lines: Vec, + pub(super) pending_math_start: Option, /// Source prefix containing only completed top-level markdown blocks. stable_source_len: usize, /// Rendered-line boundary corresponding to `stable_source_len`. @@ -38,6 +39,7 @@ impl StreamingRender { pub(super) fn new() -> Self { Self { lines: Vec::with_capacity(64), + pending_math_start: None, stable_source_len: 0, stable_rendered_len: 0, has_reference_link_definition: false, @@ -48,6 +50,7 @@ impl StreamingRender { pub(super) fn clear(&mut self) { self.lines.clear(); + self.pending_math_start = None; self.stable_source_len = 0; self.stable_rendered_len = 0; self.has_reference_link_definition = false; @@ -68,16 +71,26 @@ impl StreamingRender { inline_visualization_context: Option<&InlineVisualizationContext>, ) { self.open_code_fence = None; + self.pending_math_start = None; self.has_inline_visualization_directive = contains_inline_visualization(source); self.lines = match (render_mode, inline_visualization_context) { (HistoryRenderMode::Rich, None) if !self.has_inline_visualization_directive => { let rendered = render_streaming_markdown_agent_with_links_and_cwd(source, width, Some(cwd)); self.has_reference_link_definition = rendered.has_reference_link_definition; + self.pending_math_start = rendered.pending_math_start; rendered.lines } _ => { self.has_reference_link_definition = false; + if render_mode == HistoryRenderMode::Rich { + self.pending_math_start = render_streaming_markdown_agent_with_links_and_cwd( + source, + width, + Some(cwd), + ) + .pending_math_start; + } render_source( source, width, @@ -149,6 +162,9 @@ impl StreamingRender { let theme_revision = syntax_theme_revision(); let pending = render_streaming_markdown_agent_with_links_and_cwd(pending_source, width, Some(cwd)); + self.pending_math_start = pending + .pending_math_start + .map(|start| self.stable_source_len + start); if pending.has_reference_link_definition { self.has_reference_link_definition = true; self.recompute( diff --git a/codex-rs/tui/src/streaming/render_tests.rs b/codex-rs/tui/src/streaming/render_tests.rs index 44ad11e55d..55aefd1c9c 100644 --- a/codex-rs/tui/src/streaming/render_tests.rs +++ b/codex-rs/tui/src/streaming/render_tests.rs @@ -447,3 +447,53 @@ fn paragraphs_after_unwrapped_table_fence_advance_stable_source() { previous_stable_source_len = render.stable_source_len; } } + +#[test] +fn shell_pid_preserves_following_equations() { + for shell in [ + "Shell examples: $HOME and echo $$.", + "Shell examples: $HOME.", + ] { + let source = format!("{shell}\n\nAfter rejected equations: $\\alpha$.\n\n$$\\beta$$"); + for width in [80, 24] { + let (_, render) = assert_rich_stream_matches_full_render( + &source.split_inclusive('$').collect::>(), + Some(width), + ); + let expected = format!("{shell}\n\nAfter rejected equations: α.\n\nβ"); + assert_eq!( + render.lines, + render_source( + &expected, + Some(width), + &test_cwd(), + HistoryRenderMode::Rich, + /*inline_visualization_context*/ None, + ) + ); + } + } +} + +#[test] +fn rejected_math_openers_allow_bounded_incremental_rendering() { + let cwd = test_cwd(); + for (open, close) in [ + ("echo $$", "$$"), + ("Equation: \\[", "\\]"), + (r"\[label\]\*", "\\]"), + ] { + let (mut source, mut render) = + assert_rich_stream_matches_full_render(&[&format!("{open}\n\n")], Some(80)); + let distant_closer = format!("{close}\n\nAfter $\\alpha$.\n\n"); + for chunk in std::iter::repeat_n( + "An ordinary paragraph that must not retain the entire response.\n\n", + /*count*/ 160, + ) + .chain(std::iter::once(distant_closer.as_str())) + { + append_rich_and_assert_matches_full(&mut render, &mut source, chunk, Some(80), &cwd); + assert!(source.len() - render.stable_source_len < 4200); + } + } +} diff --git a/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_live_source_narrow.snap b/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_live_source_narrow.snap new file mode 100644 index 0000000000..8a24860f07 --- /dev/null +++ b/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_live_source_narrow.snap @@ -0,0 +1,8 @@ +--- +source: tui/src/streaming/math_tests.rs +assertion_line: 134 +expression: "stream.current_tail_lines().iter().map(|line|\nline.line.to_string()).collect::>().join(\"\\n\")" +--- +$$ + \frac{| +a+b+c+d+e+f|}{g+h} diff --git a/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_raw_preview.snap b/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_raw_preview.snap new file mode 100644 index 0000000000..480853d92f --- /dev/null +++ b/codex-rs/tui/src/streaming/snapshots/codex_tui__streaming__controller__math_tests__unicode_math_raw_preview.snap @@ -0,0 +1,5 @@ +--- +source: tui/src/streaming/math_tests.rs +expression: "preview[0].line.to_string()" +--- +$$\frac{abcdefghijk}{lmnop}