From 7e39023ea5e71e7a54bfd66e869f9c66071b3ff3 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Wed, 8 Jan 2025 14:54:54 -0500
Subject: [PATCH] assistant2: Push logic for adding thread context down into
 the `ContextStore` (#22855)

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

Release Notes:

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

diff --git a/crates/assistant2/src/context_picker/thread_context_picker.rs b/crates/assistant2/src/context_picker/thread_context_picker.rs
index 1d86e6612da25733c25ef701adca288e21fda9ca..86578cb8f780bbeb3288d1166a30eb50f7638d12 100644
--- a/crates/assistant2/src/context_picker/thread_context_picker.rs
+++ b/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(
diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs
index 03ae83102c560e6b3d0385db67659e627ad9a433..3ae7bfb9b48c8b1b17b5455b3669f2e241a8d21a 100644
--- a/crates/assistant2/src/context_store.rs
+++ b/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);