fix(tools): failed/cancelled/rejected tool calls no longer block the execution queue (#1852)

Davidyz · web-flow · commit 48d7e8f6b2fb · 2025-07-18T23:24:23.000+01:00
Closes #1850
diff --git a/doc/extending/tools.md b/doc/extending/tools.md
@@ -590,7 +590,7 @@ output = {
 
 This will notify the user with the message: `Perform the calculation 100 multiply 50?`. The user can choose to proceed, reject or cancel. The latter will cancel any tools from running.
 
-You can also customize the output if a user rejects the approval:
+You can also customize the output if a user rejects the approval or cancels the tool execution:
 
 ```lua
 output = {
@@ -604,6 +604,15 @@ output = {
   rejected = function(self, agent, cmd)
     agent.chat:add_tool_output(self, "The user declined to run the calculator tool")
   end,
+
+  ---Cancellation message back to the LLM
+  ---@param self CodeCompanion.Tool.Calculator
+  ---@param agent CodeCompanion.Agent
+  ---@param cmd table
+  ---@return nil
+  cancelled = function(self, agent, cmd)
+    agent.chat:add_tool_output(self, "The user cancelled the execution of the calculator tool")
+  end,
 },
 ```
 
diff --git a/lua/codecompanion/config.lua b/lua/codecompanion/config.lua
@@ -164,6 +164,7 @@ local defaults = {
           folds = {
             enabled = true, -- Fold tool output in the buffer?
             failure_words = { -- Words that indicate an error in the tool output. Used to apply failure highlighting
+              "cancelled",
               "error",
               "failed",
               "invalid",
diff --git a/lua/codecompanion/strategies/chat/agents/executor/cmd.lua b/lua/codecompanion/strategies/chat/agents/executor/cmd.lua
@@ -84,12 +84,14 @@ function CmdExecutor:run(cmd, index)
               return self.executor:setup()
             end
           else
-            return self.executor:error(cmd, string.format("Failed with code %s", code))
+            self.executor:error(cmd, string.format("Failed with code %s", code))
+            return self.executor:close()
           end
         end)
 
         if not ok then
-          return self.executor:error(cmd, string.format("Error whilst running command %s: %s", cmd, output))
+          self.executor:error(cmd, string.format("Error whilst running command %s: %s", cmd, output))
+          return self.executor:close()
         end
       end)
     end,
diff --git a/lua/codecompanion/strategies/chat/agents/executor/func.lua b/lua/codecompanion/strategies/chat/agents/executor/func.lua
@@ -76,7 +76,8 @@ function FuncExecutor:run(func, action, input, callback)
     end
     tool_finished = true
     if msg.status == self.executor.agent.constants.STATUS_ERROR then
-      return self.executor:error(action, msg.data or "An error occurred")
+      self.executor:error(action, msg.data or "An error occurred")
+      return self.executor:close()
     end
 
     self.executor:success(action, msg.data)
@@ -90,7 +91,8 @@ function FuncExecutor:run(func, action, input, callback)
     return func(self.executor.agent, action, input, output_handler)
   end)
   if not ok then
-    return self.executor:error(action, output)
+    self.executor:error(action, output)
+    return self.executor:close()
   end
 
   if output ~= nil then
diff --git a/lua/codecompanion/strategies/chat/agents/executor/init.lua b/lua/codecompanion/strategies/chat/agents/executor/init.lua
@@ -4,6 +4,15 @@ local Queue = require("codecompanion.strategies.chat.agents.executor.queue")
 local log = require("codecompanion.utils.log")
 local util = require("codecompanion.utils")
 
+local fmt = string.format
+
+---Add a response to the chat buffer regarding a tool's execution
+---@param exec CodeCompanion.Agent.Executor
+---@msg string
+local send_response_to_chat = function(exec, msg)
+  exec.agent.chat:add_tool_output(exec.tool, msg)
+end
+
 ---@class CodeCompanion.Agent.Executor
 ---@field agent CodeCompanion.Agent
 ---@field current_cmd_tool table The current cmd tool that's being executed
@@ -63,16 +72,30 @@ function Executor:setup_handlers()
     rejected = function(cmd)
       if self.tool.output and self.tool.output.rejected then
         self.tool.output.rejected(self.tool, self.agent, cmd)
+      else
+        -- If no handler is set then return a default message
+        send_response_to_chat(self, fmt("User rejected `%s`", self.tool.name))
       end
     end,
     error = function(cmd)
       if self.tool.output and self.tool.output.error then
         self.tool.output.error(self.tool, self.agent, cmd, self.agent.stderr, self.agent.stdout or {})
+      else
+        send_response_to_chat(self, fmt("Error calling `%s`", self.tool.name))
+      end
+    end,
+    cancelled = function(cmd)
+      if self.tool.output and self.tool.output.cancelled then
+        self.tool.output.cancelled(self.tool, self.agent, cmd)
+      else
+        send_response_to_chat(self, fmt("Cancelled `%s`", self.tool.name))
       end
     end,
     success = function(cmd)
       if self.tool.output and self.tool.output.success then
         self.tool.output.success(self.tool, self.agent, cmd, self.agent.stdout)
+      else
+        send_response_to_chat(self, fmt("Executed `%s`", self.tool.name))
       end
     end,
   }
@@ -95,7 +118,7 @@ function Executor:setup(input)
   end
   if self.agent.status == self.agent.constants.STATUS_ERROR then
     log:debug("Executor:execute - Error")
-    return finalize_agent(self)
+    self:close()
   end
 
   -- Get the next tool to run
@@ -146,9 +169,9 @@ function Executor:setup(input)
       }, function(choice)
         if not choice or choice == "Cancel" then -- No selection or cancelled
           log:debug("Executor:execute - Tool cancelled")
