feat(lsp): migrate document_color to capability framework (#38344)

* feat(lsp): migrate `document_color` to capability framework

* feat(lsp): use `vim.Range` in `document_color` module

runtime/doc/lsp.txt

@@ -122,7 +122,7 @@ BUFFER-LOCAL DEFAULTS
 - |K| is mapped to |vim.lsp.buf.hover()| unless 'keywordprg' is customized or
   a custom keymap for `K` exists.
 - Document colors are enabled for highlighting color references in a document.
-  - To opt out call `vim.lsp.document_color.enable(false, ev.buf)` on |LspAttach|.
+  - To opt out call `vim.lsp.document_color.enable(false, { bufnr = ev.buf })` on |LspAttach|.
 
 DISABLING DEFAULTS                                      *lsp-defaults-disable*
 
@@ -140,7 +140,7 @@ To remove or override BUFFER-LOCAL defaults, define a |LspAttach| handler: >lua
         -- Unmap K
         vim.keymap.del('n', 'K', { buffer = ev.buf })
         -- Disable document colors
-        vim.lsp.document_color.enable(false, ev.buf)
+        vim.lsp.document_color.enable(false, { bufnr = ev.buf })
       end,
     })
 <
@@ -2144,9 +2144,8 @@ document. Highlighting is enabled by default.
 color_presentation()             *vim.lsp.document_color.color_presentation()*
     Select from a list of presentations for the color under the cursor.
 
-enable({enable}, {bufnr}, {opts})            *vim.lsp.document_color.enable()*
-    Enables document highlighting from the given language client in the given
-    buffer.
+enable({enable}, {filter}, {opts})           *vim.lsp.document_color.enable()*
+    Enables or disables document color highlighting for the {filter}ed scope.
 
     To "toggle", pass the inverse of `is_enabled()`: >lua
         vim.lsp.document_color.enable(not vim.lsp.document_color.is_enabled())
@@ -2155,21 +2154,29 @@ enable({enable}, {bufnr}, {opts})            *vim.lsp.document_color.enable()*
     Parameters: ~
       • {enable}  (`boolean?`) True to enable, false to disable. (default:
                   `true`)
-      • {bufnr}   (`integer?`) Buffer handle, or 0 for current. (default: 0)
+      • {filter}  (`table?`) Optional filters |kwargs|,
+                  • {bufnr}? (`integer`, default: all) Buffer number, or 0 for
+                    current buffer, or nil for all.
+                  • {client_id}? (`integer`, default: all) Client ID, or nil
+                    for all.
       • {opts}    (`table?`) A table with the following fields:
                   • {style}?
-                    (`'background'|'foreground'|'virtual'|string|fun(bufnr: integer, range: Range4, hex_code: string)`)
+                    (`'background'|'foreground'|'virtual'|string|fun(bufnr: integer, range: vim.Range, hex_code: string)`)
                     Highlight style. It can be one of the pre-defined styles,
                     a string to be used as virtual text, or a function that
                     receives the buffer handle, the range (start line, start
                     col, end line, end col) and the resolved hex color.
                     (default: `'background'`)
 
-is_enabled({bufnr})                      *vim.lsp.document_color.is_enabled()*
-    Query whether document colors are enabled in the given buffer.
+is_enabled({filter})                     *vim.lsp.document_color.is_enabled()*
+    Query whether document colors are enabled in the {filter}ed scope.
 
     Parameters: ~
-      • {bufnr}  (`integer?`) Buffer handle, or 0 for current. (default: 0)
+      • {filter}  (`table?`) Optional filters |kwargs|,
+                  • {bufnr}? (`integer`, default: all) Buffer number, or 0 for
+                    current buffer, or nil for all.
+                  • {client_id}? (`integer`, default: all) Client ID, or nil
+                    for all.
 
     Return: ~
         (`boolean`)

runtime/lua/vim/lsp/_capability.lua

@@ -2,6 +2,7 @@ local api = vim.api
 
 ---@alias vim.lsp.capability.Name
 ---| 'codelens'
+---| 'document_color'
 ---| 'semantic_tokens'
 ---| 'folding_range'
 ---| 'linked_editing_range'

runtime/lua/vim/lsp/client.lua

@@ -527,6 +527,7 @@ function Client:initialize()
   require('vim.lsp.semantic_tokens')
   require('vim.lsp._folding_range')
   require('vim.lsp.inline_completion')
