diff --git a/lua/avante/api.lua b/lua/avante/api.lua
index a37f4d0..aa67ad0 100644
--- a/lua/avante/api.lua
+++ b/lua/avante/api.lua
@@ -1,6 +1,6 @@
 local Config = require("avante.config")
 local Utils = require("avante.utils")
-local PromptInput = require("avante.prompt_input")
+local PromptInput = require("avante.ui.prompt_input")
 
 ---@class avante.ApiToggle
 ---@operator call(): boolean
diff --git a/lua/avante/llm_tools.lua b/lua/avante/llm_tools.lua
index cd18f1b..4126df3 100644
--- a/lua/avante/llm_tools.lua
+++ b/lua/avante/llm_tools.lua
@@ -21,16 +21,18 @@ end
 
 ---@param message string
 ---@param callback fun(yes: boolean)
----@return NuiPopup | nil
+---@return avante.ui.Confirm | nil
 function M.confirm(message, callback)
-  local UI = require("avante.ui")
+  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
     Utils.error("Avante sidebar not found", { title = "Avante" })
     callback(false)
     return
   end
-  return UI.confirm(message, callback, { container_winid = sidebar.input_container.winid })
+  local confirm = Confirm:new(message, callback, { container_winid = sidebar.input_container.winid })
+  confirm:open()
+  return confirm
 end
 
 ---@param abs_path string
@@ -277,7 +279,7 @@ function M.str_replace_editor(opts, on_log, on_complete)
     vim.cmd("normal! zz")
     vim.api.nvim_set_current_win(current_winid)
     local augroup = vim.api.nvim_create_augroup("avante_str_replace_editor", { clear = true })
-    local popup = M.confirm("Are you sure you want to apply this modification?", function(ok)
+    local confirm = M.confirm("Are you sure you want to apply this modification?", function(ok)
       vim.api.nvim_del_augroup_by_id(augroup)
       vim.api.nvim_set_current_win(sidebar.code.winid)
       vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes("<Esc>", true, false, true), "n", true)
@@ -299,7 +301,7 @@ function M.str_replace_editor(opts, on_log, on_complete)
         local current_lines_content = table.concat(current_lines, "\n")
         if current_lines_content:find(patch_end_line_content) then return end
         vim.api.nvim_del_augroup_by_id(augroup)
-        if popup then popup:unmount() end
+        if confirm then confirm:close() end
         if lines_content == current_lines_content then
           on_complete(false, "User canceled")
           return
diff --git a/lua/avante/selection.lua b/lua/avante/selection.lua
index 6483c67..e02c813 100644
--- a/lua/avante/selection.lua
+++ b/lua/avante/selection.lua
@@ -3,7 +3,7 @@ local Config = require("avante.config")
 local Llm = require("avante.llm")
 local Provider = require("avante.providers")
 local RepoMap = require("avante.repo_map")
-local PromptInput = require("avante.prompt_input")
+local PromptInput = require("avante.ui.prompt_input")
 
 local api = vim.api
 local fn = vim.fn
@@ -19,7 +19,7 @@ local PRIORITY = vim.highlight.priorities.user
 ---@field selected_code_extmark_id integer | nil
 ---@field augroup integer | nil
 ---@field code_winid integer | nil
----@field prompt_input PromptInput | nil
+---@field prompt_input avante.ui.PromptInput | nil
 local Selection = {}
 Selection.__index = Selection
 
diff --git a/lua/avante/ui.lua b/lua/avante/ui/confirm.lua
similarity index 87%
rename from lua/avante/ui.lua
rename to lua/avante/ui/confirm.lua
index e2dca7c..ac58da3 100644
--- a/lua/avante/ui.lua
+++ b/lua/avante/ui/confirm.lua
@@ -3,13 +3,31 @@ local NuiText = require("nui.text")
 local Highlights = require("avante.highlights")
 local Utils = require("avante.utils")
 
+---@class avante.ui.Confirm
+---@field message string
+---@field callback fun(yes: boolean)
+---@field opts { container_winid: number }
+---@field _popup NuiPopup | nil
 local M = {}
+M.__index = M
 
 ---@param message string
 ---@param callback fun(yes: boolean)
 ---@param opts { container_winid: number }
----@return NuiPopup
-function M.confirm(message, callback, opts)
+---@return avante.ui.Confirm
+function M:new(message, callback, opts)
+  local this = setmetatable({}, M)
+  this.message = message
+  this.callback = callback
+  this.opts = opts
+  return this
+end
+
+function M:open()
+  local message = self.message
+  local callback = self.callback
+  local opts = self.opts
+
   local focus_index = 2 -- 1 = Yes, 2 = No
   local yes_button_pos = { 23, 28 }
   local no_button_pos = { 33, 37 }
@@ -161,7 +179,16 @@ function M.confirm(message, callback, opts)
 
   popup:mount()
   render_buttons()
-  return popup
+  self._popup = popup
+end
+
+function M:close()
+  if self._popup then
+    self._popup:unmount()
+    self._popup = nil
+    return true
+  end
+  return false
 end
 
 return M
diff --git a/lua/avante/prompt_input.lua b/lua/avante/ui/prompt_input.lua
similarity index 98%
rename from lua/avante/prompt_input.lua
rename to lua/avante/ui/prompt_input.lua
index 7f4c7ae..54f3a9b 100644
--- a/lua/avante/prompt_input.lua
+++ b/lua/avante/ui/prompt_input.lua
@@ -3,7 +3,7 @@ local fn = vim.fn
 local Config = require("avante.config")
 local Utils = require("avante.utils")
 
----@class PromptInput
+---@class avante.ui.PromptInput
 ---@field bufnr integer | nil
 ---@field winid integer | nil
 ---@field win_opts table
@@ -21,7 +21,7 @@ local Utils = require("avante.utils")
 local PromptInput = {}
 PromptInput.__index = PromptInput
 
----@class PromptInputOptions
+---@class avante.ui.PromptInputOptions
 ---@field start_insert? boolean
 ---@field submit_callback? fun(input: string):nil
 ---@field cancel_callback? fun():nil
@@ -29,7 +29,7 @@ PromptInput.__index = PromptInput
 ---@field win_opts? table
 ---@field default_value? string
 
----@param opts? PromptInputOptions
+---@param opts? avante.ui.PromptInputOptions
 function PromptInput:new(opts)
   opts = opts or {}
   local obj = setmetatable({}, PromptInput)