assistant: Fix completions for slash commands provided by context servers (#20423)

Marshall Bowers created 1 year ago

This PR fixes an issue introduced in #20372 that was causing slash
commands provided by context servers to not show up in the completions
menu.

Release Notes:

- N/A

Change summary

crates/assistant/src/assistant_panel.rs           |  1 
crates/assistant/src/prompt_library.rs            |  7 +++++
crates/assistant/src/slash_command.rs             | 12 +++++++---
crates/assistant/src/slash_command_working_set.rs | 19 ++++++++++++++--
4 files changed, 31 insertions(+), 8 deletions(-)

Detailed changes

crates/assistant/src/assistant_panel.rs

@@ -1547,6 +1547,7 @@ impl ContextEditor {
         cx: &mut ViewContext<Self>,
     ) -> Self {
         let completion_provider = SlashCommandCompletionProvider::new(
+            context.read(cx).slash_commands.clone(),
             Some(cx.view().downgrade()),
             Some(workspace.clone()),
         );

crates/assistant/src/prompt_library.rs

@@ -1,3 +1,4 @@
+use crate::SlashCommandWorkingSet;
 use crate::{slash_command::SlashCommandCompletionProvider, AssistantPanel, InlineAssistant};
 use anyhow::{anyhow, Result};
 use chrono::{DateTime, Utc};
@@ -522,7 +523,11 @@ impl PromptLibrary {
                             editor.set_use_modal_editing(false);
                             editor.set_current_line_highlight(Some(CurrentLineHighlight::None));
                             editor.set_completion_provider(Some(Box::new(
-                                SlashCommandCompletionProvider::new(None, None),
+                                SlashCommandCompletionProvider::new(
+                                    Arc::new(SlashCommandWorkingSet::default()),
+                                    None,
+                                    None,
+                                ),
                             )));
                             if focus {
                                 editor.focus(cx);

crates/assistant/src/slash_command.rs

@@ -1,4 +1,5 @@
 use crate::assistant_panel::ContextEditor;
+use crate::SlashCommandWorkingSet;
 use anyhow::Result;
 use assistant_slash_command::AfterCompletion;
 pub use assistant_slash_command::{SlashCommand, SlashCommandOutput, SlashCommandRegistry};
@@ -39,6 +40,7 @@ pub mod terminal_command;
 
 pub(crate) struct SlashCommandCompletionProvider {
     cancel_flag: Mutex<Arc<AtomicBool>>,
+    slash_commands: Arc<SlashCommandWorkingSet>,
     editor: Option<WeakView<ContextEditor>>,
     workspace: Option<WeakView<Workspace>>,
 }
@@ -52,11 +54,13 @@ pub(crate) struct SlashCommandLine {
 
 impl SlashCommandCompletionProvider {
     pub fn new(
+        slash_commands: Arc<SlashCommandWorkingSet>,
         editor: Option<WeakView<ContextEditor>>,
         workspace: Option<WeakView<Workspace>>,
     ) -> Self {
         Self {
             cancel_flag: Mutex::new(Arc::new(AtomicBool::new(false))),
+            slash_commands,
             editor,
             workspace,
         }
@@ -69,9 +73,9 @@ impl SlashCommandCompletionProvider {
         name_range: Range<Anchor>,
         cx: &mut WindowContext,
     ) -> Task<Result<Vec<project::Completion>>> {
-        let commands = SlashCommandRegistry::global(cx);
-        let candidates = commands
-            .command_names()
+        let slash_commands = self.slash_commands.clone();
+        let candidates = slash_commands
+            .command_names(cx)
             .into_iter()
             .enumerate()
             .map(|(ix, def)| StringMatchCandidate {
@@ -98,7 +102,7 @@ impl SlashCommandCompletionProvider {
                 matches
                     .into_iter()
                     .filter_map(|mat| {
-                        let command = commands.command(&mat.string)?;
+                        let command = slash_commands.command(&mat.string, cx)?;
                         let mut new_text = mat.string.clone();
                         let requires_argument = command.requires_argument();
                         let accepts_arguments = command.accepts_arguments();

crates/assistant/src/slash_command_working_set.rs

@@ -4,7 +4,7 @@ use gpui::AppContext;
 use parking_lot::Mutex;
 use std::sync::Arc;
 
-#[derive(Copy, Clone, PartialEq, Eq, Hash, Default)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
 pub struct SlashCommandId(usize);
 
 /// A working set of slash commands for use in one instance of the Assistant Panel.
@@ -16,7 +16,7 @@ pub struct SlashCommandWorkingSet {
 #[derive(Default)]
 struct WorkingSetState {
     context_server_commands_by_id: HashMap<SlashCommandId, Arc<dyn SlashCommand>>,
-    context_server_commands_by_name: HashMap<String, Arc<dyn SlashCommand>>,
+    context_server_commands_by_name: HashMap<Arc<str>, Arc<dyn SlashCommand>>,
     next_command_id: SlashCommandId,
 }
 
@@ -30,6 +30,19 @@ impl SlashCommandWorkingSet {
             .or_else(|| SlashCommandRegistry::global(cx).command(name))
     }
 
+    pub fn command_names(&self, cx: &AppContext) -> Vec<Arc<str>> {
+        let mut command_names = SlashCommandRegistry::global(cx).command_names();
+        command_names.extend(
+            self.state
+                .lock()
+                .context_server_commands_by_name
+                .keys()
+                .cloned(),
+        );
+
+        command_names
+    }
+
     pub fn featured_command_names(&self, cx: &AppContext) -> Vec<Arc<str>> {
         SlashCommandRegistry::global(cx).featured_command_names()
     }
@@ -60,7 +73,7 @@ impl WorkingSetState {
         self.context_server_commands_by_name.extend(
             self.context_server_commands_by_id
                 .values()
-                .map(|command| (command.name(), command.clone())),
+                .map(|command| (command.name().into(), command.clone())),
         );
     }
 }