Retain completed turn height index

Author: DinahK-2SO

tools/wta/src/app/tab_state.rs

@@ -12,6 +12,14 @@ use super::{TabAutofixState, TurnState};
 pub(crate) const DEFAULT_TAB_ID: &str = "0";
 
 static NEXT_COMPLETED_TURN_LAYOUT_NAMESPACE: AtomicU64 = AtomicU64::new(1);
+const MAX_COMPLETED_TURN_LAYOUT_CHANGES: usize = 2048;
+
+pub(crate) struct CompletedTurnLayoutChanges {
+    pub namespace: u64,
+    pub generation: u64,
+    pub len: usize,
+    pub dirty_indices: Option<Vec<usize>>,
+}
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub enum NoticeKind {
@@ -407,6 +415,8 @@ pub struct TabSession {
     pub(crate) completed_turn_layout_ids: Vec<u64>,
     pub(crate) completed_turn_layout_revisions: Vec<u64>,
     pub(crate) next_completed_turn_layout_id: u64,
+    pub(crate) completed_turn_layout_generation: u64,
+    pub(crate) completed_turn_layout_changes: VecDeque<(u64, usize)>,
     /// Tab/Shift+Tab selects a past turn (most recent first). Enter then
     /// toggles `CompletedTurn.expanded`. None means no selection — Enter
     /// goes to the input/prompt path as before.
@@ -523,22 +533,50 @@ pub struct TabSession {
 }
 
 impl TabSession {
-    pub(crate) fn completed_turn_layout_metadata(&mut self) -> (u64, Vec<(u64, u64)>) {
+    fn record_completed_turn_layout_change(&mut self, index: usize) {
+        if self.completed_turn_layout_generation == u64::MAX {
+            self.completed_turn_layout_namespace =
+                NEXT_COMPLETED_TURN_LAYOUT_NAMESPACE.fetch_add(1, Ordering::Relaxed);
+            self.completed_turn_layout_generation = 1;
+            self.completed_turn_layout_changes.clear();
+        } else {
+            self.completed_turn_layout_generation += 1;
+        }
+        self.completed_turn_layout_changes
+            .push_back((self.completed_turn_layout_generation, index));
+        while self.completed_turn_layout_changes.len() > MAX_COMPLETED_TURN_LAYOUT_CHANGES {
+            self.completed_turn_layout_changes.pop_front();
+        }
+    }
+
+    fn sync_completed_turn_layout_metadata(&mut self) {
         if self.completed_turn_layout_namespace == 0 {
             self.completed_turn_layout_namespace =
                 NEXT_COMPLETED_TURN_LAYOUT_NAMESPACE.fetch_add(1, Ordering::Relaxed);
         }
 
         let len = self.completed_turns.len();
+        if self.completed_turn_layout_ids.len() > len {
+            self.completed_turn_layout_namespace =
+                NEXT_COMPLETED_TURN_LAYOUT_NAMESPACE.fetch_add(1, Ordering::Relaxed);
+            self.completed_turn_layout_generation = 0;
+            self.completed_turn_layout_changes.clear();
+        }
         self.completed_turn_layout_ids.truncate(len);
         self.completed_turn_layout_revisions.truncate(len);
         while self.completed_turn_layout_ids.len() < len {
+            let index = self.completed_turn_layout_ids.len();
             self.next_completed_turn_layout_id =
                 self.next_completed_turn_layout_id.wrapping_add(1).max(1);
             self.completed_turn_layout_ids
                 .push(self.next_completed_turn_layout_id);
             self.completed_turn_layout_revisions.push(0);
+            self.record_completed_turn_layout_change(index);
         }
+    }
+
+    pub(crate) fn completed_turn_layout_metadata(&mut self) -> (u64, Vec<(u64, u64)>) {
+        self.sync_completed_turn_layout_metadata();
 
         (
             self.completed_turn_layout_namespace,
@@ -550,17 +588,98 @@ impl TabSession {
         )
     }
 
+    pub(crate) fn completed_turn_layout_changes_since(
+        &mut self,
+        previous: Option<(u64, u64)>,
+    ) -> CompletedTurnLayoutChanges {
+        self.sync_completed_turn_layout_metadata();
+        let namespace = self.completed_turn_layout_namespace;
+        let generation = self.completed_turn_layout_generation;
+        let len = self.completed_turns.len();
+        let Some((previous_namespace, previous_generation)) = previous else {
+            return CompletedTurnLayoutChanges {
+                namespace,
+                generation,
+                len,
+                dirty_indices: None,
+            };
+        };
+        if previous_namespace != namespace {
+            return CompletedTurnLayoutChanges {
+                namespace,
+                generation,
+                len,
+                dirty_indices: None,
+            };
+        }
+        if previous_generation == generation {
+            return CompletedTurnLayoutChanges {
+                namespace,
+                generation,
+                len,
+                dirty_indices: Some(Vec::new()),
+            };
+        }
+
+        let Some(first_retained_generation) = self
+            .completed_turn_layout_changes
+            .front()
+            .map(|(generation, _)| *generation)
+        else {
+            return CompletedTurnLayoutChanges {
+                namespace,
+                generation,
+                len,
+                dirty_indices: None,
+            };
+        };
+        if previous_generation.saturating_add(1) < first_retained_generation {
+            return CompletedTurnLayoutChanges {
+                namespace,
+                generation,
+                len,
+                dirty_indices: None,
+            };
+        }
+
+        let mut dirty_indices = self
+            .completed_turn_layout_changes
+            .iter()
+            .filter(|(change_generation, _)| *change_generation > previous_generation)
+            .map(|(_, index)| *index)
+            .collect::<Vec<_>>();
+        dirty_indices.sort_unstable();
+        dirty_indices.dedup();
+        CompletedTurnLayoutChanges {
+            namespace,
+            generation,
+            len,
+            dirty_indices: Some(dirty_indices),
+        }
+    }
+
+    pub(crate) fn completed_turn_layout_item(&self, index: usize) -> Option<(u64, u64)> {
+        Some((
+            *self.completed_turn_layout_ids.get(index)?,
+            *self.completed_turn_layout_revisions.get(index)?,
+        ))
+    }
+
     pub(crate) fn mark_completed_turn_layout_dirty(&mut self, index: usize) {
-        let _ = self.completed_turn_layout_metadata();
+        self.sync_completed_turn_layout_metadata();
         if let Some(revision) = self.completed_turn_layout_revisions.get_mut(index) {
             *revision = revision.wrapping_add(1);
+            self.record_completed_turn_layout_change(index);
         }
     }
 
     pub fn clear_completed_turns(&mut self) {
         self.completed_turns.clear();
+        self.completed_turn_layout_namespace = 0;
         self.completed_turn_layout_ids.clear();
         self.completed_turn_layout_revisions.clear();
+        self.completed_turn_layout_generation = 0;
+        self.completed_turn_layout_changes.clear();
     }
 
     pub fn set_last_completed_turn_trailing_marker(&mut self, marker: String) -> bool {

tools/wta/src/app_tests.rs

@@ -7966,7 +7966,7 @@ fn completed_turn_prompt_rows_expose_state_aware_action_links() {
             && link.end_column > triangle.start_column
     }));
 
-    app.current_tab_mut().completed_turns[0].expanded = false;
+    assert!(app.current_tab_mut().set_completed_turn_expanded(0, false));
     render_to_text(&mut app, 80, 16);
     assert_eq!(app.completed_turn_action_links.len(), 1);
     assert_eq!(
@@ -10116,34 +10116,40 @@ fn render_chat_keeps_keyboard_selected_completed_turn_visible() {
 
     app.current_tab_mut().toggle_selected_completed_turn();
     crate::ui::chat::reset_completed_turn_line_build_count();
+    crate::ui::chat::reset_completed_turn_descriptor_lookup_count();
     let collapsed = render_to_text(&mut app, 80, 16);
     assert_eq!(
         crate::ui::chat::completed_turn_line_build_count(),
         1,
         "collapsing one completed turn must invalidate only that item",
     );
+    assert_eq!(crate::ui::chat::completed_turn_descriptor_lookup_count(), 1);
     assert_ne!(collapsed, retained);
 
     app.current_tab_mut().toggle_selected_completed_turn();
     crate::ui::chat::reset_completed_turn_line_build_count();
+    crate::ui::chat::reset_completed_turn_descriptor_lookup_count();
     let expanded = render_to_text(&mut app, 80, 16);
     assert_eq!(
         crate::ui::chat::completed_turn_line_build_count(),
         1,
         "expanding one completed turn must invalidate only that item",
     );
+    assert_eq!(crate::ui::chat::completed_turn_descriptor_lookup_count(), 1);
     assert!(expanded.contains("ACK_SELECT_SCROLL_TURN_00"));
 
     assert!(app
         .current_tab_mut()
         .set_last_completed_turn_trailing_marker("CACHED_MARKER".into()));
     crate::ui::chat::reset_completed_turn_line_build_count();
+    crate::ui::chat::reset_completed_turn_descriptor_lookup_count();
     let _ = render_to_text(&mut app, 80, 16);
     assert_eq!(
         crate::ui::chat::completed_turn_line_build_count(),
         1,
         "updating one trailing marker must invalidate only that item",
     );
+    assert_eq!(crate::ui::chat::completed_turn_descriptor_lookup_count(), 1);
 
     for _ in 0..11 {
         app.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE));
@@ -10178,10 +10184,38 @@ fn clearing_completed_turns_renews_retained_layout_identity() {
     tab.completed_turns.push(turn);
     let (same_namespace, after) = tab.completed_turn_layout_metadata();
 
-    assert_eq!(same_namespace, namespace);
+    assert_ne!(same_namespace, namespace);
     assert_ne!(after[0].0, before[0].0);
 }
 
+#[test]
+fn completed_turn_layout_change_log_wrap_and_overflow_force_safe_rebuilds() {
+    let mut tab = TabSession::default();
+    tab.completed_turns.push(CompletedTurn {
+        prompt: "turn".into(),
+        details: Vec::new(),
+        expanded: false,
+        trailing_marker: None,
+    });
+    let (namespace, _) = tab.completed_turn_layout_metadata();
+
+    tab.completed_turn_layout_generation = u64::MAX;
+    tab.mark_completed_turn_layout_dirty(0);
+    let wrapped = tab.completed_turn_layout_changes_since(Some((namespace, u64::MAX)));
+    assert_ne!(wrapped.namespace, namespace);
+    assert!(wrapped.dirty_indices.is_none());
+
+    let baseline = (wrapped.namespace, wrapped.generation);
+    for _ in 0..=2048 {
+        tab.mark_completed_turn_layout_dirty(0);
+    }
+    let overflowed = tab.completed_turn_layout_changes_since(Some(baseline));
+    assert!(
+        overflowed.dirty_indices.is_none(),
+        "falling behind the bounded change log must force a full rebuild",
+    );
+}
+
 #[test]
 fn render_chat_materializes_only_viewport_plus_overscan_turns() {
     let mut app = test_app();
@@ -10204,6 +10238,30 @@ fn render_chat_materializes_only_viewport_plus_overscan_turns() {
         (1..=48).contains(&materialized),
         "a 16-row viewport plus 32-row overscan should materialize at most 48 turns; materialized {materialized}",
     );
+
+    crate::ui::chat::reset_completed_turn_descriptor_lookup_count();
+    let retained = render_to_text(&mut app, 80, 16);
+    assert_eq!(retained, rendered);
+    assert_eq!(
+        crate::ui::chat::completed_turn_descriptor_lookup_count(),
+        0,
+        "an unchanged bottom frame must reuse the retained height index without scanning history",
+    );
+
+    app.current_tab_mut().completed_turns.push(CompletedTurn {
+        prompt: "RETAINED_VIEWPORT_TURN_200".into(),
+        details: vec![ChatMessage::Agent("ACK_200".into())],
+        expanded: false,
+        trailing_marker: None,
+    });
+    crate::ui::chat::reset_completed_turn_descriptor_lookup_count();
+    let appended = render_to_text(&mut app, 80, 16);
+    assert!(appended.contains("RETAINED_VIEWPORT_TURN_200"));
+    assert_eq!(
+        crate::ui::chat::completed_turn_descriptor_lookup_count(),
+        1,
+        "appending one turn must update only the new height-index item",
+    );
 }
 
 #[test]