assistant2: Add support for referencing other threads as context (#22092)

This PR adds the ability to reference other threads as context:

<img width="1159" alt="Screenshot 2024-12-16 at 11 29 54 AM"
src="https://github.com/user-attachments/assets/bb8a24ff-56d3-4406-ab8c-6657e65d8c70"
/>

<img width="1159" alt="Screenshot 2024-12-16 at 11 29 35 AM"
src="https://github.com/user-attachments/assets/7a02ebda-a2f5-40e9-9dd4-1bb029cb1c43"
/>


Release Notes:

- N/A

crates/assistant2/src/assistant_panel.rs

@@ -94,7 +94,9 @@ impl AssistantPanel {
                     cx,
                 )
             }),
-            message_editor: cx.new_view(|cx| MessageEditor::new(workspace, thread.clone(), cx)),
+            message_editor: cx.new_view(|cx| {
+                MessageEditor::new(workspace, thread_store.downgrade(), thread.clone(), cx)
+            }),
             tools,
             local_timezone: UtcOffset::from_whole_seconds(
                 chrono::Local::now().offset().local_minus_utc(),
@@ -123,8 +125,14 @@ impl AssistantPanel {
                 cx,
             )
         });
-        self.message_editor =
+        self.message_editor = cx.new_view(|cx| {
-            cx.new_view(|cx| MessageEditor::new(self.workspace.clone(), thread, cx));
+            MessageEditor::new(
+                self.workspace.clone(),
+                self.thread_store.downgrade(),
+                thread,
+                cx,
+            )
+        });
         self.message_editor.focus_handle(cx).focus(cx);
     }
 
@@ -146,8 +154,14 @@ impl AssistantPanel {
                 cx,
             )
         });
-        self.message_editor =
+        self.message_editor = cx.new_view(|cx| {
-            cx.new_view(|cx| MessageEditor::new(self.workspace.clone(), thread, cx));
+            MessageEditor::new(
+                self.workspace.clone(),
+                self.thread_store.downgrade(),
+                thread,
+                cx,
+            )
+        });
         self.message_editor.focus_handle(cx).focus(cx);
     }
 

crates/assistant2/src/context.rs

@@ -24,4 +24,5 @@ pub struct Context {
 pub enum ContextKind {
     File,
     FetchedUrl,
+    Thread,
 }

crates/assistant2/src/context_picker.rs

@@ -1,11 +1,12 @@
 mod fetch_context_picker;
 mod file_context_picker;
+mod thread_context_picker;
 
 use std::sync::Arc;
 
 use gpui::{
     AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, SharedString, Task, View,
-    WeakView,
+    WeakModel, WeakView,
 };
 use picker::{Picker, PickerDelegate};
 use ui::{prelude::*, ListItem, ListItemSpacing, Tooltip};
@@ -14,13 +15,16 @@ use workspace::Workspace;
 
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
+use crate::context_picker::thread_context_picker::ThreadContextPicker;
 use crate::message_editor::MessageEditor;
+use crate::thread_store::ThreadStore;
 
 #[derive(Debug, Clone)]
 enum ContextPickerMode {
     Default,
     File(View<FileContextPicker>),
     Fetch(View<FetchContextPicker>),
+    Thread(View<ThreadContextPicker>),
 }
 
 pub(super) struct ContextPicker {
@@ -31,13 +35,15 @@ pub(super) struct ContextPicker {
 impl ContextPicker {
     pub fn new(
         workspace: WeakView<Workspace>,
+        thread_store: WeakModel<ThreadStore>,
         message_editor: WeakView<MessageEditor>,
         cx: &mut ViewContext<Self>,
     ) -> Self {
         let delegate = ContextPickerDelegate {
             context_picker: cx.view().downgrade(),
-            workspace: workspace.clone(),
+            workspace,
-            message_editor: message_editor.clone(),
+            thread_store,
+            message_editor,
             entries: vec![
                 ContextPickerEntry {
                     name: "directory".into(),
@@ -54,6 +60,11 @@ impl ContextPicker {
                     description: "Fetch content from URL".into(),
                     icon: IconName::Globe,
                 },
+                ContextPickerEntry {
+                    name: "thread".into(),
+                    description: "Insert any thread".into(),
+                    icon: IconName::MessageBubbles,
+                },
             ],
             selected_ix: 0,
         };
@@ -81,6 +92,7 @@ impl FocusableView for ContextPicker {
             ContextPickerMode::Default => self.picker.focus_handle(cx),
             ContextPickerMode::File(file_picker) => file_picker.focus_handle(cx),
             ContextPickerMode::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),
+            ContextPickerMode::Thread(thread_picker) => thread_picker.focus_handle(cx),
         }
     }
 }
@@ -94,6 +106,7 @@ impl Render for ContextPicker {
                 ContextPickerMode::Default => parent.child(self.picker.clone()),
                 ContextPickerMode::File(file_picker) => parent.child(file_picker.clone()),
                 ContextPickerMode::Fetch(fetch_picker) => parent.child(fetch_picker.clone()),
+                ContextPickerMode::Thread(thread_picker) => parent.child(thread_picker.clone()),
             })
     }
 }
