feat(tui/history): persist assistant stream records

.github/workflows/upstream-merge.yml

@@ -694,6 +694,9 @@ jobs:
           - Treat warnings as failures during the build-fast step.
           PROMPT
           } | env -i PATH="$SAFE_PATH" HOME="$SAFE_HOME" \
+            OPENAI_API_KEY="x" \
+            OPENAI_BASE_URL="http://127.0.0.1:5055/v1" \
+            OPENAI_API_BASE="http://127.0.0.1:5055/v1" \
             GH_TOKEN="$GH_TOKEN" \
             RUSTC_WRAPPER="sccache" SCCACHE_GHA_ENABLED="true" SCCACHE_IDLE_TIMEOUT="1800" SCCACHE_CACHE_SIZE="5G" \
             npm_config_cache="$SAFE_HOME/.npm" \

codex-rs/tui/HISTORY_CELLS_PLAN.md

@@ -125,14 +125,14 @@ Status legend: ✅ complete (semantic deterministic state ready), ⏳ still need
 
 ## Step 2 – Exec / Streaming / Diff Family *(Pending)*
 2.1 **Exec state extraction & module split** – Create `ExecRecord { command, parsed, action, stdout_chunks, stderr_chunks, exit_code, wait_notes }`, move exec rendering/state into `history_cell/exec.rs`, and drop layout caches from the shared module.
-2.2 **Streaming assistant module** – Capture raw markdown deltas + metadata in `AssistantStreamState`, relocate the streaming renderer to `history_cell/stream.rs`, and support merge/upsert by stream id. **Status:** renderer now lives in `history_cell/stream.rs`; state struct work (raw markdown + metadata) still pending.
-2.3 **Finalized assistant markdown module** – Use `AssistantMessageState` storing markdown + citations + token metadata. **Status:** rendering has been split into `history_cell/assistant.rs`; semantic state refactor remains outstanding.
+2.2 **Streaming assistant module** – Capture raw markdown deltas + metadata in `AssistantStreamState`, relocate the streaming renderer to `history_cell/stream.rs`, and support merge/upsert by stream id. **Status:** renderer now lives in `history_cell/stream.rs`; `HistoryState` records per-stream deltas/preview snapshots; **follow-up:** flow citations/token usage from streaming events, persist provenance metadata, and remove redundant layout caches.
+2.3 **Finalized assistant markdown module** – Use `AssistantMessageState` storing markdown + citations + token metadata. **Status:** rendering has been split into `history_cell/assistant.rs`; finalized answers now land in `HistoryState`; **follow-up:** carry citation/token metadata through `InsertFinalAnswer`, expose state to resume snapshots, and drop raw `AssistantMarkdownCell` caching once render adapters exist.
 2.4 **Diff module breakout** – Persist diff hunks (`DiffHunk`, `DiffLine`) and patch metadata, move diff rendering into `history_cell/diff.rs`, and have the renderer apply styling based on line kind.
 2.5 **Merged exec views** – Rebuild aggregated exec cells from `Vec<ExecRecord>` snapshots rather than cached text blocks once the exec module exists.
 
 ## Step 3 – HistoryState Manager *(Pending)*
 3.1 **HistoryRecord enum** – ✅ complete: `history/state.rs` defines `HistoryRecord`, per-cell state structs, and `HistoryState` scaffolding with ID management helpers.
-3.2 **ChatWidget incremental adoption** – Introduce `HistoryState`/`HistoryRenderState` alongside the legacy `history_cells` vector, migrate low-risk cell types (plain, loading, wait status) to the new state, then remove the legacy vector once the path is stable.
+3.2 **ChatWidget incremental adoption** – Introduce `HistoryState`/`HistoryRenderState` alongside the legacy `history_cells` vector, migrate low-risk cell types (plain, loading, wait status) to the new state, then remove the legacy vector once the path is stable. **Status:** assistant streaming/final answer flows now write into `HistoryState`; next migrate remaining cell constructors and have render adaptors consume the recorded state.
 3.3 **Apply-event pipeline** – Implement `HistoryState::apply_event(&mut self, event: &EventMsg)` covering all core/TUI event types (exec lifecycle, tool updates, background notices, resume snapshots, undo) and route migrated cells through it.
 3.4 **Undo/resume hooks** – Expose `snapshot`, `restore`, and `truncate_after(id)` to support /undo and resume flows.
 3.5 **ChatWidget full integration** – Delete `history_cells: Vec<Box<dyn HistoryCell>>`, wire all helper methods (`history_push`, `history_replace`, etc.) into `HistoryState`, and treat Step 4 as the follow-up for centralized rendering.

