refactor(completions): work with auto-pairs (#2400)

Co-authored-by: Oli Morris <[email protected]>

Author: olimorris

lua/codecompanion/strategies/inline/completion.lua

@@ -68,7 +68,7 @@ local function extract_next_line(text)
   return string.sub(text, 1, newline_pos)
 end
 
----Validate that the completion is not stale (cursor matches range end)
+---Validate that cursor is within the completion range
 ---@param item table The completion item
 ---@param cursor_row number The current cursor row (0-indexed)
 ---@param cursor_col number The current cursor column (0-indexed)
@@ -78,11 +78,19 @@ local function validate_completion(item, cursor_row, cursor_col)
     return true
   end
 
+  local start_row = item.range.start.row
+  local start_col = item.range.start.col
   local end_row = item.range.end_.row
   local end_col = item.range.end_.col
 
-  if end_row ~= cursor_row or end_col ~= cursor_col then
-    log:debug("Ignoring stale completion (cursor: %d,%d; range end: %d,%d)", cursor_row, cursor_col, end_row, end_col)
+  -- Allow cursor anywhere within range to support auto-pairs
+  if cursor_row < start_row or cursor_row > end_row then
+    return false
+  end
+  if cursor_row == start_row and cursor_col < start_col then
+    return false
+  end
+  if cursor_row == end_row and cursor_col > end_col then
     return false
   end
 
@@ -110,34 +118,34 @@ end
 ---Get the new text by removing the existing buffer text prefix
 ---@param item table The completion item
 ---@param text string The full completion text
+---@param cursor_row number The current cursor row (0-indexed)
+---@param cursor_col number The current cursor column (0-indexed)
 ---@return string new_text The text after removing existing prefix
-local function get_new_text(item, text)
+local function get_new_text(item, text, cursor_row, cursor_col)
   local bufnr = api.nvim_get_current_buf()
   local start_row = item.range.start.row
   local start_col = item.range.start.col
-  local end_row = item.range.end_.row
-  local end_col = item.range.end_.col
 
-  local existing_lines = api.nvim_buf_get_text(bufnr, start_row, start_col, end_row, end_col, {})
+  -- Use cursor position instead of range.end_ to handle auto-pairs
+  local existing_lines = api.nvim_buf_get_text(bufnr, start_row, start_col, cursor_row, cursor_col, {})
   local existing_text = table.concat(existing_lines, "\n")
 
-  -- Remove existing text prefix to get only new text
   if vim.startswith(text, existing_text) then
     return text:sub(#existing_text + 1)
   else
-    log:warn("Suggestion doesn't start with existing text")
     return text
   end
 end
 
----Insert text at cursor and move cursor to the end
+---Insert text at cursor (replacing to range end) and move cursor to the end
 ---@param text string The text to insert
 ---@param cursor_row number The current cursor row (0-indexed)
 ---@param cursor_col number The current cursor column (0-indexed)
-local function insert_and_move_cursor(text, cursor_row, cursor_col)
-  -- Split the text if it contains newlines
+---@param end_row number The range end row (0-indexed)
+---@param end_col number The range end column (0-indexed)
+local function insert_and_move_cursor(text, cursor_row, cursor_col, end_row, end_col)
   local lines = vim.split(text, "\n", { plain = true })
-  api.nvim_buf_set_text(0, cursor_row, cursor_col, cursor_row, cursor_col, lines)
+  api.nvim_buf_set_text(0, cursor_row, cursor_col, end_row, end_col, lines)
 
   -- Move cursor to the end of the inserted text
   local new_row, new_col
@@ -176,15 +184,17 @@ local function accept_partial(extract_fn)
         return
       end
 
-      local new_text = get_new_text(item, text)
+      local new_text = get_new_text(item, text, cursor_row, cursor_col)
 
       local extracted = extract_fn(new_text)
       if not extracted then
         log:debug("No text could be extracted from completion")
         return
       end
 
-      insert_and_move_cursor(extracted, cursor_row, cursor_col)
+      local end_row = item.range.end_.row
+      local end_col = item.range.end_.col
+      insert_and_move_cursor(extracted, cursor_row, cursor_col, end_row, end_col)
       accepted = true
     end,
   })

tests/strategies/inline/test_completion.lua

@@ -53,10 +53,45 @@ local T = new_set({
   },
 })
 
--- Helper to enter insert mode and set cursor
-local function set_cursor_insert(row, col)
-  child.cmd("startinsert!")
-  child.api.nvim_win_set_cursor(0, { row, col })
+---Helper to set buffer with cursor marker and enter insert mode
+---Accepts string "text|" or table {"line1", "line|2"} where | marks cursor position
+---@param text string|table
+local function set_buffer_text(text)
+  if type(text) == "string" then
+    local cursor_pos = text:find("|", 1, true)
+    if not cursor_pos then
+      error("No cursor marker '|' found in text")
+    end
+
+    local line = text:sub(1, cursor_pos - 1) .. text:sub(cursor_pos + 1)
+    child.api.nvim_buf_set_lines(0, 0, -1, true, { line })
+    child.cmd("startinsert!")
+    child.api.nvim_win_set_cursor(0, { 1, cursor_pos - 1 })
+  elseif type(text) == "table" then
+    local cursor_row, cursor_col
+    local lines = {}
+
+    for i, line in ipairs(text) do
+      local pos = line:find("|", 1, true)
+      if pos then
+        cursor_row = i
+        cursor_col = pos - 1
+        table.insert(lines, line:sub(1, pos - 1) .. line:sub(pos + 1))
+      else
+        table.insert(lines, line)
+      end
+    end
+
+    if not cursor_row then
+      error("No cursor marker '|' found in text")
+    end
+
+    child.api.nvim_buf_set_lines(0, 0, -1, true, lines)
+    child.cmd("startinsert!")
+    child.api.nvim_win_set_cursor(0, { cursor_row, cursor_col })
+  else
+    error("Expected string or table")
+  end
 end
 
 T["accept_word()"] = new_set()
