feat(lsp): add per-client commands (#16101)

mjlbach · web-flow · commit 519d8deb08f3 · 2021-11-01T03:14:59.000-07:00
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
@@ -239,6 +239,7 @@ local function validate_client_config(config)
     on_exit         = { config.on_exit, "f", true };
     on_init         = { config.on_init, "f", true };
     settings        = { config.settings, "t", true };
+    commands        = { config.commands, 't', true };
     before_init     = { config.before_init, "f", true };
     offset_encoding = { config.offset_encoding, "s", true };
     flags           = { config.flags, "t", true };
@@ -590,6 +591,11 @@ end
 --- returned to the language server if requested via `workspace/configuration`.
 --- Keys are case-sensitive.
 ---
+---@param commands table Table that maps string of clientside commands to user-defined functions.
+--- Commands passed to start_client take precedence over the global command registry. Each key
+--- must be a unique comand name, and the value is a function which is called if any LSP action
+--- (code action, code lenses, ...) triggers the command.
+---
 ---@param init_options Values to pass in the initialization request
 --- as `initializationOptions`. See `initialize` in the LSP spec.
 ---
@@ -772,8 +778,9 @@ function lsp.start_client(config)
     attached_buffers = {};
 
     handlers = handlers;
-    requests = {};
+    commands = config.commands or {};
 
+    requests = {};
     -- for $/progress report
     messages = { name = name, messages = {}, progress = {}, status = {} };
   }
diff --git a/runtime/lua/vim/lsp/buf.lua b/runtime/lua/vim/lsp/buf.lua
@@ -480,7 +480,7 @@ local function on_code_action_results(results, ctx)
     end
     if action.command then
       local command = type(action.command) == 'table' and action.command or action
-      local fn = vim.lsp.commands[command.command]
+      local fn = client.commands[command.command] or vim.lsp.commands[command.command]
       if fn then
         local enriched_ctx = vim.deepcopy(ctx)
         enriched_ctx.client_id = client.id
diff --git a/runtime/lua/vim/lsp/codelens.lua b/runtime/lua/vim/lsp/codelens.lua
@@ -31,15 +31,15 @@ local function execute_lens(lens, bufnr, client_id)
   local line = lens.range.start.line
   api.nvim_buf_clear_namespace(bufnr, namespaces[client_id], line, line + 1)
 
+  local client = vim.lsp.get_client_by_id(client_id)
+  assert(client, 'Client is required to execute lens, client_id=' .. client_id)
   local command = lens.command
-  local fn = vim.lsp.commands[command.command]
+  local fn = client.commands[command.command] or vim.lsp.commands[command.command]
   if fn then
     fn(command, { bufnr = bufnr, client_id = client_id })
     return
   end
   -- Need to use the client that returned the lens → must not use buf_request
-  local client = vim.lsp.get_client_by_id(client_id)
-  assert(client, 'Client is required to execute lens, client_id=' .. client_id)
   local command_provider = client.server_capabilities.executeCommandProvider
   local commands = type(command_provider) == 'table' and command_provider.commands or {}
   if not vim.tbl_contains(commands, command.command) then
diff --git a/test/functional/fixtures/fake-lsp-server.lua b/test/functional/fixtures/fake-lsp-server.lua
@@ -671,6 +671,20 @@ function tests.code_action_with_resolve()
   }
 end
 
+function tests.clientside_commands()
+  skeleton {
+    on_init = function()
+      return {
+        capabilities = {}
+      }
+    end;
+    body = function()
+      notify('start')
+      notify('shutdown')
+    end;
+  }
+end
+
 -- Tests will be indexed by TEST_NAME
 
 local kill_timer = vim.loop.new_timer()
diff --git a/test/functional/plugin/lsp/codelens_spec.lua b/test/functional/plugin/lsp/codelens_spec.lua
@@ -60,31 +60,4 @@ describe('vim.lsp.codelens', function()
     eq({[1] = {'Lens1', 'LspCodeLens'}}, virtual_text_chunks)
 
   end)
-  it('codelens uses client commands', function()
-    local fake_uri = "file:///fake/uri"
-    local cmd = exec_lua([[
-      fake_uri = ...
-      local bufnr = vim.uri_to_bufnr(fake_uri)
-      vim.fn.bufload(bufnr)
-      vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {'One line'})
-      local lenses = {
-        {
-          range = {
-            start = { line = 0, character = 0, },
-            ['end'] = { line = 0, character = 8 }
-          },
-          command = { title = 'Lens1', command = 'Dummy' }
-        },
-      }
-      vim.lsp.codelens.on_codelens(nil, lenses, {method='textDocument/codeLens', client_id=1, bufnr=bufnr})
-      local cmd_called = nil
-      vim.lsp.commands['Dummy'] = function(command)
-        cmd_called = command
-      end
-      vim.api.nvim_set_current_buf(bufnr)
-      vim.lsp.codelens.run()
-      return cmd_called
-    ]], fake_uri)
-    eq({ command = 'Dummy', title = 'Lens1' }, cmd)
-  end)
 end)
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
@@ -2531,4 +2531,57 @@ describe('LSP', function()
       )
     end)
   end)
+  describe('vim.lsp.codelens', function()
+    it('uses client commands', function()
+      local client
+      local expected_handlers = {
+        {NIL, {}, {method="shutdown", client_id=1}};
+        {NIL, {}, {method="start", client_id=1}};
+      }
+      test_rpc_server {
+        test_name = 'clientside_commands',
+        on_init = function(client_)
+          client = client_
+        end,
+        on_setup = function()
+        end,
+        on_exit = function(code, signal)
+          eq(0, code, "exit code", fake_lsp_logfile)
+          eq(0, signal, "exit signal", fake_lsp_logfile)
+        end,
+        on_handler = function(err, result, ctx)
+          eq(table.remove(expected_handlers), {err, result, ctx})
+          if ctx.method == 'start' then
+            local fake_uri = "file:///fake/uri"
+            local cmd = exec_lua([[
+              fake_uri = ...
+              local bufnr = vim.uri_to_bufnr(fake_uri)
+              vim.fn.bufload(bufnr)
+              vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {'One line'})
+              local lenses = {
+                {
+                  range = {
+                    start = { line = 0, character = 0, },
+                    ['end'] = { line = 0, character = 8 }
+                  },
+                  command = { title = 'Lens1', command = 'Dummy' }
+                },
+              }
+              vim.lsp.codelens.on_codelens(nil, lenses, {method='textDocument/codeLens', client_id=1, bufnr=bufnr})
+              local cmd_called = nil
+              vim.lsp.commands['Dummy'] = function(command)
+                cmd_called = command
+              end
+              vim.api.nvim_set_current_buf(bufnr)
+              vim.lsp.codelens.run()
+              return cmd_called
+            ]], fake_uri)
+         eq({ command = 'Dummy', title = 'Lens1' }, cmd)
+         elseif ctx.method == 'shutdown' then
+           client.stop()
+          end
+        end
+      }
+    end)
+  end)
 end)
