fix(lsp): retrigger diagnostics request on server cancellation (#31345)

Co-authored-by: Jesse <github@jessebakker.com>

7 files changed, 120 insertions, 5 deletions

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index 1253c01547..350edc068f 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -1534,12 +1534,13 @@ get_namespace({client_id}, {is_pull})
                      client. Defaults to push
 
                                           *vim.lsp.diagnostic.on_diagnostic()*
-on_diagnostic({_}, {result}, {ctx})
+on_diagnostic({error}, {result}, {ctx})
     |lsp-handler| for the method "textDocument/diagnostic"
 
     See |vim.diagnostic.config()| for configuration options.
 
     Parameters: ~
+      • {error}   (`lsp.ResponseError?`)
       • {result}  (`lsp.DocumentDiagnosticReport`)
       • {ctx}     (`lsp.HandlerContext`)
 
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index 12984f8c26..3f135d84f3 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -246,10 +246,18 @@ end
 ---
 --- See |vim.diagnostic.config()| for configuration options.
 ---
----@param _ lsp.ResponseError?
+---@param error lsp.ResponseError?
 ---@param result lsp.DocumentDiagnosticReport
 ---@param ctx lsp.HandlerContext
-function M.on_diagnostic(_, result, ctx)
+function M.on_diagnostic(error, result, ctx)
+  if error ~= nil and error.code == protocol.ErrorCodes.ServerCancelled then
+    if error.data == nil or error.data.retriggerRequest ~= false then
+      local client = assert(vim.lsp.get_client_by_id(ctx.client_id))
+      client:request(ctx.method, ctx.params)
+    end
+    return
+  end
+
   if result == nil or result.kind == 'unchanged' then
     return
   end
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index 5c28d88b38..1945040bda 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -659,7 +659,8 @@ for k, fn in pairs(M) do
       })
     end
 
-    if err then
+    -- ServerCancelled errors should be propagated to the request handler
+    if err and err.code ~= protocol.ErrorCodes.ServerCancelled then
       -- LSP spec:
       -- interface ResponseError:
       --  code: integer;
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 7db48b0c06..3d29dad90a 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -174,6 +174,7 @@ local constants = {
     -- Defined by the protocol.
     RequestCancelled = -32800,
     ContentModified = -32801,
+    ServerCancelled = -32802,
   },
 
   -- Describes the content type that a client supports in various
diff --git a/test/functional/fixtures/fake-lsp-server.lua b/test/functional/fixtures/fake-lsp-server.lua
index f813927f77..5d7ab2ad12 100644
--- a/test/functional/fixtures/fake-lsp-server.lua
+++ b/test/functional/fixtures/fake-lsp-server.lua
@@ -386,6 +386,21 @@ function tests.check_forward_content_modified()
   }
 end
 
+function tests.check_forward_server_cancelled()
+  skeleton {
+    on_init = function()
+      return { capabilities = {} }
+    end,
+    body = function()
+      expect_request('error_code_test', function()
+        return { code = -32802 }, nil, { method = 'error_code_test', client_id = 1 }
+      end)
+      expect_notification('finish')
+      notify('finish')
+    end,
+  }
+end
+
 function tests.check_pending_request_tracked()
   skeleton {
     on_init = function(_)
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index 5afbe22793..ca9196562c 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -209,10 +209,16 @@ describe('vim.lsp.diagnostic', function()
     before_each(function()
       exec_lua(create_server_definition)
       exec_lua(function()
+        _G.requests = 0
         _G.server = _G._create_server({
           capabilities = {
             diagnosticProvider = {},
           },
+          handlers = {
+            [vim.lsp.protocol.Methods.textDocument_diagnostic] = function()
+              _G.requests = _G.requests + 1
+            end,
+          },
         })
 
         function _G.get_extmarks(bufnr, client_id0)
@@ -373,5 +379,56 @@ describe('vim.lsp.diagnostic', function()
         end)
       )
     end)
+
+    it('handles server cancellation', function()
+      eq(
+        1,
+        exec_lua(function()
+          vim.lsp.diagnostic.on_diagnostic({
+            code = vim.lsp.protocol.ErrorCodes.ServerCancelled,
+            -- Empty data defaults to retriggering request
+            data = {},
+            message = '',
+          }, {}, {
+            method = vim.lsp.protocol.Methods.textDocument_diagnostic,
+            client_id = client_id,
+          })
+
+          return _G.requests
+        end)
+      )
+
+      eq(
+        2,
+        exec_lua(function()
+          vim.lsp.diagnostic.on_diagnostic({
+            code = vim.lsp.protocol.ErrorCodes.ServerCancelled,
+            data = { retriggerRequest = true },
+            message = '',
+          }, {}, {
+            method = vim.lsp.protocol.Methods.textDocument_diagnostic,
+            client_id = client_id,
+          })
+
+          return _G.requests
+        end)
+      )
+
+      eq(
+        2,
+        exec_lua(function()
+          vim.lsp.diagnostic.on_diagnostic({
+            code = vim.lsp.protocol.ErrorCodes.ServerCancelled,
+            data = { retriggerRequest = false },
+            message = '',
+          }, {}, {
+            method = vim.lsp.protocol.Methods.textDocument_diagnostic,
+            client_id = client_id,
+          })
+
+          return _G.requests
+        end)
+      )
+    end)
   end)
 end)
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 332a1a48bb..e30d1ba411 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -1066,6 +1066,39 @@ describe('LSP', function()
       }
     end)
 
+    it('should forward ServerCancelled to callback', function()
+      local expected_handlers = {
+        { NIL, {}, { method = 'finish', client_id = 1 } },
+        {
+          { code = -32802 },
+          NIL,
+          { method = 'error_code_test', bufnr = 1, client_id = 1, version = 0 },
+        },
+      }
+      local client --- @type vim.lsp.Client
+      test_rpc_server {
+        test_name = 'check_forward_server_cancelled',
+        on_init = function(_client)
+          _client:request('error_code_test')
+          client = _client
+        end,
+        on_exit = function(code, signal)
+          eq(0, code, 'exit code')
+          eq(0, signal, 'exit signal')
+          eq(0, #expected_handlers, 'did not call expected handler')
+        end,
+        on_handler = function(err, _, ctx)
+          eq(table.remove(expected_handlers), { err, _, ctx }, 'expected handler')
+          if ctx.method ~= 'finish' then
+            client:notify('finish')
+          end
+          if ctx.method == 'finish' then
+            client:stop()
+          end
+        end,
+      }
+    end)
+
     it('should forward ContentModified to callback', function()
       local expected_handlers = {
         { NIL, {}, { method = 'finish', client_id = 1 } },
@@ -1089,7 +1122,6 @@ describe('LSP', function()
         end,
         on_handler = function(err, _, ctx)
           eq(table.remove(expected_handlers), { err, _, ctx }, 'expected handler')
-          -- if ctx.method == 'error_code_test' then client.notify("finish") end
           if ctx.method ~= 'finish' then
             client:notify('finish')
           end