feat: tab flow (#1077)

README.md

@@ -544,8 +544,8 @@ If you have the following structure:
 - [x] Edit the selected block
 - [x] Smart Tab (Cursor Flow)
 - [x] Chat with project (You can use `@codebase` to chat with the whole project)
+- [x] Chat with selected files
 - [ ] CoT
-- [ ] Chat with selected files
 
 ## Roadmap
 

lua/avante/config.lua

@@ -246,6 +246,10 @@ M._defaults = {
     -- Options override for custom providers
     provider_opts = {},
   },
+  suggestion = {
+    debounce = 600,
+    throttle = 600,
+  },
 }
 
 ---@type avante.Config

lua/avante/highlights.lua

@@ -15,6 +15,7 @@ local Highlights = {
   ANNOTATION = { name = "AvanteAnnotation", link = "Comment" },
   POPUP_HINT = { name = "AvantePopupHint", link = "NormalFloat" },
   INLINE_HINT = { name = "AvanteInlineHint", link = "Keyword" },
+  TO_BE_DELETED = { name = "AvanteToBeDeleted", bg = "#ffcccc", strikethrough = true },
 }
 
 Highlights.conflict = {
@@ -46,7 +47,11 @@ M.setup = function()
       end)
       :each(function(_, hl)
         if not has_set_colors(hl.name) then
-          api.nvim_set_hl(0, hl.name, { fg = hl.fg or nil, bg = hl.bg or nil, link = hl.link or nil })
+          api.nvim_set_hl(
+            0,
+            hl.name,
+            { fg = hl.fg or nil, bg = hl.bg or nil, link = hl.link or nil, strikethrough = hl.strikethrough }
+          )
         end
       end)
   end

lua/avante/suggestion.lua

@@ -21,11 +21,11 @@ local SUGGESTION_NS = api.nvim_create_namespace("avante_suggestion")
 ---@class avante.Suggestion
 ---@field id number
 ---@field augroup integer
----@field extmark_id integer
 ---@field ignore_patterns table
 ---@field negate_patterns table
 ---@field _timer? table
 ---@field _contexts table
+---@field is_on_throttle boolean
 local Suggestion = {}
 Suggestion.__index = Suggestion
 
@@ -37,11 +37,11 @@ function Suggestion:new(id)
   local gitignore_patterns, gitignore_negate_patterns = Utils.parse_gitignore(gitignore_path)
 
   instance.id = id
-  instance.extmark_id = 1
   instance._timer = nil
   instance._contexts = {}
   instance.ignore_patterns = gitignore_patterns
   instance.negate_patterns = gitignore_negate_patterns
+  instance.is_on_throttle = false
   if Config.behaviour.auto_suggestions then
     if not vim.g.avante_login or vim.g.avante_login == false then
       api.nvim_exec_autocmds("User", { pattern = Providers.env.REQUEST_LOGIN_PATTERN })
@@ -80,7 +80,13 @@ function Suggestion:suggest()
       role = "user",
       content = [[
 <filepath>a.py</filepath>
-<code>def fib</code>
+<code>
+L1: def fib
+L2:
+L3: if __name__ == "__main__":
+L4:    # just pass
+L5:    pass
+</code>
       ]],
     },
     {
@@ -95,11 +101,30 @@ function Suggestion:suggest()
       role = "assistant",
       content = [[
 [
-  {
+  [
-    "row": 1,
+    {
-    "col": 8,
+      "start_row": 1,
-    "content": "(n):\n    if n < 2:\n        return n\n    return fib(n - 1) + fib(n - 2)"
+      "end_row": 1,
-  }
+      "content": "def fib(n):\n    if n < 2:\n        return n\n    return fib(n - 1) + fib(n - 2)"
+    },
+    {
+      "start_row": 4,
+      "end_row": 5,
+      "content": "    fib(int(input()))"
+    },
+  ],
+  [
+    {
+      "start_row": 1,
+      "end_row": 1,
+      "content": "def fib(n):\n    a, b = 0, 1\n    for _ in range(n):\n        yield a\n        a, b = b, a + b"
+    },
+    {
+      "start_row": 4,
+      "end_row": 5,
+      "content": "    list(fib(int(input())))"
+    },
+  ]
 ]
       ]],
     },