-          finalize_agent(self)
           self:close()
-          return
+          self.output.cancelled(cmd)
+          return self:setup()
         elseif choice == "Yes" then -- Selected yes
           log:debug("Executor:execute - Tool approved")
           self:execute(cmd, input)
@@ -190,8 +213,7 @@ function Executor:error(action, error)
     log:warn("Tool %s: %s", self.tool.name, error)
   end
   self.output.error(action)
-  finalize_agent(self)
-  self:close()
+  self:setup()
 end
 
 ---Handle a successful completion of a tool
@@ -210,9 +232,14 @@ end
 ---Close the execution of the tool
 ---@return nil
 function Executor:close()
-  log:debug("Executor:close")
-  self.handlers.on_exit()
-  util.fire("ToolFinished", { id = self.id, name = self.tool.name, bufnr = self.agent.bufnr })
+  --TODO: This is a workaround that avoids the close method being called more than once
+  if self.tool then
+    log:debug("Executor:close")
+    self.handlers.on_exit()
+    util.fire("ToolFinished", { id = self.id, name = self.tool.name, bufnr = self.agent.bufnr })
+    self.tool = nil
+    self.current_cmd_tool = {}
+  end
 end
 
 return Executor
diff --git a/lua/codecompanion/strategies/chat/agents/init.lua b/lua/codecompanion/strategies/chat/agents/init.lua
@@ -127,11 +127,6 @@ function Agent:execute(chat, tools)
   ---@param executor CodeCompanion.Agent.Executor The executor instance
   ---@param tool table The tool to run
   local function enqueue_tool(executor, tool)
-    -- If an error occurred, don't run any more tools
-    if self.status == CONSTANTS.STATUS_ERROR then
-      return
-    end
-
     local name = tool["function"].name
     local tool_config = self.tools_config[name]
     local function handle_missing_tool(tool_call, err_message)
@@ -190,7 +185,6 @@ function Agent:execute(chat, tools)
       local env = type(self.tool.env) == "function" and self.tool.env(vim.deepcopy(self.tool)) or {}
       util.replace_placeholders(self.tool.cmds, env)
     end
-
     return executor.queue:push(self.tool)
   end
 
diff --git a/lua/codecompanion/strategies/chat/init.lua b/lua/codecompanion/strategies/chat/init.lua
@@ -1275,7 +1275,11 @@ function Chat:add_tool_output(tool, for_llm, for_user)
   -- Ensure that tool output is merged if it has the same tool call ID
   local existing = find_tool_call(tool_call.id, self.messages)
   if existing then
-    existing.content = existing.content .. "\n\n" .. output.content
+    if existing.content ~= "" then
+      existing.content = existing.content .. "\n\n" .. output.content
+    else
+      existing.content = output.content
+    end
   else
     table.insert(self.messages, output)
   end
