feat(lsp): support willSave & willSaveWaitUntil capability (#21315)

`willSaveWaitUntil` allows servers to respond with text edits before
saving a document. That is used by some language servers to format a
document or apply quick fixes like removing unused imports.

4 files changed, 114 insertions, 13 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 4896cf19d8..047973242f 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -39,6 +39,11 @@ NEW FEATURES                                                    *news-features*
 
 The following new APIs or features were added.
 
+• Added support for the `willSave` and `willSaveWaitUntil` capabilities to the
+  LSP client. `willSaveWaitUntil` allows a server to modify a document before it
+  gets saved. Example use-cases by language servers include removing unused
+  imports, or formatting the file.
+
 • Treesitter syntax highlighting for `help` files now supports highlighted
   code examples. To enable, create a `.config/nvim/ftplugin/help.lua` with
   the contents >lua
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 9595f0b12c..9c42e9df52 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -1611,9 +1611,37 @@ function lsp.buf_attach_client(bufnr, client_id)
     all_buffer_active_clients[bufnr] = buffer_client_ids
 
     local uri = vim.uri_from_bufnr(bufnr)
-    local augroup = ('lsp_c_%d_b_%d_did_save'):format(client_id, bufnr)
+    local augroup = ('lsp_c_%d_b_%d_save'):format(client_id, bufnr)
+    local group = api.nvim_create_augroup(augroup, { clear = true })
+    api.nvim_create_autocmd('BufWritePre', {
+      group = group,
+      buffer = bufnr,
+      desc = 'vim.lsp: textDocument/willSave',
+      callback = function(ctx)
+        for_each_buffer_client(ctx.buf, function(client)
+          local params = {
+            textDocument = {
+              uri = uri,
+            },
+            reason = protocol.TextDocumentSaveReason.Manual,
+          }
+          if vim.tbl_get(client.server_capabilities, 'textDocumentSync', 'willSave') then
+            client.notify('textDocument/willSave', params)
+          end
+          if vim.tbl_get(client.server_capabilities, 'textDocumentSync', 'willSaveWaitUntil') then
+            local result, err =
+              client.request_sync('textDocument/willSaveWaitUntil', params, 1000, ctx.buf)
+            if result and result.result then
+              util.apply_text_edits(result.result, ctx.buf, client.offset_encoding)
+            elseif err then
+              log.error(vim.inspect(err))
+            end
+          end
+        end)
+      end,
+    })
     api.nvim_create_autocmd('BufWritePost', {
-      group = api.nvim_create_augroup(augroup, { clear = true }),
+      group = group,
       buffer = bufnr,
       desc = 'vim.lsp: textDocument/didSave handler',
       callback = function(ctx)
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 8dc93b3b67..925115d056 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -151,6 +151,7 @@ local constants = {
   },
 
   -- Represents reasons why a text document is saved.
+  ---@enum lsp.TextDocumentSaveReason
   TextDocumentSaveReason = {
     -- Manually triggered, e.g. by the user pressing save, by starting debugging,
     -- or by an API call.
@@ -631,11 +632,8 @@ function protocol.make_client_capabilities()
       synchronization = {
         dynamicRegistration = false,
 
-        -- TODO(ashkan) Send textDocument/willSave before saving (BufWritePre)
-        willSave = false,
-
-        -- TODO(ashkan) Implement textDocument/willSaveWaitUntil
-        willSaveWaitUntil = false,
+        willSave = true,
+        willSaveWaitUntil = true,
 
         -- Send textDocument/didSave after saving (BufWritePost)
         didSave = true,
@@ -870,8 +868,8 @@ function protocol._resolve_capabilities_compat(server_capabilities)
       text_document_sync_properties = {
         text_document_open_close = if_nil(textDocumentSync.openClose, false),
         text_document_did_change = if_nil(textDocumentSync.change, TextDocumentSyncKind.None),
-        text_document_will_save = if_nil(textDocumentSync.willSave, false),
-        text_document_will_save_wait_until = if_nil(textDocumentSync.willSaveWaitUntil, false),
+        text_document_will_save = if_nil(textDocumentSync.willSave, true),
+        text_document_will_save_wait_until = if_nil(textDocumentSync.willSaveWaitUntil, true),
         text_document_save = if_nil(textDocumentSync.save, false),
         text_document_save_include_text = if_nil(
           type(textDocumentSync.save) == 'table' and textDocumentSync.save.includeText,
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 79c83af5d9..f38a7ad044 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -65,9 +65,9 @@ local create_server_definition = [[
         })
         local handler = handlers[method]
         if handler then
-          local response = handler(method, params)
+          local response, err = handler(params)
           if response then
-            callback(nill, response)
+            callback(err, response)
           end
         elseif method == 'initialize' then
           callback(nil, {
@@ -76,9 +76,18 @@ local create_server_definition = [[
         elseif method == 'shutdown' then
           callback(nil, nil)
         end
+        local request_id = #server.messages
+        return true, request_id
       end
 
       function srv.notify(method, params)
+        table.insert(server.messages, {
+          method = method,
+          params = params
+        })
+        if method == 'exit' then
+          dispatchers.on_exit(0, 15)
+        end
       end
 
       function srv.is_closing()
@@ -612,6 +621,67 @@ describe('LSP', function()
       }
     end)
 
+    it('BufWritePre does not send notifications if server lacks willSave capabilities', function()
+      exec_lua(create_server_definition)
+      local messages = exec_lua([[
+        local server = _create_server({
+          capabilities = {
+            textDocumentSync = {
+              willSave = false,
+              willSaveWaitUntil = false,
+            }
+          },
+        })
+        local client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+        local buf = vim.api.nvim_get_current_buf()
+        vim.api.nvim_exec_autocmds('BufWritePre', { buffer = buf, modeline = false })
+        vim.lsp.stop_client(client_id)
+        return server.messages
+      ]])
+      eq(#messages, 4)
+      eq(messages[1].method, 'initialize')
+      eq(messages[2].method, 'initialized')
+      eq(messages[3].method, 'shutdown')
+      eq(messages[4].method, 'exit')
+    end)
+    it('BufWritePre sends willSave / willSaveWaitUntil, applies textEdits', function()
+      exec_lua(create_server_definition)
+      local result = exec_lua([[
+        local server = _create_server({
+          capabilities = {
+            textDocumentSync = {
+              willSave = true,
+              willSaveWaitUntil = true,
+            }
+          },
+          handlers = {
+            ['textDocument/willSaveWaitUntil'] = function()
+              local text_edit = {
+                range = {
+                  start = { line = 0, character = 0 },
+                  ['end'] = { line = 0, character = 0 },
+                },
+                newText = 'Hello'
+              }
+              return { text_edit, }
+            end
+          },
+        })
+        local buf = vim.api.nvim_get_current_buf()
+        local client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+        vim.api.nvim_exec_autocmds('BufWritePre', { buffer = buf, modeline = false })
+        vim.lsp.stop_client(client_id)
+        return {
+          messages = server.messages,
+          lines = vim.api.nvim_buf_get_lines(buf, 0, -1, true)
+        }
+      ]])
+      local messages = result.messages
+      eq('textDocument/willSave', messages[3].method)
+      eq('textDocument/willSaveWaitUntil', messages[4].method)
+      eq({'Hello'}, result.lines)
+    end)
+
     it('saveas sends didOpen if filename changed', function()
       local expected_handlers = {
         { NIL, {}, { method = 'shutdown', client_id = 1 } },
@@ -3517,12 +3587,12 @@ describe('LSP', function()
         vim.lsp.buf.format({ bufnr = bufnr, false })
         return server.messages
       ]])
-      eq("textDocument/rangeFormatting", result[2].method)
+      eq("textDocument/rangeFormatting", result[3].method)
       local expected_range = {
         start = { line = 0, character = 0 },
         ['end'] = { line = 1, character = 4 },
       }
-      eq(expected_range, result[2].params.range)
+      eq(expected_range, result[3].params.range)
     end)
   end)
   describe('cmd', function()