lsp: only send buf requests to servers that support the request (#12764)

Refactors how required capabilities are detected and validated, and make
sure requests are only sent to clients that support it (and only fail if
no clients support the provided method).

The validation happens at the buf_request level, because we assume that
if someone is sending the request directly through the client, they know
what they're doing. Also, let unknown methods go through.

This is extracted from #12518 and closes #12755.

Co-authored-by: francisco souza <fsouza@users.noreply.github.com>

4 files changed, 160 insertions, 33 deletions

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 95d9c585ee..fad213212a 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -25,6 +25,27 @@ local lsp = {
   -- format_rpc_error = lsp_rpc.format_rpc_error;
 }
 
+-- maps request name to the required resolved_capability in the client.
+lsp._request_name_to_capability = {
+  ['textDocument/hover'] = 'hover';
+  ['textDocument/signatureHelp'] = 'signature_help';
+  ['textDocument/definition'] = 'goto_definition';
+  ['textDocument/implementation'] = 'implementation';
+  ['textDocument/declaration'] = 'declaration';
+  ['textDocument/typeDefinition'] = 'type_definition';
+  ['textDocument/documentSymbol'] = 'document_symbol';
+  ['textDocument/workspaceSymbol'] = 'workspace_symbol';
+  ['textDocument/prepareCallHierarchy'] = 'call_hierarchy';
+  ['textDocument/rename'] = 'rename';
+  ['textDocument/codeAction'] = 'code_action';
+  ['workspace/executeCommand'] = 'execute_command';
+  ['textDocument/references'] = 'find_references';
+  ['textDocument/rangeFormatting'] = 'document_range_formatting';
+  ['textDocument/formatting'] = 'document_formatting';
+  ['textDocument/completion'] = 'completion';
+  ['textDocument/documentHighlight'] = 'document_highlight';
+}
+
 -- TODO improve handling of scratch buffers with LSP attached.
 
 --@private
@@ -51,6 +72,16 @@ local function resolve_bufnr(bufnr)
 end
 
 --@private
+--- callback called by the client when trying to call a method that's not
+--- supported in any of the servers registered for the current buffer.
+--@param method (string) name of the method
+function lsp._unsupported_method(method)
+  local msg = string.format("method %s is not supported by any of the servers registered for the current buffer", method)
+  log.warn(msg)
+  return lsp.rpc_response_error(protocol.ErrorCodes.MethodNotFound, msg)
+end
+
+--@private
 --- Checks whether a given path is a directory.
 ---
 --@param filename (string) path to check
@@ -575,6 +606,15 @@ function lsp.start_client(config)
       -- These are the cleaned up capabilities we use for dynamically deciding
       -- when to send certain events to clients.
       client.resolved_capabilities = protocol.resolve_capabilities(client.server_capabilities)
+      client.supports_method = function(method)
+        local required_capability = lsp._request_name_to_capability[method]
+        -- if we don't know about the method, assume that the client supports it.
+        if not required_capability then
+          return true
+        end
+
+        return client.resolved_capabilities[required_capability]
+      end
       if config.on_init then
         local status, err = pcall(config.on_init, client, result)
         if not status then
@@ -598,19 +638,6 @@ function lsp.start_client(config)
   end
 
   --@private
-  --- Throws error for a method that is not supported by the current LSP
-  --- server.
-  ---
-  --@param method (string) an LSP method name not supported by the LSP server.
-  --@returns (error) a 'MethodNotFound' JSON-RPC error response.
-  local function unsupported_method(method)
-    local msg = "server doesn't support "..method
-    local _ = log.warn() and log.warn(msg)
-    err_message(msg)
-    return lsp.rpc_response_error(protocol.ErrorCodes.MethodNotFound, msg)
-  end
-
-  --@private
   --- Sends a request to the server.
   ---
   --- This is a thin wrapper around {client.rpc.request} with some additional
@@ -637,20 +664,6 @@ function lsp.start_client(config)
         or error(string.format("not found: %q request callback for client %q.", method, client.name))
     end
     local _ = log.debug() and log.debug(log_prefix, "client.request", client_id, method, params, callback, bufnr)
-    -- TODO keep these checks or just let it go anyway?
-    if (not client.resolved_capabilities.hover and method == 'textDocument/hover')
-      or (not client.resolved_capabilities.signature_help and method == 'textDocument/signatureHelp')
-      or (not client.resolved_capabilities.goto_definition and method == 'textDocument/definition')
-      or (not client.resolved_capabilities.implementation and method == 'textDocument/implementation')
-      or (not client.resolved_capabilities.declaration and method == 'textDocument/declaration')
-      or (not client.resolved_capabilities.type_definition and method == 'textDocument/typeDefinition')
-      or (not client.resolved_capabilities.document_symbol and method == 'textDocument/documentSymbol')
-      or (not client.resolved_capabilities.workspace_symbol and method == 'textDocument/workspaceSymbol')
-      or (not client.resolved_capabilities.call_hierarchy and method == 'textDocument/prepareCallHierarchy')
-    then
-      callback(unsupported_method(method), method, nil, client_id, bufnr)
-      return
-    end
     return rpc.request(method, params, function(err, result)
       callback(err, method, result, client_id, bufnr)
     end)
@@ -997,16 +1010,32 @@ function lsp.buf_request(bufnr, method, params, callback)
     callback = { callback, 'f', true };
   }
   local client_request_ids = {}
-  for_each_buffer_client(bufnr, function(client, client_id, resolved_bufnr)
-    local request_success, request_id = client.request(method, params, callback, resolved_bufnr)
 
-    -- This could only fail if the client shut down in the time since we looked
-    -- it up and we did the request, which should be rare.
-    if request_success then
-      client_request_ids[client_id] = request_id
+  local method_supported = false
+  for_each_buffer_client(bufnr, function(client, client_id, resolved_bufnr)
+    if client.supports_method(method) then
+      method_supported = true
+      local request_success, request_id = client.request(method, params, callback, resolved_bufnr)
+
+      -- This could only fail if the client shut down in the time since we looked
+      -- it up and we did the request, which should be rare.
+      if request_success then
+        client_request_ids[client_id] = request_id
+      end
     end
   end)
 
+  -- if no clients support the given method, call the callback with the proper
+  -- error message.
+  if not method_supported then
+    local unsupported_err = lsp._unsupported_method(method)
+    local cb = callback or lsp.callbacks['method']
+    if cb then
+      cb(unsupported_err, method, bufnr)
+    end
+    return
+  end
+
   local function _cancel_all_requests()
     for client_id, request_id in pairs(client_request_ids) do
       local client = active_clients[client_id]
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 4e926381e0..2773f59b45 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -703,6 +703,10 @@ function protocol.make_client_capabilities()
         };
         hierarchicalDocumentSymbolSupport = true;
       };
+      rename = {
+        dynamicRegistration = false;
+        prepareSupport = true;
+      };
     };
     workspace = {
       symbol = {
@@ -914,6 +918,7 @@ function protocol.resolve_capabilities(server_capabilities)
       return nil, string.format("Invalid type for textDocumentSync: %q", type(textDocumentSync))
     end
   end
+  general_properties.completion = server_capabilities.completionProvider ~= nil
   general_properties.hover = server_capabilities.hoverProvider or false
   general_properties.goto_definition = server_capabilities.definitionProvider or false
   general_properties.find_references = server_capabilities.referencesProvider or false
@@ -923,6 +928,15 @@ function protocol.resolve_capabilities(server_capabilities)
   general_properties.document_formatting = server_capabilities.documentFormattingProvider or false
   general_properties.document_range_formatting = server_capabilities.documentRangeFormattingProvider or false
   general_properties.call_hierarchy = server_capabilities.callHierarchyProvider or false
+  general_properties.execute_command = server_capabilities.executeCommandProvider ~= nil
+
+  if server_capabilities.renameProvider == nil then
+    general_properties.rename = false
+  elseif type(server_capabilities.renameProvider) == 'boolean' then
+    general_properties.rename = server_capabilities.renameProvider
+  else
+    general_properties.rename = true
+  end
 
   if server_capabilities.codeActionProvider == nil then
     general_properties.code_action = false
diff --git a/test/functional/fixtures/fake-lsp-server.lua b/test/functional/fixtures/fake-lsp-server.lua
index dca7f35923..a30eb748d0 100644
--- a/test/functional/fixtures/fake-lsp-server.lua
+++ b/test/functional/fixtures/fake-lsp-server.lua
@@ -125,6 +125,26 @@ function tests.basic_check_capabilities()
   }
 end
 
+function tests.capabilities_for_client_supports_method()
+  skeleton {
+    on_init = function(params)
+      local expected_capabilities = protocol.make_client_capabilities()
+      assert_eq(params.capabilities, expected_capabilities)
+      return {
+        capabilities = {
+          textDocumentSync = protocol.TextDocumentSyncKind.Full;
+          completionProvider = true;
+          hoverProvider = true;
+          definitionProvider = false;
+          referencesProvider = false;
+        }
+      }
+    end;
+    body = function()
+    end;
+  }
+end
+
 function tests.basic_finish()
   skeleton {
     on_init = function(params)
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index f07a2d18a2..067a13ce68 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -270,6 +270,70 @@ describe('LSP', function()
         test_name = "basic_check_capabilities";
         on_init = function(client)
           client.stop()
+          local full_kind = exec_lua("return require'vim.lsp.protocol'.TextDocumentSyncKind.Full")
+          eq(full_kind, client.resolved_capabilities().text_document_did_change)
+        end;
+        on_exit = function(code, signal)
+          eq(0, code, "exit code", fake_lsp_logfile)
+          eq(0, signal, "exit signal", fake_lsp_logfile)
+        end;
+        on_callback = function(...)
+          eq(table.remove(expected_callbacks), {...}, "expected callback")
+        end;
+      }
+    end)
+
+    it('client.supports_methods() should validate capabilities', function()
+      local expected_callbacks = {
+        {NIL, "shutdown", {}, 1};
+      }
+      test_rpc_server {
+        test_name = "capabilities_for_client_supports_method";
+        on_init = function(client)
+          client.stop()
+          local full_kind = exec_lua("return require'vim.lsp.protocol'.TextDocumentSyncKind.Full")
+          eq(full_kind, client.resolved_capabilities().text_document_did_change)
+          eq(true, client.resolved_capabilities().completion)
+          eq(true, client.resolved_capabilities().hover)
+          eq(false, client.resolved_capabilities().goto_definition)
+          eq(false, client.resolved_capabilities().rename)
+
+          -- known methods for resolved capabilities
+          eq(true, client.supports_method("textDocument/hover"))
+          eq(false, client.supports_method("textDocument/definition"))
+
+          -- unknown methods are assumed to be supported.
+          eq(true, client.supports_method("unknown-method"))
+        end;
+        on_exit = function(code, signal)
+          eq(0, code, "exit code", fake_lsp_logfile)
+          eq(0, signal, "exit signal", fake_lsp_logfile)
+        end;
+        on_callback = function(...)
+          eq(table.remove(expected_callbacks), {...}, "expected callback")
+        end;
+      }
+    end)
+
+    it('should call unsupported_method when trying to call an unsupported method', function()
+      local expected_callbacks = {
+        {NIL, "shutdown", {}, 1};
+      }
+      test_rpc_server {
+        test_name = "capabilities_for_client_supports_method";
+        on_setup = function()
+            exec_lua([=[
+              vim.lsp._unsupported_method = function(method)
+                vim.lsp._last_unsupported_method = method
+                return 'fake-error'
+              end
+              vim.lsp.buf.hover()
+            ]=])
+        end;
+        on_init = function(client)
+          client.stop()
+          local method = exec_lua("return vim.lsp._last_unsupported_method")
+          eq("textDocument/hover", method)
         end;
         on_exit = function(code, signal)
           eq(0, code, "exit code", fake_lsp_logfile)