From 5c20cc177966be153e1a9fd668f7c2a4bade9b2d Mon Sep 17 00:00:00 2001
From: Fernando Freire <dogonthehorizon@gmail.com>
Date: Fri, 13 Dec 2024 06:57:42 -0800
Subject: [PATCH] feat(context): add current buffer to selected file ctx (#941)

---
 lua/avante/config.lua        |  3 +++
 lua/avante/file_selector.lua | 34 +++++++++++++++++++++++++++++++++-
 lua/avante/sidebar.lua       | 15 +++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/lua/avante/config.lua b/lua/avante/config.lua
index 97ac4e0..da820da 100644
--- a/lua/avante/config.lua
+++ b/lua/avante/config.lua
@@ -189,6 +189,9 @@ M.defaults = {
       remove_file = "d",
       add_file = "@",
     },
+    files = {
+      add_current = "<leader>ac", -- Add current buffer to selected files
+    },
   },
   windows = {
     ---@alias AvantePosition "right" | "left" | "top" | "bottom" | "smart"
diff --git a/lua/avante/file_selector.lua b/lua/avante/file_selector.lua
index 794e3ea..2ea6c9a 100644
--- a/lua/avante/file_selector.lua
+++ b/lua/avante/file_selector.lua
@@ -32,7 +32,39 @@ function FileSelector:reset()
   self.event_handlers = {}
 end
 
-function FileSelector:add_selected_file(filepath) table.insert(self.selected_filepaths, Utils.uniform_path(filepath)) end
+function FileSelector:add_selected_file(filepath)
+    local uniform_path = Utils.uniform_path(filepath)
+    -- Avoid duplicates
+    if not vim.tbl_contains(self.selected_filepaths, uniform_path) then
+        table.insert(self.selected_filepaths, uniform_path)
+        self:emit("update")
+    end
+end
+
+function FileSelector:add_current_buffer()
+    local current_buf = vim.api.nvim_get_current_buf()
+    local filepath = vim.api.nvim_buf_get_name(current_buf)
+
+    -- Only process if it's a real file buffer
+    if filepath and filepath ~= "" and not vim.startswith(filepath, "avante://") then
+        local relative_path = require("avante.utils").relative_path(filepath)
+
+        -- Check if file is already in list
+        for i, path in ipairs(self.selected_filepaths) do
+            if path == relative_path then
+                -- Remove if found
+                table.remove(self.selected_filepaths, i)
+                self:emit("update")
+                return true
+            end
+        end
+
+        -- Add if not found
+        self:add_selected_file(relative_path)
+        return true
+    end
+    return false
+end
 
 function FileSelector:on(event, callback)
   local handlers = self.event_handlers[event]
diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua
index 4f8745f..ec9c9bc 100644
--- a/lua/avante/sidebar.lua
+++ b/lua/avante/sidebar.lua
@@ -826,6 +826,21 @@ end
 function Sidebar:on_mount(opts)
   self:refresh_winids()
 
+  -- Add keymap to add current buffer while sidebar is open
+  if Config.mappings.files and Config.mappings.files.add_current then
+    vim.keymap.set("n", Config.mappings.files.add_current, function()
+      if self:is_open() and self.file_selector:add_current_buffer() then
+        vim.notify("Added current buffer to file selector", vim.log.levels.DEBUG, { title = "Avante" })
+      else
+        vim.notify("Failed to add current buffer", vim.log.levels.WARN, { title = "Avante" })
+      end
+    end, {
+      desc = "avante: add current buffer to file selector",
+      noremap = true,
+      silent = true,
+    })
+  end
+
   api.nvim_set_option_value("wrap", Config.windows.wrap, { win = self.result_container.winid })
 
   local current_apply_extmark_id = nil