lsp: client stop cleanups (#13877)

* lsp: client stop cleanups

* Add diagnostic clearing to client.stop() method used by nvim-lspconfig
* Clear diagnostic cache to prevent stale diagnostics on client restart

* lsp: Add test for vim.lsp.diagnostic.reset

3 files changed, 93 insertions, 22 deletions

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index a6f118abde..0c8d0bdc95 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -548,19 +548,13 @@ function lsp.start_client(config)
   function dispatch.on_exit(code, signal)
     active_clients[client_id] = nil
     uninitialized_clients[client_id] = nil
-    local active_buffers = {}
-    for bufnr, client_ids in pairs(all_buffer_active_clients) do
-      if client_ids[client_id] then
-        table.insert(active_buffers, bufnr)
-      end
+
+    lsp.diagnostic.reset(client_id, all_buffer_active_clients)
+    all_client_active_buffers[client_id] = nil
+    for _, client_ids in pairs(all_buffer_active_clients) do
       client_ids[client_id] = nil
     end
-    -- Buffer level cleanup
-    vim.schedule(function()
-      for _, bufnr in ipairs(active_buffers) do
-        lsp.diagnostic.clear(bufnr)
-      end
-    end)
+
     if config.on_exit then
       pcall(config.on_exit, code, signal, client_id)
     end
@@ -751,6 +745,13 @@ function lsp.start_client(config)
   ---
   --@param force (bool, optional)
   function client.stop(force)
+
+    lsp.diagnostic.reset(client_id, all_buffer_active_clients)
+    all_client_active_buffers[client_id] = nil
+    for _, client_ids in pairs(all_buffer_active_clients) do
+      client_ids[client_id] = nil
+    end
+
     local handle = rpc.handle
     if handle:is_closing() then
       return
@@ -1016,23 +1017,12 @@ end
 function lsp.stop_client(client_id, force)
   local ids = type(client_id) == 'table' and client_id or {client_id}
   for _, id in ipairs(ids) do
-    local resolved_client_id
     if type(id) == 'table' and id.stop ~= nil then
       id.stop(force)
-      resolved_client_id = id.id
     elseif active_clients[id] then
       active_clients[id].stop(force)
-      resolved_client_id = id
     elseif uninitialized_clients[id] then
       uninitialized_clients[id].stop(true)
-      resolved_client_id = id
-    end
-    if resolved_client_id then
-      local client_buffers = lsp.get_buffers_by_client_id(resolved_client_id)
-      for idx = 1, #client_buffers do
-        lsp.diagnostic.clear(client_buffers[idx], resolved_client_id)
-      end
-      all_client_active_buffers[resolved_client_id] = nil
     end
   end
 end
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index a625098bab..df25943ecd 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -1158,6 +1158,24 @@ local loclist_type_map = {
   [DiagnosticSeverity.Hint] = 'I',
 }
 
+
+--- Clear diagnotics and diagnostic cache
+---
+--- Handles saving diagnostics from multiple clients in the same buffer.
+---@param client_id number
+---@param buffer_client_map table map of buffers to active clients
+function M.reset(client_id, buffer_client_map)
+  buffer_client_map = vim.deepcopy(buffer_client_map)
+  vim.schedule(function()
+    for bufnr, client_ids in pairs(buffer_client_map) do
+      if client_ids[client_id] then
+        clear_diagnostic_cache(bufnr, client_id)
+        M.clear(bufnr, client_id)
+      end
+    end
+  end)
+end
+
 --- Sets the location list
 ---@param opts table|nil Configuration table. Keys:
 ---         - {open_loclist}: (boolean, default true)
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index 4705a76465..8c91c4ab2c 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -208,6 +208,69 @@ describe('vim.lsp.diagnostic', function()
         ]]))
       end)
 
+      describe('reset', function()
+        it('diagnostic count is 0 and displayed diagnostics are 0 after call', function()
+          -- 1 Error (1)
+          -- 1 Warning (2)
+          -- 1 Warning (2) + 1 Warning (1)
+          -- 2 highlights and 2 underlines (since error)
+          -- 1 highlight + 1 underline
+          local all_highlights = {1, 1, 2, 4, 2}
+          eq(all_highlights, exec_lua [[
+            local server_1_diags = {
+              make_error("Error 1", 1, 1, 1, 5),
+              make_warning("Warning on Server 1", 2, 1, 2, 5),
+            }
+            local server_2_diags = {
+              make_warning("Warning 1", 2, 1, 2, 5),
+            }
+
+            vim.lsp.diagnostic.on_publish_diagnostics(nil, nil, { uri = fake_uri, diagnostics = server_1_diags }, 1)
+            vim.lsp.diagnostic.on_publish_diagnostics(nil, nil, { uri = fake_uri, diagnostics = server_2_diags }, 2)
+            return {
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
+              count_of_extmarks_for_client(diagnostic_bufnr, 1),
+              count_of_extmarks_for_client(diagnostic_bufnr, 2),
+            }
+          ]])
+
+          -- Reset diagnostics from server 1
+          exec_lua([[ vim.lsp.diagnostic.reset(1, { [ diagnostic_bufnr ] = { [ 1 ] = true ; [ 2 ] = true } } )]])
+
+          -- Make sure we have the right diagnostic count
+          eq({0, 1, 1, 0, 2} , exec_lua [[
+            local diagnostic_count = {}
+            vim.wait(100, function () diagnostic_count = {
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
+              count_of_extmarks_for_client(diagnostic_bufnr, 1),
+              count_of_extmarks_for_client(diagnostic_bufnr, 2),
+            } end )
+            return diagnostic_count
+          ]])
+
+          -- Reset diagnostics from server 2
+          exec_lua([[ vim.lsp.diagnostic.reset(2, { [ diagnostic_bufnr ] = { [ 1 ] = true ; [ 2 ] = true } } )]])
+
+          -- Make sure we have the right diagnostic count
+          eq({0, 0, 0, 0, 0}, exec_lua [[
+            local diagnostic_count = {}
+            vim.wait(100, function () diagnostic_count = {
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
+              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
+              count_of_extmarks_for_client(diagnostic_bufnr, 1),
+              count_of_extmarks_for_client(diagnostic_bufnr, 2),
+            } end )
+            return diagnostic_count
+          ]])
+
+          end)
+        end)
+
       describe('get_next_diagnostic_pos', function()
         it('can find the next pos with only one client', function()
           eq({1, 1}, exec_lua [[