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

Co-authored-by: Jesse <github@jessebakker.com>
(cherry picked from commit 29c72cdf4a)

runtime/doc/lsp.txt

@@ -1472,7 +1472,7 @@ get_namespace({client_id}, {is_pull})
                      client. Defaults to push
 
                                           *vim.lsp.diagnostic.on_diagnostic()*
-on_diagnostic({_}, {result}, {ctx}, {config})
+on_diagnostic({error}, {result}, {ctx}, {config})
     |lsp-handler| for the method "textDocument/diagnostic"
 
     See |vim.diagnostic.config()| for configuration options. Handler-specific
@@ -1497,6 +1497,7 @@ on_diagnostic({_}, {result}, {ctx}, {config})
 <
 
     Parameters: ~
+      • {error}   (`lsp.ResponseError?`)
       • {result}  (`lsp.DocumentDiagnosticReport`)
       • {ctx}     (`lsp.HandlerContext`)
       • {config}  (`vim.diagnostic.Opts`) Configuration table (see

runtime/lua/vim/lsp/diagnostic.lua

@@ -315,11 +315,19 @@ end
 --- )
 --- ```
 ---
----@param _ lsp.ResponseError?
+---@param error lsp.ResponseError?
 ---@param result lsp.DocumentDiagnosticReport
 ---@param ctx lsp.HandlerContext
 ---@param config vim.diagnostic.Opts Configuration table (see |vim.diagnostic.config()|).
-function M.on_diagnostic(_, result, ctx, config)
+function M.on_diagnostic(error, result, ctx, config)
+  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

runtime/lua/vim/lsp/handlers.lua

@@ -716,7 +716,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;

runtime/lua/vim/lsp/protocol.lua

@@ -166,6 +166,7 @@ local constants = {
     -- Defined by the protocol.
     RequestCancelled = -32800,
     ContentModified = -32801,
+    ServerCancelled = -32802,
   },
 
   -- Describes the content type that a client supports in various

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(_)

test/functional/plugin/lsp/diagnostic_spec.lua

@@ -264,11 +264,16 @@ describe('vim.lsp.diagnostic', function()
     before_each(function()
       exec_lua(create_server_definition)
       exec_lua([[
+        _G.requests = 0
         server = _create_server({
           capabilities = {
-            diagnosticProvider = {
-            }
-          }
+            diagnosticProvider = {},
+          },
+          handlers = {
+            [vim.lsp.protocol.Methods.textDocument_diagnostic] = function()
+              _G.requests = _G.requests + 1
+            end,
+          },
         })
 
         function get_extmarks(bufnr, client_id)
@@ -409,5 +414,56 @@ describe('vim.lsp.diagnostic', function()
       diags = exec_lua([[return vim.diagnostic.get(diagnostic_bufnr)]])
       eq(1, #diags)
     end)
+
+    it('handles server cancellation', function()
+      eq(
+        1,
+        exec_lua([[
+          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
+        ]])
+      )
+
+      eq(
+        2,
+        exec_lua([[
+          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
+        ]])
+      )
+
+      eq(
+        2,
+        exec_lua([[
+          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)

test/functional/plugin/lsp_spec.lua

@@ -945,6 +945,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 },
+        },
+      }
+      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 } },
@@ -964,7 +997,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