codex-rs/tui/src/chatwidget.rs

@@ -176,6 +176,7 @@ use crate::history_cell::HistoryCell;
 use crate::history_cell::HistoryCellType;
 use crate::history_cell::PatchEventType;
 use crate::history_cell::PlainHistoryCell;
+use crate::history::state::{AssistantStreamDelta, HistoryState};
 use crate::slash_command::SlashCommand;
 use crate::live_wrap::RowBuilder;
 use crate::streaming::StreamKind;
@@ -366,6 +367,7 @@ pub(crate) struct ChatWidget<'a> {
     active_exec_cell: Option<ExecCell>,
     history_cells: Vec<Box<dyn HistoryCell>>, // Store all history in memory
     history_render: HistoryRenderState,
+    history_state: HistoryState,
     config: Config,
     latest_upgrade_version: Option<String>,
     initial_user_message: Option<UserMessage>,
@@ -2632,6 +2634,7 @@ impl ChatWidget<'_> {
             agents_terminal: AgentsTerminalState::new(),
             pending_upgrade_notice: None,
             history_render: HistoryRenderState::new(),
+            history_state: HistoryState::new(),
             height_manager: RefCell::new(HeightManager::new(
                 crate::height_manager::HeightManagerConfig::default(),
             )),
@@ -2863,6 +2866,7 @@ impl ChatWidget<'_> {
             agents_terminal: AgentsTerminalState::new(),
             pending_upgrade_notice: None,
             history_render: HistoryRenderState::new(),
+            history_state: HistoryState::new(),
             height_manager: RefCell::new(HeightManager::new(
                 crate::height_manager::HeightManagerConfig::default(),
             )),
@@ -10467,6 +10471,9 @@ impl ChatWidget<'_> {
                 );
                 // Any incoming Answer means reasoning is no longer bottom-most
                 self.clear_reasoning_in_progress();
+                if let Some(ref want) = id {
+                    self.ensure_answer_stream_state(want);
+                }
                 // Keep a single StreamingContentCell and append to it
                 if let Some(last) = self.history_cells.last_mut() {
                     if let Some(stream_cell) = last
@@ -10642,6 +10649,48 @@ impl ChatWidget<'_> {
         self.request_redraw();
     }
 
+    fn ensure_answer_stream_state(&mut self, stream_id: &str) {
+        let preview = self
+            .stream
+            .preview_source_for_kind(StreamKind::Answer)
+            .unwrap_or_default();
+        self.history_state
+            .upsert_assistant_stream_state(stream_id, preview, None, None);
+    }
+
+    fn track_answer_stream_delta(&mut self, stream_id: &str, delta: &str, seq: Option<u64>) {
+        let preview = self
+            .stream
+            .preview_source_for_kind(StreamKind::Answer)
+            .unwrap_or_default();
+        let delta = if delta.is_empty() {
+            None
+        } else {
+            Some(AssistantStreamDelta {
+                delta: delta.to_string(),
+                sequence: seq,
+                received_at: SystemTime::now(),
+            })
+        };
+        self
+            .history_state
+            .upsert_assistant_stream_state(stream_id, preview, delta, None);
+    }
+
+    fn finalize_answer_stream_state(&mut self, stream_id: Option<&str>, source: &str) {
+        self.history_state.finalize_assistant_stream_state(
+            stream_id,
+            source.to_string(),
+            None,
+            None,
+        );
+    }
+
+    #[cfg(test)]
+    pub(crate) fn history_state(&self) -> &HistoryState {
+        &self.history_state
+    }
+
     /// Replace the in-progress streaming assistant cell with a final markdown cell that
     /// stores raw markdown for future re-rendering.
     pub(crate) fn insert_final_answer_with_id(
@@ -10657,6 +10706,7 @@ impl ChatWidget<'_> {
             lines.len()
         );
         tracing::info!("[order] final Answer id={:?}", id);
