assistant2: Push logic for adding thread context down into the `ContextStore` (#22855)

Marshall Bowers created 1 year ago

This PR takes the logic for adding thread context out of the
`ThreadContextPicker` and pushes it down into the `ContextStore`.

Release Notes:

- N/A

Change summary

crates/assistant2/src/context_picker/thread_context_picker.rs | 12 +---
crates/assistant2/src/context_store.rs                        | 10 +++
2 files changed, 12 insertions(+), 10 deletions(-)

Detailed changes

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -167,13 +167,7 @@ impl PickerDelegate for ThreadContextPickerDelegate {
         };
 
         self.context_store
-            .update(cx, |context_store, cx| {
-                if let Some(context_id) = context_store.included_thread(&entry.id) {
-                    context_store.remove_context(&context_id);
-                } else {
-                    context_store.insert_thread(thread.read(cx));
-                }
-            })
+            .update(cx, |context_store, cx| context_store.add_thread(thread, cx))
             .ok();
 
         match self.confirm_behavior {
@@ -199,8 +193,8 @@ impl PickerDelegate for ThreadContextPickerDelegate {
     ) -> Option<Self::ListItem> {
         let thread = &self.matches[ix];
 
-        let added = self.context_store.upgrade().map_or(false, |ctx_store| {
-            ctx_store.read(cx).included_thread(&thread.id).is_some()
+        let added = self.context_store.upgrade().map_or(false, |context_store| {
+            context_store.read(cx).included_thread(&thread.id).is_some()
         });
 
         Some(

crates/assistant2/src/context_store.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, bail, Result};
 use collections::{HashMap, HashSet};
-use gpui::{ModelContext, SharedString, Task, WeakView};
+use gpui::{Model, ModelContext, SharedString, Task, WeakView};
 use language::Buffer;
 use project::{ProjectPath, Worktree};
 use workspace::Workspace;
@@ -229,6 +229,14 @@ impl ContextStore {
         });
     }
 
+    pub fn add_thread(&mut self, thread: Model<Thread>, cx: &mut ModelContext<Self>) {
+        if let Some(context_id) = self.included_thread(&thread.read(cx).id()) {
+            self.remove_context(&context_id);
+        } else {
+            self.insert_thread(thread.read(cx));
+        }
+    }
+
     pub fn insert_thread(&mut self, thread: &Thread) {
         let context_id = self.next_context_id.post_inc();
         self.threads.insert(thread.id().clone(), context_id);