From 06d8fee048ab7964fe2a29d260bd96252e7b102b Mon Sep 17 00:00:00 2001
From: Schmiddiii <git@schmidhuberj.de>
Date: Sat, 6 Jul 2024 17:34:33 +0200
Subject: [PATCH] Allow multiple language server with lsp-workspace-command
 (#10176)

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 ab12fab10..3a8ed1461 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -1370,37 +1370,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)))
                 },
             ));
@@ -1409,21 +1423,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 6a4655fde..0a65b12b5 100644
--- a/helix-term/src/ui/mod.rs
+++ b/helix-term/src/ui/mod.rs
@@ -364,14 +364,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()