+  require('vim.lsp.document_color')
 
   local init_params = {
     -- The process Id of the parent process that started the server. Is null if
@@ -1137,10 +1138,6 @@ function Client:on_attach(bufnr)
   self:_text_document_did_open_handler(bufnr)
 
   lsp._set_defaults(self, bufnr)
-  -- `enable(true)` cannot be called from `_set_defaults` for features with dynamic registration,
-  -- because it overrides the state every time `client/registerCapability` is received.
-  -- To allow disabling it once in `LspAttach`, we enable it once here instead.
-  lsp.document_color.enable(true, bufnr)
 
   api.nvim_exec_autocmds('LspAttach', {
     buffer = bufnr,

runtime/lua/vim/lsp/document_color.lua

@@ -4,41 +4,32 @@
 local api = vim.api
 local lsp = vim.lsp
 local util = lsp.util
-local Range = vim.treesitter._range
-
-local document_color_ns = api.nvim_create_namespace('nvim.lsp.document_color')
-local document_color_augroup = api.nvim_create_augroup('nvim.lsp.document_color', {})
+local Capability = require('vim.lsp._capability')
 
 local M = {}
 
 --- @class (private) vim.lsp.document_color.HighlightInfo
 --- @field lsp_info lsp.ColorInformation Unprocessed LSP color information
 --- @field hex_code string Resolved HEX color
---- @field range Range4 Range of the highlight
+--- @field range vim.Range Range of the highlight
 --- @field hl_group? string Highlight group name. Won't be present if the style is a custom function.
 
---- @class (private) vim.lsp.document_color.BufState
---- @field enabled boolean Whether document_color is enabled for the current buffer
---- @field processed_version table<integer, integer?> (client_id -> buffer version) Buffer version for which the color ranges correspond to
---- @field applied_version table<integer, integer?> (client_id -> buffer version) Last buffer version for which we applied color ranges
---- @field hl_info table<integer, vim.lsp.document_color.HighlightInfo[]?> (client_id -> color highlights) Processed highlight information
-
---- @type table<integer, vim.lsp.document_color.BufState?>
-local bufstates = {}
-
---- @type table<integer, integer> (client_id -> namespace ID) documentColor namespace ID for each client.
-local client_ns = {}
+--- @class (private) vim.lsp.document_color.ClientState
+--- @field namespace integer Extmark namespace for this client
+--- @field hl_info vim.lsp.document_color.HighlightInfo[] Processed highlight information
+--- @field processed_version? integer Buffer version for which the color ranges correspond to
+--- @field applied_version? integer Last buffer version for which we applied color ranges
 
 --- @inlinedoc
---- @class vim.lsp.document_color.enable.Opts
+--- @class vim.lsp.document_color.Opts
 ---
 --- Highlight style. It can be one of the pre-defined styles, a string to be used as virtual text, or a
 --- function that receives the buffer handle, the range (start line, start col, end line, end col) and
 --- the resolved hex color. (default: `'background'`)
---- @field style? 'background'|'foreground'|'virtual'|string|fun(bufnr: integer, range: Range4, hex_code: string)
+--- @field style? 'background'|'foreground'|'virtual'|string|fun(bufnr: integer, range: vim.Range, hex_code: string)
 
 -- Default options.
---- @type vim.lsp.document_color.enable.Opts
+--- @type vim.lsp.document_color.Opts
 local document_color_opts = { style = 'background' }
 
 --- @param color string
@@ -100,55 +91,115 @@ local function get_hl_group(hex_code, style)
   return hl_name
 end
 
+--- @class (private) vim.lsp.document_color.Provider : vim.lsp.Capability
+--- @field active table<integer, vim.lsp.document_color.Provider?>
+--- @field client_state table<integer, vim.lsp.document_color.ClientState?>
+local Provider = {
+  name = 'document_color',
+  method = 'textDocument/documentColor',
+  active = {},
+}
+Provider.__index = Provider
+setmetatable(Provider, Capability)
+Capability.all[Provider.name] = Provider
+
+--- @package
 --- @param bufnr integer
---- @param enabled boolean
-local function reset_bufstate(bufnr, enabled)
-  bufstates[bufnr] = {
-    enabled = enabled,
-    processed_version = {},
-    applied_version = {},
+--- @return vim.lsp.document_color.Provider
+function Provider:new(bufnr)
+  --- @type vim.lsp.document_color.Provider
+  self = Capability.new(self, bufnr)
+
+  api.nvim_buf_attach(bufnr, false, {
+    on_lines = function(_, buf)
+      local provider = Provider.active[buf]
+      if not provider then
+        return true
+      end
+      provider:request()
+    end,
+    on_reload = function(_, buf)
+      local provider = Provider.active[buf]
+      if provider then
+        provider:clear()
+        provider:request()
+      end
+    end,
+    on_detach = function(_, buf)
+      local provider = Provider.active[buf]
+      if provider then
+        provider:destroy()
+      end
+    end,
+  })
+
+  api.nvim_create_autocmd('ColorScheme', {
+    group = self.augroup,
+    desc = 'Refresh document_color',
+    callback = function()
+      color_cache = {}
+      local provider = Provider.active[bufnr]
+      if provider then
+        provider:clear()
+        provider:request()
+      end
+    end,
+  })
+
+  return self
+end
+
+--- @package
+--- @param client_id integer
+function Provider:on_attach(client_id)
+  self.client_state[client_id] = {
+    namespace = api.nvim_create_namespace('nvim.lsp.document_color:' .. client_id),
     hl_info = {},
   }
+  self:request(client_id)
+end
+
+--- @package
+--- @param client_id integer
+function Provider:on_detach(client_id)
+  local state = self.client_state[client_id]
+  if state then
+    api.nvim_buf_clear_namespace(self.bufnr, state.namespace, 0, -1)
+    self.client_state[client_id] = nil
+  end
+  api.nvim__redraw({ buf = self.bufnr, valid = true, flush = false })
 end
 
 --- |lsp-handler| for the `textDocument/documentColor` method.
 ---
+--- @package
 --- @param err? lsp.ResponseError
 --- @param result? lsp.ColorInformation[]
 --- @param ctx lsp.HandlerContext
-local function on_document_color(err, result, ctx)
+function Provider:handler(err, result, ctx)
   if err then
     lsp.log.error('document_color', err)
     return
   end
 
-  local bufnr = assert(ctx.bufnr)
-  local bufstate = assert(bufstates[bufnr])
-  local client_id = ctx.client_id
+  local state = self.client_state[ctx.client_id]
+  if not state then
+    return
+  end
 
   if
-    util.buf_versions[bufnr] ~= ctx.version
+    util.buf_versions[self.bufnr] ~= ctx.version
     or not result
-    or not api.nvim_buf_is_loaded(bufnr)
-    or not bufstate.enabled
+    or not api.nvim_buf_is_loaded(self.bufnr)
   then
     return
   end
 
-  if not client_ns[client_id] then
-    client_ns[client_id] = api.nvim_create_namespace('nvim.lsp.document_color.client_' .. client_id)
-  end
-
   local hl_infos = {} --- @type vim.lsp.document_color.HighlightInfo[]
   local style = document_color_opts.style
-  local position_encoding = assert(lsp.get_client_by_id(client_id)).offset_encoding
+  local position_encoding = assert(lsp.get_client_by_id(ctx.client_id)).offset_encoding
   for _, res in ipairs(result) do
-    local range = {
-      res.range.start.line,
-      util._get_line_byte_from_position(bufnr, res.range.start, position_encoding),
-      res.range['end'].line,
-      util._get_line_byte_from_position(bufnr, res.range['end'], position_encoding),
-    }
+    local range = vim.range.lsp(self.bufnr, res.range, position_encoding)
     local hex_code = get_hex_code(res.color)
     --- @type vim.lsp.document_color.HighlightInfo
     local hl_info = { range = range, hex_code = hex_code, lsp_info = res }
@@ -160,225 +211,102 @@ local function on_document_color(err, result, ctx)
     table.insert(hl_infos, hl_info)
   end
 
-  bufstate.hl_info[client_id] = hl_infos
-  bufstate.processed_version[client_id] = ctx.version
+  state.hl_info = hl_infos
+  state.processed_version = ctx.version
 
-  api.nvim__redraw({ buf = bufnr, valid = true, flush = false })
+  api.nvim__redraw({ buf = self.bufnr, valid = true, flush = false })
 end
 
---- @param bufnr integer
-local function buf_clear(bufnr)
-  local bufstate = bufstates[bufnr]
-  if not bufstate then
-    return
-  end
-
-  local client_ids = vim.tbl_keys(bufstate.hl_info) --- @type integer[]
-
-  for _, client_id in ipairs(client_ids) do
-    bufstate.hl_info[client_id] = {}
-    api.nvim_buf_clear_namespace(bufnr, client_ns[client_id], 0, -1)
-  end
-
-  api.nvim__redraw({ buf = bufnr, valid = true, flush = false })
-end
-
---- @param bufnr integer
-local function buf_disable(bufnr)
-  buf_clear(bufnr)
-  reset_bufstate(bufnr, false)
-  api.nvim_clear_autocmds({
-    buffer = bufnr,
-    group = document_color_augroup,
-  })
-end
-
---- @param bufnr integer
-local function buf_enable(bufnr)
-  reset_bufstate(bufnr, true)
-  api.nvim_clear_autocmds({
-    buffer = bufnr,
-    group = document_color_augroup,
-  })
-
-  api.nvim_buf_attach(bufnr, false, {
-    on_reload = function(_, buf)
-      buf_clear(buf)
-      if assert(bufstates[buf]).enabled then
-        M._buf_refresh(buf)
-      end
-    end,
-    on_detach = function(_, buf)
-      buf_disable(buf)
-    end,
-  })
-
-  api.nvim_create_autocmd('LspNotify', {
-    buffer = bufnr,
-    group = document_color_augroup,
-    desc = 'Refresh document_color on document changes',
-    callback = function(ev)
-      local method = ev.data.method --- @type string
-
-      if
-        (method == 'textDocument/didChange' or method == 'textDocument/didOpen')
-        and assert(bufstates[ev.buf]).enabled
-      then
-        M._buf_refresh(ev.buf, ev.data.client_id)
-      end
-    end,
-  })
-
-  api.nvim_create_autocmd('LspDetach', {
-    buffer = bufnr,
-    group = document_color_augroup,
-    desc = 'Disable document_color if all supporting clients detach',
-    callback = function(ev)
-      local clients = lsp.get_clients({ bufnr = ev.buf, method = 'textDocument/documentColor' })
-
-      if
-        not vim.iter(clients):any(function(c)
-          return c.id ~= ev.data.client_id
-        end)
-      then
-        -- There are no clients left in the buffer that support document color, so turn it off.
-        buf_disable(ev.buf)
-      end
-    end,
-  })
-
-  M._buf_refresh(bufnr)
-end
-
---- @param bufnr integer
+--- @package
 --- @param client_id? integer
-function M._buf_refresh(bufnr, client_id)
-  for _, client in
-    ipairs(lsp.get_clients({
-      bufnr = bufnr,
-      id = client_id,
-      method = 'textDocument/documentColor',
-    }))
-  do
-    ---@type lsp.DocumentColorParams
-    local params = { textDocument = util.make_text_document_params(bufnr) }
-    client:request('textDocument/documentColor', params, on_document_color, bufnr)
-  end
-end
-
---- Query whether document colors are enabled in the given buffer.
----
---- @param bufnr? integer Buffer handle, or 0 for current. (default: 0)
---- @return boolean
-function M.is_enabled(bufnr)
-  vim.validate('bufnr', bufnr, 'number', true)
-
-  bufnr = vim._resolve_bufnr(bufnr)
-
-  if not bufstates[bufnr] then
-    reset_bufstate(bufnr, false)
-  end
-
-  return assert(bufstates[bufnr]).enabled
-end
-
---- Enables document highlighting from the given language client in the given buffer.
----
---- To "toggle", pass the inverse of `is_enabled()`:
----
---- ```lua
---- vim.lsp.document_color.enable(not vim.lsp.document_color.is_enabled())
---- ```
----
---- @param enable? boolean True to enable, false to disable. (default: `true`)
---- @param bufnr? integer Buffer handle, or 0 for current. (default: 0)
---- @param opts? vim.lsp.document_color.enable.Opts
-function M.enable(enable, bufnr, opts)
-  vim.validate('enable', enable, 'boolean', true)
-  vim.validate('bufnr', bufnr, 'number', true)
-  vim.validate('opts', opts, 'table', true)
-
-  enable = enable == nil or enable
-  bufnr = vim._resolve_bufnr(bufnr)
-  document_color_opts = vim.tbl_extend('keep', opts or {}, document_color_opts)
-
-  if enable then
-    buf_enable(bufnr)
-  else
-    buf_disable(bufnr)
-  end
-end
-
-api.nvim_create_autocmd('ColorScheme', {
-  pattern = '*',
-  group = document_color_augroup,
-  desc = 'Refresh document_color',
-  callback = function()
-    color_cache = {}
-
-    for _, bufnr in ipairs(api.nvim_list_bufs()) do
-      buf_clear(bufnr)
-      if api.nvim_buf_is_loaded(bufnr) and vim.tbl_get(bufstates, bufnr, 'enabled') then
-        M._buf_refresh(bufnr)
-      else
-        reset_bufstate(bufnr, false)
-      end
+function Provider:request(client_id)
+  for id in pairs(self.client_state) do
+    if not client_id or client_id == id then
+      local client = assert(lsp.get_client_by_id(id))
+      ---@type lsp.DocumentColorParams
+      local params = { textDocument = util.make_text_document_params(self.bufnr) }
+      client:request('textDocument/documentColor', params, function(...)
+        self:handler(...)
+      end, self.bufnr)
     end
-  end,
-})
+  end
+end
 
+--- @package
+function Provider:clear()
+  for _, state in pairs(self.client_state) do
+    state.hl_info = {}
+    state.applied_version = nil
+    api.nvim_buf_clear_namespace(self.bufnr, state.namespace, 0, -1)
+  end
+  api.nvim__redraw({ buf = self.bufnr, valid = true, flush = false })
+end
+
+local document_color_ns = api.nvim_create_namespace('nvim.lsp.document_color')
 api.nvim_set_decoration_provider(document_color_ns, {
   on_win = function(_, _, bufnr)
-    if not bufstates[bufnr] then
-      reset_bufstate(bufnr, false)
+    local provider = Provider.active[bufnr]
+    if not provider then
+      return
     end
-    local bufstate = assert(bufstates[bufnr])
 
     local style = document_color_opts.style
 
-    for client_id, client_hls in pairs(bufstate.hl_info) do
+    for _, state in pairs(provider.client_state) do
       if
-        bufstate.processed_version[client_id] == util.buf_versions[bufnr]
-        and bufstate.processed_version[client_id] ~= bufstate.applied_version[client_id]
+        state.processed_version == util.buf_versions[bufnr]
+        and state.processed_version ~= state.applied_version
       then
-        api.nvim_buf_clear_namespace(bufnr, client_ns[client_id], 0, -1)
+        api.nvim_buf_clear_namespace(bufnr, state.namespace, 0, -1)
 
-        for _, hl in ipairs(client_hls) do
+        for _, hl in ipairs(state.hl_info) do
           if type(style) == 'function' then
             style(bufnr, hl.range, hl.hex_code)
           elseif style == 'foreground' or style == 'background' then
-            api.nvim_buf_set_extmark(bufnr, client_ns[client_id], hl.range[1], hl.range[2], {
-              end_row = hl.range[3],
-              end_col = hl.range[4],
-              hl_group = hl.hl_group,
-              strict = false,
-            })
+            api.nvim_buf_set_extmark(
+              bufnr,
+              state.namespace,
+              hl.range.start.row,
+              hl.range.start.col,
+              {
+                end_row = hl.range.end_.row,
+                end_col = hl.range.end_.col,
+                hl_group = hl.hl_group,
+                strict = false,
+              }
+            )
           else
             -- Default swatch: \uf0c8
-            local swatch = style == 'virtual' and ' ' or style
-            api.nvim_buf_set_extmark(bufnr, client_ns[client_id], hl.range[1], hl.range[2], {
-              virt_text = { { swatch, hl.hl_group } },
-              virt_text_pos = 'inline',
-            })
+            local swatch = style == 'virtual' and ' ' or style
+            api.nvim_buf_set_extmark(
+              bufnr,
+              state.namespace,
+              hl.range.start.row,
+              hl.range.start.col,
+              {
+                virt_text = { { swatch, hl.hl_group } },
+                virt_text_pos = 'inline',
+              }
+            )
           end
         end
 
-        bufstate.applied_version[client_id] = bufstate.processed_version[client_id]
+        state.applied_version = state.processed_version
       end
     end
   end,
 })
 
---- @param bufstate vim.lsp.document_color.BufState
+--- @param provider vim.lsp.document_color.Provider
 --- @return vim.lsp.document_color.HighlightInfo?, integer?
-local function get_hl_info_under_cursor(bufstate)
+local function get_hl_info_under_cursor(provider)
   local cursor_row, cursor_col = unpack(api.nvim_win_get_cursor(0)) --- @type integer, integer
   cursor_row = cursor_row - 1 -- Convert to 0-based index
-  local cursor_range = { cursor_row, cursor_col, cursor_row, cursor_col } --- @type Range4
+  local cursor_pos = vim.pos(cursor_row, cursor_col)
 
-  for client_id, hls in pairs(bufstate.hl_info) do
-    for _, hl in ipairs(hls) do
-      if Range.contains(hl.range, cursor_range) then
+  for client_id, state in pairs(provider.client_state) do
+    for _, hl in ipairs(state.hl_info) do
+      if vim.Range.has(hl.range, cursor_pos) then
         return hl, client_id
       end
     end
@@ -388,13 +316,13 @@ end
 --- Select from a list of presentations for the color under the cursor.
 function M.color_presentation()
   local bufnr = api.nvim_get_current_buf()
-  local bufstate = bufstates[bufnr]
-  if not bufstate then
+  local provider = Provider.active[bufnr]
+  if not provider then
     vim.notify('documentColor is not enabled for this buffer.', vim.log.levels.WARN)
     return
   end
 
-  local hl_info, client_id = get_hl_info_under_cursor(bufstate)
+  local hl_info, client_id = get_hl_info_under_cursor(provider)
   if not hl_info or not client_id then
     vim.notify('No color information under cursor.', vim.log.levels.WARN)
     return
@@ -407,10 +335,7 @@ function M.color_presentation()
   local params = {
     textDocument = { uri = uri },
     color = hl_info.lsp_info.color,
-    range = {
-      start = { line = hl_info.range[1], character = hl_info.range[2] },
-      ['end'] = { line = hl_info.range[3], character = hl_info.range[4] },
-    },
+    range = hl_info.range:to_lsp(client.offset_encoding),
   }
 
   --- @param result lsp.ColorPresentation[]
@@ -424,7 +349,7 @@ function M.color_presentation()
       util.buf_versions[bufnr] ~= ctx.version
       or not next(result)
       or not api.nvim_buf_is_loaded(bufnr)
-      or not bufstate.enabled
+      or not Provider.active[bufnr]
     then
       return
     end
@@ -453,4 +378,40 @@ function M.color_presentation()
   end, bufnr)
 end
 
+--- Query whether document colors are enabled in the {filter}ed scope.
+---
+---@param filter? vim.lsp.capability.enable.Filter
+---@return boolean
+function M.is_enabled(filter)
+  return Capability.is_enabled('document_color', filter)
+end
+
+--- Enables or disables document color highlighting for the {filter}ed scope.
+---
+--- To "toggle", pass the inverse of `is_enabled()`:
+---
+--- ```lua
+--- vim.lsp.document_color.enable(not vim.lsp.document_color.is_enabled())
+--- ```
+---
+---@param enable? boolean True to enable, false to disable. (default: `true`)
+---@param filter? vim.lsp.capability.enable.Filter
+---@param opts? vim.lsp.document_color.Opts
+function M.enable(enable, filter, opts)
+  vim.validate('opts', opts, 'table', true)
+
+  if opts then
+    document_color_opts = vim.tbl_extend('keep', opts, document_color_opts)
+    -- Re-process highlights with new style and refresh active providers.
+    for _, provider in pairs(Provider.active) do
+      provider:clear()
+      provider:request()
+    end
+  end
+
+  Capability.enable('document_color', enable, filter)
+end
+
+Capability.enable('document_color', true)
+
 return M

runtime/lua/vim/lsp/handlers.lua

@@ -151,10 +151,18 @@ RSC['client/registerCapability'] = function(_, params, ctx)
     vim.lsp._set_defaults(client, bufnr)
   end
   for _, reg in ipairs(params.registrations) do
-    if reg.method == 'textDocument/documentColor' then
-      for bufnr in pairs(client.attached_buffers) do
-        if vim.lsp.document_color.is_enabled(bufnr) then
-          vim.lsp.document_color._buf_refresh(bufnr, client.id)
+    -- The capability framework's attach loop only runs during initial Client:on_attach.
+    -- When a client dynamically registers a new method, we need to manually trigger
+    -- on_attach for capabilities that are now supported.
+    for _, Cap in pairs(vim.lsp._capability.all) do
+      if reg.method == Cap.method then
+        for bufnr in pairs(client.attached_buffers) do
+          if vim.lsp._capability.is_enabled(Cap.name, { bufnr = bufnr, client_id = client.id }) then
+            local capability = Cap.active[bufnr] or Cap:new(bufnr)
+            if not capability.client_state[client.id] then
+              capability:on_attach(client.id)
+            end
+          end
         end
       end
     end

test/functional/plugin/lsp/document_color_spec.lua

@@ -64,6 +64,8 @@ body {
     }
 
     bufnr = n.api.nvim_get_current_buf()
+    insert(text)
+
     client_id = exec_lua(function()
       _G.server = _G._create_server({
         capabilities = {
@@ -94,12 +96,6 @@ body {
       return vim.lsp.start({ name = 'dummy', cmd = _G.server.cmd })
     end)
 
-    insert(text)
-
-    exec_lua(function()
-      vim.lsp.document_color.enable(true, bufnr)
-    end)
-
     screen:expect({ grid = grid_with_colors })
   end)
 
@@ -159,7 +155,7 @@ body {
   end)
 
   it('does not clear document colors when one of several clients detaches', function()
-    local client_id2 = exec_lua(function()
+    exec_lua(function()
       _G.server2 = _G._create_server({
         capabilities = {
           colorProvider = true,
@@ -170,16 +166,14 @@ body {
           end,
         },
       })
-      local client_id2 = vim.lsp.start({ name = 'dummy2', cmd = _G.server2.cmd })
-      vim.lsp.document_color.enable(true, bufnr)
-      return client_id2
+      vim.lsp.start({ name = 'dummy2', cmd = _G.server2.cmd })
     end)
 
     exec_lua(function()
-      vim.lsp.get_client_by_id(client_id2):stop()
+      vim.lsp.get_client_by_id(client_id):stop()
     end)
 
-    screen:expect({ grid = grid_with_colors, unchanged = true })
+    screen:expect({ grid = grid_without_colors })
   end)
 
   describe('is_enabled()', function()
@@ -187,27 +181,27 @@ body {
       eq(
         true,
         exec_lua(function()
-          return vim.lsp.document_color.is_enabled(bufnr)
+          return vim.lsp.document_color.is_enabled({ bufnr = bufnr })
         end)
       )
 
       exec_lua(function()
-        vim.lsp.get_client_by_id(client_id):stop()
+        vim.lsp.document_color.enable(false, { bufnr = bufnr })
       end)
 
       eq(
         false,
         exec_lua(function()
-          return vim.lsp.document_color.is_enabled(bufnr)
+          return vim.lsp.document_color.is_enabled({ bufnr = bufnr })
         end)
       )
     end)
 
     it('does not error when called on a new unattached buffer', function()
       eq(
-        false,
+        true,
         exec_lua(function()
-          return vim.lsp.document_color.is_enabled(vim.api.nvim_create_buf(false, true))
+          return vim.lsp.document_color.is_enabled({ bufnr = vim.api.nvim_create_buf(false, true) })
         end)
       )
     end)
@@ -226,7 +220,7 @@ body {                                               |
       ]]
 
       exec_lua(function()
-        vim.lsp.document_color.enable(true, bufnr, { style = 'foreground' })
+        vim.lsp.document_color.enable(true, nil, { style = 'foreground' })
       end)
 
       screen:expect({ grid = grid_with_fg_colors })
@@ -244,7 +238,7 @@ body {                                               |
       ]]
 
       exec_lua(function()
-        vim.lsp.document_color.enable(true, bufnr, { style = ' :) ' })
+        vim.lsp.document_color.enable(true, nil, { style = ' :) ' })
       end)
 
       screen:expect({ grid = grid_with_swatches })
@@ -252,9 +246,7 @@ body {                                               |
 
     it('will not create highlights with custom style function', function()
       exec_lua(function()
-        vim.lsp.document_color.enable(true, bufnr, {
-          style = function() end,
-        })
+        vim.lsp.document_color.enable(true, nil, { style = function() end })
       end)
 
       screen:expect({ grid = grid_without_colors })