fix: confirm ux (#1653)

2025-03-20 19:00:23 +08:00
parent 968d5fbd52
commit af372d7328
3 changed files with 118 additions and 37 deletions
--- a/lua/avante/llm_tools.lua
+++ b/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
--- a/lua/avante/sidebar.lua
+++ b/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
--- a/lua/avante/ui/confirm.lua
+++ b/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 button_line_num = 2 + #vim.split(message, "\n")
+  local keybindings_content = "<C-w>f: focus; c: code; r: resp; i: input"
+  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")),
-      "",
-      button_line,
+      keybindings_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 },
-    enter = true,
+    size = { width = win_width, height = #content + 3 },
+    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(popup.bufnr, 0, no_style, button_line_num, no_button_pos[1], no_button_pos[2])
-    focus_button(button_line_num + 1)
+    vim.api.nvim_buf_add_highlight(
+      popup.bufnr,
+      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 col >= yes_button_pos[1] and col <= yes_button_pos[2] then
-          focus_index = 1
-          render_buttons()
-        elseif col >= no_button_pos[1] and col <= no_button_pos[2] then
-          focus_index = 2
-          render_buttons()
-        end
+      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
+        focus_index = 1
+        render_content()
+      elseif col >= no_button_pos[1] and col <= no_button_pos[2] then
+        focus_index = 2
+        render_content()
      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