+        let final_source = source.clone();
         if self.is_review_flow_active() {
             if let Some(ref want) = id {
                 if let Some(idx) = self.history_cells.iter().rposition(|c| {
@@ -10678,7 +10728,8 @@ impl ChatWidget<'_> {
             }) {
                 self.history_remove_at(idx);
             }
-            self.last_assistant_message = Some(source);
+            self.last_assistant_message = Some(final_source.clone());
+            self.finalize_answer_stream_state(id.as_deref(), &final_source);
             return;
         }
         // Debug: list last few history cell kinds so we can see what's present
@@ -10809,6 +10860,7 @@ impl ChatWidget<'_> {
                     .closed_answer_ids
                     .insert(StreamId(want.clone()));
             }
+            self.finalize_answer_stream_state(id.as_deref(), &final_source);
             self.autoscroll_if_near_bottom();
             return;
         }
@@ -10842,6 +10894,7 @@ impl ChatWidget<'_> {
                         .closed_answer_ids
                         .insert(StreamId(want.clone()));
                 }
+                self.finalize_answer_stream_state(id.as_deref(), &final_source);
                 self.autoscroll_if_near_bottom();
                 return;
             }
@@ -10888,6 +10941,7 @@ impl ChatWidget<'_> {
                     &self.config,
                 );
                 self.history_replace_at(idx, Box::new(cell));
+                self.finalize_answer_stream_state(id.as_deref(), &final_source);
                 self.autoscroll_if_near_bottom();
                 return;
             }
@@ -10926,6 +10980,7 @@ impl ChatWidget<'_> {
                 .closed_answer_ids
                 .insert(StreamId(want.clone()));
         }
+        self.finalize_answer_stream_state(id.as_deref(), &final_source);
     }
 
     // Assign or fetch a stable sequence for a stream kind+id within its originating turn

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

@@ -27,10 +27,14 @@ pub(super) fn begin(chat: &mut ChatWidget<'_>, kind: StreamKind, id: Option<Stri
 pub(super) fn delta_text(chat: &mut ChatWidget<'_>, kind: StreamKind, id: String, delta: String, seq: Option<u64>) {
     chat.stream_state.current_kind = Some(kind);
     let sink = AppEventHistorySink(chat.app_event_tx.clone());
+    let stream_id = id.clone();
     chat.stream.begin_with_id(kind, Some(id), &sink);
     chat.stream.set_last_sequence_number(kind, seq);
     
     chat.stream.push_and_maybe_commit(&delta, &sink);
+    if matches!(kind, StreamKind::Answer) {
+        chat.track_answer_stream_delta(&stream_id, &delta, seq);
+    }
 }
 
 // New facade: finalize a specific kind and optionally follow bottom

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

@@ -3,6 +3,7 @@
 use super::*;
 use crate::app_event::{AppEvent, BackgroundPlacement};
 use crate::app_event_sender::AppEventSender;
+use crate::history::state::{HistoryRecord, HistoryState};
 use crate::slash_command::SlashCommand;
 use codex_core::config::Config;
 use codex_core::config::ConfigOverrides;
@@ -100,6 +101,63 @@ fn final_answer_without_newline_is_flushed_immediately() {
     );
 }
 