@@ -108,6 +121,7 @@ struct ContextPickerEntry {
 pub(crate) struct ContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
+    thread_store: WeakModel<ThreadStore>,
     message_editor: WeakView<MessageEditor>,
     entries: Vec<ContextPickerEntry>,
     selected_ix: usize,
@@ -162,6 +176,16 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
+                        "thread" => {
+                            this.mode = ContextPickerMode::Thread(cx.new_view(|cx| {
+                                ThreadContextPicker::new(
+                                    self.thread_store.clone(),
+                                    self.context_picker.clone(),
+                                    self.message_editor.clone(),
+                                    cx,
+                                )
+                            }));
+                        }
                         _ => {}
                     }
 
@@ -175,7 +199,9 @@ impl PickerDelegate for ContextPickerDelegate {
         self.context_picker
             .update(cx, |this, cx| match this.mode {
                 ContextPickerMode::Default => cx.emit(DismissEvent),
-                ContextPickerMode::File(_) | ContextPickerMode::Fetch(_) => {}
+                ContextPickerMode::File(_)
+                | ContextPickerMode::Fetch(_)
+                | ContextPickerMode::Thread(_) => {}
             })
             .log_err();
     }

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -0,0 +1,209 @@
+use std::sync::Arc;
+
+use fuzzy::StringMatchCandidate;
+use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView};
+use picker::{Picker, PickerDelegate};
+use ui::{prelude::*, ListItem};
+
+use crate::context::ContextKind;
+use crate::context_picker::ContextPicker;
+use crate::message_editor::MessageEditor;
+use crate::thread::ThreadId;
+use crate::thread_store::ThreadStore;
+
+pub struct ThreadContextPicker {
+    picker: View<Picker<ThreadContextPickerDelegate>>,
+}
+
+impl ThreadContextPicker {
+    pub fn new(
+        thread_store: WeakModel<ThreadStore>,
+        context_picker: WeakView<ContextPicker>,
+        message_editor: WeakView<MessageEditor>,
+        cx: &mut ViewContext<Self>,
+    ) -> Self {
+        let delegate =
+            ThreadContextPickerDelegate::new(thread_store, context_picker, message_editor);
+        let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
+
+        ThreadContextPicker { picker }
+    }
+}
+
+impl FocusableView for ThreadContextPicker {
+    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
+        self.picker.focus_handle(cx)
+    }
+}
+
+impl Render for ThreadContextPicker {
+    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+        self.picker.clone()
+    }
+}
+
+#[derive(Debug, Clone)]
+struct ThreadContextEntry {
+    id: ThreadId,
+    summary: SharedString,
+}
+
+pub struct ThreadContextPickerDelegate {
+    thread_store: WeakModel<ThreadStore>,
+    context_picker: WeakView<ContextPicker>,
+    message_editor: WeakView<MessageEditor>,
+    matches: Vec<ThreadContextEntry>,
+    selected_index: usize,
+}
+
+impl ThreadContextPickerDelegate {
+    pub fn new(
+        thread_store: WeakModel<ThreadStore>,
+        context_picker: WeakView<ContextPicker>,
+        message_editor: WeakView<MessageEditor>,
+    ) -> Self {
+        ThreadContextPickerDelegate {
+            thread_store,
+            context_picker,
+            message_editor,
+            matches: Vec::new(),
+            selected_index: 0,
+        }
+    }
+}
+
+impl PickerDelegate for ThreadContextPickerDelegate {
+    type ListItem = ListItem;
+
+    fn match_count(&self) -> usize {
+        self.matches.len()
+    }
+
+    fn selected_index(&self) -> usize {
+        self.selected_index
+    }
+
+    fn set_selected_index(&mut self, ix: usize, _cx: &mut ViewContext<Picker<Self>>) {
+        self.selected_index = ix;
+    }
+
+    fn placeholder_text(&self, _cx: &mut WindowContext) -> Arc<str> {
+        "Search threads…".into()
+    }
+
+    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)
+                .into_iter()
+                .map(|thread| {
+                    const DEFAULT_SUMMARY: SharedString = SharedString::new_static("New Thread");
+
+                    let id = thread.read(cx).id().clone();
+                    let summary = thread.read(cx).summary().unwrap_or(DEFAULT_SUMMARY);
+                    ThreadContextEntry { id, summary }
+                })
+                .collect::<Vec<_>>()
+        }) else {
+            return Task::ready(());
+        };
+
+        let executor = cx.background_executor().clone();
+        let search_task = cx.background_executor().spawn(async move {
+            if query.is_empty() {
+                threads
+            } else {
+                let candidates = threads
+                    .iter()
+                    .enumerate()
+                    .map(|(id, thread)| StringMatchCandidate::new(id, &thread.summary))
+                    .collect::<Vec<_>>();
+                let matches = fuzzy::match_strings(
+                    &candidates,
+                    &query,
+                    false,
+                    100,
+                    &Default::default(),
+                    executor,
+                )
+                .await;
+
+                matches
+                    .into_iter()
+                    .map(|mat| threads[mat.candidate_id].clone())
+                    .collect()
+            }
+        });
+
+        cx.spawn(|this, mut cx| async move {
+            let matches = search_task.await;
+            this.update(&mut cx, |this, cx| {
+                this.delegate.matches = matches;
+                this.delegate.selected_index = 0;
+                cx.notify();
+            })
+            .ok();
+        })
+    }
+
+    fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
+        let entry = &self.matches[self.selected_index];
+
+        let Some(thread_store) = self.thread_store.upgrade() else {
+            return;
+        };
+
+        let Some(thread) = thread_store.update(cx, |this, cx| this.open_thread(&entry.id, cx))
+        else {
+            return;
+        };
+
+        self.message_editor
+            .update(cx, |message_editor, cx| {
+                let text = thread.update(cx, |thread, _cx| {
+                    let mut text = String::new();
+
+                    for message in thread.messages() {
+                        text.push_str(match message.role {
+                            language_model::Role::User => "User:",
+                            language_model::Role::Assistant => "Assistant:",
+                            language_model::Role::System => "System:",
+                        });
+                        text.push('\n');
+
+                        text.push_str(&message.text);
+                        text.push('\n');
+                    }
+
+                    text
+                });
+
+                message_editor.insert_context(ContextKind::Thread, entry.summary.clone(), text);
+            })
+            .ok();
+    }
+
+    fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
+        self.context_picker
+            .update(cx, |this, cx| {
+                this.reset_mode();
+                cx.emit(DismissEvent);
+            })
+            .ok();
+    }
+
+    fn render_match(
+        &self,
+        ix: usize,
+        selected: bool,
+        _cx: &mut ViewContext<Picker<Self>>,
+    ) -> Option<Self::ListItem> {
+        let thread = &self.matches[ix];
+
+        Some(
+            ListItem::new(ix)
+                .inset(true)
+                .toggle_state(selected)
+                .child(thread.summary.clone()),
+        )
+    }
+}

