From 93f117b21a07cb0e5c41682658339a473794d2da Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Mon, 13 Jan 2025 17:25:58 -0500 Subject: [PATCH] Improve registration for Assistant code action providers (#23099) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/assistant/src/inline_assistant.rs | 45 ++++++++++++------ crates/assistant2/src/inline_assistant.rs | 56 +++++++++++++++-------- crates/editor/src/editor.rs | 22 ++++++++- 3 files changed, 89 insertions(+), 34 deletions(-) diff --git a/crates/assistant/src/inline_assistant.rs b/crates/assistant/src/inline_assistant.rs index 4f43f92ede..8acdaec50f 100644 --- a/crates/assistant/src/inline_assistant.rs +++ b/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::({ |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, telemetry: Arc, fs: Arc, + 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::(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, } +const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant"; + impl CodeActionProvider for AssistantCodeActionProvider { + fn id(&self) -> Arc { + ASSISTANT_CODE_ACTION_PROVIDER_ID.into() + } + fn code_actions( &self, buffer: &Model, diff --git a/crates/assistant2/src/inline_assistant.rs b/crates/assistant2/src/inline_assistant.rs index 4367b3641a..73e539473f 100644 --- a/crates/assistant2/src/inline_assistant.rs +++ b/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::({ |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, telemetry: Arc, fs: Arc, + 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::(cx) { editor.update(cx, |editor, cx| { - let thread_store = workspace - .read(cx) - .panel::(cx) - .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade()); + if is_assistant2_enabled { + let thread_store = workspace + .read(cx) + .panel::(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, - ); + 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>, } +const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant2"; + impl CodeActionProvider for AssistantCodeActionProvider { + fn id(&self) -> Arc { + ASSISTANT_CODE_ACTION_PROVIDER_ID.into() + } + fn code_actions( &self, buffer: &Model, diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs index 4170f70d2b..0ca6796b37 100644 --- a/crates/editor/src/editor.rs +++ b/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, cx: &mut ViewContext, ) { + 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, cx: &mut ViewContext) { + self.code_action_providers + .retain(|provider| provider.id() != id); + self.refresh_code_actions(cx); + } + fn refresh_code_actions(&mut self, cx: &mut ViewContext) -> 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; + fn code_actions( &self, buffer: &Model, @@ -13609,6 +13625,10 @@ pub trait CodeActionProvider { } impl CodeActionProvider for Model { + fn id(&self) -> Arc { + "project".into() + } + fn code_actions( &self, buffer: &Model,