feat(lsp): track pending+cancel requests on client object #15949

2 files changed, 45 insertions, 7 deletions

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 56ac1cbc66..3a067373d0 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -772,8 +772,10 @@ function lsp.start_client(config)
     attached_buffers = {};
 
     handlers = handlers;
+    requests = {};
+
     -- for $/progress report
-    messages = { name = name, messages = {}, progress = {}, status = {} }
+    messages = { name = name, messages = {}, progress = {}, status = {} };
   }
 
   -- Store the uninitialized_clients for cleanup in case we exit before initialize finishes.
@@ -906,11 +908,21 @@ function lsp.start_client(config)
     end
     -- Ensure pending didChange notifications are sent so that the server doesn't operate on a stale state
     changetracking.flush(client)
-
+    bufnr = resolve_bufnr(bufnr)
     local _ = log.debug() and log.debug(log_prefix, "client.request", client_id, method, params, handler, bufnr)
-    return rpc.request(method, params, function(err, result)
+    local success, request_id = rpc.request(method, params, function(err, result)
       handler(err, result, {method=method, client_id=client_id, bufnr=bufnr, params=params})
+    end, function(request_id)
+      client.requests[request_id] = nil
+      nvim_command("doautocmd <nomodeline> User LspRequest")
     end)
+
+    if success then
+      client.requests[request_id] = { type='pending', bufnr=bufnr, method=method }
+      nvim_command("doautocmd <nomodeline> User LspRequest")
+    end
+
+    return success, request_id
   end
 
   ---@private
@@ -970,6 +982,11 @@ function lsp.start_client(config)
   ---@see |vim.lsp.client.notify()|
   function client.cancel_request(id)
     validate{id = {id, 'n'}}
+    local request = client.requests[id]
+    if request and request.type == 'pending' then
+      request.type = 'cancel'
+      nvim_command("doautocmd <nomodeline> User LspRequest")
+    end
     return rpc.notify("$/cancelRequest", { id = id })
   end
 
diff --git a/runtime/lua/vim/lsp/rpc.lua b/runtime/lua/vim/lsp/rpc.lua
index d9a684a738..bce1e9f35d 100644
--- a/runtime/lua/vim/lsp/rpc.lua
+++ b/runtime/lua/vim/lsp/rpc.lua
@@ -297,6 +297,7 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
 
   local message_index = 0
   local message_callbacks = {}
+  local notify_reply_callbacks = {}
 
   local handle, pid
   do
@@ -309,8 +310,9 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
       stdout:close()
       stderr:close()
       handle:close()
-      -- Make sure that message_callbacks can be gc'd.
+      -- Make sure that message_callbacks/notify_reply_callbacks can be gc'd.
       message_callbacks = nil
+      notify_reply_callbacks = nil
       dispatchers.on_exit(code, signal)
     end
     local spawn_params = {
@@ -375,10 +377,12 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
   ---@param method (string) The invoked LSP method
   ---@param params (table) Parameters for the invoked LSP method
   ---@param callback (function) Callback to invoke
+  ---@param notify_reply_callback (function) Callback to invoke as soon as a request is no longer pending
   ---@returns (bool, number) `(true, message_id)` if request could be sent, `false` if not
-  local function request(method, params, callback)
+  local function request(method, params, callback, notify_reply_callback)
     validate {
       callback = { callback, 'f' };
+      notify_reply_callback = { notify_reply_callback, 'f', true };
     }
     message_index = message_index + 1
     local message_id = message_index
@@ -388,8 +392,15 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
       method = method;
       params = params;
     }
-    if result and message_callbacks then
-      message_callbacks[message_id] = schedule_wrap(callback)
+    if result then
+      if message_callbacks then
+        message_callbacks[message_id] = schedule_wrap(callback)
+      else
+        return false
+      end
+      if notify_reply_callback and notify_reply_callbacks then
+        notify_reply_callbacks[message_id] = schedule_wrap(notify_reply_callback)
+      end
       return result, message_id
     else
       return false
@@ -466,6 +477,16 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
       -- We sent a number, so we expect a number.
       local result_id = tonumber(decoded.id)
 
+      -- Notify the user that a response was received for the request
+      local notify_reply_callback = notify_reply_callbacks and notify_reply_callbacks[result_id]
+      if notify_reply_callback then
+        validate {
+          notify_reply_callback = { notify_reply_callback, 'f' };
+        }
+        notify_reply_callback(result_id)
+        notify_reply_callbacks[result_id] = nil
+      end
+
       -- Do not surface RequestCancelled to users, it is RPC-internal.
       if decoded.error then
         local mute_error = false