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

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

3 files changed, 105 insertions, 1 deletions

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