agent: Thread history update improvements (#30415)

- Try to preserve previously selected item on update
- Do not clear list items while updating to avoid a frame with no items
rendered

Release Notes:

- agent: Preserve previously selected item in Thread History on update

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
Co-authored-by: Richard Feldman <oss@rtfeldman.com>

crates/agent/src/history_store.rs

@@ -1,4 +1,4 @@
-use std::{collections::VecDeque, path::Path};
+use std::{collections::VecDeque, path::Path, sync::Arc};
 
 use anyhow::{Context as _, anyhow};
 use assistant_context_editor::{AssistantContext, SavedContextMetadata};
@@ -34,6 +34,20 @@ impl HistoryEntry {
             HistoryEntry::Context(context) => context.mtime.to_utc(),
         }
     }
+
+    pub fn id(&self) -> HistoryEntryId {
+        match self {
+            HistoryEntry::Thread(thread) => HistoryEntryId::Thread(thread.id.clone()),
+            HistoryEntry::Context(context) => HistoryEntryId::Context(context.path.clone()),
+        }
+    }
+}
+
+/// Generic identifier for a history entry.
+#[derive(Clone, PartialEq, Eq)]
+pub enum HistoryEntryId {
+    Thread(ThreadId),
+    Context(Arc<Path>),
 }
 
 #[derive(Clone, Debug)]

crates/agent/src/thread_history.rs

@@ -117,40 +117,21 @@ impl ThreadHistory {
     }
 
     fn update_all_entries(&mut self, cx: &mut Context<Self>) {
-        self.all_entries = self
+        let new_entries: Arc<Vec<HistoryEntry>> = self
             .history_store
             .update(cx, |store, cx| store.entries(cx))
             .into();
 
-        self.set_selected_entry_index(0, cx);
-        self.update_separated_items(cx);
-
-        match &self.search_state {
-            SearchState::Empty => {}
-            SearchState::Searching { query, .. } | SearchState::Searched { query, .. } => {
-                self.search(query.clone(), cx);
-            }
-        }
-        cx.notify();
-    }
-
-    fn update_separated_items(&mut self, cx: &mut Context<Self>) {
         self._separated_items_task.take();
-        let all_entries = self.all_entries.clone();
 
-        let mut items = std::mem::take(&mut self.separated_items);
+        let mut items = Vec::with_capacity(new_entries.len() + 1);
-        let mut indexes = std::mem::take(&mut self.separated_item_indexes);
+        let mut indexes = Vec::with_capacity(new_entries.len() + 1);
-        items.clear();
-        indexes.clear();
-        // We know there's going to be at least one bucket separator
-        items.reserve(all_entries.len() + 1);
-        indexes.reserve(all_entries.len() + 1);
 
         let bg_task = cx.background_spawn(async move {
             let mut bucket = None;
             let today = Local::now().naive_local().date();
 
-            for (index, entry) in all_entries.iter().enumerate() {
+            for (index, entry) in new_entries.iter().enumerate() {
                 let entry_date = entry
                     .updated_at()
                     .with_timezone(&Local)
@@ -169,14 +150,41 @@ impl ThreadHistory {
                     format: entry_bucket.into(),
                 });
             }
-            (items, indexes)
+            (new_entries, items, indexes)
         });
 
         let task = cx.spawn(async move |this, cx| {
-            let (items, indexes) = bg_task.await;
+            let (new_entries, items, indexes) = bg_task.await;
             this.update(cx, |this, cx| {
+                let previously_selected_entry =
+                    this.all_entries.get(this.selected_index).map(|e| e.id());
+
+                this.all_entries = new_entries;
                 this.separated_items = items;
                 this.separated_item_indexes = indexes;
+
+                match &this.search_state {
+                    SearchState::Empty => {
+                        if this.selected_index >= this.all_entries.len() {
+                            this.set_selected_entry_index(
+                                this.all_entries.len().saturating_sub(1),
+                                cx,
+                            );
+                        } else if let Some(prev_id) = previously_selected_entry {
+                            if let Some(new_ix) = this
+                                .all_entries
+                                .iter()
+                                .position(|probe| probe.id() == prev_id)
+                            {
+                                this.set_selected_entry_index(new_ix, cx);
+                            }
+                        }
+                    }
+                    SearchState::Searching { query, .. } | SearchState::Searched { query, .. } => {
+                        this.search(query.clone(), cx);
+                    }
+                }
+
                 cx.notify();
             })
             .log_err();