feat: next edit suggestion (nes)

Fixes zbirenbaum#544

Author: AntoineGS

README.md

@@ -1,6 +1,7 @@
 # copilot.lua
 
 This plugin is the pure lua replacement for [github/copilot.vim](https://github.com/github/copilot.vim).
+A huge thank you to @tris203 for the code behind the nes functionality ([copilot-lsp](https://github.com/copilotlsp-nvim/copilot-lsp)).
 
 <details>
 <summary>Motivation behind `copilot.lua`</summary>
@@ -48,7 +49,11 @@ Install the plugin with your preferred plugin manager.
 For example, with [packer.nvim](https://github.com/wbthomason/packer.nvim):
 
 ```lua
-use { "zbirenbaum/copilot.lua" }
+use { "zbirenbaum/copilot.lua" 
+  requires = {
+    "copilotlsp-nvim/copilot-lsp", -- (optional) for NES functionality
+  },
+}
 ```
 
 ### Authentication
@@ -93,6 +98,9 @@ For example:
 ```lua
 use {
   "zbirenbaum/copilot.lua",
+  requires = {
+    "copilotlsp-nvim/copilot-lsp", -- (optional) for NES functionality
+  },
   cmd = "Copilot",
   event = "InsertEnter",
   config = function()
@@ -136,16 +144,14 @@ require('copilot').setup({
       dismiss = "<C-]>",
     },
   },
-  filetypes = {
-    yaml = false,
-    markdown = false,
-    help = false,
-    gitcommit = false,
-    gitrebase = false,
-    hgcommit = false,
-    svn = false,
-    cvs = false,
-    ["."] = false,
+  nes = {
+    enabled = false, -- requires copilot-lsp as a dependency
+    auto_trigger = false,
+    keymap = {
+      accept_and_goto = false,
+      accept = false,
+      dismiss = false,
+    },
   },
   auth_provider_url = nil, -- URL to authentication provider, if not "https://github.com/"
   logger = {
@@ -270,6 +276,43 @@ require("copilot.suggestion").toggle_auto_trigger()
 ```
 These can also be accessed through the `:Copilot suggestion <function>` command (eg. `:Copilot suggestion accept`).
 
+### nes (next edit suggestion)
+
+>[!WARNING]
+> This feature is still experimental and may not work as expected in all scenarios, please report any issues you encounter.
+
+When `enabled` is `true`, copilot will provide suggestions based on the next edit you are likely to make, through [copilot-lsp](https://github.com/copilotlsp-nvim/copilot-lsp).
+If there is no suggestion, the keymaps will pass through the original keymap.
+
+`copilot-lsp` has a few configurations built-in as well, for additional configurations, please refer to the [copilot-lsp documentation](https://github.com/copilotlsp-nvim/copilot-lsp/blob/main/README.md).
+These configurations should be set in the `init` function of the `copilot-lsp` dependency.
+
+```lua
+use {
+  "zbirenbaum/copilot.lua",
+  requires = {
+    "copilotlsp-nvim/copilot-lsp",
+    init = function()
+      vim.g.copilot_nes_debounce = 500
+    end,
+  },
+  cmd = "Copilot",
+  event = "InsertEnter",
+  config = function()
+    require("copilot").setup({
+      nes = {
+        enabled = true,
+        keymap = {
+          accept_and_goto = "<leader>p",
+          accept = false,
+          dismiss = "<Esc>",
+        },
+      },
+    })
+  end,
+}
+```
+
 ### filetypes
 
 Specify filetypes for attaching copilot.

lua/copilot/client/config.lua

@@ -122,6 +122,8 @@ function M.prepare_client_config(overrides, client)
         if token_env_set then
           require("copilot.auth").signin()
         end
+
+        require("copilot.nes").setup(lsp_client)
       end)
     end,
     on_exit = function(code, _, client_id)

lua/copilot/config/init.lua

@@ -5,6 +5,7 @@ local logger = require("copilot.logger")
 ---@field suggestion SuggestionConfig
 ---@field logger LoggerConfig
 ---@field server ServerConfig
+---@field nes NesConfig
 ---@field filetypes table<string, boolean> Filetypes to enable Copilot for
 ---@field auth_provider_url string|nil URL for the authentication provider
 ---@field workspace_folders string[] Workspace folders to enable Copilot for
@@ -23,6 +24,7 @@ local M = {
   suggestion = require("copilot.config.suggestion").default,
   logger = require("copilot.config.logger").default,
   server = require("copilot.config.server").default,
+  nes = require("copilot.config.nes").default,
   root_dir = require("copilot.config.root_dir").default,
   should_attach = require("copilot.config.should_attach").default,
   filetypes = {},
@@ -59,6 +61,7 @@ end
 function M.validate(config)
   vim.validate("panel", config.panel, "table")
   vim.validate("suggestion", config.suggestion, "table")
+  vim.validate("nes", config.nes, "table")
   vim.validate("logger", config.logger, "table")
   vim.validate("server", config.server, "table")
   vim.validate("filetypes", config.filetypes, "table")
@@ -76,6 +79,7 @@ function M.validate(config)
   require("copilot.config.server").validate(config.server)
   require("copilot.config.root_dir").validate(config.root_dir)
   require("copilot.config.should_attach").validate(config.should_attach)
+  require("copilot.config.nes").validate(config.nes)
 end
 
 return M

lua/copilot/config/nes.lua

@@ -0,0 +1,58 @@
+local logger = require("copilot.logger")
+local nes_api = require("copilot.nes.api")
+---@class NesKeymap
+---@field accept_and_goto string|false Keymap to accept the suggestion and go to the end of the suggestion
+---@field accept string|false Keymap to accept the suggestion
+---@field dismiss string|false Keymap to dismiss the suggestion
+
+---@class NesConfig
+---@field enabled boolean Whether to enable nes (next edit suggestions)
+---@field auto_trigger boolean Whether to automatically trigger next edit suggestions
+---@field keymap NesKeymap Keymaps for nes actions
+
+local M = {
+  ---@type NesConfig
+  default = {
+    enabled = false,
+    auto_trigger = false,
+    keymap = {
+      accept_and_goto = false,
+      accept = false,
+      dismiss = false,
+    },
+  },
+}
+
+---@type NesConfig
+M.config = vim.deepcopy(M.default)
+
+---@param opts? NesConfig
+function M.setup(opts)
+  opts = opts or {}
+  M.config = vim.tbl_deep_extend("force", M.default, opts)
+end
+
+---@param config NesConfig
+function M.validate(config)
+  vim.validate("enabled", config.enabled, "boolean")
+  vim.validate("auto_trigger", config.auto_trigger, "boolean")
+  vim.validate("keymap", config.keymap, "table")
+  vim.validate("keymap.accept_and_goto", config.keymap.accept_and_goto, { "string", "boolean" })
+  vim.validate("keymap.accept", config.keymap.accept, { "string", "boolean" })
+  vim.validate("keymap.dismiss", config.keymap.dismiss, { "string", "boolean" })
+
+  if config.enabled then
+    local has_nes, _ = pcall(function()
+      nes_api.test()
+    end)
+
+    if not has_nes then
+      logger.error(
+        "copilot-lsp is not available, disabling nes.\nPlease refer to the documentation and ensure it is installed."
+      )
+      config.enabled = false
+    end
+  end
+end
+
+return M

lua/copilot/highlight.lua

@@ -1,3 +1,7 @@
+local logger = require("copilot.logger")
+local config = require("copilot.config")
+local nes_api = require("copilot.nes.api")
+
 local M = {
   group = {
     CopilotAnnotation = "CopilotAnnotation",
@@ -19,6 +23,16 @@ function M.setup()
         vim.api.nvim_set_hl(0, from_group, { link = to_group })
       end
     end
+
+    if config.nes.enabled then
+      local ok, err = pcall(function()
+        nes_api.set_hl()
+      end)
+
+      if not ok then
+        logger.error("Error setting copilot-lsp highlights: ", err)
+      end
+    end
   end)
 end
 

lua/copilot/keymaps/init.lua

@@ -0,0 +1,89 @@
+local logger = require("copilot.logger")
+local config = require("copilot.config")
+
+local M = {}
+
+local previous_keymaps = {}
+
+---@param mode string
+---@param key string
+---@param action function
+---@param desc string
+function M.register_keymap(mode, key, action, desc)
+  if not mode or not key or not action then
+    return
+  end
+
+  vim.keymap.set(mode, key, function()
+    action()
+  end, {
+    desc = desc,
+    silent = true,
+  })
+end
+
+---@param mode string
+---@param key string
+---@param action function: boolean
+---@param desc string
+function M.register_keymap_with_passthrough(mode, key, action, desc)
+  if not mode or not key or not action then
+    return
+  end
+
+  local keymap_key = mode .. ":" .. key
+  -- Save any existing mapping for this key
+  local existing = vim.fn.maparg(key, mode, false, true)
+  if existing and existing.rhs and existing.rhs ~= "" then
+    previous_keymaps[keymap_key] = existing.rhs
+  else
+    previous_keymaps[keymap_key] = nil
+  end
+
+  vim.keymap.set(mode, key, function()
+    local action_ran = action()
+    if not action_ran then
+      -- If there was a previous mapping, execute it
+      local prev = previous_keymaps[keymap_key]
+      if prev then
+        vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes(prev, true, false, true), mode, true)
+      end
+    end
+  end, {
+    desc = desc,
+    silent = true,
+  })
+end
+
+---@param mode string
+---@param key string|false
+function M.unset_keymap_if_exists(mode, key)
+  if not key then
+    return
+  end
+
+  local ok, err = pcall(vim.api.nvim_del_keymap, mode, key)
+
+  if not ok then
+    local suggestion_keymaps = config.suggestion.keymap or {}
+    local panel_keymaps = config.panel.keymap or {}
+    local found = false
+
+    for _, tbl in ipairs({ suggestion_keymaps, panel_keymaps }) do
+      for _, v in pairs(tbl) do
+        if v == key then
+          if found then
+            logger.error("Keymap " .. key .. " is used for two different actions, please review your configuration.")
+            return
+          else
+            found = true
+          end
+        end
+      end
+    end
+
+    logger.error("Could not unset keymap for " .. mode .. " " .. key .. ": " .. err)
+  end
+end
+
+return M

lua/copilot/logger/init.lua

@@ -171,6 +171,20 @@ function M.handle_lsp_progress(_, result, _)
   M.trace(string.format("LSP progress - token %s", result.token), result.value)
 end
 
+-- Known noisy errors that we do not want to show as they seem to be out of our control
+---@param msg string
+---@return boolean
+local function force_log_to_trace(msg)
+  if
+    msg:match(".*Request textDocument/copilotInlineEdit: AbortError: The operation was aborted.*")
+    or msg:match(".*AsyncCompletionManager.*Request errored with AbortError: The operation was aborted.*")
+  then
+    return true
+  end
+
+  return false
+end
+
 function M.handle_log_lsp_messages(_, result, _)
   if not result then
     return
@@ -179,7 +193,7 @@ function M.handle_log_lsp_messages(_, result, _)
   local message = string.format("LSP message: %s", result.message)
   local message_type = result.type --[[@as integer]]
 
-  if message_type == 1 then
+  if message_type == 1 and not force_log_to_trace(message) then
     M.error(message)
   elseif message_type == 2 then
     M.warn(message)

lua/copilot/nes/api.lua

@@ -0,0 +1,46 @@
+-- Abstraction to the copilot-lsp module
+
+local M = {}
+
+function M.nes_set_auto_trigger(value)
+  local config = require("copilot-lsp.config")
+  config.config.require("copilot-lsp.nes").auto_trigger = value
+end
+
+function M.nes_lsp_on_init(client, au)
+  require("copilot-lsp.nes").lsp_on_init(client, au)
+end
+
+function M.set_hl()
+  local util = require("copilot-lsp.util")
+  util.set_hl()
+end
+
+---@param bufnr? integer
+---@return boolean --if the cursor walked
+function M.nes_walk_cursor_start_edit(bufnr)
+  return require("copilot-lsp.nes").walk_cursor_start_edit(bufnr)
+end
+
+---@param bufnr? integer
+---@return boolean --if the cursor walked
+function M.nes_walk_cursor_end_edit(bufnr)
+  return require("copilot-lsp.nes").walk_cursor_end_edit(bufnr)
+end
+
+---@param bufnr? integer
+---@return boolean --if the nes was applied
+function M.nes_apply_pending_nes(bufnr)
+  return require("copilot-lsp.nes").apply_pending_nes(bufnr)
+end
+
+---@return boolean -- true if a suggestion was cleared, false if no suggestion existed
+function M.nes_clear()
+  return require("copilot-lsp.nes").clear()
+end
+
+function M.test()
+  require("copilot-lsp.nes")
+end
+
+return M