fix(lsp): only disable inlay hints / diagnostics if no other clients are connected (#24535)

This fixes the issue where the LspNotify handlers for inlay_hint /
diagnostics would end up refreshing all attached clients.

The handler would call util._refresh, which called
vim.lsp.buf_request, which calls the method on all attached clients.

Now util._refresh takes an optional client_id parameter, which is used
to specify a specific client to update.

This commit also fixes util._refresh's handling of the `only_visible`
flag. Previously if `only_visible` was false, two requests would be made
to the server: one for the visible region, and one for the entire file.

Co-authored-by: Stanislav Asunkin <1353637+stasjok@users.noreply.github.com>
Co-authored-by: Mathias Fußenegger <mfussenegger@users.noreply.github.com>

6 files changed, 239 insertions, 29 deletions

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 2a16bafbfc..1990c09561 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -2118,7 +2118,7 @@ api.nvim_create_autocmd('VimLeavePre', {
 ---@param bufnr (integer) Buffer handle, or 0 for current.
 ---@param method (string) LSP method name
 ---@param params table|nil Parameters to send to the server
----@param handler lsp-handler See |lsp-handler|
+---@param handler? lsp-handler See |lsp-handler|
 ---       If nil, follows resolution strategy defined in |lsp-handler-configuration|
 ---
 ---@return table<integer, integer> client_request_ids Map of client-id:request-id pairs
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index a0568bc09c..2a77992c4d 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -408,6 +408,16 @@ local function disable(bufnr)
   clear(bufnr)
 end
 
+--- Refresh diagnostics, only if we have attached clients that support it
+---@param bufnr (integer) buffer number
+---@param opts? table Additional options to pass to util._refresh
+---@private
+local function _refresh(bufnr, opts)
+  opts = opts or {}
+  opts['bufnr'] = bufnr
+  util._refresh(ms.textDocument_diagnostic, opts)
+end
+
 --- Enable pull diagnostics for a buffer
 ---@param bufnr (integer) Buffer handle, or 0 for current
 ---@private
@@ -429,7 +439,7 @@ function M._enable(bufnr)
           return
         end
         if bufstates[bufnr] and bufstates[bufnr].enabled then
-          util._refresh(ms.textDocument_diagnostic, { bufnr = bufnr, only_visible = true })
+          _refresh(bufnr, { only_visible = true, client_id = opts.data.client_id })
         end
       end,
       group = augroup,
@@ -438,7 +448,7 @@ function M._enable(bufnr)
     api.nvim_buf_attach(bufnr, false, {
       on_reload = function()
         if bufstates[bufnr] and bufstates[bufnr].enabled then
-          util._refresh(ms.textDocument_diagnostic, { bufnr = bufnr })
+          _refresh(bufnr)
         end
       end,
       on_detach = function()
@@ -448,8 +458,16 @@ function M._enable(bufnr)
 
     api.nvim_create_autocmd('LspDetach', {
       buffer = bufnr,
-      callback = function()
-        disable(bufnr)
+      callback = function(args)
+        local clients = vim.lsp.get_clients({ bufnr = bufnr, method = ms.textDocument_diagnostic })
+
+        if
+          not vim.iter(clients):any(function(c)
+            return c.id ~= args.data.client_id
+          end)
+        then
+          disable(bufnr)
+        end
       end,
       group = augroup,
     })
diff --git a/runtime/lua/vim/lsp/inlay_hint.lua b/runtime/lua/vim/lsp/inlay_hint.lua
index 8407105d47..7b58188c53 100644
--- a/runtime/lua/vim/lsp/inlay_hint.lua
+++ b/runtime/lua/vim/lsp/inlay_hint.lua
@@ -131,6 +131,16 @@ local function disable(bufnr)
   end
 end
 
+--- Refresh inlay hints, only if we have attached clients that support it
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@param opts? table Additional options to pass to util._refresh
+---@private
+local function _refresh(bufnr, opts)
+  opts = opts or {}
+  opts['bufnr'] = bufnr
+  util._refresh(ms.textDocument_inlayHint, opts)
+end
+
 --- Enable inlay hints for a buffer
 ---@param bufnr (integer) Buffer handle, or 0 for current
 local function enable(bufnr)
@@ -150,18 +160,18 @@ local function enable(bufnr)
           return
         end
         if bufstates[bufnr] and bufstates[bufnr].enabled then
-          util._refresh(ms.textDocument_inlayHint, { bufnr = bufnr })
+          _refresh(bufnr, { client_id = opts.data.client_id })
         end
       end,
       group = augroup,
     })
-    util._refresh(ms.textDocument_inlayHint, { bufnr = bufnr })
+    _refresh(bufnr)
     api.nvim_buf_attach(bufnr, false, {
       on_reload = function(_, cb_bufnr)
         clear(cb_bufnr)
         if bufstates[cb_bufnr] and bufstates[cb_bufnr].enabled then
           bufstates[cb_bufnr].applied = {}
-          util._refresh(ms.textDocument_inlayHint, { bufnr = cb_bufnr })
+          _refresh(cb_bufnr)
         end
       end,
       on_detach = function(_, cb_bufnr)
@@ -170,14 +180,22 @@ local function enable(bufnr)
     })
     api.nvim_create_autocmd('LspDetach', {
       buffer = bufnr,
-      callback = function()
-        disable(bufnr)
+      callback = function(args)
+        local clients = vim.lsp.get_clients({ bufnr = bufnr, method = ms.textDocument_inlayHint })
+
+        if
+          not vim.iter(clients):any(function(c)
+            return c.id ~= args.data.client_id
+          end)
+        then
+          disable(bufnr)
+        end
       end,
       group = augroup,
     })
   else
     bufstate.enabled = true
-    util._refresh(ms.textDocument_inlayHint, { bufnr = bufnr })
+    _refresh(bufnr)
   end
 end
 
diff --git a/runtime/lua/vim/lsp/util.lua b/runtime/lua/vim/lsp/util.lua
index 633cddca2d..2e376f9093 100644
--- a/runtime/lua/vim/lsp/util.lua
+++ b/runtime/lua/vim/lsp/util.lua
@@ -2122,6 +2122,7 @@ end
 function M.make_workspace_params(added, removed)
   return { event = { added = added, removed = removed } }
 end
+
 --- Returns indentation size.
 ---
 ---@see 'shiftwidth'
@@ -2192,32 +2193,46 @@ end
 ---@private
 --- Request updated LSP information for a buffer.
 ---
+---@class lsp.util.RefreshOptions
+---@field bufnr integer? Buffer to refresh (default: 0)
+---@field only_visible? boolean Whether to only refresh for the visible regions of the buffer (default: false)
+---@field client_id? integer Client ID to refresh (default: all clients)
+--
 ---@param method string LSP method to call
----@param opts (nil|table) Optional arguments
----  - bufnr (integer, default: 0): Buffer to refresh
----  - only_visible (boolean, default: false): Whether to only refresh for the visible regions of the buffer
+---@param opts? lsp.util.RefreshOptions Options table
 function M._refresh(method, opts)
   opts = opts or {}
   local bufnr = opts.bufnr
   if bufnr == nil or bufnr == 0 then
     bufnr = api.nvim_get_current_buf()
   end
+
+  local clients = vim.lsp.get_clients({ bufnr = bufnr, method = method, id = opts.client_id })
+
+  if #clients == 0 then
+    return
+  end
+
   local only_visible = opts.only_visible or false
-  for _, window in ipairs(api.nvim_list_wins()) do
-    if api.nvim_win_get_buf(window) == bufnr then
-      local first = vim.fn.line('w0', window)
-      local last = vim.fn.line('w$', window)
-      local params = {
-        textDocument = M.make_text_document_params(bufnr),
-        range = {
-          start = { line = first - 1, character = 0 },
-          ['end'] = { line = last, character = 0 },
-        },
-      }
-      vim.lsp.buf_request(bufnr, method, params)
+
+  if only_visible then
+    for _, window in ipairs(api.nvim_list_wins()) do
+      if api.nvim_win_get_buf(window) == bufnr then
+        local first = vim.fn.line('w0', window)
+        local last = vim.fn.line('w$', window)
+        local params = {
+          textDocument = M.make_text_document_params(bufnr),
+          range = {
+            start = { line = first - 1, character = 0 },
+            ['end'] = { line = last, character = 0 },
+          },
+        }
+        for _, client in ipairs(clients) do
+          client.request(method, params, nil, bufnr)
+        end
+      end
     end
-  end
-  if not only_visible then
+  else
     local params = {
       textDocument = M.make_text_document_params(bufnr),
       range = {
@@ -2225,7 +2240,9 @@ function M._refresh(method, opts)
         ['end'] = { line = api.nvim_buf_line_count(bufnr), character = 0 },
       },
     }
-    vim.lsp.buf_request(bufnr, method, params)
+    for _, client in ipairs(clients) do
+      client.request(method, params, nil, bufnr)
+    end
   end
 end
 
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index d1c3fd6b1e..1da0222114 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -84,6 +84,7 @@ describe('vim.lsp.diagnostic', function()
       local lines = {"1st line of text", "2nd line of text", "wow", "cool", "more", "lines"}
       vim.fn.bufload(diagnostic_bufnr)
       vim.api.nvim_buf_set_lines(diagnostic_bufnr, 0, 1, false, lines)
+      vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
       return diagnostic_bufnr
     ]], fake_uri)
   end)