@@ -128,21 +153,49 @@ function Suggestion:suggest()
         full_response = full_response:gsub("(.-)\n```\n?$", "%1")
         -- Remove everything before the first '[' to ensure we get just the JSON array
         full_response = full_response:gsub("^.-(%[.*)", "%1")
-        local ok, suggestions = pcall(vim.json.decode, full_response)
+        local ok, suggestions_list = pcall(vim.json.decode, full_response)
         if not ok then
           Utils.error("Error while decoding suggestions: " .. full_response, { once = true, title = "Avante" })
           return
         end
-        if not suggestions then
+        if not suggestions_list then
           Utils.info("No suggestions found", { once = true, title = "Avante" })
           return
         end
-        suggestions = vim
+        local current_lines = Utils.get_buf_lines(0, -1, bufnr)
-          .iter(suggestions)
+        suggestions_list = vim
-          :map(function(s) return { row = s.row, col = s.col, content = Utils.trim_all_line_numbers(s.content) } end)
+          .iter(suggestions_list)
+          :map(function(suggestions)
+            local new_suggestions = vim
+              .iter(suggestions)
+              :map(function(s)
+                local lines = vim.split(s.content, "\n")
+                local new_start_row = s.start_row
+                local new_content_lines = lines
+                for i = s.start_row, s.start_row + #lines - 1 do
+                  if current_lines[i] == lines[1] then
+                    new_start_row = i + 1
+                    new_content_lines = vim.list_slice(new_content_lines, 2)
+                  else
+                    break
+                  end
+                end
+                return {
+                  id = s.start_row,
+                  original_start_row = s.start_row,
+                  start_row = new_start_row,
+                  end_row = s.end_row,
+                  content = Utils.trim_all_line_numbers(table.concat(new_content_lines, "\n")),
+                }
+              end)
+              :totable()
+            --- sort the suggestions by start_row
+            table.sort(new_suggestions, function(a, b) return a.start_row < b.start_row end)
+            return new_suggestions
+          end)
           :totable()
-        ctx.suggestions = suggestions
+        ctx.suggestions_list = suggestions_list
-        ctx.current_suggestion_idx = 1
+        ctx.current_suggestions_idx = 1
         self:show()
       end)
     end,
@@ -155,74 +208,92 @@ function Suggestion:show()
   if not fn.mode():match("^[iR]") then return end
 
   local ctx = self:ctx()
-  local suggestion = ctx.suggestions[ctx.current_suggestion_idx]
-  if not suggestion then return end
-
-  local cursor_row, cursor_col = Utils.get_cursor_pos()
-
-  if suggestion.row < cursor_row then return end
 
   local bufnr = api.nvim_get_current_buf()
-  local row = suggestion.row
-  local col = suggestion.col
-  local content = suggestion.content
 
-  local lines = vim.split(content, "\n")
+  local suggestions = ctx.suggestions_list and ctx.suggestions_list[ctx.current_suggestions_idx] or nil
 
-  local extmark_col = cursor_col
+  if not suggestions then return end
 
-  if cursor_row < row then extmark_col = 0 end
+  for _, suggestion in ipairs(suggestions) do
+    local start_row = suggestion.start_row
+    local end_row = suggestion.end_row
+    local content = suggestion.content
 
-  local current_lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
+    local lines = vim.split(content, "\n")
 
