fix: prevent partial accepts from triggering a new suggestion (zbirenbaum#469)

Allows users to quickly accept a series of words or lines without waiting for the suggestion every time
Fixes zbirenbaum#467

Author: AntoineGS

lua/copilot/api/init.lua

@@ -130,7 +130,7 @@ function M.notify_shown(client, params, callback)
   return M.request(client, "notifyShown", params, callback)
 end
 
----@alias copilot_get_completions_data_completion { displayText: string, position: { character: integer, line: integer }, range: { ['end']: { character: integer, line: integer }, start: { character: integer, line: integer } }, text: string, uuid: string }
+---@alias copilot_get_completions_data_completion { displayText: string, position: { character: integer, line: integer }, range: { ['end']: { character: integer, line: integer }, start: { character: integer, line: integer } }, text: string, uuid: string, partial_text: string }
 ---@alias copilot_get_completions_data { completions: copilot_get_completions_data_completion[] }
 
 ---@return any|nil err

lua/copilot/suggestion/init.lua

@@ -9,7 +9,7 @@ local utils = require("copilot.client.utils")
 
 local M = {}
 
----@alias copilot_suggestion_context { first?: integer, cycling?: integer, cycling_callbacks?: (fun(ctx: copilot_suggestion_context):nil)[], params?: table, suggestions?: copilot_get_completions_data_completion[], choice?: integer, shown_choices?: table<string, true> }
+---@alias copilot_suggestion_context { first?: integer, cycling?: integer, cycling_callbacks?: (fun(ctx: copilot_suggestion_context):nil)[], params?: table, suggestions?: copilot_get_completions_data_completion[], choice?: integer, shown_choices?: table<string, true>, accepted_partial?: boolean }
 
 local copilot = {
   setup_done = false,
@@ -26,6 +26,8 @@ local copilot = {
   debounce = 75,
 }
 
+local ignore_next_cursor_moved = false
+
 local function with_client(fn)
   local client = c.get()
   if client then
@@ -58,6 +60,28 @@ local function get_ctx(bufnr)
   return ctx
 end
 
+---@param idx integer
+---@param new_line integer
+---@param new_end_col integer
+---@param bufnr? integer
+local function update_ctx_suggestion_position(idx, new_line, new_end_col, bufnr)
+  bufnr = bufnr or vim.api.nvim_get_current_buf()
+
+  if not copilot.context[bufnr] then
+    return
+  end
+
+  if not copilot.context[bufnr].suggestions[idx] then
+    return
+  end
+
+  local suggestion = copilot.context[bufnr].suggestions[idx]
+  suggestion.range["start"].line = new_line
+  suggestion.range["start"].character = 0
+  suggestion.range["end"].line = new_line
+  suggestion.range["end"].character = new_end_col
+end
+
 ---@param idx integer
 ---@param text string
 ---@param bufnr? integer
@@ -89,6 +113,7 @@ local function reset_ctx(ctx)
   ctx.suggestions = nil
   ctx.choice = nil
   ctx.shown_choices = nil
+  ctx.accepted_partial = nil
 end
 
 local function set_keymap(keymap)
@@ -249,10 +274,10 @@ local function get_current_suggestion(ctx)
       return nil
     end
 
-    if choice.range.start.character ~= 0 then
-      -- unexpected range
-      return nil
-    end
+    -- if choice.range.start.character ~= 0 then
+    --   -- unexpected range
+    --   return nil
+    -- end
 
     return choice
   end)
@@ -475,7 +500,7 @@ local function schedule(ctx)
     stop_timer()
   end
 
-  update_preview(ctx)
+  -- update_preview(ctx)
   local bufnr = vim.api.nvim_get_current_buf()
   copilot._copilot_timer = vim.fn.timer_start(copilot.debounce, function(timer)
     logger.trace("suggestion schedule timer", bufnr)
@@ -487,6 +512,10 @@ function M.next()
   local ctx = get_ctx()
   logger.trace("suggestion next", ctx)
 
+  if ctx.accepted_partial then
+    reset_ctx(ctx)
+  end
+
   -- no suggestion request yet
   if not ctx.first then
     logger.trace("suggestion next, no first request")
@@ -503,6 +532,10 @@ function M.prev()
   local ctx = get_ctx()
   logger.trace("suggestion prev", ctx)
 
+  if ctx.accepted_partial then
+    reset_ctx(ctx)
+  end
+
   -- no suggestion request yet
   if not ctx.first then
     logger.trace("suggestion prev, no first request", ctx)
@@ -532,13 +565,17 @@ function M.accept(modifier)
     return
   end
 
-  cancel_inflight_requests(ctx)
-  reset_ctx(ctx)
-
   if type(modifier) == "function" then
     suggestion = modifier(suggestion)
   end
 
+  local accepted_partial = suggestion.partial_text and suggestion.partial_text ~= ""
+
+  if not accepted_partial then
+    cancel_inflight_requests(ctx)
+    reset_ctx(ctx)
+  end
+
   with_client(function(client)
     local ok, _ = pcall(function()
       api.notify_accepted(
@@ -552,34 +589,62 @@ function M.accept(modifier)
     end
   end)
 
-  clear_preview()
+  local newText
+
+  if accepted_partial then
+    newText = suggestion.partial_text
+    ctx.accepted_partial = true
+    ignore_next_cursor_moved = true
+  else
+    clear_preview()
+    newText = suggestion.text
+  end
 
-  local range, newText = suggestion.range, suggestion.text
+  local range = suggestion.range
   local cursor = vim.api.nvim_win_get_cursor(0)
   local line, character = cursor[1] - 1, cursor[2]
   if range["end"].line == line and range["end"].character < character then
     range["end"].character = character
   end
 
-  -- Hack for 'autoindent', makes the indent persist. Check `:help 'autoindent'`.
   vim.schedule_wrap(function()
     -- Create an undo breakpoint
     vim.cmd("let &undolevels=&undolevels")
+    -- Hack for 'autoindent', makes the indent persist. Check `:help 'autoindent'`.
     vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes("<Space><Left><Del>", true, false, true), "n", false)
     local bufnr = vim.api.nvim_get_current_buf()
+
     local encoding = vim.api.nvim_get_option_value("fileencoding", { buf = bufnr }) ~= ""
         and vim.api.nvim_get_option_value("fileencoding", { buf = bufnr })
       or vim.api.nvim_get_option_value("encoding", { scope = "global" })
-    vim.lsp.util.apply_text_edits({ { range = range, newText = newText } }, bufnr, encoding)
 
-    -- instead of calling <End>, go to the pos of the row after the last \n of inserted text
-    -- local cursor_keys = string.rep("<Down>", #vim.split(newText, "\n", { plain = true }) - 1) .. "<End>"
-    -- vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes(cursor_keys, true, false, true), "n", false)
     local lines = vim.split(newText, "\n", { plain = true })
-    local last_line = lines[#lines]
-    local cursor_keys = string.rep("<Down>", #lines - 1)
+    local lines_count = #lines
+    local last_col = #lines[lines_count]
+
+    -- apply_text_edits will remove the last \n if the last line is empty,
+    -- so we trick it by adding an extra one
+    if last_col == 0 then
+      newText = newText .. "\n"
+    end
+
+    vim.lsp.util.apply_text_edits({ { range = range, newText = newText } }, bufnr, encoding)
+
     -- Position cursor at the end of the last inserted line
-    vim.api.nvim_win_set_cursor(0, { range["start"].line + #lines, #last_line })
+    local new_cursor_line = range["start"].line + #lines
+    vim.api.nvim_win_set_cursor(0, { new_cursor_line, last_col })
+
+    if accepted_partial then
+      suggestion.partial_text = nil
+
+      for _ = 1, lines_count - 1 do
+        suggestion.text = suggestion.text:sub(suggestion.text:find("\n") + 1)
+        suggestion.displayText = suggestion.displayText:sub(suggestion.displayText:find("\n") + 1)
+      end
+
+      update_ctx_suggestion_position(ctx.choice, new_cursor_line - 1, last_col, bufnr)
+      update_preview(ctx)
+    end
   end)()
 end
 
@@ -592,29 +657,30 @@ function M.accept_word()
 
     local _, char_idx = string.find(text, "%s*%p*[^%s%p]*%s*", character + 1)
     if char_idx then
-      suggestion.text = string.sub(text, 1, char_idx)
-
-      range["end"].line = range["start"].line
+      suggestion.partial_text = string.sub(text, 1, char_idx)
       range["end"].character = char_idx
     end
 
+    range["end"].line = range["start"].line
     return suggestion
   end)
 end
 
 function M.accept_line()
   M.accept(function(suggestion)
-    local text = suggestion.text
+    local range, text = suggestion.range, suggestion.text
 
     local cursor = vim.api.nvim_win_get_cursor(0)
     local _, character = cursor[1], cursor[2]
 
     local next_char = string.sub(text, character + 1, character + 1)
     local _, char_idx = string.find(text, next_char == "\n" and "\n%s*[^\n]*\n%s*" or "\n%s*", character)
     if char_idx then
-      suggestion.text = string.sub(text, 1, char_idx)
+      suggestion.partial_text = string.sub(text, 1, char_idx)
+      range["end"].character = char_idx
     end
 
+    range["end"].line = range["start"].line
     return suggestion
   end)
 end
@@ -659,6 +725,11 @@ local function on_buf_enter()
 end
 
 local function on_cursor_moved_i()
+  if ignore_next_cursor_moved then
+    ignore_next_cursor_moved = false
+    return
+  end
+
   local ctx = get_ctx()
   if copilot._copilot_timer or ctx.params or should_auto_trigger() then
     logger.trace("suggestion on cursor moved insert")

tests/child_helper.lua

@@ -21,7 +21,7 @@ function M.new_child_neovim(test_name)
       panel = "",
       suggestion = [[
         suggestion_notification = function(virt_text, _)
-          if (#virt_text > 0) and (#virt_text[1] > 0) and (virt_text[1][1] == "89") then
+          if (#virt_text > 0) and (#virt_text[1] > 0) then
             M.suggested = true
           end
         end,

tests/screenshots/tests-test_suggestion.lua---suggestion()---accept_line,-1-line,-then-accept

@@ -0,0 +1,103 @@
+--|---------|---------|---------|---------|
+01|# Numbers in a 3x3 grid, up to 63       
+02|{                                       
+03|  1,2,3                                 
+04|  4,5,6                                 
+05|  7,8,9                                 
+06|}                                       
+07|{                                       
+08|  10,11,12                              
+09|  13,14,15                              
+10|  16,17,18                              
+11|}                                       
+12|{                                       
+13|  19,20,21                              
+14|  22,23,24                              
+15|  25,26,27                              
+16|}                                       
+17|{                                       
+18|  28,29,30                              
+19|  31,32,33                              
+20|  34,35,36                              
+21|}                                       
+22|{                                       
+23|  37,38,39                              
+24|  40,41,42                              
+25|  43,44,45                              
+26|}                                       
+27|{                                       
+28|  46,47,48                              
+29|  49,50,51                              
+30|  52,53,54                              
+31|}                                       
+32|{                                       
+33|  55,56,57                              
+34|  58,59,60                              
+35|  61,62,63                              
+36|}                                       
+37|~                                       
+38|~                                       
+39|~                                       
+40|~                                       
+41|~                                       
+42|~                                       
+43|~                                       
+44|~                                       
+45|~                                       
+46|~                                       
+47|~                                       
+48|~                                       
+49|[No Name] [+]         36,2           All
+50|-- INSERT --                            
+
+--|---------|---------|---------|---------|
+01|0000000000000000000000000000000000000000
+02|0000000000000000000000000000000000000000
+03|0000000000000000000000000000000000000000
+04|0000000000000000000000000000000000000000
+05|0000000000000000000000000000000000000000
+06|0000000000000000000000000000000000000000
+07|0000000000000000000000000000000000000000
+08|0000000000000000000000000000000000000000
+09|0000000000000000000000000000000000000000
+10|0000000000000000000000000000000000000000
+11|0000000000000000000000000000000000000000
+12|0000000000000000000000000000000000000000
+13|0000000000000000000000000000000000000000
+14|0000000000000000000000000000000000000000
+15|0000000000000000000000000000000000000000
+16|0000000000000000000000000000000000000000
+17|0000000000000000000000000000000000000000
+18|0000000000000000000000000000000000000000
+19|0000000000000000000000000000000000000000
+20|0000000000000000000000000000000000000000
+21|0000000000000000000000000000000000000000
+22|0000000000000000000000000000000000000000
+23|0000000000000000000000000000000000000000
+24|0000000000000000000000000000000000000000
+25|0000000000000000000000000000000000000000
+26|0000000000000000000000000000000000000000
+27|0000000000000000000000000000000000000000
+28|0000000000000000000000000000000000000000
+29|0000000000000000000000000000000000000000
+30|0000000000000000000000000000000000000000
+31|0000000000000000000000000000000000000000
+32|1000000000000000000000000000000000000000
+33|0000000000000000000000000000000000000000
+34|0000000000000000000000000000000000000000
+35|0000000000000000000000000000000000000000
+36|1000000000000000000000000000000000000000
+37|2222222222222222222222222222222222222222
+38|2222222222222222222222222222222222222222
+39|2222222222222222222222222222222222222222
+40|2222222222222222222222222222222222222222
+41|2222222222222222222222222222222222222222
+42|2222222222222222222222222222222222222222
+43|2222222222222222222222222222222222222222
+44|2222222222222222222222222222222222222222
+45|2222222222222222222222222222222222222222
+46|2222222222222222222222222222222222222222
+47|2222222222222222222222222222222222222222
+48|2222222222222222222222222222222222222222
+49|3333333333333333333333333333333333333333
+50|4444444444445555555555555555555555555555