fix(ui): set selection per buffer and remove spinner (closes #32) (#62)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
2024-08-18 05:36:30 -04:00 · 2024-08-18 05:36:30 -04:00 · d885bd9680
commit d885bd9680
parent c19ea9a48a
4 changed files with 102 additions and 85 deletions
--- a/lua/avante/ai_bot.lua
+++ b/lua/avante/ai_bot.lua
@ -282,8 +282,6 @@ local function call_claude_api_stream(question, code_lang, code_content, selecte

  local url = Utils.trim_suffix(Config.claude.endpoint, "/") .. "/v1/messages"

-  -- print("Sending request to Claude API...")
-
  curl.post(url, {
    ---@diagnostic disable-next-line: unused-local
    stream = function(err, data, job)
@ -390,8 +388,6 @@ local function call_openai_api_stream(question, code_lang, code_content, selecte
    }
  end

-  -- print("Sending request to " .. (config.get().provider == "azure" and "Azure OpenAI" or "OpenAI") .. " API...")
-
  curl.post(url, {
    ---@diagnostic disable-next-line: unused-local
    stream = function(err, data, job)
--- a/lua/avante/init.lua
+++ b/lua/avante/init.lua
@ -1,22 +1,21 @@
 local api = vim.api

-local Tiktoken = require("avante.tiktoken")
 local Sidebar = require("avante.sidebar")
-local Config = require("avante.config")
-local Diff = require("avante.diff")
-local AiBot = require("avante.ai_bot")
 local Selection = require("avante.selection")
+local Config = require("avante.config")

 ---@class Avante
 local M = {
  ---@type avante.Sidebar[] we use this to track chat command across tabs
  sidebars = {},
-  ---@type avante.Sidebar
-  current = nil,
-  selection = nil,
-  _once = false,
+  ---@type avante.Selection[]
+  selections = {},
+  ---@type {sidebar?: avante.Sidebar, selection?: avante.Selection}
+  current = { sidebar = nil, selection = nil },
 }

+M.did_setup = false
+
 local H = {}

 H.commands = function()
@ -52,7 +51,7 @@ H.autocmds = function()
    local name = "avante.nvim"
    local load_path = function()
      require("tiktoken_lib").load()
-      Tiktoken.setup("gpt-4o")
+      require("avante.tiktoken").setup("gpt-4o")
    end

    if LazyConfig.plugins[name] and LazyConfig.plugins[name]._.loaded then
@ -80,15 +79,24 @@ H.autocmds = function()
    callback = function(ev)
      local tab = tonumber(ev.file)
      local s = M.sidebars[tab]
+      local sl = M.selections[tab]
      if s then
        s:destroy()
      end
+      if sl then
+        sl:delete_autocmds()
+      end
      if tab ~= nil then
        M.sidebars[tab] = nil
      end
    end,
  })

+  vim.schedule(function()
+    M._init(api.nvim_get_current_tabpage())
+    M.current.selection:setup_autocmds()
+  end)
+
  -- automatically setup Avante filetype to markdown
  vim.treesitter.language.register("markdown", "Avante")
 end
@ -98,24 +106,33 @@ end
 function M._get(current)
  local tab = api.nvim_get_current_tabpage()
  local sidebar = M.sidebars[tab]
+  local selection = M.selections[tab]
  if current ~= false then
-    M.current = sidebar
+    M.current.sidebar = sidebar
+    M.current.selection = selection
  end
  return sidebar
 end

-M.open = function()
-  local tab = api.nvim_get_current_tabpage()
-  local sidebar = M.sidebars[tab]
+---@param id integer
+function M._init(id)
+  local sidebar = M.sidebars[id]
+  local selection = M.selections[id]

  if not sidebar then
-    sidebar = Sidebar:new(tab)
-    M.sidebars[tab] = sidebar
+    sidebar = Sidebar:new(id)
+    M.sidebars[id] = sidebar
  end
+  if not selection then
+    selection = Selection:new(id)
+    M.selections[id] = selection
+  end
+  M.current = { sidebar = sidebar, selection = selection }
+  return M
+end

-  M.current = sidebar
-
-  return sidebar:open()
+M.open = function()
+  M._init(api.nvim_get_current_tabpage())._get(false):open()
 end

 M.toggle = function()
@ -160,20 +177,19 @@ function M.setup(opts)
  ---but most of the other functionality will only be called once from lazy.nvim
  Config.setup(opts)

-  if M._once then
+  if M.did_setup then
    return
  end

-  Diff.setup()
-  AiBot.setup()
-  M.selection = Selection:new():setup()
+  require("avante.diff").setup()
+  require("avante.ai_bot").setup()

  -- setup helpers
  H.autocmds()
  H.commands()
  H.keymaps()

-  M._once = true
+  M.did_setup = true
 end

 return M
--- a/lua/avante/selection.lua
+++ b/lua/avante/selection.lua
@ -6,13 +6,17 @@ local fn = vim.fn
 local NAMESPACE = api.nvim_create_namespace("avante_selection")
 local PRIORITY = vim.highlight.priorities.user

+---@class avante.Selection
 local Selection = {}

-function Selection:new()
+Selection.did_setup = false
+
+---@param id integer the tabpage id retrieved from vim.api.nvim_get_current_tabpage()
+function Selection:new(id)
  return setmetatable({
    hints_popup_extmark_id = nil,
    edit_popup_renderer = nil,
-    augroup = api.nvim_create_augroup("avante_selection", { clear = true }),
+    augroup = api.nvim_create_augroup("avante_selection_" .. id, { clear = true }),
  }, { __index = self })
 end

@ -56,8 +60,9 @@ function Selection:close_hints_popup()
  end
 end

-function Selection:setup()
-  vim.api.nvim_create_autocmd({ "ModeChanged" }, {
+function Selection:setup_autocmds()
+  Selection.did_setup = true
+  api.nvim_create_autocmd({ "ModeChanged" }, {
    group = self.augroup,
    pattern = { "n:v", "n:V", "n:" }, -- Entering Visual mode from Normal mode
    callback = function()
@ -91,6 +96,7 @@ function Selection:delete_autocmds()
    vim.api.nvim_del_augroup_by_id(self.augroup)
  end
  self.augroup = nil
+  Selection.did_setup = false
 end

 return Selection
--- a/lua/avante/sidebar.lua
+++ b/lua/avante/sidebar.lua
@ -219,28 +219,34 @@ function Sidebar:intialize()
  return self
 end

+function Sidebar:is_focused()
+  return self.view:is_open() and self.code.win == api.nvim_get_current_win()
+end
+
 ---@param content string concatenated content of the buffer
---@param focus? boolean whether to focus the result view
-function Sidebar:update_content(content, focus, callback)
-  focus = focus or false
+---@param opts? {focus?: boolean, scroll?: boolean, callback?: fun(): nil} whether to focus the result view
+function Sidebar:update_content(content, opts)
+  opts = vim.tbl_deep_extend("force", { focus = true, scroll = true, callback = nil }, opts or {})
  vim.defer_fn(function()
    api.nvim_set_option_value("modifiable", true, { buf = self.view.buf })
    api.nvim_buf_set_lines(self.view.buf, 0, -1, false, vim.split(content, "\n"))
    api.nvim_set_option_value("modifiable", false, { buf = self.view.buf })
    api.nvim_set_option_value("filetype", "Avante", { buf = self.view.buf })
-    if callback ~= nil then
-      callback()
+    if opts.callback ~= nil then
+      opts.callback()
    end
-    if focus then
+    if opts.focus and not self:is_focused() then
      xpcall(function()
        --- set cursor to bottom of result view
        api.nvim_set_current_win(self.winid.result)
      end, function(err)
-        -- XXX: omit error for now, but should fix me why it can't jump here.
        return err
      end)
+    end
+
+    if opts.scroll then
      xpcall(function()
-        api.nvim_win_set_cursor(self.winid.result, { api.nvim_buf_line_count(self.view.buf), 0 })
+        api.nvim_win_set_cursor(self.winid.result, { api.nvim_buf_line_count(self.bufnr.result), 0 })
      end, function(err)
        return err
      end)
@ -411,7 +417,7 @@ function Sidebar:update_content_with_history(history)
    content = content .. entry.response .. "\n\n"
    content = content .. "---\n\n"
  end
-  self:update_content(content, true)
+  self:update_content(content)
 end

 local function get_conflict_content(content, snippets)
@ -582,10 +588,15 @@ function Sidebar:render()
  self:update_content_with_history(chat_history)

  local function handle_submit()
+    signal.is_loading = true
    local state = signal:get_value()
    local user_input = state.text

    local timestamp = get_timestamp()
+    --- HACK: we need to set focus to true and scroll to false to
+    --- prevent the cursor from jumping to the bottom of the
+    --- buffer at the beginning
+    self:update_content("", { focus = true, scroll = false })
    self:update_content(
      "## "
        .. timestamp
@ -607,8 +618,6 @@ function Sidebar:render()

    local full_response = ""

-    signal.is_loading = true
-
    local filetype = api.nvim_get_option_value("filetype", { buf = self.code.buf })

    AiBot.call_ai_api_stream(
@ -617,11 +626,9 @@ function Sidebar:render()
      content_with_line_numbers,
      selected_code_content_with_line_numbers,
      function(chunk)
+        signal.is_loading = true
        full_response = full_response .. chunk
-        self:update_content(
-          "## " .. timestamp .. "\n\n> " .. user_input:gsub("\n", "\n> ") .. "\n\n" .. full_response,
-          true
-        )
+        self:update_content("## " .. timestamp .. "\n\n> " .. user_input:gsub("\n", "\n> ") .. "\n\n" .. full_response)
        vim.schedule(function()
          vim.cmd("redraw")
        end)
@ -638,8 +645,7 @@ function Sidebar:render()
              .. "\n\n"
              .. full_response
              .. "\n\n🚨 Error: "
-              .. vim.inspect(err),
-            true
+              .. vim.inspect(err)
          )
          return
        end
@ -653,10 +659,11 @@ function Sidebar:render()
            .. "\n\n"
            .. full_response
            .. "\n\n🎉🎉🎉 **Generation complete!** Please review the code suggestions above.\n\n\n\n",
-          true,
-          function()
-            api.nvim_exec_autocmds("User", { pattern = VIEW_BUFFER_UPDATED_PATTERN })
-          end
+          {
+            callback = function()
+              api.nvim_exec_autocmds("User", { pattern = VIEW_BUFFER_UPDATED_PATTERN })
+            end,
+          }
        )

        -- Save chat history
@ -676,11 +683,11 @@ function Sidebar:render()
    local code_file_fullpath = api.nvim_buf_get_name(self.code.buf)
    local code_filename = fn.fnamemodify(code_file_fullpath, ":t")

-    local input_label = string.format(" 🙋 with %s %s (<Tab> switch focus): ", icon, code_filename)
+    local input_label = string.format(" 🙋 with %s %s (<Tab>: switch focus): ", icon, code_filename)

    if self.code.selection ~= nil then
      input_label = string.format(
-        " 🙋 with selected code in %s %s(%d:%d) (<Tab> switch focus): ",
+        " 🙋 with %s %s(%d:%d) (<Tab>: switch focus): ",
        icon,
        code_filename,
        self.code.selection.range.start.line,
@ -747,37 +754,29 @@ function Sidebar:render()
          right = 1,
        },
      }),
-      N.columns(
-        { flex = 0 },
-        N.text_input({
-          id = "input",
-          border_label = {
-            text = input_label,
-          },
-          placeholder = "Enter your question",
-          autofocus = true,
-          wrap = true,
-          flex = 1,
-          on_change = function(value)
-            local state = signal:get_value()
-            local is_enter = value:sub(-1) == "\n" and #state.text < #value
-            if is_enter then
-              value = value:sub(1, -2)
-            end
-            signal.text = value
-            if is_enter and #value > 0 then
-              handle_submit()
-            end
-          end,
-          padding = { left = 1, right = 1 },
-        }),
-        N.gap(1),
-        N.spinner({
-          is_loading = signal.is_loading,
-          padding = { top = 1, right = 1 },
-          hidden = signal.is_loading:negate(),
-        })
-      )
+      N.gap(1),
+      N.text_input({
+        id = "input",
+        border_label = {
+          text = input_label,
+        },
+        placeholder = "Enter your question",
+        autofocus = true,
+        wrap = true,
+        flex = 1,
+        on_change = function(value)
+          local state = signal:get_value()
+          local is_enter = value:sub(-1) == "\n" and #state.text < #value
+          if is_enter then
+            value = value:sub(1, -2)
+          end
+          signal.text = value
+          if is_enter and #value > 0 then
+            handle_submit()
+          end
+        end,
+        padding = { left = 1, right = 1 },
+      })
    )
  end