codex: address PR review feedback (#20252)

codex-rs/tui/src/app/resize_reflow.rs

@@ -29,7 +29,9 @@ use super::App;
 use super::InitialHistoryReplayBuffer;
 use crate::history_cell;
 use crate::history_cell::HistoryCell;
+use crate::markdown::append_markdown;
 use crate::render::line_utils::prefix_lines;
+use crate::style::proposed_plan_style;
 use crate::transcript_reflow::TRANSCRIPT_REFLOW_DEBOUNCE;
 use crate::tui;
 
@@ -44,6 +46,12 @@ struct AgentMessageStreamSource {
     is_first_line: bool,
 }
 
+struct ProposedPlanStreamSource {
+    source: String,
+    cwd: PathBuf,
+    include_bottom_padding: bool,
+}
+
 /// Rendered transcript lines ready to be replayed into terminal scrollback.
 ///
 /// This is intentionally line-oriented rather than cell-oriented because the terminal only accepts
@@ -399,22 +407,34 @@ impl App {
         let row_cap = self.resize_reflow_max_rows();
         let mut cell_displays = VecDeque::new();
         let mut rendered_rows = 0usize;
-        let stream_run_start =
-            trailing_run_start::<history_cell::AgentMessageCell>(&self.transcript_cells);
-        let stream_source = self
-            .trailing_agent_message_stream_source(stream_run_start, self.transcript_cells.len());
-        let mut start = stream_source
-            .as_ref()
-            .map(|_| stream_run_start)
-            .unwrap_or(self.transcript_cells.len());
+        let mut start = self.transcript_cells.len();
 
-        if let Some(stream_source) = stream_source {
-            let lines = render_agent_message_stream_source(&stream_source, width);
-            rendered_rows += lines.len();
-            cell_displays.push_front(ReflowCellDisplay {
-                lines,
-                is_stream_continuation: !stream_source.is_first_line,
-            });
+        let plan_stream_run_start =
+            trailing_run_start::<history_cell::ProposedPlanStreamCell>(&self.transcript_cells);
+        if let Some(displays) = self.trailing_proposed_plan_stream_displays(
+            plan_stream_run_start,
+            self.transcript_cells.len(),
+            width,
+        ) {
+            start = plan_stream_run_start;
+            for display in displays.into_iter().rev() {
+                rendered_rows += display.lines.len();
+                cell_displays.push_front(display);
+            }
+        } else {
+            let agent_stream_run_start =
+                trailing_run_start::<history_cell::AgentMessageCell>(&self.transcript_cells);
+            if let Some(displays) = self.trailing_agent_message_stream_displays(
+                agent_stream_run_start,
+                self.transcript_cells.len(),
+                width,
+            ) {
+                start = agent_stream_run_start;
+                for display in displays.into_iter().rev() {
+                    rendered_rows += display.lines.len();
+                    cell_displays.push_front(display);
+                }
+            }
         }
 
         while start > 0 {
@@ -470,11 +490,12 @@ impl App {
         }
     }
 
-    fn trailing_agent_message_stream_source(
+    fn trailing_agent_message_stream_displays(
         &self,
         start: usize,
         end: usize,
-    ) -> Option<AgentMessageStreamSource> {
+        width: u16,
+    ) -> Option<Vec<ReflowCellDisplay>> {
         if start == end {
             return None;
         }
@@ -485,7 +506,7 @@ impl App {
         let is_first_line = first_agent_cell.is_first_line();
 
         let mut latest_source = None;
-        let mut latest_source_offset = None;
+        let mut latest_source_offset = 0;
         for (offset, cell) in self.transcript_cells[start..end].iter().enumerate() {
             let agent_cell = cell
                 .as_any()
@@ -496,13 +517,66 @@ impl App {
                     cwd: cwd.to_path_buf(),
                     is_first_line,
                 });
-                latest_source_offset = Some(offset);
+                latest_source_offset = offset;
             }
         }
 
-        (latest_source_offset == Some(end - start - 1))
-            .then_some(latest_source)
-            .flatten()
+        let latest_source = latest_source?;
+        let mut displays = vec![ReflowCellDisplay {
+            lines: render_agent_message_stream_source(&latest_source, width),
+            is_stream_continuation: !latest_source.is_first_line,
+        }];
+
+        for cell in &self.transcript_cells[start + latest_source_offset + 1..end] {
+            displays.push(ReflowCellDisplay {
+                lines: cell.display_lines(width),
+                is_stream_continuation: cell.is_stream_continuation(),
+            });
+        }
+
+        Some(displays)
+    }
+
+    fn trailing_proposed_plan_stream_displays(
+        &self,
+        start: usize,
+        end: usize,
+        width: u16,
+    ) -> Option<Vec<ReflowCellDisplay>> {
+        if start == end {
+            return None;
+        }
+
+        let mut latest_source = None;
+        let mut latest_source_offset = 0;
+        for (offset, cell) in self.transcript_cells[start..end].iter().enumerate() {
+            let plan_cell = cell
+                .as_any()
+                .downcast_ref::<history_cell::ProposedPlanStreamCell>()?;
+            if let Some((source, cwd, include_bottom_padding)) = plan_cell.markdown_source() {
+                latest_source = Some(ProposedPlanStreamSource {
+                    source: source.to_string(),
+                    cwd: cwd.to_path_buf(),
+                    include_bottom_padding,
+                });
+                latest_source_offset = offset;
+            }
+        }
+
+        let latest_source = latest_source?;
+        let mut displays = vec![ReflowCellDisplay {
+            lines: render_proposed_plan_stream_source(&latest_source, width),
+            is_stream_continuation: false,
+        }];
+
+        for cell in &self.transcript_cells[start + latest_source_offset + 1..end] {
+            displays.push(ReflowCellDisplay {
+                lines: cell.display_lines(width),
+                is_stream_continuation: cell.is_stream_continuation(),
+            });
+        }
+
+        Some(displays)
     }
 
     /// Return whether current transcript state should be treated as stream-time resize state.
@@ -554,3 +628,33 @@ fn render_agent_message_stream_source(
         "  ".into(),
     )
 }
+
+fn render_proposed_plan_stream_source(
+    stream_source: &ProposedPlanStreamSource,
+    width: u16,
+) -> Vec<Line<'static>> {
+    let mut lines: Vec<Line<'static>> = vec![vec!["• ".dim(), "Proposed Plan".bold()].into()];
+    lines.push(Line::from(" "));
+
+    let mut body = Vec::new();
+    let wrap_width = width.saturating_sub(4).max(1) as usize;
+    append_markdown(
+        &stream_source.source,
+        Some(wrap_width),
+        Some(stream_source.cwd.as_path()),
+        &mut body,
+    );
+    if body.is_empty() {
+        body.push(Line::from("(empty)".dim().italic()));
+    }
+
+    let plan_style = proposed_plan_style();
+    let mut plan_lines: Vec<Line<'static>> = vec![Line::from(" ")];
+    plan_lines.extend(prefix_lines(body, "  ".into(), "  ".into()));
+    if stream_source.include_bottom_padding {
+        plan_lines.push(Line::from(" "));
+    }
+
+    lines.extend(plan_lines.into_iter().map(|line| line.style(plan_style)));
+    lines
+}

codex-rs/tui/src/app/tests.rs

@@ -18,6 +18,9 @@ use crate::history_cell::AgentMessageCell;
 use crate::history_cell::HistoryCell;
 use crate::history_cell::PlainHistoryCell;
 use crate::history_cell::UserHistoryCell;
+use crate::history_cell::new_proposed_plan;
+use crate::history_cell::new_proposed_plan_stream;
+use crate::history_cell::new_proposed_plan_stream_with_markdown_source;
 use crate::history_cell::new_session_info;
 use crate::multi_agents::AgentPickerThreadEntry;
 use assert_matches::assert_matches;
@@ -4024,7 +4027,87 @@ async fn table_resize_lifecycle_stream_reflow_uses_markdown_source_not_transient
 }
 
 #[tokio::test]
