diff --git a/crates/agent/src/history_store.rs b/crates/agent/src/history_store.rs index 2c03c81d41..85fdac2bab 100644 --- a/crates/agent/src/history_store.rs +++ b/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), } #[derive(Clone, Debug)] diff --git a/crates/agent/src/thread_history.rs b/crates/agent/src/thread_history.rs index 73c5b47001..353291f4e3 100644 --- a/crates/agent/src/thread_history.rs +++ b/crates/agent/src/thread_history.rs @@ -117,40 +117,21 @@ impl ThreadHistory { } fn update_all_entries(&mut self, cx: &mut Context) { - self.all_entries = self + let new_entries: Arc> = 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._separated_items_task.take(); - let all_entries = self.all_entries.clone(); - let mut items = std::mem::take(&mut self.separated_items); - let mut indexes = std::mem::take(&mut self.separated_item_indexes); - 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 mut items = Vec::with_capacity(new_entries.len() + 1); + let mut indexes = Vec::with_capacity(new_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();