crates/assistant2/src/message_editor.rs

@@ -1,7 +1,7 @@
 use std::rc::Rc;
 
 use editor::{Editor, EditorElement, EditorStyle};
-use gpui::{AppContext, FocusableView, Model, TextStyle, View, WeakView};
+use gpui::{AppContext, FocusableView, Model, TextStyle, View, WeakModel, WeakView};
 use language_model::{LanguageModelRegistry, LanguageModelRequestTool};
 use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
 use settings::Settings;
@@ -15,6 +15,7 @@ use workspace::Workspace;
 use crate::context::{Context, ContextId, ContextKind};
 use crate::context_picker::ContextPicker;
 use crate::thread::{RequestKind, Thread};
+use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
 use crate::{Chat, ToggleModelSelector};
 
@@ -32,6 +33,7 @@ pub struct MessageEditor {
 impl MessageEditor {
     pub fn new(
         workspace: WeakView<Workspace>,
+        thread_store: WeakModel<ThreadStore>,
         thread: Model<Thread>,
         cx: &mut ViewContext<Self>,
     ) -> Self {
@@ -46,7 +48,9 @@ impl MessageEditor {
             }),
             context: Vec::new(),
             next_context_id: ContextId(0),
-            context_picker: cx.new_view(|cx| ContextPicker::new(workspace.clone(), weak_self, cx)),
+            context_picker: cx.new_view(|cx| {
+                ContextPicker::new(workspace.clone(), thread_store.clone(), weak_self, cx)
+            }),
             context_picker_handle: PopoverMenuHandle::default(),
             language_model_selector: cx.new_view(|cx| {
                 LanguageModelSelector::new(

crates/assistant2/src/thread.rs

@@ -194,6 +194,7 @@ impl Thread {
             if let Some(context) = self.context_for_message(message.id) {
                 let mut file_context = String::new();
                 let mut fetch_context = String::new();
+                let mut thread_context = String::new();
 
                 for context in context.iter() {
                     match context.kind {
@@ -207,6 +208,12 @@ impl Thread {
                             fetch_context.push_str(&context.text);
                             fetch_context.push('\n');
                         }
+                        ContextKind::Thread => {
+                            thread_context.push_str(&context.name);
+                            thread_context.push('\n');
+                            thread_context.push_str(&context.text);
+                            thread_context.push('\n');
+                        }
                     }
                 }
 
@@ -221,6 +228,12 @@ impl Thread {
                     context_text.push_str(&fetch_context);
                 }
 
+                if !thread_context.is_empty() {
+                    context_text
+                        .push_str("The following previous conversation threads are available\n");
+                    context_text.push_str(&thread_context);
+                }
+
                 request_message
                     .content
                     .push(MessageContent::Text(context_text))