feat(context): add a ui for selecting and adding files to the sidebar as context (#912)

* feat(sidebar): supports select files

chore (context) update add type annotations to context functions

chore (sidebar) remove unused notify function call

refactor (sidebar) remove setting search file to file path

chore (sidebar) remove nvim_notify debugging api call

* feat (files) allow selecting a file by string via cmp suggestion menu

* chore (context) refactor to allow context using @file with a context view

* refactor (context) refactor seletected file types as an array of path and content

* refactor (config) remove unused configuration options

* refactor (sidebar) remove unused unbild key

* refactor (context) remove unused imports

* refactor (mentions) update mentions to support items with callback functions and removal of the underlying selection.

* fix (sidebar) add file context as a window that is visitable via the tab key

* refactor (file_content) remove file content as an input to llm

* feat (sidebar) support suggesting and applying code in all languages that are in the context

* feat (sidebar) configurable mapping for removing a file from the context.

* feat (context_view) configure hints for the context view for adding and deleting a file.

* feat (context) add hints for the context view.

* fix (sidebar) type when scrolling the results buffer.

* refactor (selected files) refactor llm stream to accept an array of selected file metadata

* refactor: context => selected_files

---------

Co-authored-by: yetone <yetoneful@gmail.com>

crates/avante-templates/src/lib.rs

@@ -15,13 +15,19 @@ impl<'a> State<'a> {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+struct SelectedFile {
+    path: String,
+    content: String,
+    file_type: String,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 struct TemplateContext {
     use_xml_format: bool,
     ask: bool,
     code_lang: String,
-    filepath: String,
+    selected_files: Vec<SelectedFile>,
-    file_content: String,
     selected_code: Option<String>,
     project_context: Option<String>,
     diagnostics: Option<String>,
@@ -44,8 +50,7 @@ fn render(state: &State, template: &str, context: TemplateContext) -> LuaResult<
                   use_xml_format => context.use_xml_format,
                   ask => context.ask,
                   code_lang => context.code_lang,
-                  filepath => context.filepath,
+                  selected_files => context.selected_files,
-                  file_content => context.file_content,
                   selected_code => context.selected_code,
                   project_context => context.project_context,
                   diagnostics => context.diagnostics,

lua/avante/api.lua

@@ -161,7 +161,7 @@ M.refresh = function(opts)
   if not sidebar:is_open() then return end
   local curbuf = vim.api.nvim_get_current_buf()
 
-  local focused = sidebar.result.bufnr == curbuf or sidebar.input.bufnr == curbuf
+  local focused = sidebar.result_container.bufnr == curbuf or sidebar.input_container.bufnr == curbuf
   if focused or not sidebar:is_open() then return end
   local listed = vim.api.nvim_get_option_value("buflisted", { buf = curbuf })
 
@@ -185,19 +185,19 @@ M.focus = function(opts)
   local curwin = vim.api.nvim_get_current_win()
 
   if sidebar:is_open() then
-    if curbuf == sidebar.input.bufnr then
+    if curbuf == sidebar.input_container.bufnr then
       if sidebar.code.winid and sidebar.code.winid ~= curwin then vim.api.nvim_set_current_win(sidebar.code.winid) end
-    elseif curbuf == sidebar.result.bufnr then
+    elseif curbuf == sidebar.result_container.bufnr then
       if sidebar.code.winid and sidebar.code.winid ~= curwin then vim.api.nvim_set_current_win(sidebar.code.winid) end
     else
-      if sidebar.input.winid and sidebar.input.winid ~= curwin then
+      if sidebar.input_container.winid and sidebar.input_container.winid ~= curwin then
-        vim.api.nvim_set_current_win(sidebar.input.winid)
+        vim.api.nvim_set_current_win(sidebar.input_container.winid)
       end
     end
   else
     if sidebar.code.winid then vim.api.nvim_set_current_win(sidebar.code.winid) end
     sidebar:open(opts)
-    if sidebar.input.winid then vim.api.nvim_set_current_win(sidebar.input.winid) end
+    if sidebar.input_container.winid then vim.api.nvim_set_current_win(sidebar.input_container.winid) end
   end
 end
 

lua/avante/config.lua

@@ -186,6 +186,8 @@ M.defaults = {
       apply_cursor = "a",
       switch_windows = "<Tab>",
       reverse_switch_windows = "<S-Tab>",
+      remove_file = "d",
+      add_file = "@",
     },
   },
   windows = {

lua/avante/diff.lua

@@ -2,7 +2,6 @@ local api = vim.api
 
 local Config = require("avante.config")
 local Utils = require("avante.utils")
-local Highlights = require("avante.highlights")
 
 local H = {}
 local M = {}
@@ -558,7 +557,7 @@ end
 ---@param enable_autojump boolean
 function M.process_position(bufnr, side, position, enable_autojump)
   local lines = {}
-  if vim.tbl_contains({ SIDES.OURS, SIDES.THEIRS, SIDES.BASE }, side) then
+  if vim.tbl_contains({ SIDES.OURS, SIDES.THEIRS }, side) then
     local data = position[name_map[side]]
     lines = Utils.get_buf_lines(data.content_start, data.content_end + 1)
   elseif side == SIDES.BOTH then
@@ -569,7 +568,7 @@ function M.process_position(bufnr, side, position, enable_autojump)
     lines = {}
   elseif side == SIDES.CURSOR then
     local cursor_line = Utils.get_cursor_pos()
-    for _, pos in ipairs({ SIDES.OURS, SIDES.THEIRS, SIDES.BASE }) do
+    for _, pos in ipairs({ SIDES.OURS, SIDES.THEIRS }) do
       local data = position[name_map[pos]] or {}
       if data.range_start and data.range_start + 1 <= cursor_line and data.range_end + 1 >= cursor_line then
         side = pos

lua/avante/file_selector.lua

@@ -0,0 +1,146 @@
+local Utils = require("avante.utils")
+local Path = require("plenary.path")
+local scan = require("plenary.scandir")
+
+--- @class FileSelector
+local FileSelector = {}
+
+--- @class FileSelector
+--- @field id integer
+--- @field selected_filepaths string[]
+--- @field file_cache string[]
+--- @field event_handlers table<string, function[]>
+
+---@param id integer
+---@return FileSelector
+function FileSelector:new(id)
+  return setmetatable({
+    id = id,
+    selected_files = {},
+    file_cache = {},
+    event_handlers = {},
+  }, { __index = self })
+end
+
+function FileSelector:reset()
+  self.selected_filepaths = {}
+  self.event_handlers = {}
+end
+
+function FileSelector:add_selected_file(filepath) table.insert(self.selected_filepaths, Utils.uniform_path(filepath)) end
+
+function FileSelector:on(event, callback)
+  local handlers = self.event_handlers[event]
+  if not handlers then
+    handlers = {}
+    self.event_handlers[event] = handlers
+  end
+
+  table.insert(handlers, callback)
+end
+
+function FileSelector:emit(event, ...)
+  local handlers = self.event_handlers[event]
+  if not handlers then return end
+
+  for _, handler in ipairs(handlers) do
+    handler(...)
+  end
+end
+
+function FileSelector:off(event, callback)
+  if not callback then
+    self.event_handlers[event] = {}
+    return
+  end
+  local handlers = self.event_handlers[event]
+  if not handlers then return end
+
+  for i, handler in ipairs(handlers) do
+    if handler == callback then
+      table.remove(handlers, i)
+      break
+    end
+  end
+end
+
+---@return nil
+function FileSelector:open()
+  self:update_file_cache()
+  self:show_select_ui()
+end
+
+---@return nil
+function FileSelector:update_file_cache()
+  local project_root = Path:new(Utils.get_project_root()):absolute()
+
+  local filepaths = scan.scan_dir(project_root, {
+    respect_gitignore = true,
+  })
+
+  -- Sort buffer names alphabetically
+  table.sort(filepaths, function(a, b) return a < b end)
+
+  self.file_cache = vim
+    .iter(filepaths)
+    :map(function(filepath) return Path:new(filepath):make_relative(project_root) end)
+    :totable()
+end
+
+---@return nil
+function FileSelector:show_select_ui()
+  vim.schedule(function()
+    local filepaths = vim
+      .iter(self.file_cache)
+      :filter(function(filepath) return not vim.tbl_contains(self.selected_filepaths, filepath) end)
+      :totable()
+
+    vim.ui.select(filepaths, {
+      prompt = "(Avante) Add a file:",
+      format_item = function(item) return item end,
+    }, function(filepath)
+      if not filepath then return end
+      table.insert(self.selected_filepaths, Utils.uniform_path(filepath))
+      self:emit("update")
+    end)
+  end)
+
+  -- unlist the current buffer as vim.ui.select will be listed
+  local winid = vim.api.nvim_get_current_win()
+  local bufnr = vim.api.nvim_win_get_buf(winid)
+  vim.api.nvim_set_option_value("buflisted", false, { buf = bufnr })
+  vim.api.nvim_set_option_value("bufhidden", "wipe", { buf = bufnr })
+end
+
+---@param idx integer
+---@return boolean
+function FileSelector:remove_selected_filepaths(idx)
+  if idx > 0 and idx <= #self.selected_filepaths then
+    table.remove(self.selected_filepaths, idx)
+    self:emit("update")
+    return true
+  end
+  return false
+end
+
+---@return { path: string, content: string, file_type: string }[]
+function FileSelector:get_selected_files_contents()
+  local contents = {}
+  for _, file_path in ipairs(self.selected_filepaths) do
+    local file = io.open(file_path, "r")
+    if file then
+      local content = file:read("*all")
+      file:close()
+
+      -- Detect the file type
+      local filetype = vim.filetype.match({ filename = file_path, contents = contents }) or "unknown"
+
+      table.insert(contents, { path = file_path, content = content, file_type = filetype })
+    end
+  end
+  return contents
+end
+
+function FileSelector:get_selected_filepaths() return vim.deepcopy(self.selected_filepaths) end
+
+return FileSelector

lua/avante/llm.lua

@@ -44,14 +44,11 @@ M._stream = function(opts, Provider)
 
   Path.prompts.initialize(Path.prompts.get(opts.bufnr))
 
-  local filepath = Utils.relative_path(api.nvim_buf_get_name(opts.bufnr))
-
   local template_opts = {
     use_xml_format = Provider.use_xml_format,
     ask = opts.ask, -- TODO: add mode without ask instruction
     code_lang = opts.code_lang,
-    filepath = filepath,
+    selected_files = opts.selected_files,
-    file_content = opts.file_content,
     selected_code = opts.selected_code,
     project_context = opts.project_context,
     diagnostics = opts.diagnostics,
@@ -335,14 +332,19 @@ end
 
 ---@alias LlmMode "planning" | "editing" | "suggesting"
 ---
+---@class SelectedFiles
+---@field path string
+---@field content string
+---@field file_type string
+---
 ---@class TemplateOptions
 ---@field use_xml_format boolean
 ---@field ask boolean
 ---@field question string
 ---@field code_lang string
----@field file_content string
 ---@field selected_code string | nil
 ---@field project_context string | nil
+---@field selected_files SelectedFiles[] | nil
 ---@field diagnostics string | nil
 ---@field history_messages AvanteLLMMessage[]
 ---
@@ -357,8 +359,22 @@ end
 
 ---@param opts StreamOptions
 M.stream = function(opts)
-  if opts.on_chunk ~= nil then opts.on_chunk = vim.schedule_wrap(opts.on_chunk) end
+  local is_completed = false
-  if opts.on_complete ~= nil then opts.on_complete = vim.schedule_wrap(opts.on_complete) end
+  if opts.on_chunk ~= nil then
+    local original_on_chunk = opts.on_chunk
+    opts.on_chunk = vim.schedule_wrap(function(chunk)
+      if is_completed then return end
+      return original_on_chunk(chunk)
+    end)
+  end
+  if opts.on_complete ~= nil then
+    local original_on_complete = opts.on_complete
+    opts.on_complete = vim.schedule_wrap(function(err)
+      if is_completed then return end
+      is_completed = true
+      return original_on_complete(err)
+    end)
+  end
   local Provider = opts.provider or P[Config.provider]
   if Config.dual_boost.enabled then
     M._dual_boost_stream(opts, Provider, P[Config.dual_boost.first_provider], P[Config.dual_boost.second_provider])

lua/avante/path.lua

@@ -15,6 +15,7 @@ local Config = require("avante.config")
 ---@field selected_file {filepath: string}?
 ---@field selected_code {filetype: string, content: string}?
 ---@field reset_memory boolean?
+---@field selected_filepaths string[] | nil
 
 ---@class avante.Path
 ---@field history_path Path

lua/avante/selection.lua

@@ -210,7 +210,7 @@ function Selection:create_editing_input()
       ask = true,
       project_context = vim.json.encode(project_context),
       diagnostics = vim.json.encode(diagnostics),
-      file_content = code_content,
+      selected_files = { { content = code_content, file_type = filetype, path = "" } },
       code_lang = filetype,
       selected_code = self.selection.content,
       instructions = input,

lua/avante/suggestion.lua

@@ -71,7 +71,7 @@ function Suggestion:suggest()
     provider = provider,
     bufnr = bufnr,
     ask = true,
-    file_content = code_content,
+    selected_files = { { content = code_content, file_type = filetype, path = "" } },
     code_lang = filetype,
     instructions = vim.json.encode(doc),
     mode = "suggesting",

lua/avante/templates/_context.avanterules

@@ -1,12 +1,15 @@
 {%- if use_xml_format -%}
-<filepath>{{filepath}}</filepath>
 
 {% if selected_code -%}
+{% for file in selected_files %}
+<filepath>{{file.path}}</filepath>
+
 <context>
-```{{code_lang}}
+```{{file.file_type}}
-{{file_content}}
+{{file.content}}
 ```
 </context>
+{% endfor %}
 
 <code>
 ```{{code_lang}}
@@ -14,28 +17,38 @@
 ```
 </code>
 {%- else -%}
+{% for file in selected_files %}
+<filepath>{{file.path}}</filepath>
+
 <code>
-```{{code_lang}}
+```{{file.file_type}}
-{{file_content}}
+{{file.content}}
 ```
 </code>
+{% endfor %}
 {%- endif %}
 {% else %}
-FILEPATH: {{filepath}}
-
 {% if selected_code -%}
+{% for file in selected_files %}
+FILEPATH: {{file.path}}
+
 CONTEXT:
-```{{code_lang}}
+```{{file.file_type}}
-{{file_content}}
+{{file.content}}
 ```
+{% endfor %}
 
 CODE:
 ```{{code_lang}}
 {{selected_code}}
 ```
 {%- else -%}
+{% for file in selected_files %}
+FILEPATH: {{file.path}}
+
 CODE:
-```{{code_lang}}
+```{{file.file_type}}
-{{file_content}}
+{{file.content}}
 ```
+{% endfor %}
 {%- endif %}{%- endif %}

lua/avante/utils/init.lua

@@ -58,17 +58,17 @@ M.shell_run = function(input_cmd)
   -- powershell then we can just run the cmd
   if shell:match("powershell") or shell:match("pwsh") then
     cmd = input_cmd
-  elseif vim.fn.has("wsl") > 0 then
+  elseif fn.has("wsl") > 0 then
     -- wsl: powershell.exe -Command 'command "/path"'
     cmd = "powershell.exe -NoProfile -Command '" .. input_cmd:gsub("'", '"') .. "'"
-  elseif vim.fn.has("win32") > 0 then
+  elseif fn.has("win32") > 0 then
     cmd = 'powershell.exe -NoProfile -Command "' .. input_cmd:gsub('"', "'") .. '"'
   else
     -- linux and macos we wil just do sh -c
-    cmd = "sh -c " .. vim.fn.shellescape(input_cmd)
+    cmd = "sh -c " .. fn.shellescape(input_cmd)
   end
 
-  local output = vim.fn.system(cmd)
+  local output = fn.system(cmd)
   local code = vim.v.shell_error
 
   return { stdout = output, code = code }
@@ -562,10 +562,10 @@ function M.debounce(func, delay)
 end
 
 function M.winline(winid)
-  local current_win = vim.api.nvim_get_current_win()
+  local current_win = api.nvim_get_current_win()
-  vim.api.nvim_set_current_win(winid)
+  api.nvim_set_current_win(winid)
-  local line = vim.fn.winline()
+  local line = fn.winline()
-  vim.api.nvim_set_current_win(current_win)
+  api.nvim_set_current_win(current_win)
   return line
 end
 
@@ -725,7 +725,7 @@ function M.get_or_create_buffer_with_filepath(filepath)
   api.nvim_set_current_buf(buf)
 
   -- Use the edit command to load the file content and set the buffer name
-  vim.cmd("edit " .. vim.fn.fnameescape(filepath))
+  vim.cmd("edit " .. fn.fnameescape(filepath))
 
   return buf
 end
@@ -823,4 +823,13 @@ function M.get_current_selection_diagnostics(bufnr, selection)
   return selection_diagnostics
 end
 
+function M.uniform_path(path)
+  local project_root = M.get_project_root()
+  local abs_path = Path:new(project_root):joinpath(path):absolute()
+  local relative_path = Path:new(abs_path):make_relative(project_root)
+  return relative_path
+end
+
+function M.is_same_file(filepath_a, filepath_b) return M.uniform_path(filepath_a) == M.uniform_path(filepath_b) end
+
 return M

lua/cmp_avante/mentions.lua

@@ -44,4 +44,37 @@ function mentions_source:complete(_, callback)
   })
 end
 
+---@param completion_item table
+---@param callback fun(response: {behavior: number})
+function mentions_source:execute(completion_item, callback)
+  local current_line = api.nvim_get_current_line()
+  local label = completion_item.label:match("^@(%S+)") -- Extract mention command without '@' and space
+
+  -- Find the corresponding mention
+  local selected_mention
+  for _, mention in ipairs(self.mentions) do
+    if mention.command == label then
+      selected_mention = mention
+      break
+    end
+  end
+
+  -- Execute the mention's callback if it exists
+  if selected_mention and type(selected_mention.callback) == "function" then
+    selected_mention.callback(selected_mention)
+    -- Get the current cursor position
+    local row, col = unpack(api.nvim_win_get_cursor(0))
+
+    -- Replace the current line with the new line (removing the mention)
+    local new_line = current_line:gsub(vim.pesc(completion_item.label), "")
+    api.nvim_buf_set_lines(0, row - 1, row, false, { new_line })
+
+    -- Adjust the cursor position if needed
+    local new_col = math.min(col, #new_line)
+    api.nvim_win_set_cursor(0, { row, new_col })
+  end
+
+  callback({ behavior = require("cmp").ConfirmBehavior.Insert })
+end
+
 return mentions_source