Improve registration for Assistant code action providers (#23099)

Marshall Bowers created 1 year ago

This PR is a follow-up to
https://github.com/zed-industries/zed/pull/22911 to further improve the
registration of code action providers for the Assistant in order to
prevent duplicates.

The `CodeActionProvider` trait now has an `id` method that is used to
return a unique ID for a code action provider. We use this to prevent
registering duplicates of the same provider.

The registration of the code action providers for Assistant1 and
Assistant2 have also been reworked. Previously we were not call the
registration function—and thus setting up the subscriptions—until we
resolved the feature flags. However, this could lead to the registration
happening too late for existing workspace items.

We now perform the registration right away and then remove the undesired
code action providers once the feature flags have been resolved.

Release Notes:

- N/A

Change summary

crates/assistant/src/inline_assistant.rs  | 45 ++++++++++++------
crates/assistant2/src/inline_assistant.rs | 58 ++++++++++++++++--------
crates/editor/src/editor.rs               | 22 +++++++++
3 files changed, 90 insertions(+), 35 deletions(-)

Detailed changes

crates/assistant/src/inline_assistant.rs

@@ -72,16 +72,16 @@ pub fn init(
 ) {
     cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
     cx.observe_new_views(|_, cx| {
+        let workspace = cx.view().clone();
+        InlineAssistant::update_global(cx, |inline_assistant, cx| {
+            inline_assistant.register_workspace(&workspace, cx)
+        });
+
         cx.observe_flag::<Assistant2FeatureFlag, _>({
             |is_assistant2_enabled, _view, cx| {
-                if is_assistant2_enabled {
-                    // Assistant2 enabled, nothing to do for Assistant1.
-                } else {
-                    let workspace = cx.view().clone();
-                    InlineAssistant::update_global(cx, |inline_assistant, cx| {
-                        inline_assistant.register_workspace(&workspace, cx)
-                    })
-                }
+                InlineAssistant::update_global(cx, |inline_assistant, _cx| {
+                    inline_assistant.is_assistant2_enabled = is_assistant2_enabled;
+                });
             }
         })
         .detach();
@@ -102,6 +102,7 @@ pub struct InlineAssistant {
     prompt_builder: Arc<PromptBuilder>,
     telemetry: Arc<Telemetry>,
     fs: Arc<dyn Fs>,
+    is_assistant2_enabled: bool,
 }
 
 impl Global for InlineAssistant {}
@@ -123,6 +124,7 @@ impl InlineAssistant {
             prompt_builder,
             telemetry,
             fs,
+            is_assistant2_enabled: false,
         }
     }
 
@@ -183,15 +185,22 @@ impl InlineAssistant {
         item: &dyn ItemHandle,
         cx: &mut WindowContext,
     ) {
+        let is_assistant2_enabled = self.is_assistant2_enabled;
+
         if let Some(editor) = item.act_as::<Editor>(cx) {
             editor.update(cx, |editor, cx| {
-                editor.push_code_action_provider(
-                    Rc::new(AssistantCodeActionProvider {
-                        editor: cx.view().downgrade(),
-                        workspace: workspace.downgrade(),
-                    }),
-                    cx,
-                );
+                if is_assistant2_enabled {
+                    editor
+                        .remove_code_action_provider(ASSISTANT_CODE_ACTION_PROVIDER_ID.into(), cx);
+                } else {
+                    editor.add_code_action_provider(
+                        Rc::new(AssistantCodeActionProvider {
+                            editor: cx.view().downgrade(),
+                            workspace: workspace.downgrade(),
+                        }),
+                        cx,
+                    );
+                }
             });
         }
     }
@@ -3437,7 +3446,13 @@ struct AssistantCodeActionProvider {
     workspace: WeakView<Workspace>,
 }
 
+const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant";
+
 impl CodeActionProvider for AssistantCodeActionProvider {
+    fn id(&self) -> Arc<str> {
+        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
+    }
+
     fn code_actions(
         &self,
         buffer: &Model<Buffer>,

crates/assistant2/src/inline_assistant.rs

@@ -51,14 +51,16 @@ pub fn init(
 ) {
     cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
     cx.observe_new_views(|_workspace: &mut Workspace, cx| {
+        let workspace = cx.view().clone();
+        InlineAssistant::update_global(cx, |inline_assistant, cx| {
+            inline_assistant.register_workspace(&workspace, cx)
+        });
+
         cx.observe_flag::<Assistant2FeatureFlag, _>({
             |is_assistant2_enabled, _view, cx| {
-                if is_assistant2_enabled {
-                    let workspace = cx.view().clone();
-                    InlineAssistant::update_global(cx, |inline_assistant, cx| {
-                        inline_assistant.register_workspace(&workspace, cx)
-                    })
-                }
+                InlineAssistant::update_global(cx, |inline_assistant, _cx| {
+                    inline_assistant.is_assistant2_enabled = is_assistant2_enabled;
+                });
             }
         })
         .detach();
@@ -84,6 +86,7 @@ pub struct InlineAssistant {
     prompt_builder: Arc<PromptBuilder>,
     telemetry: Arc<Telemetry>,
     fs: Arc<dyn Fs>,
+    is_assistant2_enabled: bool,
 }
 
 impl Global for InlineAssistant {}
@@ -105,6 +108,7 @@ impl InlineAssistant {
             prompt_builder,
             telemetry,
             fs,
+            is_assistant2_enabled: false,
         }
     }
 
@@ -165,21 +169,31 @@ impl InlineAssistant {
         item: &dyn ItemHandle,
         cx: &mut WindowContext,
     ) {
+        let is_assistant2_enabled = self.is_assistant2_enabled;
+
         if let Some(editor) = item.act_as::<Editor>(cx) {
             editor.update(cx, |editor, cx| {
-                let thread_store = workspace
-                    .read(cx)
-                    .panel::<AssistantPanel>(cx)
-                    .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
-
-                editor.push_code_action_provider(
-                    Rc::new(AssistantCodeActionProvider {
-                        editor: cx.view().downgrade(),
-                        workspace: workspace.downgrade(),
-                        thread_store,
-                    }),
-                    cx,
-                );
+                if is_assistant2_enabled {
+                    let thread_store = workspace
+                        .read(cx)
+                        .panel::<AssistantPanel>(cx)
+                        .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
+
+                    editor.add_code_action_provider(
+                        Rc::new(AssistantCodeActionProvider {
+                            editor: cx.view().downgrade(),
+                            workspace: workspace.downgrade(),
+                            thread_store,
+                        }),
+                        cx,
+                    );
+
+                    // Remove the Assistant1 code action provider, as it still might be registered.
+                    editor.remove_code_action_provider("assistant".into(), cx);
+                } else {
+                    editor
+                        .remove_code_action_provider(ASSISTANT_CODE_ACTION_PROVIDER_ID.into(), cx);
+                }
             });
         }
     }
@@ -1581,7 +1595,13 @@ struct AssistantCodeActionProvider {
     thread_store: Option<WeakModel<ThreadStore>>,
 }
 
+const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant2";
+
 impl CodeActionProvider for AssistantCodeActionProvider {
+    fn id(&self) -> Arc<str> {
+        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
+    }
+
     fn code_actions(
         &self,
         buffer: &Model<Buffer>,

crates/editor/src/editor.rs

@@ -4299,15 +4299,29 @@ impl Editor {
         self.available_code_actions.take();
     }
 
-    pub fn push_code_action_provider(
+    pub fn add_code_action_provider(
         &mut self,
         provider: Rc<dyn CodeActionProvider>,
         cx: &mut ViewContext<Self>,
     ) {
+        if self
+            .code_action_providers
+            .iter()
+            .any(|existing_provider| existing_provider.id() == provider.id())
+        {
+            return;
+        }
+
         self.code_action_providers.push(provider);
         self.refresh_code_actions(cx);
     }
 
+    pub fn remove_code_action_provider(&mut self, id: Arc<str>, cx: &mut ViewContext<Self>) {
+        self.code_action_providers
+            .retain(|provider| provider.id() != id);
+        self.refresh_code_actions(cx);
+    }
+
     fn refresh_code_actions(&mut self, cx: &mut ViewContext<Self>) -> Option<()> {
         let buffer = self.buffer.read(cx);
         let newest_selection = self.selections.newest_anchor().clone();
@@ -13591,6 +13605,8 @@ pub trait CompletionProvider {
 }
 
 pub trait CodeActionProvider {
+    fn id(&self) -> Arc<str>;
+
     fn code_actions(
         &self,
         buffer: &Model<Buffer>,
@@ -13609,6 +13625,10 @@ pub trait CodeActionProvider {
 }
 
 impl CodeActionProvider for Model<Project> {
+    fn id(&self) -> Arc<str> {
+        "project".into()
+    }
+
     fn code_actions(
         &self,
         buffer: &Model<Buffer>,