assistant2: Add thread persistence (#23582)

This PR adds persistence for threads in Assistant2.

Threads are now persisted to an LMDB database.

Release Notes:

- N/A

crates/assistant2/Cargo.toml

@@ -39,6 +39,7 @@ fs.workspace = true
 futures.workspace = true
 fuzzy.workspace = true
 gpui.workspace = true
+heed.workspace = true
 html_to_markdown.workspace = true
 http_client.workspace = true
 itertools.workspace = true
@@ -52,6 +53,7 @@ markdown.workspace = true
 menu.workspace = true
 multi_buffer.workspace = true
 parking_lot.workspace = true
+paths.workspace = true
 picker.workspace = true
 project.workspace = true
 prompt_library.workspace = true
@@ -71,7 +73,6 @@ theme.workspace = true
 time.workspace = true
 time_format.workspace = true
 ui.workspace = true
-unindent.workspace = true
 util.workspace = true
 uuid.workspace = true
 workspace.workspace = true

crates/assistant2/src/active_thread.rs

@@ -16,12 +16,14 @@ use ui::prelude::*;
 use workspace::Workspace;
 
 use crate::thread::{MessageId, Thread, ThreadError, ThreadEvent};
+use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
 
 pub struct ActiveThread {
     workspace: WeakView<Workspace>,
     language_registry: Arc<LanguageRegistry>,
     tools: Arc<ToolWorkingSet>,
+    thread_store: Model<ThreadStore>,
     thread: Model<Thread>,
     messages: Vec<MessageId>,
     list_state: ListState,
@@ -33,6 +35,7 @@ pub struct ActiveThread {
 impl ActiveThread {
     pub fn new(
         thread: Model<Thread>,
+        thread_store: Model<ThreadStore>,
         workspace: WeakView<Workspace>,
         language_registry: Arc<LanguageRegistry>,
         tools: Arc<ToolWorkingSet>,
@@ -47,6 +50,7 @@ impl ActiveThread {
             workspace,
             language_registry,
             tools,
+            thread_store,
             thread: thread.clone(),
             messages: Vec::new(),
             rendered_messages_by_id: HashMap::default(),
@@ -192,8 +196,13 @@ impl ActiveThread {
             ThreadEvent::ShowError(error) => {
                 self.last_error = Some(error.clone());
             }
-            ThreadEvent::StreamedCompletion => {}
-            ThreadEvent::SummaryChanged => {}
+            ThreadEvent::StreamedCompletion | ThreadEvent::SummaryChanged => {
+                self.thread_store
+                    .update(cx, |thread_store, cx| {
+                        thread_store.save_thread(&self.thread, cx)
+                    })
+                    .detach_and_log_err(cx);
+            }
             ThreadEvent::StreamedAssistantText(message_id, text) => {
                 if let Some(markdown) = self.rendered_messages_by_id.get_mut(&message_id) {
                     markdown.update(cx, |markdown, cx| {
@@ -211,6 +220,12 @@ impl ActiveThread {
                     self.push_message(message_id, message_text, cx);
                 }
 
+                self.thread_store
+                    .update(cx, |thread_store, cx| {
+                        thread_store.save_thread(&self.thread, cx)
+                    })
+                    .detach_and_log_err(cx);
+
                 cx.notify();
             }
             ThreadEvent::UsePendingTools => {

crates/assistant2/src/assistant_panel.rs

@@ -179,6 +179,7 @@ impl AssistantPanel {
             thread: cx.new_view(|cx| {
                 ActiveThread::new(
                     thread.clone(),
+                    thread_store.clone(),
                     workspace,
                     language_registry,
                     tools.clone(),
@@ -239,6 +240,7 @@ impl AssistantPanel {
         self.thread = cx.new_view(|cx| {
             ActiveThread::new(
                 thread.clone(),
+                self.thread_store.clone(),
                 self.workspace.clone(),
                 self.language_registry.clone(),
                 self.tools.clone(),
@@ -361,34 +363,41 @@ impl AssistantPanel {
         })
     }
 
-    pub(crate) fn open_thread(&mut self, thread_id: &ThreadId, cx: &mut ViewContext<Self>) {
-        let Some(thread) = self
+    pub(crate) fn open_thread(
+        &mut self,
+        thread_id: &ThreadId,
+        cx: &mut ViewContext<Self>,
+    ) -> Task<Result<()>> {
+        let open_thread_task = self
             .thread_store
-            .update(cx, |this, cx| this.open_thread(thread_id, cx))
-        else {
-            return;
-        };
+            .update(cx, |this, cx| this.open_thread(thread_id, cx));
 
-        self.active_view = ActiveView::Thread;
-        self.thread = cx.new_view(|cx| {
-            ActiveThread::new(
-                thread.clone(),
-                self.workspace.clone(),
-                self.language_registry.clone(),
-                self.tools.clone(),
-                cx,
-            )
-        });
-        self.message_editor = cx.new_view(|cx| {
-            MessageEditor::new(
-                self.fs.clone(),
-                self.workspace.clone(),
-                self.thread_store.downgrade(),
-                thread,
-                cx,
-            )
-        });
-        self.message_editor.focus_handle(cx).focus(cx);
+        cx.spawn(|this, mut cx| async move {
+            let thread = open_thread_task.await?;
+            this.update(&mut cx, |this, cx| {
+                this.active_view = ActiveView::Thread;
+                this.thread = cx.new_view(|cx| {
+                    ActiveThread::new(
+                        thread.clone(),
+                        this.thread_store.clone(),
+                        this.workspace.clone(),
+                        this.language_registry.clone(),
+                        this.tools.clone(),
+                        cx,
+                    )
+                });
+                this.message_editor = cx.new_view(|cx| {
+                    MessageEditor::new(
+                        this.fs.clone(),
+                        this.workspace.clone(),
+                        this.thread_store.downgrade(),
+                        thread,
+                        cx,
+                    )
+                });
+                this.message_editor.focus_handle(cx).focus(cx);
+            })
+        })
     }
 
     pub(crate) fn open_configuration(&mut self, cx: &mut ViewContext<Self>) {
@@ -437,7 +446,8 @@ impl AssistantPanel {
 
     pub(crate) fn delete_thread(&mut self, thread_id: &ThreadId, cx: &mut ViewContext<Self>) {
         self.thread_store
-            .update(cx, |this, cx| this.delete_thread(thread_id, cx));
+            .update(cx, |this, cx| this.delete_thread(thread_id, cx))
+            .detach_and_log_err(cx);
     }
 }
 
@@ -655,7 +665,7 @@ impl AssistantPanel {
     fn render_thread_empty_state(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let recent_threads = self
             .thread_store
-            .update(cx, |this, cx| this.recent_threads(3, cx));
+            .update(cx, |this, _cx| this.recent_threads(3));
 
         v_flex()
             .gap_2()

crates/assistant2/src/context_picker.rs

@@ -6,10 +6,12 @@ mod thread_context_picker;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use anyhow::{anyhow, Result};
 use editor::Editor;
 use file_context_picker::render_file_context_entry;
 use gpui::{
-    AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, View, WeakModel, WeakView,
+    AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, Task, View, WeakModel,
+    WeakView,
 };
 use project::ProjectPath;
 use thread_context_picker::{render_thread_context_entry, ThreadContextEntry};
@@ -237,7 +239,8 @@ impl ContextPicker {
                     },
                     move |cx| {
                         context_picker.update(cx, |this, cx| {
-                            this.add_recent_thread(thread.clone(), cx);
+                            this.add_recent_thread(thread.clone(), cx)
+                                .detach_and_log_err(cx);
                         })
                     },
                 )
@@ -260,25 +263,32 @@ impl ContextPicker {
         cx.notify();
     }
 
-    fn add_recent_thread(&self, thread: ThreadContextEntry, cx: &mut ViewContext<Self>) {
+    fn add_recent_thread(
+        &self,
+        thread: ThreadContextEntry,
+        cx: &mut ViewContext<Self>,
+    ) -> Task<Result<()>> {
         let Some(context_store) = self.context_store.upgrade() else {
-            return;
+            return Task::ready(Err(anyhow!("context store not available")));
         };
 
-        let Some(thread) = self
+        let Some(thread_store) = self
             .thread_store
-            .clone()
-            .and_then(|this| this.upgrade())
-            .and_then(|this| this.update(cx, |this, cx| this.open_thread(&thread.id, cx)))
+            .as_ref()
+            .and_then(|thread_store| thread_store.upgrade())
         else {
-            return;
+            return Task::ready(Err(anyhow!("thread store not available")));
         };
 
-        context_store.update(cx, |context_store, cx| {
-            context_store.add_thread(thread, cx);
-        });
+        let open_thread_task = thread_store.update(cx, |this, cx| this.open_thread(&thread.id, cx));
+        cx.spawn(|this, mut cx| async move {
+            let thread = open_thread_task.await?;
+            context_store.update(&mut cx, |context_store, cx| {
+                context_store.add_thread(thread, cx);
+            })?;
 
-        cx.notify();
+            this.update(&mut cx, |_this, cx| cx.notify())
+        })
     }
 
     fn recent_entries(&self, cx: &mut WindowContext) -> Vec<RecentEntry> {
@@ -332,19 +342,17 @@ impl ContextPicker {
             return recent;
         };
 
-        thread_store.update(cx, |thread_store, cx| {
+        thread_store.update(cx, |thread_store, _cx| {
             recent.extend(
                 thread_store
-                    .threads(cx)
+                    .threads()
                     .into_iter()
-                    .filter(|thread| !current_threads.contains(thread.read(cx).id()))
+                    .filter(|thread| !current_threads.contains(&thread.id))
                     .take(2)
                     .map(|thread| {
-                        let thread = thread.read(cx);
-
                         RecentEntry::Thread(ThreadContextEntry {
-                            id: thread.id().clone(),
-                            summary: thread.summary_or_default(),
+                            id: thread.id,
+                            summary: thread.summary,
                         })
                     }),
             )

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -99,13 +99,12 @@ impl PickerDelegate for ThreadContextPickerDelegate {
     }
 
     fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
-        let Ok(threads) = self.thread_store.update(cx, |this, cx| {
-            this.threads(cx)
+        let Ok(threads) = self.thread_store.update(cx, |this, _cx| {
+            this.threads()
                 .into_iter()
-                .map(|thread| {
-                    let id = thread.read(cx).id().clone();
-                    let summary = thread.read(cx).summary_or_default();
-                    ThreadContextEntry { id, summary }
+                .map(|thread| ThreadContextEntry {
+                    id: thread.id,
+                    summary: thread.summary,
                 })
                 .collect::<Vec<_>>()
         }) else {
@@ -159,19 +158,23 @@ impl PickerDelegate for ThreadContextPickerDelegate {
             return;
         };
 
-        let Some(thread) = thread_store.update(cx, |this, cx| this.open_thread(&entry.id, cx))
-        else {
-            return;
-        };
+        let open_thread_task = thread_store.update(cx, |this, cx| this.open_thread(&entry.id, cx));
 
-        self.context_store
-            .update(cx, |context_store, cx| context_store.add_thread(thread, cx))
-            .ok();
-
-        match self.confirm_behavior {
-            ConfirmBehavior::KeepOpen => {}
-            ConfirmBehavior::Close => self.dismissed(cx),
-        }
+        cx.spawn(|this, mut cx| async move {
+            let thread = open_thread_task.await?;
+            this.update(&mut cx, |this, cx| {
+                this.delegate
+                    .context_store
+                    .update(cx, |context_store, cx| context_store.add_thread(thread, cx))
+                    .ok();
+
+                match this.delegate.confirm_behavior {
+                    ConfirmBehavior::KeepOpen => {}
+                    ConfirmBehavior::Close => this.delegate.dismissed(cx),
+                }
+            })
+        })
+        .detach_and_log_err(cx);
     }
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {