fix(lsp): improve dynamic registration handling #37161

Work on #37166 

- Dynamic Registration Tracking via Provider
- Supports_Method
    - Multiple Registrations
    - RegistrationOptions may dictate support for a method

runtime/lua/vim/lsp.lua

@@ -35,13 +35,6 @@ local changetracking = lsp._changetracking
 ---@nodoc
 lsp.rpc_response_error = lsp.rpc.rpc_response_error
 
-lsp._resolve_to_request = {
-  ['codeAction/resolve'] = 'textDocument/codeAction',
-  ['codeLens/resolve'] = 'textDocument/codeLens',
-  ['documentLink/resolve'] = 'textDocument/documentLink',
-  ['inlayHint/resolve'] = 'textDocument/inlayHint',
-}
-
 -- TODO improve handling of scratch buffers with LSP attached.
 
 --- Called by the client when trying to call a method that's not

runtime/lua/vim/lsp/client.lua

@@ -438,13 +438,13 @@ function Client.create(config)
       return self:_unregister_dynamic(unregistrations)
     end,
     get = function(_, method, opts)
-      return self:_get_registration(method, opts and opts.bufnr)
+      return self:_get_registrations(method, opts and opts.bufnr)
     end,
     supports_registration = function(_, method)
       return self:_supports_registration(method)
     end,
     supports = function(_, method, opts)
-      return self:_get_registration(method, opts and opts.bufnr) ~= nil
+      return self:_get_registrations(method, opts and opts.bufnr) ~= nil
     end,
   }
 
@@ -917,17 +917,24 @@ function Client:_supports_registration(method)
   return type(capability) == 'table' and capability.dynamicRegistration
 end
 
+--- Get provider for a method to be registered dyanamically.
+--- @param method vim.lsp.protocol.Method | vim.lsp.protocol.Method.Registration
+function Client:_registration_provider(method)
+  local capability_path = lsp.protocol._request_name_to_server_capability[method]
+  return capability_path and capability_path[1] or method
+end
+
 --- @private
 --- @param registrations lsp.Registration[]
 function Client:_register_dynamic(registrations)
   -- remove duplicates
   self:_unregister_dynamic(registrations)
   for _, reg in ipairs(registrations) do
-    local method = reg.method
+    local provider = self:_registration_provider(reg.method)
-    if not self.registrations[method] then
+    if not self.registrations[provider] then
-      self.registrations[method] = {}
+      self.registrations[provider] = {}
     end
-    table.insert(self.registrations[method], reg)
+    table.insert(self.registrations[provider], reg)
   end
 end
 
@@ -958,7 +965,8 @@ end
 --- @param unregistrations lsp.Unregistration[]
 function Client:_unregister_dynamic(unregistrations)
   for _, unreg in ipairs(unregistrations) do
-    local sreg = self.registrations[unreg.method]
+    local provider = self:_registration_provider(unreg.method)
+    local sreg = self.registrations[provider]
     -- Unegister dynamic capability
     for i, reg in ipairs(sreg or {}) do
       if reg.id == unreg.id then
@@ -984,12 +992,13 @@ function Client:_get_language_id(bufnr)
   return self.get_language_id(bufnr, vim.bo[bufnr].filetype)
 end
 
---- @param method vim.lsp.protocol.Method | vim.lsp.protocol.Method.Registration
+--- @param provider string
 --- @param bufnr? integer
---- @return lsp.Registration?
+--- @return lsp.Registration[]?
-function Client:_get_registration(method, bufnr)
+function Client:_get_registrations(provider, bufnr)
   bufnr = vim._resolve_bufnr(bufnr)
-  for _, reg in ipairs(self.registrations[method] or {}) do
+  local matched_regs = {} --- @type lsp.Registration[]
+  for _, reg in ipairs(self.registrations[provider] or {}) do
     local regoptions = reg.registerOptions --[[@as {documentSelector:lsp.DocumentSelector|lsp.null}]]
     if
       not regoptions
@@ -997,22 +1006,24 @@ function Client:_get_registration(method, bufnr)
       or not regoptions.documentSelector
       or regoptions.documentSelector == vim.NIL
     then
-      return reg
+      matched_regs[#matched_regs + 1] = reg
-    end
+    else
-    local language = self:_get_language_id(bufnr)
+      local language = self:_get_language_id(bufnr)
-    local uri = vim.uri_from_bufnr(bufnr)
+      local uri = vim.uri_from_bufnr(bufnr)
-    local fname = vim.uri_to_fname(uri)
+      local fname = vim.uri_to_fname(uri)
-    for _, filter in ipairs(regoptions.documentSelector) do
+      for _, filter in ipairs(regoptions.documentSelector) do
-      local flang, fscheme, fpat = filter.language, filter.scheme, filter.pattern
+        local flang, fscheme, fpat = filter.language, filter.scheme, filter.pattern
-      if
+        if
-        not (flang and language ~= flang)
+          not (flang and language ~= flang)
-        and not (fscheme and not vim.startswith(uri, fscheme .. ':'))
+          and not (fscheme and not vim.startswith(uri, fscheme .. ':'))
-        and not (type(fpat) == 'string' and not vim.glob.to_lpeg(fpat):match(fname))
+          and not (type(fpat) == 'string' and not vim.glob.to_lpeg(fpat):match(fname))
-      then
+        then
-        return reg
+          matched_regs[#matched_regs + 1] = reg
+        end
       end
     end
   end
+  return #matched_regs > 0 and matched_regs or nil
 end
 
 --- Checks whether a client is stopped.
@@ -1166,17 +1177,24 @@ function Client:supports_method(method, bufnr)
     return true
   end
 
-  local rmethod = lsp._resolve_to_request[method]
+  local provider = self:_registration_provider(method)
-  if rmethod then
+  local regs = self:_get_registrations(provider, bufnr)
-    if self:_supports_registration(rmethod) then
+  if lsp.protocol._request_name_allows_registration[method] and not regs then
-      local reg = self:_get_registration(rmethod, bufnr)
+    return false
-      return vim.tbl_get(reg or {}, 'registerOptions', 'resolveProvider') or false
-    end
-  else
-    if self:_supports_registration(method) then
-      return self:_get_registration(method, bufnr) ~= nil
-    end
   end
+  if regs then
+    for _, reg in ipairs(regs or {}) do
+      if required_capability and #required_capability > 1 then
+        if vim.tbl_get(reg, 'registerOptions', unpack(required_capability, 2)) then
+          return self:_supports_registration(reg.method)
+        end
+      else
+        return self:_supports_registration(reg.method)
+      end
+    end
+    return false
+  end
+
   -- if we don't know about the method, assume that the client supports it.
   -- This needs to be at the end, so that dynamic_capabilities are checked first
   return required_capability == nil

runtime/lua/vim/lsp/protocol.lua

@@ -1299,4 +1299,64 @@ protocol._request_name_to_server_capability = {
 }
 -- stylua: ignore end
 
+-- stylua: ignore start
+-- Generated by gen_lsp.lua, keep at end of file.
+--- Maps method names to the required client capability
+protocol._request_name_allows_registration = {
+  ['notebookDocument/didChange'] = true,
+  ['notebookDocument/didClose'] = true,
+  ['notebookDocument/didOpen'] = true,
+  ['notebookDocument/didSave'] = true,
+  ['textDocument/codeAction'] = true,
+  ['textDocument/codeLens'] = true,
+  ['textDocument/colorPresentation'] = true,
+  ['textDocument/completion'] = true,
+  ['textDocument/declaration'] = true,
+  ['textDocument/definition'] = true,
+  ['textDocument/diagnostic'] = true,
+  ['textDocument/didChange'] = true,
+  ['textDocument/didClose'] = true,
+  ['textDocument/didOpen'] = true,
+  ['textDocument/didSave'] = true,
+  ['textDocument/documentColor'] = true,
+  ['textDocument/documentHighlight'] = true,
+  ['textDocument/documentLink'] = true,
+  ['textDocument/documentSymbol'] = true,
+  ['textDocument/foldingRange'] = true,
+  ['textDocument/formatting'] = true,
+  ['textDocument/hover'] = true,
+  ['textDocument/implementation'] = true,
+  ['textDocument/inlayHint'] = true,
+  ['textDocument/inlineCompletion'] = true,
+  ['textDocument/inlineValue'] = true,
+  ['textDocument/linkedEditingRange'] = true,
+  ['textDocument/moniker'] = true,
+  ['textDocument/onTypeFormatting'] = true,
+  ['textDocument/prepareCallHierarchy'] = true,
+  ['textDocument/prepareTypeHierarchy'] = true,
+  ['textDocument/rangeFormatting'] = true,
+  ['textDocument/rangesFormatting'] = true,
+  ['textDocument/references'] = true,
+  ['textDocument/rename'] = true,
+  ['textDocument/selectionRange'] = true,
+  ['textDocument/semanticTokens/full'] = true,
+  ['textDocument/semanticTokens/full/delta'] = true,
+  ['textDocument/signatureHelp'] = true,
+  ['textDocument/typeDefinition'] = true,
+  ['textDocument/willSave'] = true,
+  ['textDocument/willSaveWaitUntil'] = true,
+  ['workspace/didChangeConfiguration'] = true,
+  ['workspace/didChangeWatchedFiles'] = true,
+  ['workspace/didCreateFiles'] = true,
+  ['workspace/didDeleteFiles'] = true,
+  ['workspace/didRenameFiles'] = true,
+  ['workspace/executeCommand'] = true,
+  ['workspace/symbol'] = true,
+  ['workspace/textDocumentContent'] = true,
+  ['workspace/willCreateFiles'] = true,
+  ['workspace/willDeleteFiles'] = true,
+  ['workspace/willRenameFiles'] = true,
+}
+-- stylua: ignore end
+
 return protocol

src/gen/gen_lsp.lua

@@ -282,6 +282,23 @@ local function write_to_vim_protocol(protocol)
 
     output[#output + 1] = '}'
     output[#output + 1] = '-- stylua: ignore end'
+
+    vim.list_extend(output, {
+      '',
+      '-- stylua: ignore start',
+      '-- Generated by gen_lsp.lua, keep at end of file.',
+      '--- Maps method names to the required client capability',
+      'protocol._request_name_allows_registration = {',
+    })
+
+    for _, item in ipairs(all) do
+      if item.registrationOptions then
+        output[#output + 1] = ("  ['%s'] = %s,"):format(item.method, true)
+      end
+    end
+
+    output[#output + 1] = '}'
+    output[#output + 1] = '-- stylua: ignore end'
   end
 
   output[#output + 1] = ''

test/functional/plugin/lsp_spec.lua

@@ -5930,10 +5930,73 @@ describe('LSP', function()
         check('workspace/didChangeWatchedFiles')
         check('workspace/didChangeWatchedFiles', tmpfile)
 
+        -- Initial support false
+        check('workspace/diagnostic')
+
+        vim.lsp.handlers['client/registerCapability'](nil, {
+          registrations = {
+            {
+              id = 'diag1',
+              method = 'textDocument/diagnostic',
+              registerOptions = {
+                -- workspaceDiagnostics field omitted
+              },
+            },
+          },
+        }, { client_id = client_id })
+
+        -- Checks after registering without worspaceDiagnostics support
+        -- Returns false
+        check('workspace/diagnostic')
+
+        vim.lsp.handlers['client/registerCapability'](nil, {
+          registrations = {
+            {
+              id = 'diag2',
+              method = 'textDocument/diagnostic',
+              registerOptions = {
+                workspaceDiagnostics = true,
+              },
+            },
+          },
+        }, { client_id = client_id })
+
+        -- Check after second registration with support
+        -- Returns true
+        check('workspace/diagnostic')
+
+        vim.lsp.handlers['client/unregisterCapability'](nil, {
+          unregisterations = {
+            { id = 'diag2', method = 'textDocument/diagnostic' },
+          },
+        }, { client_id = client_id })
+
+        -- Check after unregistering
+        -- Returns false
+        check('workspace/diagnostic')
+
+        check('textDocument/codeAction')
+        check('codeAction/resolve')
+
+        vim.lsp.handlers['client/registerCapability'](nil, {
+          registrations = {
+            {
+              id = 'codeAction',
+              method = 'textDocument/codeAction',
+              registerOptions = {
+                resolveProvider = true,
+              },
+            },
+          },
+        }, { client_id = client_id })
+
+        check('textDocument/codeAction')
+        check('codeAction/resolve')
+
         return result
       end)
 
-      eq(9, #result)
+      eq(17, #result)
       eq({ method = 'textDocument/formatting', supported = false }, result[1])
       eq({ method = 'textDocument/formatting', supported = true, fname = tmpfile }, result[2])
       eq({ method = 'textDocument/rangeFormatting', supported = true }, result[3])
@@ -5949,6 +6012,14 @@ describe('LSP', function()
         { method = 'workspace/didChangeWatchedFiles', supported = true, fname = tmpfile },
         result[9]
       )
+      eq({ method = 'workspace/diagnostic', supported = false }, result[10])
+      eq({ method = 'workspace/diagnostic', supported = false }, result[11])
+      eq({ method = 'workspace/diagnostic', supported = true }, result[12])
+      eq({ method = 'workspace/diagnostic', supported = false }, result[13])
+      eq({ method = 'textDocument/codeAction', supported = false }, result[14])
+      eq({ method = 'codeAction/resolve', supported = false }, result[15])
+      eq({ method = 'textDocument/codeAction', supported = true }, result[16])
+      eq({ method = 'codeAction/resolve', supported = true }, result[17])
     end)
 
     it('identifies client dynamic registration capability', function()
@@ -6011,7 +6082,7 @@ describe('LSP', function()
         true,
         exec_lua(function()
           local client = assert(vim.lsp.get_client_by_id(client_id))
-          return client.dynamic_capabilities:get('textDocument/documentColor') ~= nil
+          return client.dynamic_capabilities:get('colorProvider') ~= nil
         end)
       )
     end)