From 883b74f03f18b53e27038c5f5de71a5db3e4f676 Mon Sep 17 00:00:00 2001
From: j-hui <j-hui@cs.columbia.edu>
Date: Mon, 27 Nov 2023 00:45:29 -0500
Subject: [PATCH] feat(progress): Clear notification group on LspDetach

---
 doc/fidget.md           | 29 +++++++++++++++++++++
 lua/fidget/progress.lua | 56 +++++++++++++++++++++++++++++++++++------
 2 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/doc/fidget.md b/doc/fidget.md
index de4e8bd..78cbc7e 100644
--- a/doc/fidget.md
+++ b/doc/fidget.md
@@ -30,6 +30,11 @@ The following table shows the default options for this plugin:
     suppress_on_insert = false,   -- Suppress new messages while in insert mode
     ignore_done_already = false,  -- Ignore new tasks that are already complete
     ignore_empty_message = false, -- Ignore new tasks that don't contain a message
+    clear_on_detach =             -- Clear notification group when LSP server detaches
+      function(client_id)
+        local client = vim.lsp.get_client_by_id(client_id)
+        return client and client.name or nil
+      end,
     notification_group =          -- How to get a progress message's notification group key
       function(msg) return msg.lsp_client.name end,
     ignore = {},                  -- List of LSP servers to ignore
@@ -174,6 +179,30 @@ end
 
 Type: `fun(msg: ProgressMessage): NotificationKey` (default: `msg.lsp_client.name`)
 
+progress.clear_on_detach
+: Clear notification group when LSP server detaches
+
+This option should be set to a function that, given a client ID number,
+returns the notification group to clear. No group will be cleared if the
+the function returns `nil`.
+
+The default setting looks up and returns the LSP client name, which is also used
+by `progress.notification_group`.
+
+Set this option to `nil` to disable this feature entirely (no `LspDetach`
+callback will be registered).
+
+Default value:
+
+```lua
+clear_on_detach = function(client_id)
+  local client = vim.lsp.get_client_by_id(client_id)
+  return client and client.name or nil
+end
+```
+
+Type: `(fun(client_id: number): NotificationKey)?` (default: `client.name`)
+
 progress.ignore
 : List of LSP servers to ignore
 
diff --git a/lua/fidget/progress.lua b/lua/fidget/progress.lua
index 849e7c5..bfdb96f 100644
--- a/lua/fidget/progress.lua
+++ b/lua/fidget/progress.lua
@@ -6,9 +6,8 @@ local poll         = require("fidget.poll")
 local notification = require("fidget.notification")
 local logger       = require("fidget.logger")
 
---- Used to ensure only a single autocmd callback exists.
----@type number?
-local autocmd_id   = nil
+--- Table of progress-related autocmds, used to ensure setup() re-entrancy.
+local autocmds     = {}
 
 --- Options related to LSP progress notification subsystem
 require("fidget.options").declare(M, "progress", {
@@ -81,6 +80,25 @@ require("fidget.options").declare(M, "progress", {
     return msg.lsp_client.name
   end,
 
+  --- Clear notification group when LSP server detaches
+  ---
+  --- This option should be set to a function that, given a client ID number,
+  --- returns the notification group to clear. No group will be cleared if the
+  --- the function returns `nil`.
+  ---
+  --- The default setting looks up and returns the LSP client name, which is
+  --- also used by `progress.notification_group`.
+  ---
+  --- Set this option to `nil` to disable this feature entirely (no `LspDetach`
+  --- callback will be registered).
+  ---
+  ---
+  ---@type (fun(client_id: number): NotificationKey)?
+  clear_on_detach = function(client_id)
+    local client = vim.lsp.get_client_by_id(client_id)
+    return client and client.name or nil
+  end,
+
   --- List of LSP servers to ignore
   ---
   --- Example:
@@ -95,12 +113,14 @@ require("fidget.options").declare(M, "progress", {
   display = M.display,
   lsp = M.lsp,
 }, function()
-  if autocmd_id ~= nil then
-    vim.api.nvim_del_autocmd(autocmd_id)
-    autocmd_id = nil
+  -- Ensure setup() reentrancy
+  for _, autocmd in pairs(autocmds) do
+    vim.api.nvim_del_autocmd(autocmd)
   end
+  autocmds = {}
+
   if M.options.poll_rate ~= false then
-    autocmd_id = M.lsp.on_progress_message(function()
+    autocmds["LspProgress"] = M.lsp.on_progress_message(function()
       if M.options.poll_rate > 0 then
         M.poller:start_polling(M.options.poll_rate)
       else
@@ -108,6 +128,15 @@ require("fidget.options").declare(M, "progress", {
       end
     end)
   end
+
+  if M.options.clear_on_detach then
+    autocmds["LspDetach"] = vim.api.nvim_create_autocmd("LspDetach", {
+      desc = "Fidget LSP detach handler",
+      callback = function(args)
+        M.on_detach(args.data.client_id)
+      end,
+    })
+  end
 end)
 
 --- Whether progress message updates are suppressed.
@@ -210,4 +239,17 @@ function M.suppress(suppress)
   end
 end
 
+--- Called upon `LspDetach` event.
+---
+--- Clears notification group given by `options.clear_on_detach`.
+---
+---@param client_id number
+function M.on_detach(client_id)
+  local group_key = M.options.clear_on_detach(client_id)
+  if group_key == nil then
+    return
+  end
+  notification.clear(group_key)
+end
+
 return M