From e74e2863f9028da8c58ac0f889dbda9ec5c81731 Mon Sep 17 00:00:00 2001 From: "gcp-cherry-pick-bot[bot]" <98988430+gcp-cherry-pick-bot[bot]@users.noreply.github.com> Date: Tue, 14 Jan 2025 18:45:52 +0200 Subject: [PATCH] Improve registration for Assistant code action providers (cherry-pick #23099) (#23136) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cherry-picked Improve registration for Assistant code action providers (#23099) 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 Co-authored-by: Marshall Bowers --- 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 73230a48460ccca0fb6ce7c517875cb5bfc2eedf..6c142d376a08238e3e4235c1786b412b769a0493 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 cff145f25aa9e6e7dbd6b001c84461f1fa95b632..651269e3d995c6bfeb0965e0d4d1f7a38c0e7c63 100644 --- a/crates/editor/src/editor.rs +++ b/crates/editor/src/editor.rs @@ -4283,15 +4283,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(); @@ -13529,6 +13543,8 @@ pub trait CompletionProvider { } pub trait CodeActionProvider { + fn id(&self) -> Arc; + fn code_actions( &self, buffer: &Model, @@ -13547,6 +13563,10 @@ pub trait CodeActionProvider { } impl CodeActionProvider for Model { + fn id(&self) -> Arc { + "project".into() + } + fn code_actions( &self, buffer: &Model,