Skip to content

Commit

Permalink
Improve registration for Assistant code action providers (#23099)
Browse files Browse the repository at this point in the history 
This PR is a follow-up to
#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
  • Loading branch information
@maxdeviant
maxdeviant authored Jan 13, 2025
1 parent 830f45e commit 93f117b
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 35 deletions.
45 changes: 30 additions & 15 deletions crates/assistant/src/inline_assistant.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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 {}
Expand All @@ -123,6 +124,7 @@ impl InlineAssistant {
prompt_builder,
telemetry,
fs,
is_assistant2_enabled: false,
}
}

Expand Down Expand Up @@ -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,
);
}
});
}
}
Expand Down Expand Up @@ -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>,
Expand Down
58 changes: 39 additions & 19 deletions crates/assistant2/src/inline_assistant.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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 {}
Expand All @@ -105,6 +108,7 @@ impl InlineAssistant {
prompt_builder,
telemetry,
fs,
is_assistant2_enabled: false,
}
}

Expand Down Expand Up @@ -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);
}
});
}
}
Expand Down Expand Up @@ -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>,
Expand Down
22 changes: 21 additions & 1 deletion crates/editor/src/editor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -13591,6 +13605,8 @@ pub trait CompletionProvider {
}

pub trait CodeActionProvider {
fn id(&self) -> Arc<str>;

fn code_actions(
&self,
buffer: &Model<Buffer>,
Expand All @@ -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>,
Expand Down

0 comments on commit 93f117b

Please sign in to comment.