@@ -66,9 +101,7 @@ T["accept_word()"]["works with simple word completion"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  -- Set up: buffer has "-- Create a fib" and cursor is at end
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "-- Create a fib" })
-  set_cursor_insert(1, 15) -- After "-- Create a fib"
+  set_buffer_text("-- Create a fib|")
 
   -- Set up the mock completion
   child.lua([[
@@ -94,8 +127,7 @@ T["accept_word()"]["works with punctuation in completion"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "local x" })
-  set_cursor_insert(1, 7)
+  set_buffer_text("local x|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -118,8 +150,7 @@ T["accept_word()"]["works with newline in word"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "function test()" })
-  set_cursor_insert(1, 15)
+  set_buffer_text("function test()|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -143,9 +174,7 @@ T["accept_word()"]["ignores stale completion (cursor before range end)"] = funct
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  -- Make a longer buffer so cursor doesn't get clamped
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "test some long text here" })
-  set_cursor_insert(1, 10) -- Cursor at position 10
+  set_buffer_text("test some |long text here")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -168,8 +197,7 @@ T["accept_word()"]["handles empty buffer"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "" })
-  set_cursor_insert(1, 0)
+  set_buffer_text("|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -187,13 +215,59 @@ T["accept_word()"]["handles empty buffer"] = function()
   h.eq(child.api.nvim_win_get_cursor(0), { 1, 6 })
 end
 
+T["accept_word()"]["handles auto-pairs with cursor inside brackets"] = function()
+  if vim.fn.has("nvim-0.12") == 0 then
+    MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
+  end
+
+  set_buffer_text("local function hello_world(|)")
+
+  child.lua([[
+    _G.mock_completion_item = {
+      insert_text = "local function hello_world(arg)",
+      range = {
+        start = { row = 0, col = 0 },
+        end_ = { row = 0, col = 27 }
+      }
+    }
+
+    _G.completion.accept_word()
+  ]])
+
+  -- Should insert "arg" without being confused by the trailing ")"
+  h.eq(child.api.nvim_buf_get_lines(0, 0, -1, true), { "local function hello_world(arg)" })
+  h.eq(child.api.nvim_win_get_cursor(0), { 1, 30 })
+end
+
+T["accept_word()"]["handles auto-pairs with cursor inside quotes"] = function()
+  if vim.fn.has("nvim-0.12") == 0 then
+    MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
+  end
+
+  set_buffer_text('print("|")') -- print("|")
+
+  child.lua([[
+    _G.mock_completion_item = {
+      insert_text = "print(\"hello world\")",
+      range = {
+        start = { row = 0, col = 0 },
+        end_ = { row = 0, col = 7 }
+      }
+    }
+
+    _G.completion.accept_word()
+  ]])
+
+  h.eq(child.api.nvim_buf_get_lines(0, 0, -1, true), { 'print("hello ")' })
+  h.eq(child.api.nvim_win_get_cursor(0), { 1, 13 })
+end
+
 T["accept_word()"]["handles completion that doesn't start with existing text"] = function()
   if vim.fn.has("nvim-0.12") == 0 then
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "foo bar" })
-  set_cursor_insert(1, 7)
+  set_buffer_text("foo bar|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -218,8 +292,7 @@ T["accept_line()"]["works with single line completion"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "-- Comment" })
-  set_cursor_insert(1, 10)
+  set_buffer_text("-- Comment|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -242,8 +315,7 @@ T["accept_line()"]["works with multi-line completion"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "function test()" })
-  set_cursor_insert(1, 15)
+  set_buffer_text("function test()|")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -257,7 +329,6 @@ T["accept_line()"]["works with multi-line completion"] = function()
     _G.completion.accept_line()
   ]])
 
-  -- Should insert newline + next line content + newline (avoiding the double-press issue)
   h.eq(child.api.nvim_buf_get_lines(0, 0, -1, true), { "function test()", "  return 42", "" })
   h.eq(child.api.nvim_win_get_cursor(0), { 3, 0 })
 end
@@ -267,9 +338,7 @@ T["accept_line()"]["ignores stale completion"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  -- Make a longer buffer so cursor doesn't get clamped
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "test some long text here" })
-  set_cursor_insert(1, 10)
+  set_buffer_text("test some |long text here")
 
   child.lua([[
     _G.mock_completion_item = {
@@ -292,8 +361,7 @@ T["accept_line()"]["handles completion without newline"] = function()
     MiniTest.skip("Requires Neovim 0.12+ for vim.lsp.inline_completion")
   end
 
-  child.api.nvim_buf_set_lines(0, 0, -1, true, { "hello" })
-  set_cursor_insert(1, 5)
+  set_buffer_text("hello|")
 
   child.lua([[
     _G.mock_completion_item = {