Use the same context store for all inline assists in a project (#29953)

Max Brunsfeld and Michael Sloan created 7 months ago

Release Notes:

- Made context attachments in inline assist prompts persist across
inline assist invocations.

Co-authored-by: Michael Sloan <[e-mail]>

Change summary

crates/agent/src/assistant_panel.rs  | 20 ++++++++++++++++++--
crates/agent/src/inline_assistant.rs | 24 ++++++++++++------------
2 files changed, 30 insertions(+), 14 deletions(-)

Detailed changes

crates/agent/src/assistant_panel.rs

@@ -54,8 +54,8 @@ use crate::thread::{Thread, ThreadError, ThreadId, TokenUsageRatio};
 use crate::thread_history::{PastContext, PastThread, ThreadHistory};
 use crate::thread_store::{TextThreadStore, ThreadStore};
 use crate::{
-    AddContextServer, AgentDiffPane, DeleteRecentlyOpenThread, ExpandMessageEditor, Follow,
-    InlineAssistant, NewTextThread, NewThread, OpenActiveThreadAsMarkdown, OpenAgentDiff,
+    AddContextServer, AgentDiffPane, ContextStore, DeleteRecentlyOpenThread, ExpandMessageEditor,
+    Follow, InlineAssistant, NewTextThread, NewThread, OpenActiveThreadAsMarkdown, OpenAgentDiff,
     OpenHistory, ThreadEvent, ToggleContextPicker, ToggleNavigationMenu, ToggleOptionsMenu,
 };
 
@@ -315,6 +315,7 @@ pub struct AssistantPanel {
     _default_model_subscription: Subscription,
     context_store: Entity<TextThreadStore>,
     prompt_store: Option<Entity<PromptStore>>,
+    inline_assist_context_store: Entity<crate::context_store::ContextStore>,
     configuration: Option<Entity<AssistantConfiguration>>,
     configuration_subscription: Option<Subscription>,
     local_timezone: UtcOffset,
@@ -438,6 +439,12 @@ impl AssistantPanel {
                 Some(thread_store.downgrade()),
             )
         });
+        let inline_assist_context_store = cx.new(|_cx| {
+            crate::context_store::ContextStore::new(
+                project.downgrade(),
+                Some(thread_store.downgrade()),
+            )
+        });
 
         let message_editor = cx.new(|cx| {
             MessageEditor::new(
@@ -640,6 +647,7 @@ impl AssistantPanel {
                 chrono::Local::now().offset().local_minus_utc(),
             )
             .unwrap(),
+            inline_assist_context_store,
             previous_view: None,
             history_store: history_store.clone(),
             history: cx.new(|cx| ThreadHistory::new(weak_self, history_store, window, cx)),
@@ -674,6 +682,12 @@ impl AssistantPanel {
         &self.prompt_store
     }
 
+    pub(crate) fn inline_assist_context_store(
+        &self,
+    ) -> &Entity<crate::context_store::ContextStore> {
+        &self.inline_assist_context_store
+    }
+
     pub(crate) fn thread_store(&self) -> &Entity<ThreadStore> {
         &self.thread_store
     }
@@ -2436,9 +2450,11 @@ impl rules_library::InlineAssistDelegate for PromptLibraryInlineAssist {
             let prompt_store = None;
             let thread_store = None;
             let text_thread_store = None;
+            let context_store = cx.new(|_| ContextStore::new(project.clone(), None));
             assistant.assist(
                 &prompt_editor,
                 self.workspace.clone(),
+                context_store,
                 project,
                 prompt_store,
                 thread_store,

crates/agent/src/inline_assistant.rs

@@ -251,15 +251,15 @@ impl InlineAssistant {
                 .map_or(false, |model| model.provider.is_authenticated(cx))
         };
 
-        let assistant_panel = workspace
-            .panel::<AssistantPanel>(cx)
-            .map(|assistant_panel| assistant_panel.read(cx));
-        let prompt_store = assistant_panel
-            .and_then(|assistant_panel| assistant_panel.prompt_store().as_ref().cloned());
-        let thread_store =
-            assistant_panel.map(|assistant_panel| assistant_panel.thread_store().downgrade());
-        let text_thread_store =
-            assistant_panel.map(|assistant_panel| assistant_panel.text_thread_store().downgrade());
+        let Some(assistant_panel) = workspace.panel::<AssistantPanel>(cx) else {
+            return;
+        };
+        let assistant_panel = assistant_panel.read(cx);
+
+        let prompt_store = assistant_panel.prompt_store().as_ref().cloned();
+        let thread_store = Some(assistant_panel.thread_store().downgrade());
+        let text_thread_store = Some(assistant_panel.text_thread_store().downgrade());
+        let context_store = assistant_panel.inline_assist_context_store().clone();
 
         let handle_assist =
             |window: &mut Window, cx: &mut Context<Workspace>| match inline_assist_target {
@@ -268,6 +268,7 @@ impl InlineAssistant {
                         assistant.assist(
                             &active_editor,
                             cx.entity().downgrade(),
+                            context_store,
                             workspace.project().downgrade(),
                             prompt_store,
                             thread_store,
@@ -338,6 +339,7 @@ impl InlineAssistant {
         &mut self,
         editor: &Entity<Editor>,
         workspace: WeakEntity<Workspace>,
+        context_store: Entity<ContextStore>,
         project: WeakEntity<Project>,
         prompt_store: Option<Entity<PromptStore>>,
         thread_store: Option<WeakEntity<ThreadStore>>,
@@ -447,8 +449,6 @@ 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(|_cx| ContextStore::new(project.clone(), thread_store.clone()));
             let codegen = cx.new(|cx| {
                 BufferCodegen::new(
                     editor.read(cx).buffer().clone(),
@@ -472,7 +472,7 @@ impl InlineAssistant {
                     prompt_buffer.clone(),
                     codegen.clone(),
                     self.fs.clone(),
-                    context_store,
+                    context_store.clone(),
                     workspace.clone(),
                     thread_store.clone(),
                     text_thread_store.clone(),