-  if cursor_row == row then
+    local current_lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
-    local cursor_line_col = #current_lines[cursor_row] - 1
+
-    if cursor_col ~= cursor_line_col then
+    local virt_text_win_col = 0
-      local current_line = current_lines[cursor_row]
+
-      lines[1] = lines[1] .. current_line:sub(col + 1, -1)
+    if
+      start_row == end_row
+      and current_lines[start_row]
+      and #lines > 0
+      and vim.startswith(lines[1], current_lines[start_row])
+    then
+      virt_text_win_col = #current_lines[start_row]
+      lines[1] = string.sub(lines[1], #current_lines[start_row] + 1)
+    end
+
+    local virt_lines = {}
+
+    for _, line in ipairs(lines) do
+      table.insert(virt_lines, { { line, Highlights.SUGGESTION } })
+    end
+
+    local extmark = {
+      id = suggestion.id,
+      virt_text_win_col = virt_text_win_col,
+      virt_lines = virt_lines,
+    }
+
+    if virt_text_win_col > 0 then
+      extmark.virt_text = { { lines[1], Highlights.SUGGESTION } }
+      extmark.virt_lines = vim.list_slice(virt_lines, 2)
+    end
+
+    extmark.hl_mode = "combine"
+
+    local buf_lines = Utils.get_buf_lines(0, -1, bufnr)
+    local buf_lines_count = #buf_lines
+
+    while buf_lines_count < end_row do
+      api.nvim_buf_set_lines(bufnr, buf_lines_count, -1, false, { "" })
+      buf_lines_count = buf_lines_count + 1
+    end
+
+    if virt_text_win_col > 0 or start_row - 2 < 0 then
+      api.nvim_buf_set_extmark(bufnr, SUGGESTION_NS, start_row - 1, 0, extmark)
+    else
+      api.nvim_buf_set_extmark(bufnr, SUGGESTION_NS, start_row - 2, 0, extmark)
+    end
+
+    for i = start_row, end_row do
+      if i == start_row and virt_text_win_col > 0 then goto continue end
+      Utils.debug("add highlight", i - 1)
+      api.nvim_buf_add_highlight(bufnr, SUGGESTION_NS, Highlights.TO_BE_DELETED, i - 1, 0, -1)
+      ::continue::
     end
   end
-
-  local extmark = {
-    id = self.extmark_id,
-    virt_text_win_col = col,
-    virt_text = { { lines[1], Highlights.SUGGESTION } },
-  }
-
-  if #lines > 1 then
-    extmark.virt_lines = {}
-    for i = 2, #lines do
-      extmark.virt_lines[i - 1] = { { lines[i], Highlights.SUGGESTION } }
-    end
-  end
-
-  extmark.hl_mode = "combine"
-
-  local buf_lines = Utils.get_buf_lines(0, -1, bufnr)
-  local buf_lines_count = #buf_lines
-
-  while buf_lines_count < row do
-    api.nvim_buf_set_lines(bufnr, buf_lines_count, -1, false, { "" })
-    buf_lines_count = buf_lines_count + 1
-  end
-
-  api.nvim_buf_set_extmark(bufnr, SUGGESTION_NS, row - 1, extmark_col, extmark)
 end
 
 function Suggestion:is_visible()
-  return not not api.nvim_buf_get_extmark_by_id(0, SUGGESTION_NS, self.extmark_id, { details = false })[1]
+  local extmarks = api.nvim_buf_get_extmarks(0, SUGGESTION_NS, 0, -1, { details = false })
+  return #extmarks > 0
 end
 
-function Suggestion:hide() api.nvim_buf_del_extmark(0, SUGGESTION_NS, self.extmark_id) end
+function Suggestion:hide() api.nvim_buf_clear_namespace(0, SUGGESTION_NS, 0, -1) end
 
 function Suggestion:ctx()
   local bufnr = api.nvim_get_current_buf()
   local ctx = self._contexts[bufnr]
   if not ctx then
     ctx = {
-      suggestions = {},
+      suggestions_list = {},
-      current_suggestion_idx = 0,
+      current_suggestions_idx = 0,
       prev_doc = {},
+      internal_move = false,
     }
     self._contexts[bufnr] = ctx
   end
@@ -244,15 +315,15 @@ end
 
 function Suggestion:next()
   local ctx = self:ctx()
-  if #ctx.suggestions == 0 then return end
+  if #ctx.suggestions_list == 0 then return end
-  ctx.current_suggestion_idx = (ctx.current_suggestion_idx % #ctx.suggestions) + 1
+  ctx.current_suggestions_idx = (ctx.current_suggestions_idx % #ctx.suggestions_list) + 1
   self:show()
 end
 
 function Suggestion:prev()
   local ctx = self:ctx()
-  if #ctx.suggestions == 0 then return end
+  if #ctx.suggestions_list == 0 then return end
-  ctx.current_suggestion_idx = ((ctx.current_suggestion_idx - 2 + #ctx.suggestions) % #ctx.suggestions) + 1
+  ctx.current_suggestions_idx = ((ctx.current_suggestions_idx - 2 + #ctx.suggestions_list) % #ctx.suggestions_list) + 1
   self:show()
 end
 
@@ -262,56 +333,131 @@ function Suggestion:dismiss()
   self:reset()
 end
 
+function Suggestion:get_current_suggestion()
+  local ctx = self:ctx()
+  local suggestions = ctx.suggestions_list and ctx.suggestions_list[ctx.current_suggestions_idx] or nil
+  if not suggestions then return nil end
+  local cursor_row, _ = Utils.get_cursor_pos(0)
+  Utils.debug("cursor row", cursor_row)
+  for _, suggestion in ipairs(suggestions) do
+    if suggestion.original_start_row - 1 <= cursor_row and suggestion.end_row >= cursor_row then return suggestion end
+  end
+end
+
+function Suggestion:get_next_suggestion()
+  local ctx = self:ctx()
+  local suggestions = ctx.suggestions_list and ctx.suggestions_list[ctx.current_suggestions_idx] or nil
+  if not suggestions then return nil end
+  local cursor_row, _ = Utils.get_cursor_pos()
+  local new_suggestions = {}
+  for _, suggestion in ipairs(suggestions) do
+    table.insert(new_suggestions, suggestion)
+  end
+  --- sort the suggestions by cursor distance
+  table.sort(
+    new_suggestions,
+    function(a, b) return math.abs(a.start_row - cursor_row) < math.abs(b.start_row - cursor_row) end
+  )
+  --- get the closest suggestion to the cursor
+  return new_suggestions[1]
+end
+
 function Suggestion:accept()
   -- Llm.cancel_inflight_request()
-  api.nvim_buf_del_extmark(0, SUGGESTION_NS, self.extmark_id)
   local ctx = self:ctx()
-  local suggestion = ctx.suggestions and ctx.suggestions[ctx.current_suggestion_idx] or nil
+  local suggestions = ctx.suggestions_list and ctx.suggestions_list[ctx.current_suggestions_idx] or nil
-  if not suggestion then
+  Utils.debug("suggestions", suggestions)
+  if not suggestions then
     if Config.mappings.suggestion and Config.mappings.suggestion.accept == "<Tab>" then
       api.nvim_feedkeys(api.nvim_replace_termcodes("<Tab>", true, false, true), "n", true)
     end
     return
   end
-  local bufnr = api.nvim_get_current_buf()
+  local suggestion = self:get_current_suggestion()
-  local current_lines = Utils.get_buf_lines(0, -1, bufnr)
+  Utils.debug("current suggestion", suggestion)
-  local row = suggestion.row
+  if not suggestion then
-  local col = suggestion.col
+    suggestion = self:get_next_suggestion()
-  local content = suggestion.content
+    if suggestion then
-  local lines = vim.split(content, "\n")
+      Utils.debug("next suggestion", suggestion)
-  local cursor_row, cursor_col = Utils.get_cursor_pos()
+      local lines = api.nvim_buf_get_lines(0, 0, -1, false)
-  if row > cursor_row then api.nvim_buf_set_lines(bufnr, row - 1, row - 1, false, { "" }) end
+      local first_line_row = suggestion.start_row
-  local line_count = #lines
+      if first_line_row > 1 then first_line_row = first_line_row - 1 end
-  if line_count > 0 then
+      local line = lines[first_line_row]
-    if cursor_row == row then
+      local col = 0
-      local cursor_line_col = #current_lines[cursor_row] - 1
+      if line ~= nil then col = #line end
-      if cursor_col ~= cursor_line_col then
+      self:set_internal_move(true)
-        local current_line_ = current_lines[cursor_row]
+      api.nvim_win_set_cursor(0, { first_line_row, col })
-        lines[1] = lines[1] .. current_line_:sub(col + 1, -1)
+      vim.cmd("normal! zz")
-      end
+      vim.cmd("startinsert")
-    end
+      self:set_internal_move(false)
-    local current_line = current_lines[row] or ""
+      return
-    local current_line_max_col = #current_line - 1
-    local start_col = col
-    if start_col > current_line_max_col then
-      lines[1] = string.rep(" ", start_col - current_line_max_col - 1) .. lines[1]
-      start_col = -1
-    end
-    api.nvim_buf_set_text(bufnr, row - 1, start_col, row - 1, -1, { lines[1] })
-    if #lines > 1 then
-      local insert_lines = vim.list_slice(lines, 2)
-      api.nvim_buf_set_lines(bufnr, row, row, true, insert_lines)
     end
   end
+  if not suggestion then return end
+  api.nvim_buf_del_extmark(0, SUGGESTION_NS, suggestion.id)
+  local bufnr = api.nvim_get_current_buf()
+  local start_row = suggestion.start_row
+  local end_row = suggestion.end_row
+  local content = suggestion.content
+  local lines = vim.split(content, "\n")
+  local cursor_row, _ = Utils.get_cursor_pos()
+
+  local replaced_line_count = end_row - start_row + 1
+
+  if replaced_line_count > #lines then
+    Utils.debug("delete lines")
+    api.nvim_buf_set_lines(bufnr, start_row + #lines - 1, end_row, false, {})
+    api.nvim_buf_set_lines(bufnr, start_row - 1, start_row + #lines, false, lines)
+  else
+    Utils.debug("replace lines", start_row - 1, end_row, lines)
+    api.nvim_buf_set_lines(bufnr, start_row - 1, end_row, false, lines)
+  end
+
+  local row_diff = #lines - replaced_line_count
+
+  ctx.suggestions_list[ctx.current_suggestions_idx] = vim
+    .iter(suggestions)
+    :filter(function(s) return s.start_row ~= suggestion.start_row end)
+    :map(function(s)
+      if s.start_row > suggestion.start_row then
+        s.original_start_row = s.original_start_row + row_diff
+        s.start_row = s.start_row + row_diff
+        s.end_row = s.end_row + row_diff
+      end
+      return s
+    end)
+    :totable()
+
+  local line_count = #lines
 
   local down_count = line_count - 1
-  if row > cursor_row then down_count = down_count + 1 end
+  if start_row > cursor_row then down_count = down_count + 1 end
 
   local cursor_keys = string.rep("<Down>", down_count) .. "<End>"
-  api.nvim_feedkeys(api.nvim_replace_termcodes(cursor_keys, true, false, true), "n", false)
+  suggestions = ctx.suggestions_list and ctx.suggestions_list[ctx.current_suggestions_idx] or {}
 
-  self:hide()
+  if #suggestions > 0 then self:set_internal_move(true) end
-  self:reset()
+  api.nvim_feedkeys(api.nvim_replace_termcodes(cursor_keys, true, false, true), "n", false)
+  if #suggestions > 0 then self:set_internal_move(false) end
+end
+
+function Suggestion:is_internal_move()
+  local ctx = self:ctx()
+  Utils.debug("is internal move", ctx and ctx.internal_move)
+  return ctx and ctx.internal_move
+end
+
+function Suggestion:set_internal_move(internal_move)
+  local ctx = self:ctx()
+  if not internal_move then
+    vim.schedule(function()
+      Utils.debug("set internal move", internal_move)
+      ctx.internal_move = internal_move
+    end)
+  else
+    Utils.debug("set internal move", internal_move)
+    ctx.internal_move = internal_move
+  end
 end
 
 function Suggestion:setup_autocmds()
@@ -319,13 +465,20 @@ function Suggestion:setup_autocmds()
   local last_cursor_pos = {}
 
   local check_for_suggestion = Utils.debounce(function()
+    if self.is_on_throttle then return end
     local current_cursor_pos = api.nvim_win_get_cursor(0)
     if last_cursor_pos[1] == current_cursor_pos[1] and last_cursor_pos[2] == current_cursor_pos[2] then
+      self.is_on_throttle = true
+      vim.defer_fn(function() self.is_on_throttle = false end, Config.suggestion.throttle)
       self:suggest()
     end
-  end, 700)
+  end, Config.suggestion.debounce)
 
   local function suggest_callback()
+    if self.is_on_throttle then return end
+
+    if self:is_internal_move() then return end
+
     if not vim.bo.buflisted then return end
 
     if vim.bo.buftype ~= "" then return end

lua/avante/templates/suggesting.avanterules

@@ -5,22 +5,45 @@ Your task is to suggest code modifications at the cursor position. Follow these
   2. You must follow this JSON format when suggesting modifications:
     {% raw %}
     [
-      {
+      [
-        "row": ${row},
+        {
-        "col": ${column},
+          "start_row": ${start_row},
-        "content": "Your suggested code here"
+          "end_row": ${end_row},
-      }
+          "content": "Your suggested code here"
+        },
+        {
+          "start_row": ${start_row},
+          "end_row": ${end_row},
+          "content": "Your suggested code here"
+        }
+      ],
+      [
+        {
+          "start_row": ${start_row},
+          "end_row": ${end_row},
+          "content": "Your suggested code here"
+        },
+        {
+          "start_row": ${start_row},
+          "end_row": ${end_row},
+          "content": "Your suggested code here"
+        }
     ]
     {% endraw %}
 
+    JSON fields explanation:
+      start_row: The starting row of the code snippet you want to replace (1-indexed), inclusive
+      end_row: The ending row of the code snippet you want to replace (1-indexed), inclusive
+      content: The suggested code you want to replace the original code with
+
 Guidelines:
   1. Make sure you have maintained the user's existing whitespace and indentation. This is REALLY IMPORTANT!
-  2. DO NOT include three backticks: {%raw%}```{%endraw%} in your suggestion. Treat the suggested code AS IS.
+  2. Each segment in the returned list must be non-overlapping, and together they constitute this code modification.
-  3. Each element in the returned list is a COMPLETE and INDEPENDENT code snippet.
+  3. DO NOT include three backticks: {%raw%}```{%endraw%} in your suggestion. Treat the suggested code AS IS.
-  4. MUST be a valid JSON format. DON NOT be lazy!
+  4. Each element in the returned list is a COMPLETE code snippet.
-  5. Only return the new code to be inserted.
+  5. MUST be a valid JSON format. DO NOT be lazy!
-  6. Your returned code should not overlap with the original code in any way. Don't be lazy!
+  6. Only return the new code to be inserted. DON NOT be lazy!
-  7. Please strictly check the code around the position and ensure that the complete code after insertion is correct. Don't be lazy!
+  7. Please strictly check the code around the position and ensure that the complete code after insertion is correct. DO NOT be lazy!
   8. Do not return the entire file content or any surrounding code.
   9. Do not include any explanations, comments, or line numbers in your response.
   10. Ensure the suggested code fits seamlessly with the existing code structure and indentation.

tests/utils/init_spec.lua

@@ -120,4 +120,57 @@ describe("Utils", function()
       assert.equals("diagnostics", mentions[2].command)
     end)
   end)
+
+  describe("debounce", function()
+    it("should debounce function calls", function()
+      local count = 0
+      local debounced = Utils.debounce(function() count = count + 1 end, 100)
+
+      -- Call multiple times in quick succession
+      debounced()
+      debounced()
+      debounced()
+
+      -- Should not have executed yet
+      assert.equals(0, count)
+
+      -- Wait for debounce timeout
+      vim.wait(200, function() return false end)
+
+      -- Should have executed once
+      assert.equals(1, count)
+    end)
+
+    it("should cancel previous timer on new calls", function()
+      local count = 0
+      local debounced = Utils.debounce(function(c) count = c end, 100)
+
+      -- First call
+      debounced(1)
+
+      -- Wait partial time
+      vim.wait(50, function() return false end)
+
+      -- Second call should cancel first
+      debounced(233)
+
+      -- Wait for timeout
+      vim.wait(200, function() return false end)
+
+      -- Should only execute the latest once
+      assert.equals(233, count)
+    end)
+
+    it("should pass arguments correctly", function()
+      local result
+      local debounced = Utils.debounce(function(x, y) result = x + y end, 100)
+
+      debounced(2, 3)
+
+      -- Wait for timeout
+      vim.wait(200, function() return false end)
+
+      assert.equals(5, result)
+    end)
+  end)
 end)