@@ -360,5 +361,63 @@ describe('vim.lsp.diagnostic', function()
       eq(2, #extmarks)
       eq(expected_spacing, #extmarks[1][4].virt_text[1][1])
     end)
+
+    it('clears diagnostics when client detaches', function()
+      exec_lua([[
+        vim.lsp.diagnostic.on_diagnostic(nil,
+          {
+            kind = 'full',
+            items = {
+              make_error('Pull Diagnostic', 4, 4, 4, 4),
+            }
+          },
+          {
+            params = {
+              textDocument = { uri = fake_uri },
+            },
+            uri = fake_uri,
+            client_id = client_id,
+          },
+          {}
+        )
+      ]])
+      local diags = exec_lua([[return vim.diagnostic.get(diagnostic_bufnr)]])
+      eq(1, #diags)
+
+      exec_lua([[ vim.lsp.stop_client(client_id) ]])
+
+      diags = exec_lua([[return vim.diagnostic.get(diagnostic_bufnr)]])
+      eq(0, #diags)
+    end)
+
+    it('keeps diagnostics when one client detaches and others still are attached', function()
+      exec_lua([[
+        client_id2 = vim.lsp.start({ name = 'dummy2', cmd = server.cmd })
+
+        vim.lsp.diagnostic.on_diagnostic(nil,
+          {
+            kind = 'full',
+            items = {
+              make_error('Pull Diagnostic', 4, 4, 4, 4),
+            }
+          },
+          {
+            params = {
+              textDocument = { uri = fake_uri },
+            },
+            uri = fake_uri,
+            client_id = client_id,
+          },
+          {}
+        )
+      ]])
+      local diags = exec_lua([[return vim.diagnostic.get(diagnostic_bufnr)]])
+      eq(1, #diags)
+
+      exec_lua([[ vim.lsp.stop_client(client_id2) ]])
+
+      diags = exec_lua([[return vim.diagnostic.get(diagnostic_bufnr)]])
+      eq(1, #diags)
+    end)
   end)
 end)
diff --git a/test/functional/plugin/lsp/inlay_hint_spec.lua b/test/functional/plugin/lsp/inlay_hint_spec.lua
index b19f2ba146..eec86fdb8e 100644
--- a/test/functional/plugin/lsp/inlay_hint_spec.lua
+++ b/test/functional/plugin/lsp/inlay_hint_spec.lua
@@ -135,5 +135,103 @@ describe('inlay hints', function()
           unchanged = true
         })
       end)
+
+    it(
+      'inlay hints are cleared when the client detaches',
+      function()
+        exec_lua([[
+          bufnr = vim.api.nvim_get_current_buf()
+          vim.api.nvim_win_set_buf(0, bufnr)
+          client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+        ]])
+
+        insert(text)
+        exec_lua([[vim.lsp.inlay_hint(bufnr, true)]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b)-> int { return a + b; }    |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(a: x,b: y);                        |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]]
+        })
+        exec_lua([[vim.lsp.stop_client(client_id)]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b) { return a + b; }          |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(x,y);                              |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]],
+          unchanged = true
+        })
+      end)
+
+    it(
+      'inlay hints are not cleared when one of several clients detaches',
+      function()
+        -- Start two clients
+        exec_lua([[
+          bufnr = vim.api.nvim_get_current_buf()
+          vim.api.nvim_win_set_buf(0, bufnr)
+          server2 = _create_server({
+            capabilities = {
+              inlayHintProvider = true,
+            },
+            handlers = {
+              ['textDocument/inlayHint'] = function()
+                return {}
+              end,
+            }
+          })
+          client1 = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+          client2 = vim.lsp.start({ name = 'dummy2', cmd = server2.cmd })
+        ]])
+
+        insert(text)
+        exec_lua([[vim.lsp.inlay_hint(bufnr, true)]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b)-> int { return a + b; }    |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(a: x,b: y);                        |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]]
+        })
+
+        -- Now stop one client
+        exec_lua([[ vim.lsp.stop_client(client2) ]])
+
+        -- We should still see the hints
+        screen:expect({
+          grid = [[
+  auto add(int a, int b)-> int { return a + b; }    |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(a: x,b: y);                        |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]],
+          unchanged = true
+        })
+      end)
   end)
 end)