From 154e5f578f8925135a9dd23764d4c33a10c7ae36 Mon Sep 17 00:00:00 2001
From: yetone <yetoneful@gmail.com>
Date: Tue, 15 Jul 2025 20:46:15 +0800
Subject: [PATCH] feat: add timeout for bash tool (#2451)

---
 lua/avante/llm_tools/bash.lua |  2 +-
 lua/avante/utils/init.lua     | 43 ++++++++++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/lua/avante/llm_tools/bash.lua b/lua/avante/llm_tools/bash.lua
index 8bf7195..3e40169 100644
--- a/lua/avante/llm_tools/bash.lua
+++ b/lua/avante/llm_tools/bash.lua
@@ -251,7 +251,7 @@ function M.func(input, opts)
       Utils.shell_run_async(input.command, "bash -c", function(output, exit_code)
         local result, err = handle_result(output, exit_code)
         opts.on_complete(result, err)
-      end, abs_path)
+      end, abs_path, 1000 * 60 * 2)
     end,
     { focus = true },
     opts.session_ctx,
diff --git a/lua/avante/utils/init.lua b/lua/avante/utils/init.lua
index 919cd2c..f0fc36f 100644
--- a/lua/avante/utils/init.lua
+++ b/lua/avante/utils/init.lua
@@ -117,11 +117,31 @@ end
 ---@param shell_cmd string?
 ---@param on_complete fun(output: string, code: integer)
 ---@param cwd? string
-function M.shell_run_async(input_cmd, shell_cmd, on_complete, cwd)
+---@param timeout? integer Timeout in milliseconds
+function M.shell_run_async(input_cmd, shell_cmd, on_complete, cwd, timeout)
   local cmd = get_cmd_for_shell(input_cmd, shell_cmd)
   ---@type string[]
   local output = {}
-  fn.jobstart(cmd, {
+  local timer = nil
+  local completed = false
+
+  -- Create a wrapper for on_complete to ensure it's only called once
+  local function complete_once(out, code)
+    if completed then return end
+    completed = true
+
+    -- Clean up timer if it exists
+    if timer then
+      timer:stop()
+      timer:close()
+      timer = nil
+    end
+
+    on_complete(out, code)
+  end
+
+  -- Start the job
+  local job_id = fn.jobstart(cmd, {
     on_stdout = function(_, data)
       if not data then return end
       vim.list_extend(output, data)
@@ -130,9 +150,26 @@ function M.shell_run_async(input_cmd, shell_cmd, on_complete, cwd)
       if not data then return end
       vim.list_extend(output, data)
     end,
-    on_exit = function(_, exit_code) on_complete(table.concat(output, "\n"), exit_code) end,
+    on_exit = function(_, exit_code) complete_once(table.concat(output, "\n"), exit_code) end,
     cwd = cwd,
   })
+
+  -- Set up timeout if specified
+  if timeout and timeout > 0 then
+    timer = vim.loop.new_timer()
+    if timer then
+      timer:start(timeout, 0, function()
+        vim.schedule(function()
+          if not completed and job_id then
+            -- Kill the job
+            fn.jobstop(job_id)
+            -- Complete with timeout error
+            complete_once("Command timed out after " .. timeout .. "ms", 124)
+          end
+        end)
+      end)
+    end
+  end
 end
 
 ---@see https://github.com/LazyVim/LazyVim/blob/main/lua/lazyvim/util/toggle.lua