From 2244d11a1d9eaa389a5da4c99b746c74240abd8c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charliemarsh@openai.com>
Date: Mon, 20 Jul 2026 13:41:08 +0000
Subject: [PATCH] Track inline visualization directives during streaming
 (#34346)

## What changed

- Record when committed streaming content contains an inline visualization directive and keep using the canonical full-source renderer afterward.
- Restore directive state during full recomputes and clear it when the streaming renderer is reset.
- Cover directive tracking with and without visualization context, including a raw-to-rich render mode switch.

GitOrigin-RevId: 17786851d6dad6805240b79072fcbd0bb441874a
---
 codex-rs/tui/src/streaming/render.rs       | 10 ++++-
 codex-rs/tui/src/streaming/render_tests.rs | 45 ++++++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/codex-rs/tui/src/streaming/render.rs b/codex-rs/tui/src/streaming/render.rs
index dccc22acd628..29070c67835e 100644
--- a/codex-rs/tui/src/streaming/render.rs
+++ b/codex-rs/tui/src/streaming/render.rs
@@ -26,6 +26,8 @@ pub(super) struct StreamingRender {
     stable_rendered_len: usize,
     /// Reference-style link definitions can affect any earlier or later markdown block.
     has_reference_link_definition: bool,
+    /// Inline visualization directives require source-wide rewriting once one is committed.
+    has_inline_visualization_directive: bool,
 }
 
 impl StreamingRender {
@@ -35,6 +37,7 @@ impl StreamingRender {
             stable_source_len: 0,
             stable_rendered_len: 0,
             has_reference_link_definition: false,
+            has_inline_visualization_directive: false,
         }
     }
 
@@ -43,6 +46,7 @@ impl StreamingRender {
         self.stable_source_len = 0;
         self.stable_rendered_len = 0;
         self.has_reference_link_definition = false;
+        self.has_inline_visualization_directive = false;
     }
 
     /// Re-render the full source and reset both stable-prefix boundaries.
@@ -57,8 +61,9 @@ impl StreamingRender {
         render_mode: HistoryRenderMode,
         inline_visualization_context: Option<&InlineVisualizationContext>,
     ) {
+        self.has_inline_visualization_directive = source.contains(DIRECTIVE_PREFIX);
         self.lines = match (render_mode, inline_visualization_context) {
-            (HistoryRenderMode::Rich, None) if !source.contains(DIRECTIVE_PREFIX) => {
+            (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;
@@ -102,7 +107,8 @@ impl StreamingRender {
             return;
         }
 
-        if raw_source.contains(DIRECTIVE_PREFIX) {
+        self.has_inline_visualization_directive |= committed_source.contains(DIRECTIVE_PREFIX);
+        if self.has_inline_visualization_directive {
             self.recompute(
                 raw_source,
                 width,
diff --git a/codex-rs/tui/src/streaming/render_tests.rs b/codex-rs/tui/src/streaming/render_tests.rs
index 925cd02dd362..c7b5ecef716f 100644
--- a/codex-rs/tui/src/streaming/render_tests.rs
+++ b/codex-rs/tui/src/streaming/render_tests.rs
@@ -197,6 +197,7 @@ fn inline_visualization_context_without_directives_keeps_stable_prefix() {
     }
 
     assert!(render.stable_source_len > 0);
+    assert!(!render.has_inline_visualization_directive);
 }
 
 #[test]
@@ -230,6 +231,8 @@ fn inline_visualizations_use_canonical_full_render() {
         );
         assert_eq!(render.stable_source_len, 0);
     }
+
+    assert!(render.has_inline_visualization_directive);
 }
 
 #[test]
@@ -240,12 +243,54 @@ fn inline_visualizations_without_context_use_canonical_full_render() {
     );
 
     assert_eq!(render.stable_source_len, 0);
+    assert!(render.has_inline_visualization_directive);
     assert_debug_snapshot!(
         "inline_visualizations_without_context_use_canonical_full_render",
         render.lines
     );
 }
 
+#[test]
+fn inline_visualization_directive_survives_raw_to_rich_render_mode_switch() {
+    let cwd = test_cwd();
+    let width = Some(80);
+    let mut source = String::new();
+    let mut render = StreamingRender::new();
+
+    append(
+        &mut render,
+        &mut source,
+        "::codex-inline-vis{file=\"missing.html\"}\n",
+        width,
+        &cwd,
+        HistoryRenderMode::Raw,
+    );
+    assert!(!render.has_inline_visualization_directive);
+
+    render.recompute(
+        &source,
+        width,
+        &cwd,
+        HistoryRenderMode::Rich,
+        /*inline_visualization_context*/ None,
+    );
+
+    assert!(render.has_inline_visualization_directive);
+    assert_eq!(
+        render.lines,
+        render_source(
+            &source,
+            width,
+            &cwd,
+            HistoryRenderMode::Rich,
+            /*inline_visualization_context*/ None,
+        ),
+    );
+
+    render.clear();
+    assert!(!render.has_inline_visualization_directive);
+}
+
 #[test]
 fn reference_link_definition_recomputes_earlier_and_later_blocks() {
     let streams: &[&[&str]] = &[