diff --git a/lua/codecompanion/types.lua b/lua/codecompanion/types.lua
@@ -119,6 +119,7 @@
 ---@field output.rejected? fun(self: CodeCompanion.Agent.Tool, agent: CodeCompanion.Agent, cmd: table): any Function to call if the user rejects running a command
 ---@field output.error? fun(self: CodeCompanion.Agent.Tool, agent: CodeCompanion.Agent, cmd: table, stderr: table, stdout?: table): any The function to call if an error occurs
 ---@field output.success? fun(self: CodeCompanion.Agent.Tool, agent: CodeCompanion.Agent, cmd: table, stdout: table): any Function to call if the tool is successful
+---@field output.cancelled? fun(self: CodeCompanion.Agent.Tool, agent: CodeCompanion.Agent, cmd: table): any Function to call if the tool is cancelled
 ---@field args table The arguments sent over by the LLM when making the request
 ---@field tool table The tool configuration from the config file
 
diff --git a/tests/config.lua b/tests/config.lua
@@ -180,6 +180,10 @@ return {
           callback = vim.fn.getcwd() .. "/tests/strategies/chat/agents/tools/stubs/cmd_queue.lua",
           description = "Cmd tool",
         },
+        ["cmd_queue_error"] = {
+          callback = vim.fn.getcwd() .. "/tests/strategies/chat/agents/tools/stubs/cmd_queue_error.lua",
+          description = "Cmd tool",
+        },
         ["mock_cmd_runner"] = {
           callback = vim.fn.getcwd() .. "/tests/strategies/chat/agents/tools/stubs/mock_cmd_runner.lua",
           description = "Cmd tool",
diff --git a/tests/strategies/chat/agents/executor/test_queue.lua b/tests/strategies/chat/agents/executor/test_queue.lua
@@ -68,5 +68,41 @@ T["Agent"]["queue"]["can queue functions and commands"] = function()
   -- Test that the function was called
   h.eq("Data 1 Data 2", child.lua_get([[_G._test_func]]))
 end
+T["Agent"]["queue"]["can proceed on error"] = function()
+  h.eq(vim.NIL, child.lua_get([[_G._test_order]]))
+
+  child.lua([[
+    local tools = {
+      {
+        ["function"] = {
+          arguments = { data = "Data 1" },
+          name = "func_queue",
+        },
+      },
+      {
+        ["function"] = {
+          name = "cmd_queue_error",
+        },
+      },
+      {
+        ["function"] = {
+          arguments = { data = "Data 2" },
+          name = "func_queue_2",
+        },
+      },
+    }
+    agent:execute(chat, tools)
+    vim.wait(1000)
+  ]])
+
+  -- Test order
+  h.eq(
+    "Func[Setup]->Func[Success]->Func[Exit]->Cmd[Setup]->Cmd[Error]->Cmd[Exit]->Func2[Setup]->Func2[Success]->Func2[Exit]",
+    child.lua_get([[_G._test_order]])
+  )
+
+  -- Test that the function was called
+  h.eq("Data 1 Data 2", child.lua_get([[_G._test_func]]))
+end
 
 return T
diff --git a/tests/strategies/chat/agents/tools/stubs/cmd_queue_error.lua b/tests/strategies/chat/agents/tools/stubs/cmd_queue_error.lua
@@ -0,0 +1,26 @@
+return {
+  name = "cmd_queue_error",
+  system_prompt = "my cmd system prompt",
+  cmds = {
+    { "echofanweoufqwefvcergv", "0.5" },
+  },
+  handlers = {
+    -- Should only be called once
+    setup = function(self)
+      _G._test_order = (_G._test_order or "") .. "->Cmd[Setup]"
+      _G._test_setup = (_G._test_setup or "") .. "Setup"
+    end,
+    -- Should only be called once
+    on_exit = function(self)
+      _G._test_order = (_G._test_order or "") .. "->Cmd[Exit]"
+      _G._test_exit = (_G._test_exit or "") .. "Exited"
+    end,
+  },
+  output = {
+    -- Should only be called once
+    error = function(self, cmd, output)
+      _G._test_order = (_G._test_order or "") .. "->Cmd[Error]"
+      _G._test_output = _G._test_output or {}
+    end,
+  },
+}