+#[test]
+fn assistant_history_state_tracks_stream_and_final() {
+    let (mut chat, rx, _op_rx) = make_chatwidget_manual();
+
+    chat.handle_codex_event(Event {
+        id: "turn-1".into(),
+        msg: EventMsg::AgentMessageDelta(AgentMessageDeltaEvent {
+            delta: "Hello world\n".into(),
+        }),
+    });
+    flush_stream_events(&mut chat, &rx);
+
+    let stream_records: Vec<_> = chat
+        .history_state()
+        .records
+        .iter()
+        .filter_map(|rec| match rec {
+            HistoryRecord::AssistantStream(state) => Some(state.clone()),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(stream_records.len(), 1, "expected single assistant stream state");
+    let stream_state = &stream_records[0];
+    assert_eq!(stream_state.stream_id, "turn-1");
+    assert!(stream_state.in_progress);
+    assert_eq!(stream_state.deltas.len(), 1);
+    assert_eq!(stream_state.deltas[0].delta, "Hello world\n");
+
+    chat.handle_codex_event(Event {
+        id: "turn-1".into(),
+        msg: EventMsg::AgentMessage(AgentMessageEvent {
+            message: "Hello world\n".into(),
+        }),
+    });
+    flush_stream_events(&mut chat, &rx);
+
+    let mut has_stream = false;
+    let mut has_final = false;
+    for record in &chat.history_state().records {
+        match record {
+            HistoryRecord::AssistantStream(state) if state.stream_id == "turn-1" => {
+                has_stream = true;
+            }
+            HistoryRecord::AssistantMessage(state)
+                if state.stream_id.as_deref() == Some("turn-1")
+                    && state.markdown.contains("Hello world") =>
+            {
+                has_final = true;
+            }
+            _ => {}
+        }
+    }
+
+    assert!(has_final, "expected finalized assistant message state");
+    assert!(!has_stream, "stream state should be removed after finalization");
+}
+
 fn cell_texts(chat: &ChatWidget<'_>) -> Vec<String> {
     chat
         .history_cells
@@ -293,6 +351,7 @@ fn make_chatwidget_manual() -> (
         app_event_tx,
         codex_op_tx: op_tx,
         bottom_pane: bottom,
+        history_state: HistoryState::new(),
         active_exec_cell: None,
         config: cfg.clone(),
         latest_upgrade_version: None,
@@ -356,6 +415,21 @@ fn lines_to_single_string(lines: &[ratatui::text::Line<'static>]) -> String {
     s
 }
 
+fn flush_stream_events(chat: &mut ChatWidget<'_>, rx: &std::sync::mpsc::Receiver<AppEvent>) {
+    while let Ok(event) = rx.try_recv() {
+        match event {
+            AppEvent::InsertHistory(lines) => chat.insert_history_lines(lines),
+            AppEvent::InsertHistoryWithKind { id, kind, lines } => {
+                chat.insert_history_lines_with_kind(kind, id, lines);
+            }
+            AppEvent::InsertFinalAnswer { id, lines, source } => {
+                chat.insert_final_answer_with_id(id, lines, source);
+            }
+            _ => {}
+        }
+    }
+}
+
 #[derive(Clone, Copy)]
 enum ScriptStep {
     Key(KeyCode, KeyModifiers),

codex-rs/tui/src/history/state.rs

@@ -316,21 +316,34 @@ pub struct ExecWaitNote {
     pub timestamp: SystemTime,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AssistantStreamDelta {
+    pub delta: String,
+    pub sequence: Option<u64>,
+    pub received_at: SystemTime,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct AssistantStreamState {
     pub id: HistoryId,
     pub stream_id: String,
-    pub markdown: String,
+    pub preview_markdown: String,
+    pub deltas: Vec<AssistantStreamDelta>,
     pub citations: Vec<String>,
+    pub metadata: Option<MessageMetadata>,
     pub in_progress: bool,
+    pub last_updated_at: SystemTime,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct AssistantMessageState {
     pub id: HistoryId,
+    pub stream_id: Option<String>,
     pub markdown: String,
     pub citations: Vec<String>,
     pub metadata: Option<MessageMetadata>,
+    pub token_usage: Option<TokenUsage>,
+    pub created_at: SystemTime,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -461,6 +474,98 @@ impl HistoryState {
         id
     }
 
+    pub fn upsert_assistant_stream_state(
+        &mut self,
+        stream_id: &str,
+        preview_markdown: String,
+        delta: Option<AssistantStreamDelta>,
+        metadata: Option<&MessageMetadata>,
+    ) -> HistoryId {
+        let now = SystemTime::now();
+        if let Some(idx) = self
+            .records
+            .iter()
+            .position(|record| matches!(record,
+                HistoryRecord::AssistantStream(state) if state.stream_id == stream_id))
+        {
+            if let Some(HistoryRecord::AssistantStream(state)) = self.records.get_mut(idx) {
+                if let Some(delta) = delta {
+                    state.deltas.push(delta);
+                }
+                state.preview_markdown = preview_markdown;
+                if let Some(meta) = metadata {
+                    state.citations = meta.citations.clone();
+                    state.metadata = Some(meta.clone());
+                }
+                state.in_progress = true;
+                state.last_updated_at = now;
+                return state.id;
+            }
+        }
+
+        let mut deltas = Vec::new();
+        if let Some(delta) = delta {
+            deltas.push(delta);
+        }
+        let citations = metadata
+            .map(|meta| meta.citations.clone())
+            .unwrap_or_default();
+        let metadata = metadata.cloned();
+        let state = AssistantStreamState {
+            id: HistoryId(0),
+            stream_id: stream_id.to_string(),
+            preview_markdown,
+            deltas,
+            citations,
+            metadata,
+            in_progress: true,
+            last_updated_at: now,
+        };
+        self.push(HistoryRecord::AssistantStream(state))
+    }
+
+    pub fn finalize_assistant_stream_state(
+        &mut self,
+        stream_id: Option<&str>,
+        markdown: String,
+        metadata: Option<&MessageMetadata>,
+        token_usage: Option<&TokenUsage>,
+    ) -> HistoryId {
+        let mut carried_citations: Vec<String> = Vec::new();
+        let mut carried_metadata: Option<MessageMetadata> = None;
+        if let Some(stream_id) = stream_id {
+            if let Some(idx) = self.records.iter().position(|record| match record {
+                HistoryRecord::AssistantStream(state) => state.stream_id == stream_id,
+                _ => false,
+            }) {
+                if let Some(HistoryRecord::AssistantStream(state)) = self.remove(idx) {
+                    if !state.citations.is_empty() {
+                        carried_citations = state.citations;
+                    }
+                    if carried_metadata.is_none() {
+                        carried_metadata = state.metadata;
+                    }
+                }
+            }
+        }
+
+        let citations = metadata
+            .map(|meta| meta.citations.clone())
+            .unwrap_or(carried_citations);
+        let metadata = metadata.cloned().or(carried_metadata);
+
+        let state = AssistantMessageState {
+            id: HistoryId(0),
+            stream_id: stream_id.map(|s| s.to_string()),
+            markdown,
+            citations,
+            metadata,
+            token_usage: token_usage.cloned(),
+            created_at: SystemTime::now(),
+        };
+        self.push(HistoryRecord::AssistantMessage(state))
+    }
+
     pub fn insert(&mut self, index: usize, record: HistoryRecord) -> HistoryId {
         let id = self.next_history_id();
         self.records.insert(index, record.with_id(id));

codex-rs/tui/src/history_cell/diff.rs

@@ -189,7 +189,9 @@ pub(crate) fn diff_record_from_string(title: String, diff: &str) -> DiffRecord {
             continue;
         }
 
-        let (kind, content) = if let Some(rest) = raw_line.strip_prefix('+') {
+        let (kind, content) = if raw_line.starts_with("+++") || raw_line.starts_with("---") {
+            (DiffLineKind::Context, raw_line.to_string())
+        } else if let Some(rest) = raw_line.strip_prefix('+') {
             (DiffLineKind::Addition, rest.to_string())
         } else if let Some(rest) = raw_line.strip_prefix('-') {
             (DiffLineKind::Removal, rest.to_string())

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

@@ -106,6 +106,13 @@ pub(crate) fn set_last_sequence_number(&mut self, kind: StreamKind, seq: Option<
     self.state_mut(kind).last_sequence_number = seq;
 }
 
+    pub(crate) fn preview_source_for_kind(&self, kind: StreamKind) -> Option<String> {
+        if self.current_stream != Some(kind) {
+            return None;
+        }
+        Some(self.state(kind).collector.full_render_source_preview())
+    }
+
     fn emit_header_if_needed(&mut self, kind: StreamKind, out_lines: &mut Lines) -> bool {
         let emitted = self.header.maybe_emit(kind, out_lines);
         if emitted {