fix(picker.lsp): use LspRequest to track completed and cancelled requests. Fixes #2364

lua/snacks/picker/source/lsp/init.lua

@@ -98,27 +98,52 @@ end
 
 ---@class snacks.picker.lsp.Requester
 ---@field async snacks.picker.Async
----@field requests {client_id:number, request_id:number}[]
----@field completed number
+---@field requests table<string, {client_id:number, request_id:number, done:boolean}>
 ---@field pending integer
+---@field autocmd_id? number
 local R = {}
 R.__index = R
+R._id = 0
 
 function R.new()
   local self = setmetatable({}, R)
   self.async = Async.running()
   self.requests = {}
   self.pending = 0
-  self.completed = 0
-  self.async:on(
-    "abort",
-    vim.schedule_wrap(function()
-      self:cancel()
-    end)
-  )
+  R._id = R._id + 1
+
+  self.async
+    :on(
+      "abort",
+      vim.schedule_wrap(function()
+        self:cancel()
+      end)
+    )
+    :on(
+      "done",
+      vim.schedule_wrap(function()
+        pcall(vim.api.nvim_del_autocmd, self.autocmd_id)
+      end)
+    )
   return self
 end
 
+---@param client_id number
+---@param request_id number
+---@param completed? boolean
+function R:track(client_id, request_id, completed)
+  local key = ("%d:%d"):format(client_id, request_id)
+  if completed and self.requests[key] and not self.requests[key].done then
+    self.requests[key].done = true
+    self.pending = self.pending - 1
+    self.async:resume()
+    return
+  elseif not completed then
+    self.requests[key] = { client_id = client_id, request_id = request_id, done = false }
+    self.pending = self.pending + 1
+  end
+end
+
 function R:cancel()
   while #self.requests > 0 do
     local req = table.remove(self.requests)
@@ -129,6 +154,20 @@ function R:cancel()
   end
 end
 
+function R:autocmd()
+  if self.autocmd_id then
+    return
+  end
+  self.autocmd_id = vim.api.nvim_create_autocmd("LspRequest", {
+    group = vim.api.nvim_create_augroup("snacks.picker.lsp." .. R._id, { clear = true }),
+    callback = function(ev)
+      if ev.data.request.type == "complete" or ev.data.request.type == "cancel" then
+        self:track(ev.data.client_id, ev.data.request_id, true)
+      end
+    end,
+  })
+end
+
 ---@param buf number|vim.lsp.Client
 ---@param method string
 ---@param params fun(client:vim.lsp.Client):table
@@ -137,21 +176,23 @@ end
 function R:request(buf, method, params, cb)
   self.pending = self.pending + 1
   vim.schedule(function()
-    local clients = type(buf) == "number" and M.get_clients(buf, method)
-      or {
-        wrap(buf --[[@as vim.lsp.Client]]),
-      }
+    self:autocmd() -- setup autocmd here, since this must be called in the main loop
+
+    ---@diagnostic disable-next-line: param-type-mismatch
+    local clients = type(buf) == "number" and M.get_clients(buf, method) or { wrap(buf) }
+
     for _, client in ipairs(clients) do
-      local p = params(client)
+      local p, done = params(client), false
       local status, request_id = client:request(method, p, function(err, result)
-        if not err and result and not self.async._aborted then
+        done = true
+        if not err and result and not self.async:aborted() then
           cb(client, result, p)
         end
-        self.completed = self.completed + 1
-        self.async:resume()
       end)
-      if status and request_id then
-        table.insert(self.requests, { client_id = client.id, request_id = request_id })
+      -- skip tracking if the request failed
+      -- or is already done (in-process syncronous response)
+      if status and request_id and not done then
+        self:track(client.id, request_id)
       end
     end
     self.pending = self.pending - 1
@@ -161,7 +202,7 @@ function R:request(buf, method, params, cb)
 end
 
 function R:wait()
-  while self.pending > 0 or self.completed < #self.requests do
+  while self.pending > 0 do
     self.async:suspend()
   end
 end