-async fn table_resize_lifecycle_stream_reflow_ignores_stale_markdown_source() {
+async fn table_resize_lifecycle_stream_reflow_preserves_source_prefix_before_unsourced_tail() {
+    let (mut app, _rx, _op_rx) = make_test_app_with_channels().await;
+    enable_terminal_resize_reflow(&mut app);
+    app.config.terminal_resize_reflow.max_rows = TerminalResizeReflowMaxRows::Disabled;
+    let cwd = std::env::temp_dir();
+    let source = markdown_table_source();
+
+    app.transcript_cells
+        .push(Arc::new(AgentMessageCell::new_with_markdown_source(
+            vec![Line::from("│ Area │ Result │")],
+            /*is_first_line*/ true,
+            source.to_string(),
+            cwd.as_path(),
+        )) as Arc<dyn HistoryCell>);
+    app.transcript_cells.push(Arc::new(AgentMessageCell::new(
+        vec![Line::from("newer emitted stream row")],
+        /*is_first_line*/ false,
+    )) as Arc<dyn HistoryCell>);
+
+    let reflowed = app.render_transcript_lines_for_reflow(/*width*/ 44);
+    let lines = reflowed
+        .lines
+        .iter()
+        .map(rendered_line_text)
+        .collect::<Vec<_>>();
+    let mut expected = AgentMarkdownCell::new(source.to_string(), cwd.as_path())
+        .display_lines(/*width*/ 44)
+        .iter()
+        .map(rendered_line_text)
+        .collect::<Vec<_>>();
+    expected.push("  newer emitted stream row".to_string());
+
+    assert_eq!(
+        lines, expected,
+        "resize reflow must rebuild the source-backed prefix and append newer unsourced stream cells",
+    );
+}
+
+#[tokio::test]
+async fn table_resize_lifecycle_plan_stream_reflow_uses_markdown_source_before_unsourced_tail() {
+    let (mut app, _rx, _op_rx) = make_test_app_with_channels().await;
+    enable_terminal_resize_reflow(&mut app);
+    app.config.terminal_resize_reflow.max_rows = TerminalResizeReflowMaxRows::Disabled;
+    let cwd = std::env::temp_dir();
+    let source = markdown_table_source();
+
+    app.transcript_cells
+        .push(Arc::new(new_proposed_plan_stream_with_markdown_source(
+            vec![Line::from("│ Area │ Result │")],
+            /*is_stream_continuation*/ false,
+            source.to_string(),
+            cwd.as_path(),
+            /*include_bottom_padding*/ false,
+        )) as Arc<dyn HistoryCell>);
+    app.transcript_cells.push(Arc::new(new_proposed_plan_stream(
+        vec![Line::from("  newer emitted plan row")],
+        /*is_stream_continuation*/ true,
+    )) as Arc<dyn HistoryCell>);
+
+    let reflowed = app.render_transcript_lines_for_reflow(/*width*/ 44);
+    let lines = reflowed
+        .lines
+        .iter()
+        .map(rendered_line_text)
+        .collect::<Vec<_>>();
+    let mut expected = new_proposed_plan(source.to_string(), cwd.as_path())
+        .display_lines(/*width*/ 44)
+        .iter()
+        .map(rendered_line_text)
+        .collect::<Vec<_>>();
+    expected.pop();
+    expected.push("  newer emitted plan row".to_string());
+
+    assert_eq!(
+        lines, expected,
+        "resize reflow must rebuild proposed-plan stream tables from source and append newer unsourced plan cells",
+    );
+}
+
+#[tokio::test]
+async fn table_resize_lifecycle_stream_reflow_does_not_drop_newer_unsourced_tail() {
     let (mut app, _rx, _op_rx) = make_test_app_with_channels().await;
     enable_terminal_resize_reflow(&mut app);
     app.config.terminal_resize_reflow.max_rows = TerminalResizeReflowMaxRows::Disabled;

codex-rs/tui/src/history_cell.rs

@@ -2620,6 +2620,25 @@ pub(crate) fn new_proposed_plan_stream(
     ProposedPlanStreamCell {
         lines,
         is_stream_continuation,
+        markdown_source: None,
+    }
+}
+
+pub(crate) fn new_proposed_plan_stream_with_markdown_source(
+    lines: Vec<Line<'static>>,
+    is_stream_continuation: bool,
+    source: String,
+    cwd: &Path,
+    include_bottom_padding: bool,
+) -> ProposedPlanStreamCell {
+    ProposedPlanStreamCell {
+        lines,
+        is_stream_continuation,
+        markdown_source: Some(ProposedPlanStreamMarkdownSource {
+            source,
+            cwd: cwd.to_path_buf(),
+            include_bottom_padding,
+        }),
     }
 }
 
@@ -2643,6 +2662,26 @@ pub(crate) struct ProposedPlanCell {
 pub(crate) struct ProposedPlanStreamCell {
     lines: Vec<Line<'static>>,
     is_stream_continuation: bool,
+    markdown_source: Option<ProposedPlanStreamMarkdownSource>,
+}
+
+#[derive(Debug)]
+pub(crate) struct ProposedPlanStreamMarkdownSource {
+    source: String,
+    cwd: PathBuf,
+    include_bottom_padding: bool,
+}
+
+impl ProposedPlanStreamCell {
+    pub(crate) fn markdown_source(&self) -> Option<(&str, &Path, bool)> {
+        self.markdown_source.as_ref().map(|source| {
+            (
+                source.source.as_str(),
+                source.cwd.as_path(),
+                source.include_bottom_padding,
+            )
+        })
+    }
 }
 
 impl HistoryCell for ProposedPlanCell {

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

@@ -450,15 +450,20 @@ impl PlanStreamController {
         }
 
         let source = std::mem::take(&mut self.core.raw_source);
-        let out = self.emit(remaining, /*include_bottom_padding*/ true);
+        let out = self.emit(
+            remaining,
+            /*include_bottom_padding*/ true,
+            Some(source.clone()),
+        );
         self.core.reset();
         (out, Some(source))
     }
 
     pub(crate) fn on_commit_tick(&mut self) -> (Option<Box<dyn HistoryCell>>, bool) {
         let step = self.core.tick();
+        let source = self.core.emitted_stable_source().map(str::to_string);
         (
-            self.emit(step, /*include_bottom_padding*/ false),
+            self.emit(step, /*include_bottom_padding*/ false, source),
             self.core.is_idle(),
         )
     }
@@ -468,8 +473,9 @@ impl PlanStreamController {
         max_lines: usize,
     ) -> (Option<Box<dyn HistoryCell>>, bool) {
         let step = self.core.tick_batch(max_lines);
+        let source = self.core.emitted_stable_source().map(str::to_string);
         (
-            self.emit(step, /*include_bottom_padding*/ false),
+            self.emit(step, /*include_bottom_padding*/ false, source),
             self.core.is_idle(),
         )
     }
@@ -494,6 +500,7 @@ impl PlanStreamController {
         &mut self,
         lines: Vec<Line<'static>>,
         include_bottom_padding: bool,
+        markdown_source: Option<String>,
     ) -> Option<Box<dyn HistoryCell>> {
         if lines.is_empty() && !include_bottom_padding {
             return None;
@@ -524,10 +531,22 @@ impl PlanStreamController {
             .collect::<Vec<_>>();
         out_lines.extend(plan_lines);
 
-        Some(Box::new(history_cell::new_proposed_plan_stream(
-            out_lines,
-            is_stream_continuation,
-        )))
+        if let Some(source) = markdown_source {
+            Some(Box::new(
+                history_cell::new_proposed_plan_stream_with_markdown_source(
+                    out_lines,
+                    is_stream_continuation,
+                    source,
+                    self.core.cwd.as_path(),
+                    include_bottom_padding,
+                ),
+            ))
+        } else {
+            Some(Box::new(history_cell::new_proposed_plan_stream(
+                out_lines,
+                is_stream_continuation,
+            )))
+        }
     }
 }
 
@@ -758,6 +777,26 @@ mod tests {
         );
     }
 
+    #[test]
+    fn plan_stream_cells_carry_stable_source_snapshots() {
+        let mut ctrl = plan_stream_controller(/*width*/ Some(80));
+
+        assert!(ctrl.push(&format!("{}\nAfter table.\n", table_source())));
+
+        let (cell, _idle) = ctrl.on_commit_tick_batch(usize::MAX);
+        let cell = cell.expect("expected queued stable plan table to emit");
+        let plan_cell = cell
+            .as_any()
+            .downcast_ref::<history_cell::ProposedPlanStreamCell>()
+            .expect("expected proposed plan stream cell");
+        let (source, _cwd, include_bottom_padding) = plan_cell
+            .markdown_source()
+            .expect("stable plan stream cell should carry markdown source");
+
+        assert_eq!(source, table_source());
+        assert!(!include_bottom_padding);
+    }
+
     #[test]
     fn table_resize_lifecycle_streaming_resize_updates_active_tail_width() {
         let mut ctrl = stream_controller(Some(36));