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

2 files changed, 30 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)