Render streamed Markdown incrementally

codex-rs/tui/src/inline_visualization.rs

@@ -24,7 +24,7 @@ use uuid::Uuid;
 
 use self::viewer::materialize_document;
 
-const DIRECTIVE_PREFIX: &str = "::codex-inline-vis{";
+pub(crate) const DIRECTIVE_PREFIX: &str = "::codex-inline-vis{";
 const MAX_FRAGMENT_BYTES: u64 = 2 * 1024 * 1024;
 
 #[derive(Clone, Debug)]

codex-rs/tui/src/markdown.rs

@@ -106,6 +106,34 @@ pub(crate) fn render_markdown_agent_with_links_cwd_and_visualizations(
     lines
 }
 
+/// Render an agent message and collect the block metadata needed for incremental rendering.
+///
+/// Block offsets are mapped back to `markdown_source` after Markdown table fences are unwrapped.
+/// If a normalized boundary cannot be expressed as a raw-source suffix, it is discarded so the
+/// transformed block remains mutable.
+pub(crate) fn render_streaming_markdown_agent_with_links_and_cwd(
+    markdown_source: &str,
+    width: Option<usize>,
+    cwd: Option<&Path>,
+) -> crate::markdown_render::StreamingMarkdownRender {
+    let normalized = unwrap_markdown_fences(markdown_source);
+    let mut rendered = crate::markdown_render::render_streaming_markdown_lines_with_width_and_cwd(
+        &normalized,
+        width,
+        cwd,
+    );
+    if normalized != markdown_source {
+        // 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.last_top_level_block_start = rendered
+            .last_top_level_block_start
+            .and_then(|boundary| markdown_source.strip_suffix(&normalized[boundary..]))
+            .map(str::len);
+    }
+    rendered
+}
+
 /// Strip `` ```md ``/`` ```markdown `` fences that contain tables, emitting their content as bare
 /// markdown so `pulldown-cmark` parses the tables natively.
 ///

codex-rs/tui/src/markdown_render.rs

@@ -77,8 +77,12 @@ use unicode_width::UnicodeWidthChar;
 use unicode_width::UnicodeWidthStr;
 use url::Url;
 
+mod streaming;
 mod table_key_value;
 
+pub(crate) use streaming::StreamingMarkdownRender;
+pub(crate) use streaming::render_streaming_markdown_lines_with_width_and_cwd;
+
 const TABLE_COLUMN_GAP: usize = 2;
 const TABLE_CELL_PADDING: usize = 1;
 const TABLE_HEADER_SEPARATOR_CHAR: char = '━';

codex-rs/tui/src/markdown_render/streaming.rs

@@ -0,0 +1,92 @@
+//! Streaming markdown render metadata collected during the writer's single parse pass.
+//!
+//! Top-level block offsets always refer to the exact source passed to this renderer; callers that
+//! normalize source before rendering must not apply those offsets to the original source.
+
+use super::DecodedTextMerge;
+use super::Event;
+use super::HyperlinkLine;
+use super::Options;
+use super::Parser;
+use super::Tag;
+use super::Writer;
+use super::never_hide_link_destination;
+use std::ops::Range;
+use std::path::Path;
+
+/// Rendered lines and the block metadata needed to keep only the final block mutable.
+pub(crate) struct StreamingMarkdownRender {
+    /// Styled output produced by the same parser pass that collected the metadata below.
+    pub(crate) lines: Vec<HyperlinkLine>,
+    /// Byte offset of the final top-level block when at least one earlier block exists.
+    pub(crate) last_top_level_block_start: Option<usize>,
+    /// Whether a reference definition can retroactively change another block's rendering.
+    pub(crate) has_reference_link_definition: bool,
+    /// Whether the first block is raw HTML, which joins a retained prefix without a separator.
+    pub(crate) first_top_level_block_is_html: bool,
+}
+
+/// Render `input` while tracking the final mutable top-level block.
+///
+/// Every reported byte offset indexes the exact `input` passed here. Callers that transform source
+/// before rendering must map the offset back to their original source before retaining a prefix.
+pub(crate) fn render_streaming_markdown_lines_with_width_and_cwd(
+    input: &str,
+    width: Option<usize>,
+    cwd: Option<&Path>,
+) -> StreamingMarkdownRender {
+    let mut options = Options::empty();
+    options.insert(Options::ENABLE_STRIKETHROUGH);
+    options.insert(Options::ENABLE_TABLES);
+    let parser = Parser::new_ext(input, options);
+    let has_reference_link_definition = parser.reference_definitions().iter().next().is_some();
+    let parser = TopLevelBlockTracker {
+        iter: DecodedTextMerge::new(parser.into_offset_iter()),
+        depth: 0,
+        block_count: 0,
+        last_start: 0,
+        first_is_html: false,
+    };
+    let mut writer = Writer::new(input, parser, width, cwd, &never_hide_link_destination);
+    writer.run();
+    StreamingMarkdownRender {
+        lines: writer.text,
+        last_top_level_block_start: (writer.iter.block_count > 1).then_some(writer.iter.last_start),
+        has_reference_link_definition,
+        first_top_level_block_is_html: writer.iter.first_is_html,
+    }
+}
+
+/// Records top-level block boundaries without adding a second parser traversal.
+struct TopLevelBlockTracker<I> {
+    iter: I,
+    depth: usize,
+    block_count: usize,
+    last_start: usize,
+    first_is_html: bool,
+}
+
+impl<'a, I> Iterator for TopLevelBlockTracker<I>
+where
+    I: Iterator<Item = (Event<'a>, Range<usize>)>,
+{
+    type Item = (Event<'a>, Range<usize>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let (event, range) = self.iter.next()?;
+        if self.depth == 0 && matches!(&event, Event::Start(_) | Event::Rule | Event::Html(_)) {
+            self.block_count += 1;
+            self.last_start = range.start;
+            if self.block_count == 1 {
+                self.first_is_html =
+                    matches!(&event, Event::Start(Tag::HtmlBlock) | Event::Html(_));
+            }
+        }
+        match event {
+            Event::Start(_) => self.depth += 1,
+            Event::End(_) => self.depth = self.depth.saturating_sub(1),
+            _ => {}
+        }
+        Some((event, range))
+    }
+}