From 3eea89d28884f985f30e310fd13db826daa1b573 Mon Sep 17 00:00:00 2001
From: Chris Busillo <info@shinycomputers.com>
Date: Mon, 1 Jun 2026 20:35:06 -0400
Subject: [PATCH] Show final assistant replay after stale snapshot

---
 code-rs/tui/src/chatwidget.rs      | 57 ++++++++++++++++++++++++++++++
 code-rs/tui/tests/resume_replay.rs | 53 ++++++++++++++++++++++++++-
 2 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/code-rs/tui/src/chatwidget.rs b/code-rs/tui/src/chatwidget.rs
index 3f8d1d8f017..07fa26cf523 100644
--- a/code-rs/tui/src/chatwidget.rs
+++ b/code-rs/tui/src/chatwidget.rs
@@ -5142,6 +5142,59 @@ impl ChatWidget<'_> {
         }
     }
 
+    fn render_replay_assistant_item_if_missing(&mut self, item: ResponseItem) {
+        let ResponseItem::Message { id, role, content, .. } = item else {
+            return;
+        };
+        if role != "assistant" {
+            return;
+        }
+
+        let mut text = String::new();
+        for c in content {
+            match c {
+                ContentItem::OutputText { text: t } | ContentItem::InputText { text: t } => {
+                    if !text.is_empty() {
+                        text.push('\n');
+                    }
+                    text.push_str(&t);
+                }
+                _ => {}
+            }
+        }
+
+        let text = text.trim();
+        if text.is_empty() || self.history_has_assistant_text(text) {
+            return;
+        }
+
+        let mut lines: Vec<ratatui::text::Line<'static>> = Vec::new();
+        crate::markdown::append_markdown(text, &mut lines, &self.config);
+        self.insert_final_answer_with_id(id, lines, text.to_string());
+    }
+
+    fn history_has_assistant_text(&self, text: &str) -> bool {
+        let normalized = Self::normalize_text(text);
+        self.history_cells.iter().any(|cell| {
+            if let Some(existing) = cell
+                .as_any()
+                .downcast_ref::<crate::history_cell::AssistantMarkdownCell>()
+            {
+                return Self::normalize_text(existing.markdown()) == normalized;
+            }
+            if let Some(existing) = cell
+                .as_any()
+                .downcast_ref::<crate::history_cell::PlainHistoryCell>()
+            {
+                if existing.state().kind == PlainMessageKind::Assistant {
+                    let existing_text = Self::message_lines_to_plain_preview(&existing.state().lines);
+                    return Self::normalize_text(&existing_text) == normalized;
+                }
+            }
+            false
+        })
+    }
+
     fn is_auto_review_cell(item: &dyn HistoryCell) -> bool {
         item.as_any()
             .downcast_ref::<crate::history_cell::PlainHistoryCell>()
@@ -14398,6 +14451,10 @@ impl ChatWidget<'_> {
                         self.last_seen_request_index =
                             self.last_seen_request_index.max(self.current_request_index);
                     }
+                } else {
+                    for item in &items {
+                        self.render_replay_assistant_item_if_missing(item.clone());
+                    }
                 }
                 if max_req > 0 {
                     self.last_seen_request_index = self.last_seen_request_index.max(max_req);
diff --git a/code-rs/tui/tests/resume_replay.rs b/code-rs/tui/tests/resume_replay.rs
index 57d1dd8f72b..35b10cb525d 100644
--- a/code-rs/tui/tests/resume_replay.rs
+++ b/code-rs/tui/tests/resume_replay.rs
@@ -8,7 +8,9 @@ use code_core::history::{
 };
 use code_core::plan_tool::StepStatus;
 use code_core::protocol::{Event, EventMsg, ReplayHistoryEvent};
-use code_protocol::models::{ContentItem, ResponseItem};
+use code_protocol::models::{
+    ContentItem, FunctionCallOutputBody, FunctionCallOutputPayload, ResponseItem,
+};
 use code_tui::test_helpers::{render_chat_widget_to_vt100, ChatWidgetHarness};
 use serde_json::to_value;
 
@@ -135,6 +137,18 @@ fn final_reasoning_snapshot() -> HistorySnapshot {
     }
 }
 
+fn tool_only_snapshot() -> HistorySnapshot {
+    HistorySnapshot {
+        records: vec![explore_record(1)],
+        next_id: 2,
+        exec_call_lookup: Default::default(),
+        tool_call_lookup: Default::default(),
+        stream_lookup: Default::default(),
+        order: vec![OrderKeySnapshot { req: 1, out: 0, seq: 1 }],
+        order_debug: Vec::new(),
+    }
+}
+
 #[test]
 fn replay_history_duplicates_short_assistant_messages() {
     let mut harness = ChatWidgetHarness::new();
@@ -259,3 +273,40 @@ fn replay_history_keeps_spacing_before_final_reasoning() {
         "expected blank line between exploring and reasoning. screen: {screen}"
     );
 }
+
+#[test]
+fn replay_history_restores_final_assistant_after_snapshot_tail() {
+    let mut harness = ChatWidgetHarness::new();
+
+    let snapshot_json = to_value(&tool_only_snapshot()).expect("snapshot to json");
+    let items = vec![
+        ResponseItem::FunctionCallOutput {
+            call_id: "call-final-tool".to_string(),
+            output: FunctionCallOutputPayload {
+                body: FunctionCallOutputBody::Text("## fix/token-percent-remaining".to_string()),
+                success: Some(true),
+            },
+        },
+        message(
+            "assistant",
+            "Removed the local custom override and verified no tracked repo files changed.",
+        ),
+    ];
+
+    harness.handle_event(Event {
+        id: "resume-replay".to_string(),
+        event_seq: 0,
+        msg: EventMsg::ReplayHistory(ReplayHistoryEvent {
+            items,
+            history_snapshot: Some(snapshot_json),
+        }),
+        order: None,
+    });
+
+    let screen = render_chat_widget_to_vt100(&mut harness, 80, 12);
+
+    assert!(
+        screen.contains("Removed the local custom override"),
+        "screen: {screen}"
+    );
+}