Merge #38047 _provider_foreach

3 files changed, 74 insertions, 45 deletions

diff --git a/runtime/lua/vim/lsp/buf.lua b/runtime/lua/vim/lsp/buf.lua
index 1d99afd39a..14bbb279f1 100644
--- a/runtime/lua/vim/lsp/buf.lua
+++ b/runtime/lua/vim/lsp/buf.lua
@@ -1365,6 +1365,7 @@ function M.code_action(opts)
       params.context = context
     else
       local ns_push = lsp.diagnostic.get_namespace(client.id, false)
+      -- TODO(tris203): should we aggregate diagnostics from all the possible pull namespaces?
       local ns_pull = lsp.diagnostic.get_namespace(client.id, true)
       local diagnostics = {}
       local lnum = api.nvim_win_get_cursor(0)[1] - 1
diff --git a/runtime/lua/vim/lsp/client.lua b/runtime/lua/vim/lsp/client.lua
index 36a4c28686..4a5b259f3c 100644
--- a/runtime/lua/vim/lsp/client.lua
+++ b/runtime/lua/vim/lsp/client.lua
@@ -1224,38 +1224,47 @@ function Client:supports_method(method, bufnr)
   return required_capability == nil
 end
 
---- Retrieves all capability values for a given LSP method, handling both static and dynamic registrations.
---- This function abstracts over differences between capabilities declared in `server_capabilities`
---- and those registered dynamically at runtime, returning all matching capability values.
---- It also handles cases where the registration method differs from the calling method by abstracting to the Provider.
---- For example, `workspace/diagnostic` uses capabilities registered under `textDocument/diagnostic`.
---- This is useful for features like diagnostics and formatting, where servers may register multiple providers
---- with different options (such as specific filetypes or document selectors).
---- @param method vim.lsp.protocol.Method.ClientToServer | vim.lsp.protocol.Method.Registration LSP method name
---- @param ... any Additional keys to index into the capability
---- @return lsp.LSPAny[] # The capability value if it exists, empty table if not found
-function Client:_provider_value_get(method, ...)
-  local matched_regs = {} --- @type any[]
+--- Executes callback fn for all registrations for a given LSP method.
+---
+--- This handles both static capabilities (declared in server_capabilities during
+--- initialization) and dynamic registrations (registered at runtime via
+--- `client/registerCapability`).
+---
+--- Some methods may have multiple registrations (e.g., different documentSelectors
+--- or configurations). The callback is invoked once for each registration.
+---
+--- Example: Getting diagnostic identifiers from all registrations
+---     client:_provider_foreach('textDocument/diagnostic', function(cap)
+---       print(cap.identifier)  -- "static-id", "dynamic-id-1", "dynamic-id-2"
+---     end)
+---
+--- Note: Some capabilities alias to different providers. For example,
+--- `workspace/diagnostic` uses the same `diagnosticProvider` as `textDocument/diagnostic`.
+---
+---@param method vim.lsp.protocol.Method.ClientToServer | vim.lsp.protocol.Method.Registration LSP method name
+---@param fn fun(capability_value: lsp.LSPAny) Callback invoked for each matching capability
+function Client:_provider_foreach(method, fn)
   local provider = self:_registration_provider(method)
+  local required_capability = lsp.protocol._request_name_to_server_capability[method]
   local dynamic_regs = self:_get_registrations(provider)
+  local has_subcap = required_capability and #required_capability > 1
   if not provider then
-    return matched_regs
+    return
   elseif not dynamic_regs then
     -- First check static capabilities
     local static_reg = vim.tbl_get(self.server_capabilities, provider)
     if static_reg then
-      matched_regs[1] = vim.tbl_get(static_reg, ...) or vim.NIL
+      if not has_subcap or vim.tbl_get(static_reg, unpack(required_capability, 2)) then
+        fn(static_reg)
+      end
     end
   else
-    local required_capability = lsp.protocol._request_name_to_server_capability[method]
     for _, reg in ipairs(dynamic_regs) do
