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 | 58 +++++++++++++++-------- crates/editor/src/editor.rs | 22 ++++++++- 3 files changed, 90 insertions(+), 35 deletions(-) diff --git a/crates/assistant/src/inline_assistant.rs b/crates/assistant/src/inline_assistant.rs index 4f43f92edec31eaafd36f977c393156330c97b36..8acdaec50f4cfed05668ab34099979f842bb7620 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 4367b3641ac5db224b3a29f5603d8c7e68f34fa7..73e539473f9bffb52562562cf6f88a34eb863151 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()); - - 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::(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>, } +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 4170f70d2b806a571e427c4b40a7e578912e62c6..0ca6796b37156d845781ce7595be8a648257b7a0 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,