From 7a484efb8fac616f5a9e29ec7458e3dfc6e7744f Mon Sep 17 00:00:00 2001
From: Julian Schmidhuber <git@schmidhuberj.de>
Date: Sat, 6 Apr 2024 15:18:09 +0200
Subject: [PATCH] Allow multiple language server with lsp-workspace-command

This fix allows for multiple language servers at once which support
workspace commands. This was previously broken as just the first
language server supporting workspace commands was queried when listing
allowed worspace commands.

The fix is in two parts. Firstly, querying all workspace commands from
all language servers available and using them when actually running the
command in `lsp_workspace_command`. Secondly, doing the same in
`completers::lsp_workspace_command` such that completion still works as
expected.

The fix has one remaining issue, which I am unsure how to handle in the
best way possible, but which I also don't think should happen often:
Multiple language servers may register commands with the same name. This
will lead to that command being listed in the popup menu and in the
completion list multiple times, which can be possibly confusing. One
could disambigue them in the popover menu, but I am not sure the same
can be done for completion. When running `lsp-workspace-command` with
parameters, this behavior is "fixed" by displaying an error in that
case. I am unsure if this is the best fix for this issue in that case,
but could not find a better one.
---
 helix-term/src/commands/typed.rs | 95 ++++++++++++++++++++------------
 helix-term/src/ui/mod.rs         | 16 +++---
 2 files changed, 69 insertions(+), 42 deletions(-)

diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index 5d7057da6..61a6c2514 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -1367,37 +1367,51 @@ fn lsp_workspace_command(
     if event != PromptEvent::Validate {
         return Ok(());
     }
+
+    struct LsIdCommand(usize, helix_lsp::lsp::Command);
+
+    impl ui::menu::Item for LsIdCommand {
+        type Data = ();
+
+        fn format(&self, _data: &Self::Data) -> Row {
+            self.1.title.as_str().into()
+        }
+    }
+
     let doc = doc!(cx.editor);
-    let Some((language_server_id, options)) = doc
+    let ls_id_commands = doc
         .language_servers_with_feature(LanguageServerFeature::WorkspaceCommand)
-        .find_map(|ls| {
+        .flat_map(|ls| {
             ls.capabilities()
                 .execute_command_provider
-                .as_ref()
-                .map(|options| (ls.id(), options))
-        })
-    else {
-        cx.editor
-            .set_status("No active language servers for this document support workspace commands");
-        return Ok(());
-    };
+                .iter()
+                .flat_map(|options| options.commands.iter())
+                .map(|command| (ls.id(), command))
+        });
 
     if args.is_empty() {
-        let commands = options
-            .commands
-            .iter()
-            .map(|command| helix_lsp::lsp::Command {
-                title: command.clone(),
-                command: command.clone(),
-                arguments: None,
+        let commands = ls_id_commands
+            .map(|(ls_id, command)| {
+                LsIdCommand(
+                    ls_id,
+                    helix_lsp::lsp::Command {
+                        title: command.clone(),
+                        command: command.clone(),
+                        arguments: None,
+                    },
+                )
             })
             .collect::<Vec<_>>();
         let callback = async move {
             let call: job::Callback = Callback::EditorCompositor(Box::new(
                 move |_editor: &mut Editor, compositor: &mut Compositor| {
-                    let picker = ui::Picker::new(commands, (), move |cx, command, _action| {
-                        execute_lsp_command(cx.editor, language_server_id, command.clone());
-                    });
+                    let picker = ui::Picker::new(
+                        commands,
+                        (),
+                        move |cx, LsIdCommand(ls_id, command), _action| {
+                            execute_lsp_command(cx.editor, *ls_id, command.clone());
+                        },
+                    );
                     compositor.push(Box::new(overlaid(picker)))
                 },
             ));
@@ -1406,21 +1420,32 @@ fn lsp_workspace_command(
         cx.jobs.callback(callback);
     } else {
         let command = args.join(" ");
-        if options.commands.iter().any(|c| c == &command) {
-            execute_lsp_command(
-                cx.editor,
-                language_server_id,
-                helix_lsp::lsp::Command {
-                    title: command.clone(),
-                    arguments: None,
-                    command,
-                },
-            );
-        } else {
-            cx.editor.set_status(format!(
-                "`{command}` is not supported for this language server"
-            ));
-            return Ok(());
+        let matches: Vec<_> = ls_id_commands
+            .filter(|(_ls_id, c)| *c == &command)
+            .collect();
+
+        match matches.as_slice() {
+            [(ls_id, _command)] => {
+                execute_lsp_command(
+                    cx.editor,
+                    *ls_id,
+                    helix_lsp::lsp::Command {
+                        title: command.clone(),
+                        arguments: None,
+                        command,
+                    },
+                );
+            }
+            [] => {
+                cx.editor.set_status(format!(
+                    "`{command}` is not supported for any language server"
+                ));
+            }
+            _ => {
+                cx.editor.set_status(format!(
+                    "`{command}` supported by multiple language servers"
+                ));
+            }
         }
     }
     Ok(())
diff --git a/helix-term/src/ui/mod.rs b/helix-term/src/ui/mod.rs
index b5969818c..57ac13f60 100644
--- a/helix-term/src/ui/mod.rs
+++ b/helix-term/src/ui/mod.rs
@@ -366,14 +366,16 @@ pub mod completers {
     }
 
     pub fn lsp_workspace_command(editor: &Editor, input: &str) -> Vec<Completion> {
-        let Some(options) = doc!(editor)
+        let commands = doc!(editor)
             .language_servers_with_feature(LanguageServerFeature::WorkspaceCommand)
-            .find_map(|ls| ls.capabilities().execute_command_provider.as_ref())
-        else {
-            return vec![];
-        };
-
-        fuzzy_match(input, &options.commands, false)
+            .flat_map(|ls| {
+                ls.capabilities()
+                    .execute_command_provider
+                    .iter()
+                    .flat_map(|options| options.commands.iter())
+            });
+
+        fuzzy_match(input, commands, false)
             .into_iter()
             .map(|(name, _)| ((0..), name.to_owned().into()))
             .collect()