Track lsp death more explicitly.

Fixes a nightly race condition on LSP detach and hopefully makes
us less flaky in general.

Author: Julian

lua/lean/infoview.lua

@@ -118,12 +118,51 @@ local autoopen_fn = nil
 
 local FOCUS_AUGROUP = vim.api.nvim_create_augroup('LeanInfoviewFocus', {})
 
----URIs whose language server has died (and not yet been restarted).
+---Track death of the Lean language server in each buffer it was attached to.
 ---
----Consulted when rendering so that asynchronous updates racing with the
----server's death can't resurrect stale contents from it.
----@type table<string, true>
-local dead_uris = {}
+---Sets a buffer-local flag which is consulted at render time (rather than
+---eagerly mutating pin contents here) so that asynchronous updates racing
+---with the server's death can't resurrect stale contents from it.
+local DEATH_AUGROUP = vim.api.nvim_create_augroup('LeanInfoviewDeath', {})
+
+---The `leanls` client for an autocmd's client ID, if that's who it is.
+---@param args table autocmd callback args
+---@return vim.lsp.Client?
+local function leanls_from(args)
+  local client = vim.lsp.get_client_by_id(args.data.client_id)
+  return client and client.name == 'leanls' and client or nil
+end
+
+vim.api.nvim_create_autocmd('LspDetach', {
+  group = DEATH_AUGROUP,
+  callback = function(args)
+    if not leanls_from(args) then
+      return
+    end
+    vim.b[args.buf].lean_lsp_died = true
+    local uri = vim.uri_from_bufnr(args.buf)
+    -- The server certainly isn't processing anything anymore.
+    progress.proc_infos[uri] = {}
+    -- LspDetach can fire in contexts where text changes aren't allowed
+    -- (e.g. during `:edit`), so delay the re-render until it's safe.
+    vim.schedule(function()
+      infoview.__update_pin_by_uri(uri)
+    end)
+  end,
+})
+
+vim.api.nvim_create_autocmd('LspAttach', {
+  group = DEATH_AUGROUP,
+  callback = function(args)
+    if not leanls_from(args) then
+      return
+    end
+    if vim.b[args.buf].lean_lsp_died then
+      vim.b[args.buf].lean_lsp_died = nil
+      infoview.__update_pin_by_uri(vim.uri_from_bufnr(args.buf))
+    end
+  end,
+})
 
 ---@class InfoviewViewOptions
 ---@field use_widgets boolean use interactive widgets (true) or plain text (false)
@@ -790,7 +829,8 @@ end
 ---@param sw std.Stopwatch
 ---@return Element
 function Infoview:render_contents(params, sw)
-  if dead_uris[params.textDocument.uri] then
+  local bufnr = vim.uri_to_bufnr(params.textDocument.uri)
+  if vim.b[bufnr].lean_lsp_died then
     return components.LSP_HAS_DIED
   end
   return self.__contents_for(params, self.view_options, sw)
@@ -941,31 +981,6 @@ function Infoview:__update()
   self.pin:request_update()
 end
 
----Directly mark that the infoview has died. What a shame.
-function Infoview:died()
-  local params = self.pin.__position_params
-  progress.proc_infos[params.textDocument.uri] = {
-    {
-      kind = progress.Kind.fatal_error,
-      range = { start = params.position, ['end'] = params.position },
-    },
-  }
-  -- Invalidate any in-flight (async) update so it can't land after us and
-  -- resurrect stale contents from the now-dead server.
-  self.pin.__tick = self.pin.__tick + 1
-  self.pin.loading = false
-  self.pin.__data_element = components.LSP_HAS_DIED
-  self.pin.__element:set_children { self.pin.__data_element }
-  -- LspDetach can fire in contexts where text changes aren't allowed
-  -- (e.g. during `:edit`), so delay the re-render until it's safe.
-  vim.schedule(function()
-    self.pin:render()
-  end)
-  if self.window then
-    self.window.o.winhighlight = 'NormalNC:leanInfoLSPDead'
-  end
-end
-
 ---Either open or close a diff window for this infoview depending on whether it has a diff pin.
 function Infoview:__refresh_diff()
   if not self.window then
@@ -1940,27 +1955,6 @@ function infoview.init(bufnr)
   -- in case we are re-entering a buffer, clear old autocmds first
   vim.api.nvim_clear_autocmds { group = FOCUS_AUGROUP, buffer = bufnr }
 
-  vim.api.nvim_create_autocmd('LspDetach', {
-    callback = function(args)
-      dead_uris[vim.uri_from_bufnr(args.buf)] = true
-      local current_infoview = infoview.get_current_infoview()
-      if not current_infoview then
-        return
-      end
-      current_infoview:died()
-    end,
-    buffer = bufnr,
-    group = FOCUS_AUGROUP,
-  })
-
-  vim.api.nvim_create_autocmd('LspAttach', {
-    callback = function(args)
-      dead_uris[vim.uri_from_bufnr(args.buf)] = nil
-    end,
-    buffer = bufnr,
-    group = FOCUS_AUGROUP,
-  })
-
   vim.api.nvim_create_autocmd('BufEnter', {
     callback = infoview_bufenter,
     buffer = bufnr,

spec/infoview/lsp_dead_spec.lua

@@ -22,6 +22,7 @@ describe('language server dead', function()
     end)
     assert.message("Couldn't kill the LSP!").is_true(succeeded)
 
+    wait:for_infoview_contents '🪦'
     assert.infoview_contents_nowait.are '🪦 The Lean language server is dead.'
     assert.are.same(
       'NormalNC:leanInfoLSPDead',
@@ -30,7 +31,7 @@ describe('language server dead', function()
 
     -- and comes back alive
     vim.cmd.edit()
-    wait:for_processing()
+    wait:for_infoview_contents 'information'
     assert.infoview_contents.are [[
       ▼ 1:1-1:6: information:
       1
@@ -56,6 +57,7 @@ describe('language server dead', function()
       return #vim.lsp.get_clients {} == 1
     end)
     assert.message("Couldn't kill the LSP!").is_true(succeeded)
+    wait:for_infoview_contents '🪦'
     assert.infoview_contents_nowait.are '🪦 The Lean language server is dead.'
 
     -- still alive for the other window
@@ -66,6 +68,7 @@ describe('language server dead', function()
     ]]
 
     vim.cmd.wincmd 'p'
+    wait:for_infoview_contents '🪦'
     assert.infoview_contents_nowait.are '🪦 The Lean language server is dead.'
   end)
 
@@ -85,6 +88,7 @@ describe('language server dead', function()
     assert.message("Couldn't kill the LSP!").is_true(succeeded)
 
     infoview.open()
+    wait:for_infoview_contents '🪦'
     assert.infoview_contents_nowait.are '🪦 The Lean language server is dead.'
   end)
 end)

spec/infoview/plain_contents_spec.lua

@@ -455,6 +455,7 @@ describe('plain infoviews', function()
         return vim.tbl_isempty(vim.lsp.get_clients { bufnr = 0 })
       end)
       assert.message("Couldn't kill the LSP!").is_true(succeeded)
+      helpers.wait:for_infoview_contents '🪦'
       assert.infoview_contents_nowait.are '🪦 The Lean language server is dead.'
     end)
   end)