From 86f5bb1cc0b982d4bbcaa34ca6b816b63d3ca2e8 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Wed, 8 Jan 2025 12:46:49 -0500 Subject: [PATCH] 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 --- .../src/context_picker/file_context_picker.rs | 74 +++++-------------- crates/assistant2/src/context_store.rs | 47 +++++++++++- crates/assistant2/src/inline_assistant.rs | 4 +- crates/assistant2/src/message_editor.rs | 2 +- .../src/terminal_inline_assistant.rs | 2 +- 5 files changed, 66 insertions(+), 63 deletions(-) diff --git a/crates/assistant2/src/context_picker/file_context_picker.rs b/crates/assistant2/src/context_picker/file_context_picker.rs index 2d885a6534df8a42370cfc0d4644b98391b4fd6f..8b7543356999890895cd28466f62b56c56059190 100644 --- a/crates/assistant2/src/context_picker/file_context_picker.rs +++ b/crates/assistant2/src/context_picker/file_context_picker.rs @@ -193,81 +193,41 @@ impl PickerDelegate for FileContextPickerDelegate { return; }; - let workspace = self.workspace.clone(); - let Some(project) = workspace - .upgrade() - .map(|workspace| workspace.read(cx).project().clone()) - else { - return; + let project_path = ProjectPath { + worktree_id: WorktreeId::from_usize(mat.worktree_id), + path: mat.path.clone(), }; - let path = mat.path.clone(); - let already_included = self + let Some(task) = self .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, - } + .update(cx, |context_store, cx| { + context_store.add_file(project_path, cx) }) - .unwrap_or(true); - if already_included { + .ok() + else { return; - } + }; - let worktree_id = WorktreeId::from_usize(mat.worktree_id); + let workspace = self.workspace.clone(); let confirm_behavior = self.confirm_behavior; cx.spawn(|this, mut cx| async move { - let Some(open_buffer_task) = project - .update(&mut cx, |project, cx| { - let project_path = ProjectPath { - 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 { + match task.await { + Ok(()) => { + this.update(&mut cx, |this, 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(()) }) diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs index f2659f53a86365e8bf9078e1c485f3dd24211631..cdb50ce5b5029a7898bb759252342d1d8529a763 100644 --- a/crates/assistant2/src/context_store.rs +++ b/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, context: Vec, next_context_id: ContextId, files: HashMap, @@ -21,8 +25,9 @@ pub struct ContextStore { } impl ContextStore { - pub fn new() -> Self { + pub fn new(workspace: WeakView) -> 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, + ) -> Task> { + 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; diff --git a/crates/assistant2/src/inline_assistant.rs b/crates/assistant2/src/inline_assistant.rs index ef1896440c414ec2f5265d668489ea7f803bc246..bd1c7bf3fb6f4fad19978ffa2b40bee7db68a6c2 100644 --- a/crates/assistant2/src/inline_assistant.rs +++ b/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( diff --git a/crates/assistant2/src/message_editor.rs b/crates/assistant2/src/message_editor.rs index c11f4056cc10226eea57c6285f59035272b50c3c..5621032b8fc9aaa18978eb0572aa5fcc0ca41704 100644 --- a/crates/assistant2/src/message_editor.rs +++ b/crates/assistant2/src/message_editor.rs @@ -47,7 +47,7 @@ impl MessageEditor { thread: Model, cx: &mut ViewContext, ) -> 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(); diff --git a/crates/assistant2/src/terminal_inline_assistant.rs b/crates/assistant2/src/terminal_inline_assistant.rs index ad8153293a7be03a4c98813d0bdc8e6dc80a72eb..d3773de0a20766a757c3b6e2b2ac03d0a42ede16 100644 --- a/crates/assistant2/src/terminal_inline_assistant.rs +++ b/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| {