fix: confirm ux (#1653)

lua/avante/llm_tools.lua

@@ -21,8 +21,9 @@ end
 
 ---@param message string
 ---@param callback fun(yes: boolean)
+---@param opts? { focus?: boolean }
 ---@return avante.ui.Confirm | nil
-function M.confirm(message, callback)
+function M.confirm(message, callback, opts)
   local Confirm = require("avante.ui.confirm")
   local sidebar = require("avante").get()
   if not sidebar or not sidebar.input_container or not sidebar.input_container.winid then
@@ -30,7 +31,8 @@ function M.confirm(message, callback)
     callback(false)
     return
   end
-  local confirm = Confirm:new(message, callback, { container_winid = sidebar.input_container.winid })
+  local confirm_opts = vim.tbl_deep_extend("force", { container_winid = sidebar.input_container.winid }, opts or {})
+  local confirm = Confirm:new(message, callback, confirm_opts)
   confirm:open()
   return confirm
 end
@@ -292,7 +294,8 @@ function M.str_replace_editor(opts, on_log, on_complete)
       vim.api.nvim_buf_set_lines(bufnr, start_line - 1, end_line, false, new_lines)
       vim.api.nvim_set_current_win(current_winid)
       on_complete(true, nil)
-    end)
+    end, { focus = false })
+    vim.api.nvim_set_current_win(sidebar.code.winid)
     vim.api.nvim_create_autocmd({ "TextChangedI", "TextChanged" }, {
       group = augroup,
       buffer = bufnr,
@@ -302,6 +305,7 @@ function M.str_replace_editor(opts, on_log, on_complete)
         if current_lines_content:find(patch_end_line_content) then return end
         vim.api.nvim_del_augroup_by_id(augroup)
         if confirm then confirm:close() end
+        if vim.api.nvim_win_is_valid(current_winid) then vim.api.nvim_set_current_win(current_winid) end
         if lines_content == current_lines_content then
           on_complete(false, "User canceled")
           return

lua/avante/sidebar.lua

@@ -37,7 +37,7 @@ local Sidebar = {}
 ---@field id integer
 ---@field augroup integer
 ---@field code avante.CodeState
----@field winids table<string, integer>
+---@field winids table<"result_container" | "selected_code_container" | "selected_files_container" | "input_container", integer>
 ---@field result_container NuiSplit | nil
 ---@field selected_code_container NuiSplit | nil
 ---@field selected_files_container NuiSplit | nil
@@ -1872,6 +1872,7 @@ function Sidebar:refresh_winids()
 
   local function reverse_switch_windows()
     local current_winid = api.nvim_get_current_win()
+    winids = vim.iter(winids):filter(function(winid) return api.nvim_win_is_valid(winid) end):totable()
     local current_idx = Utils.tbl_indexof(winids, current_winid) or 1
     if current_idx == 1 then
       current_idx = #winids

lua/avante/ui/confirm.lua

@@ -6,50 +6,62 @@ local Utils = require("avante.utils")
 ---@class avante.ui.Confirm
 ---@field message string
 ---@field callback fun(yes: boolean)
----@field opts { container_winid: number }
+---@field _container_winid number | nil
+---@field _focus boolean | nil
+---@field _group number | nil
 ---@field _popup NuiPopup | nil
+---@field _prev_winid number | nil
 local M = {}
 M.__index = M
 
 ---@param message string
 ---@param callback fun(yes: boolean)
----@param opts { container_winid: number }
+---@param opts { container_winid: number, focus?: boolean }
 ---@return avante.ui.Confirm
 function M:new(message, callback, opts)
   local this = setmetatable({}, M)
   this.message = message
   this.callback = callback
-  this.opts = opts
+  this._container_winid = opts.container_winid
+  this._focus = opts.focus
   return this
 end
 
 function M:open()
+  self._prev_winid = vim.api.nvim_get_current_win()
   local message = self.message
   local callback = self.callback
-  local opts = self.opts
+
+  local win_width = 60
 
   local focus_index = 2 -- 1 = Yes, 2 = No
-  local yes_button_pos = { 23, 28 }
-  local no_button_pos = { 33, 37 }
 
   local BUTTON_NORMAL = Highlights.BUTTON_DEFAULT
   local BUTTON_FOCUS = Highlights.BUTTON_DEFAULT_HOVER
 
-  local button_line = string.rep(" ", 23) .. " Yes       No "
+  local keybindings_content = "<C-w>f: focus; c: code; r: resp; i: input"
-  local button_line_num = 2 + #vim.split(message, "\n")
+  local keybidings_start_col = math.floor((win_width - #keybindings_content) / 2)
+  local buttons_content = " Yes       No "
+  local buttons_start_col = math.floor((win_width - #buttons_content) / 2)
+  local yes_button_pos = { buttons_start_col, buttons_start_col + 5 }
+  local no_button_pos = { buttons_start_col + 10, buttons_start_col + 14 }
+  local keybindings_line = string.rep(" ", keybidings_start_col) .. keybindings_content
+  local buttons_line = string.rep(" ", buttons_start_col) .. buttons_content
+  local keybindings_line_num = 1 + #vim.split(message, "\n")
+  local buttons_line_num = 2 + #vim.split(message, "\n")
   local content = vim
     .iter({
       "",
       vim.tbl_map(function(line) return "  " .. line end, vim.split(message, "\n")),
-      "",
+      keybindings_line,
-      button_line,
+      buttons_line,
       "",
     })
     :flatten()
     :totable()
   local button_row = #content - 1
 
-  local container_winid = opts.container_winid or vim.api.nvim_get_current_win()
+  local container_winid = self._container_winid or vim.api.nvim_get_current_win()
   local container_width = vim.api.nvim_win_get_width(container_winid)
 
   local popup = Popup({
@@ -59,10 +71,10 @@ function M:open()
     },
     position = {
       row = vim.o.lines - #content - 3,
-      col = (container_width - 60) / 2,
+      col = (container_width - win_width) / 2,
     },
-    size = { width = 60, height = #content + 3 },
+    size = { width = win_width, height = #content + 3 },
-    enter = true,
+    enter = self._focus ~= false,
     focusable = true,
     border = {
       style = "rounded",
@@ -87,7 +99,7 @@ function M:open()
     end
   end
 
-  local function render_buttons()
+  local function render_content()
     local yes_style = (focus_index == 1) and BUTTON_FOCUS or BUTTON_NORMAL
     local no_style = (focus_index == 2) and BUTTON_FOCUS or BUTTON_NORMAL
 
@@ -95,25 +107,57 @@ function M:open()
     vim.api.nvim_buf_set_lines(popup.bufnr, 0, -1, false, content)
     Utils.lock_buf(popup.bufnr)
 
-    vim.api.nvim_buf_add_highlight(popup.bufnr, 0, yes_style, button_line_num, yes_button_pos[1], yes_button_pos[2])
+    vim.api.nvim_buf_add_highlight(
-    vim.api.nvim_buf_add_highlight(popup.bufnr, 0, no_style, button_line_num, no_button_pos[1], no_button_pos[2])
+      popup.bufnr,
-    focus_button(button_line_num + 1)
+      0,
+      "Comment",
+      keybindings_line_num,
+      keybidings_start_col,
+      keybidings_start_col + #keybindings_content
+    )
+    vim.api.nvim_buf_add_highlight(popup.bufnr, 0, yes_style, buttons_line_num, yes_button_pos[1], yes_button_pos[2])
+    vim.api.nvim_buf_add_highlight(popup.bufnr, 0, no_style, buttons_line_num, no_button_pos[1], no_button_pos[2])
+    focus_button(buttons_line_num + 1)
   end
 
   local function select_button()
-    popup:unmount()
+    self:close()
     callback(focus_index == 1)
   end
 
+  vim.keymap.set("n", "c", function()
+    local sidebar = require("avante").get()
+    if not sidebar then return end
+    if sidebar.code.winid and vim.api.nvim_win_is_valid(sidebar.code.winid) then
+      vim.api.nvim_set_current_win(sidebar.code.winid)
+    end
+  end, { buffer = popup.bufnr })
+
+  vim.keymap.set("n", "r", function()
+    local sidebar = require("avante").get()
+    if not sidebar then return end
+    if sidebar.winids.result_container and vim.api.nvim_win_is_valid(sidebar.winids.result_container) then
+      vim.api.nvim_set_current_win(sidebar.winids.result_container)
+    end
+  end, { buffer = popup.bufnr })
+
+  vim.keymap.set("n", "i", function()
+    local sidebar = require("avante").get()
+    if not sidebar then return end
+    if sidebar.winids.input_container and vim.api.nvim_win_is_valid(sidebar.winids.input_container) then
+      vim.api.nvim_set_current_win(sidebar.winids.input_container)
+    end
+  end, { buffer = popup.bufnr })
+
   vim.keymap.set("n", "y", function()
     focus_index = 1
-    render_buttons()
+    render_content()
     select_button()
   end, { buffer = popup.bufnr })
 
   vim.keymap.set("n", "n", function()
     focus_index = 2
-    render_buttons()
+    render_content()
     select_button()
   end, { buffer = popup.bufnr })
 
@@ -146,11 +190,11 @@ function M:open()
       if row == button_row then
         if col >= yes_button_pos[1] and col <= yes_button_pos[2] then
           focus_index = 1
-          render_buttons()
+          render_content()
           select_button()
         elseif col >= no_button_pos[1] and col <= no_button_pos[2] then
           focus_index = 2
-          render_buttons()
+          render_content()
           select_button()
         end
       end
@@ -163,26 +207,58 @@ function M:open()
     buffer = popup.bufnr,
     callback = function()
       local row, col = unpack(vim.api.nvim_win_get_cursor(0))
-      if row == button_row then
+      if row ~= button_row then vim.api.nvim_win_set_cursor(self._popup.winid, { button_row, buttons_start_col }) end
-        if col >= yes_button_pos[1] and col <= yes_button_pos[2] then
+      if col >= yes_button_pos[1] and col <= yes_button_pos[2] then
-          focus_index = 1
+        focus_index = 1
-          render_buttons()
+        render_content()
-        elseif col >= no_button_pos[1] and col <= no_button_pos[2] then
+      elseif col >= no_button_pos[1] and col <= no_button_pos[2] then
-          focus_index = 2
+        focus_index = 2
-          render_buttons()
+        render_content()
-        end
       end
     end,
   })
 
-  -- popup:on(event.BufLeave, function() popup:unmount() end)
+  self._group = self._group and self._group or vim.api.nvim_create_augroup("AvanteConfirm", { clear = true })
+  vim.api.nvim_create_autocmd("WinClosed", {
+    group = self._group,
+    callback = function()
+      local winids = vim.api.nvim_list_wins()
+      if not vim.list_contains(winids, self._container_winid) then self:close() end
+    end,
+  })
 
   popup:mount()
-  render_buttons()
+  render_content()
   self._popup = popup
+  self:bind_window_focus_keymaps()
 end
 
+function M:window_focus_handler()
+  local current_winid = vim.api.nvim_get_current_win()
+  if
+    current_winid == self._popup.winid
+    and current_winid ~= self._prev_winid
+    and vim.api.nvim_win_is_valid(self._prev_winid)
+  then
+    vim.api.nvim_set_current_win(self._prev_winid)
+    return
+  end
+  self._prev_winid = current_winid
+  vim.api.nvim_set_current_win(self._popup.winid)
+end
+
+function M:bind_window_focus_keymaps()
+  vim.keymap.set({ "n", "i" }, "<C-w>f", function() self:window_focus_handler() end)
+end
+
+function M:unbind_window_focus_keymaps() vim.keymap.del({ "n", "i" }, "<C-w>f") end
+
 function M:close()
+  self:unbind_window_focus_keymaps()
+  if self._group then
+    vim.api.nvim_del_augroup_by_id(self._group)
+    self._group = nil
+  end
   if self._popup then
     self._popup:unmount()
     self._popup = nil