diff --git a/lua/telescope/actions/init.lua b/lua/telescope/actions/init.lua
index f9e24b6..5acc8da 100644
--- a/lua/telescope/actions/init.lua
+++ b/lua/telescope/actions/init.lua
@@ -1407,20 +1407,27 @@ actions.which_key = function(prompt_bufnr, opts)
     end
   end
 
-  -- only set up autocommand after showing preview completed
+  -- if close_with_action is true, close the which_key window when any action is triggered
+  -- otherwise close the window when the prompt buffer is closed
+  local close_event, close_pattern, close_buffer
   if opts.close_with_action then
-    vim.schedule(function()
-      vim.api.nvim_create_autocmd("User", {
-        pattern = "TelescopeKeymap",
-        once = true,
-        callback = function()
-          pcall(vim.api.nvim_win_close, km_win_id, true)
-          pcall(vim.api.nvim_win_close, km_opts.border.win_id, true)
-          require("telescope.utils").buf_delete(km_buf)
-        end,
-      })
-    end)
+    close_event, close_pattern, close_buffer = "User", "TelescopeKeymap", nil
+  else
+    close_event, close_pattern, close_buffer = "BufWinLeave", nil, prompt_bufnr
   end
+  -- only set up autocommand after showing preview completed
+  vim.schedule(function()
+    vim.api.nvim_create_autocmd(close_event, {
+      pattern = close_pattern,
+      buffer = close_buffer,
+      once = true,
+      callback = function()
+        pcall(vim.api.nvim_win_close, km_win_id, true)
+        pcall(vim.api.nvim_win_close, km_opts.border.win_id, true)
+        require("telescope.utils").buf_delete(km_buf)
+      end,
+    })
+  end)
 end
 
 --- Move from a none fuzzy search to a fuzzy one<br>