Reapply "Use Project instead of Workspace in ContextStore (#28402)" (#28441)

Michael Sloan created

Motivation for this change is to use `ContextStore` in headless
assistant, which requires it to not depend on UI entities like
`Workspace`.

This reapplies a change that was revert was in #28428, and fixes the panic.

Release Notes:

- N/A

Change summary

crates/agent/src/assistant_panel.rs                    | 24 +++++++++--
crates/agent/src/context_picker/completion_provider.rs |  2 
crates/agent/src/context_store.rs                      | 22 +++--------
crates/agent/src/inline_assistant.rs                   | 17 +++++---
crates/agent/src/terminal_inline_assistant.rs          |  5 +-
crates/agent/src/thread.rs                             |  2 
6 files changed, 41 insertions(+), 31 deletions(-)

Detailed changes

crates/agent/src/assistant_panel.rs 🔗

@@ -227,14 +227,14 @@ impl AssistantPanel {
     ) -> Self {
         let thread = thread_store.update(cx, |this, cx| this.create_thread(cx));
         let fs = workspace.app_state().fs.clone();
-        let project = workspace.project().clone();
+        let project = workspace.project();
         let language_registry = project.read(cx).languages().clone();
         let workspace = workspace.weak_handle();
         let weak_self = cx.entity().downgrade();
 
         let message_editor_context_store = cx.new(|_cx| {
             crate::context_store::ContextStore::new(
-                workspace.clone(),
+                project.downgrade(),
                 Some(thread_store.downgrade()),
             )
         });
@@ -344,7 +344,7 @@ impl AssistantPanel {
 
         let message_editor_context_store = cx.new(|_cx| {
             crate::context_store::ContextStore::new(
-                self.workspace.clone(),
+                self.project.downgrade(),
                 Some(self.thread_store.downgrade()),
             )
         });
@@ -521,7 +521,7 @@ impl AssistantPanel {
                 this.set_active_view(thread_view, window, cx);
                 let message_editor_context_store = cx.new(|_cx| {
                     crate::context_store::ContextStore::new(
-                        this.workspace.clone(),
+                        this.project.downgrade(),
                         Some(this.thread_store.downgrade()),
                     )
                 });
@@ -1624,7 +1624,21 @@ impl prompt_library::InlineAssistDelegate for PromptLibraryInlineAssist {
         cx: &mut Context<PromptLibrary>,
     ) {
         InlineAssistant::update_global(cx, |assistant, cx| {
-            assistant.assist(&prompt_editor, self.workspace.clone(), None, window, cx)
+            let Some(project) = self
+                .workspace
+                .upgrade()
+                .map(|workspace| workspace.read(cx).project().downgrade())
+            else {
+                return;
+            };
+            assistant.assist(
+                &prompt_editor,
+                self.workspace.clone(),
+                project,
+                None,
+                window,
+                cx,
+            )
         })
     }
 

crates/agent/src/context_picker/completion_provider.rs 🔗

@@ -867,7 +867,7 @@ mod tests {
                 .expect("Opened test file wasn't an editor")
         });
 
-        let context_store = cx.new(|_| ContextStore::new(workspace.downgrade(), None));
+        let context_store = cx.new(|_| ContextStore::new(project.downgrade(), None));
 
         let editor_entity = editor.downgrade();
         editor.update_in(&mut cx, |editor, window, cx| {

crates/agent/src/context_store.rs 🔗

@@ -8,11 +8,10 @@ use futures::future::join_all;
 use futures::{self, Future, FutureExt, future};
 use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};
 use language::{Buffer, File};
-use project::{ProjectItem, ProjectPath, Worktree};
+use project::{Project, ProjectItem, ProjectPath, Worktree};
 use rope::Rope;
 use text::{Anchor, BufferId, OffsetRangeExt};
 use util::{ResultExt as _, maybe};
-use workspace::Workspace;
 
 use crate::ThreadStore;
 use crate::context::{
@@ -23,7 +22,7 @@ use crate::context_strip::SuggestedContext;
 use crate::thread::{Thread, ThreadId};
 
 pub struct ContextStore {
-    workspace: WeakEntity<Workspace>,
+    project: WeakEntity<Project>,
     context: Vec<AssistantContext>,
     thread_store: Option<WeakEntity<ThreadStore>>,
     // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
@@ -40,11 +39,11 @@ pub struct ContextStore {
 
 impl ContextStore {
     pub fn new(
-        workspace: WeakEntity<Workspace>,
+        project: WeakEntity<Project>,
         thread_store: Option<WeakEntity<ThreadStore>>,
     ) -> Self {
         Self {
-            workspace,
+            project,
             thread_store,
             context: Vec::new(),
             next_context_id: ContextId(0),
@@ -81,12 +80,7 @@ impl ContextStore {
         remove_if_exists: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
-        let workspace = self.workspace.clone();
-
-        let Some(project) = workspace
-            .upgrade()
-            .map(|workspace| workspace.read(cx).project().clone())
-        else {
+        let Some(project) = self.project.upgrade() else {
             return Task::ready(Err(anyhow!("failed to read project")));
         };
 
@@ -161,11 +155,7 @@ impl ContextStore {
         remove_if_exists: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
-        let workspace = self.workspace.clone();
-        let Some(project) = workspace
-            .upgrade()
-            .map(|workspace| workspace.read(cx).project().clone())
-        else {
+        let Some(project) = self.project.upgrade() else {
             return Task::ready(Err(anyhow!("failed to read project")));
         };
 

crates/agent/src/inline_assistant.rs 🔗

@@ -28,6 +28,7 @@ use language_model::{LanguageModelRegistry, report_assistant_event};
 use multi_buffer::MultiBufferRow;
 use parking_lot::Mutex;
 use project::LspAction;
+use project::Project;
 use project::{CodeAction, ProjectTransaction};
 use prompt_store::PromptBuilder;
 use settings::{Settings, SettingsStore};
@@ -254,6 +255,7 @@ impl InlineAssistant {
                         assistant.assist(
                             &active_editor,
                             cx.entity().downgrade(),
+                            workspace.project().downgrade(),
                             thread_store,
                             window,
                             cx,
@@ -265,6 +267,7 @@ impl InlineAssistant {
                         assistant.assist(
                             &active_terminal,
                             cx.entity().downgrade(),
+                            workspace.project().downgrade(),
                             thread_store,
                             window,
                             cx,
@@ -318,6 +321,7 @@ impl InlineAssistant {
         &mut self,
         editor: &Entity<Editor>,
         workspace: WeakEntity<Workspace>,
+        project: WeakEntity<Project>,
         thread_store: Option<WeakEntity<ThreadStore>>,
         window: &mut Window,
         cx: &mut App,
@@ -425,7 +429,7 @@ impl InlineAssistant {
         for range in codegen_ranges {
             let assist_id = self.next_assist_id.post_inc();
             let context_store =
-                cx.new(|_cx| ContextStore::new(workspace.clone(), thread_store.clone()));
+                cx.new(|_cx| ContextStore::new(project.clone(), thread_store.clone()));
             let codegen = cx.new(|cx| {
                 BufferCodegen::new(
                     editor.read(cx).buffer().clone(),
@@ -519,7 +523,7 @@ impl InlineAssistant {
         initial_prompt: String,
         initial_transaction_id: Option<TransactionId>,
         focus: bool,
-        workspace: WeakEntity<Workspace>,
+        workspace: Entity<Workspace>,
         thread_store: Option<WeakEntity<ThreadStore>>,
         window: &mut Window,
         cx: &mut App,
@@ -537,8 +541,8 @@ impl InlineAssistant {
             range.end = range.end.bias_right(&snapshot);
         }
 
-        let context_store =
-            cx.new(|_cx| ContextStore::new(workspace.clone(), thread_store.clone()));
+        let project = workspace.read(cx).project().downgrade();
+        let context_store = cx.new(|_cx| ContextStore::new(project, thread_store.clone()));
 
         let codegen = cx.new(|cx| {
             BufferCodegen::new(
@@ -562,7 +566,7 @@ impl InlineAssistant {
                 codegen.clone(),
                 self.fs.clone(),
                 context_store,
-                workspace.clone(),
+                workspace.downgrade(),
                 thread_store,
                 window,
                 cx,
@@ -589,7 +593,7 @@ impl InlineAssistant {
                 end_block_id,
                 range,
                 codegen.clone(),
-                workspace.clone(),
+                workspace.downgrade(),
                 window,
                 cx,
             ),
@@ -1779,6 +1783,7 @@ impl CodeActionProvider for AssistantCodeActionProvider {
         let workspace = self.workspace.clone();
         let thread_store = self.thread_store.clone();
         window.spawn(cx, async move |cx| {
+            let workspace = workspace.upgrade().context("workspace was released")?;
             let editor = editor.upgrade().context("editor was released")?;
             let range = editor
                 .update(cx, |editor, cx| {

crates/agent/src/terminal_inline_assistant.rs 🔗

@@ -16,6 +16,7 @@ use language_model::{
     ConfiguredModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
     Role, report_assistant_event,
 };
+use project::Project;
 use prompt_store::PromptBuilder;
 use std::sync::Arc;
 use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
@@ -67,6 +68,7 @@ impl TerminalInlineAssistant {
         &mut self,
         terminal_view: &Entity<TerminalView>,
         workspace: WeakEntity<Workspace>,
+        project: WeakEntity<Project>,
         thread_store: Option<WeakEntity<ThreadStore>>,
         window: &mut Window,
         cx: &mut App,
@@ -75,8 +77,7 @@ impl TerminalInlineAssistant {
         let assist_id = self.next_assist_id.post_inc();
         let prompt_buffer =
             cx.new(|cx| MultiBuffer::singleton(cx.new(|cx| Buffer::local(String::new(), cx)), cx));
-        let context_store =
-            cx.new(|_cx| ContextStore::new(workspace.clone(), thread_store.clone()));
+        let context_store = cx.new(|_cx| ContextStore::new(project, thread_store.clone()));
         let codegen = cx.new(|_| TerminalCodegen::new(terminal, self.telemetry.clone()));
 
         let prompt_editor = cx.new(|cx| {

crates/agent/src/thread.rs 🔗

@@ -2266,7 +2266,7 @@ fn main() {{
         });
 
         let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
-        let context_store = cx.new(|_cx| ContextStore::new(workspace.downgrade(), None));
+        let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None));
 
         (
             workspace,