-      if vim.tbl_get(reg, 'registerOptions', unpack(required_capability, 2)) then
-        matched_regs[#matched_regs + 1] = vim.tbl_get(reg, 'registerOptions', ...) or vim.NIL
+      if not has_subcap or vim.tbl_get(reg, 'registerOptions', unpack(required_capability, 2)) then
+        fn(vim.tbl_get(reg, 'registerOptions') or {})
       end
     end
   end
-
-  return matched_regs
 end
 
 --- @private
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index eab69fdc62..4f6a2a346c 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -15,7 +15,7 @@ local augroup = api.nvim_create_augroup('nvim.lsp.diagnostic', {})
 
 ---@class (private) vim.lsp.diagnostic.BufState
 ---@field pull_kind 'document'|'workspace'|'disabled' Whether diagnostics are being updated via document pull, workspace pull, or disabled.
----@field client_result_id table<integer, string?> Latest responded `resultId`
+---@field client_result_id table<string, string?> Latest responded `resultId`, keyed by `client_id.identifier`
 
 ---@type table<integer, vim.lsp.diagnostic.BufState>
 local bufstates = {}
@@ -147,6 +147,13 @@ local function tags_vim_to_lsp(diagnostic)
   return tags
 end
 
+---@param client_id integer
+---@param identifier string|nil
+---@return string
+local function result_id_key(client_id, identifier)
+  return string.format('%d.%s', client_id, identifier or 'nil')
+end
+
 --- Converts the input `vim.Diagnostic`s to LSP diagnostics.
 --- @param diagnostics vim.Diagnostic[]
 --- @return lsp.Diagnostic[]
@@ -254,6 +261,8 @@ end
 ---@param params lsp.PublishDiagnosticsParams
 ---@param ctx lsp.HandlerContext
 function M.on_publish_diagnostics(_, params, ctx)
+  -- TODO(tris203): if empty array then clear diags
+  -- https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
   handle_diagnostics(params.uri, ctx.client_id, params.diagnostics, false)
 end
 
@@ -281,14 +290,15 @@ function M.on_diagnostic(error, result, ctx)
   local client_id = ctx.client_id
   local bufnr = assert(ctx.bufnr)
   local bufstate = bufstates[bufnr]
-  bufstate.client_result_id[client_id] = result.resultId
+  ---@type lsp.DocumentDiagnosticParams
+  local params = ctx.params
+  local key = result_id_key(client_id, params.identifier)
+  bufstate.client_result_id[key] = result.resultId
 
   if result.kind == 'unchanged' then
     return
   end
 
-  ---@type lsp.DocumentDiagnosticParams
-  local params = ctx.params
   handle_diagnostics(params.textDocument.uri, client_id, result.items, params.identifier or true)
 
   for uri, related_result in pairs(result.relatedDocuments or {}) do
@@ -304,7 +314,7 @@ function M.on_diagnostic(error, result, ctx)
       or { pull_kind = 'document', client_result_id = {} }
     bufstates[related_bufnr] = related_bufstate
 
-    related_bufstate.client_result_id[client_id] = related_result.resultId
+    related_bufstate.client_result_id[key] = related_result.resultId
   end
 end
 
@@ -381,12 +391,19 @@ function M._refresh(bufnr, client_id, only_visible)
     type = 'pending',
   })
   for _, client in ipairs(clients) do
-    ---@type lsp.DocumentDiagnosticParams
-    local params = {
-      textDocument = util.make_text_document_params(bufnr),
-      previousResultId = bufstate.client_result_id[client.id],
-    }
-    client:request(method, params, nil, bufnr)
+    ---@param cap lsp.DiagnosticRegistrationOptions
+    client:_provider_foreach(method, function(cap)
+      local key = result_id_key(client.id, cap.identifier)
+      ---@type lsp.DocumentDiagnosticParams
+      local params = {
+        identifier = cap.identifier,
+        textDocument = util.make_text_document_params(bufnr),
+        previousResultId = bufstate
+          and bufstate.client_result_id
+          and bufstate.client_result_id[key],
+      }
+      client:request(method, params, nil, bufnr)
+    end)
   end
 end
 
@@ -477,19 +494,20 @@ end
 
 --- Returns the result IDs from the reports provided by the given client.
 --- @return lsp.PreviousResultId[]
-local function previous_result_ids(client_id)
+--- @param client_id integer
+--- @param identifier string|nil
+local function previous_result_ids(client_id, identifier)
   local results = {} ---@type lsp.PreviousResultId[]
+  local key = result_id_key(client_id, identifier)
 
   for bufnr, state in pairs(bufstates) do
     if state.pull_kind ~= 'disabled' then
-      for buf_client_id, result_id in pairs(state.client_result_id) do
-        if buf_client_id == client_id then
-          results[#results + 1] = {
-            uri = vim.uri_from_bufnr(bufnr),
-            value = result_id,
-          }
-          break
-        end
+      local result_id = state.client_result_id[key]
+      if result_id then
+        results[#results + 1] = {
+          uri = vim.uri_from_bufnr(bufnr),
+          value = result_id,
+        }
       end
     end
   end
@@ -531,23 +549,24 @@ function M._workspace_diagnostics(opts)
         -- state if we're not pulling document diagnostics for this buffer.
         if bufstates[bufnr].pull_kind == 'workspace' and report.kind == 'full' then
           handle_diagnostics(report.uri, ctx.client_id, report.items, params.identifier or true)
-          bufstates[bufnr].client_result_id[ctx.client_id] = report.resultId
+          local key = result_id_key(ctx.client_id, params.identifier)
+          bufstates[bufnr].client_result_id[key] = report.resultId
         end
       end
     end
   end
 
   for _, client in ipairs(clients) do
-    local identifiers = client:_provider_value_get('workspace/diagnostic', 'identifier')
-    for _, id in ipairs(identifiers) do
+    ---@param cap lsp.DiagnosticRegistrationOptions
+    client:_provider_foreach('workspace/diagnostic', function(cap)
       --- @type lsp.WorkspaceDiagnosticParams
       local params = {
-        identifier = type(id) == 'string' and id or nil,
-        previousResultIds = previous_result_ids(client.id),
+        identifier = cap.identifier,
+        previousResultIds = previous_result_ids(client.id, cap.identifier),
       }
 
       client:request('workspace/diagnostic', params, handler)
-    end
+    end)
   end
 end