assistant2: Push logic for adding file context down into the ContextStore (#22846)

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

Release Notes:

- N/A

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -193,81 +193,41 @@ impl PickerDelegate for FileContextPickerDelegate {
             return;
         };
 
-        let workspace = self.workspace.clone();
+        let project_path = ProjectPath {
-        let Some(project) = workspace
+            worktree_id: WorktreeId::from_usize(mat.worktree_id),
-            .upgrade()
+            path: mat.path.clone(),
-            .map(|workspace| workspace.read(cx).project().clone())
+        };
+
+        let Some(task) = self
+            .context_store
+            .update(cx, |context_store, cx| {
+                context_store.add_file(project_path, cx)
+            })
+            .ok()
         else {
             return;
         };
-        let path = mat.path.clone();
 
-        let already_included = self
+        let workspace = self.workspace.clone();
-            .context_store
-            .update(cx, |context_store, _cx| {
-                match context_store.included_file(&path) {
-                    Some(IncludedFile::Direct(context_id)) => {
-                        context_store.remove_context(&context_id);
-                        true
-                    }
-                    Some(IncludedFile::InDirectory(_)) => true,
-                    None => false,
-                }
-            })
-            .unwrap_or(true);
-        if already_included {
-            return;
-        }
-
-        let worktree_id = WorktreeId::from_usize(mat.worktree_id);
         let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
-            let Some(open_buffer_task) = project
+            match task.await {
-                .update(&mut cx, |project, cx| {
+                Ok(()) => {
-                    let project_path = ProjectPath {
+                    this.update(&mut cx, |this, cx| match confirm_behavior {
-                        worktree_id,
-                        path: path.clone(),
-                    };
-
-                    let task = project.open_buffer(project_path, cx);
-
-                    Some(task)
-                })
-                .ok()
-                .flatten()
-            else {
-                return anyhow::Ok(());
-            };
-
-            let result = open_buffer_task.await;
-
-            this.update(&mut cx, |this, cx| match result {
-                Ok(buffer) => {
-                    this.delegate
-                        .context_store
-                        .update(cx, |context_store, cx| {
-                            context_store.insert_file(buffer.read(cx));
-                        })?;
-
-                    match confirm_behavior {
                         ConfirmBehavior::KeepOpen => {}
                         ConfirmBehavior::Close => this.delegate.dismissed(cx),
-                    }
+                    })?;
-
-                    anyhow::Ok(())
                 }
                 Err(err) => {
                     let Some(workspace) = workspace.upgrade() else {
                         return anyhow::Ok(());
                     };
 
-                    workspace.update(cx, |workspace, cx| {
+                    workspace.update(&mut cx, |workspace, cx| {
                         workspace.show_error(&err, cx);
-                    });
+                    })?;
-
-                    anyhow::Ok(())
                 }
-            })??;
+            }
 
             anyhow::Ok(())
         })

crates/assistant2/src/context_store.rs

@@ -1,9 +1,12 @@
 use std::fmt::Write as _;
 use std::path::{Path, PathBuf};
 
+use anyhow::{anyhow, Result};
 use collections::{HashMap, HashSet};
-use gpui::SharedString;
+use gpui::{ModelContext, SharedString, Task, WeakView};
 use language::Buffer;
+use project::ProjectPath;
+use workspace::Workspace;
 
 use crate::thread::Thread;
 use crate::{
@@ -12,6 +15,7 @@ use crate::{
 };
 
 pub struct ContextStore {
+    workspace: WeakView<Workspace>,
     context: Vec<Context>,
     next_context_id: ContextId,
     files: HashMap<PathBuf, ContextId>,
@@ -21,8 +25,9 @@ pub struct ContextStore {
 }
 
 impl ContextStore {
-    pub fn new() -> Self {
+    pub fn new(workspace: WeakView<Workspace>) -> Self {
         Self {
+            workspace,
             context: Vec::new(),
             next_context_id: ContextId(0),
             files: HashMap::default(),
@@ -44,6 +49,44 @@ impl ContextStore {
         self.fetched_urls.clear();
     }
 
+    pub fn add_file(
+        &mut self,
+        project_path: ProjectPath,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<()>> {
+        let workspace = self.workspace.clone();
+        let Some(project) = workspace
+            .upgrade()
+            .map(|workspace| workspace.read(cx).project().clone())
+        else {
+            return Task::ready(Err(anyhow!("failed to read project")));
+        };
+
+        let already_included = match self.included_file(&project_path.path) {
+            Some(IncludedFile::Direct(context_id)) => {
+                self.remove_context(&context_id);
+                true
+            }
+            Some(IncludedFile::InDirectory(_)) => true,
+            None => false,
+        };
+        if already_included {
+            return Task::ready(Ok(()));
+        }
+
+        cx.spawn(|this, mut cx| async move {
+            let open_buffer_task =
+                project.update(&mut cx, |project, cx| project.open_buffer(project_path, cx))?;
+
+            let buffer = open_buffer_task.await?;
+            this.update(&mut cx, |this, cx| {
+                this.insert_file(buffer.read(cx));
+            })?;
+
+            anyhow::Ok(())
+        })
+    }
+
     pub fn insert_file(&mut self, buffer: &Buffer) {
         let Some(file) = buffer.file() else {
             return;

crates/assistant2/src/inline_assistant.rs

@@ -335,7 +335,7 @@ impl InlineAssistant {
         let mut assist_to_focus = None;
         for range in codegen_ranges {
             let assist_id = self.next_assist_id.post_inc();
-            let context_store = cx.new_model(|_cx| ContextStore::new());
+            let context_store = cx.new_model(|_cx| ContextStore::new(workspace.clone()));
             let codegen = cx.new_model(|cx| {
                 BufferCodegen::new(
                     editor.read(cx).buffer().clone(),
@@ -445,7 +445,7 @@ impl InlineAssistant {
             range.end = range.end.bias_right(&snapshot);
         }
 
-        let context_store = cx.new_model(|_cx| ContextStore::new());
+        let context_store = cx.new_model(|_cx| ContextStore::new(workspace.clone()));
 
         let codegen = cx.new_model(|cx| {
             BufferCodegen::new(

crates/assistant2/src/message_editor.rs

@@ -47,7 +47,7 @@ impl MessageEditor {
         thread: Model<Thread>,
         cx: &mut ViewContext<Self>,
     ) -> Self {
-        let context_store = cx.new_model(|_cx| ContextStore::new());
+        let context_store = cx.new_model(|_cx| ContextStore::new(workspace.clone()));
         let context_picker_menu_handle = PopoverMenuHandle::default();
         let inline_context_picker_menu_handle = PopoverMenuHandle::default();
         let model_selector_menu_handle = PopoverMenuHandle::default();

crates/assistant2/src/terminal_inline_assistant.rs

@@ -78,7 +78,7 @@ impl TerminalInlineAssistant {
         let prompt_buffer = cx.new_model(|cx| {
             MultiBuffer::singleton(cx.new_model(|cx| Buffer::local(String::new(), cx)), cx)
         });
-        let context_store = cx.new_model(|_cx| ContextStore::new());
+        let context_store = cx.new_model(|_cx| ContextStore::new(workspace.clone()));
         let codegen = cx.new_model(|_| TerminalCodegen::new(terminal, self.telemetry.clone()));
 
         let prompt_editor = cx.new_view(|cx| {