orchestrator.lua

---@class CodeCompanion.Tools.Orchestrator
---@field cancelled boolean? Whether the tool execution has been cancelled
---@field handlers table<string, function>
---@field id number The id of the tools coordinator
---@field index number The index of the current command
---@field output table<string, function>
---@field queue CodeCompanion.Tools.Orchestrator.Queue
---@field status string The status of the tool execution "success" | "error"
---@field tool CodeCompanion.Tools.Tool The current tool being executed
---@field tool_handle vim.SystemObj? SystemObj of the current tool execution
---@field tool_output table? The output collected from the tool
---@field tools CodeCompanion.Tools

local Approvals = require("codecompanion.interactions.chat.tools.approvals")
local Queue = require("codecompanion.interactions.chat.tools.runtime.queue")
local Runner = require("codecompanion.interactions.chat.tools.runtime.runner")
local log = require("codecompanion.utils.log")
local os_utils = require("codecompanion.utils.os")
local ui_utils = require("codecompanion.utils.ui")
local utils = require("codecompanion.utils")

local fmt = string.format
local Orchestrator = {}

---Strip any ANSI color codes which don't render in the chat buffer
---@param tbl table
---@return table
local function strip_ansi(tbl)
  for i, v in ipairs(tbl) do
    tbl[i] = v:gsub("\027%[[0-9;]*%a", "")
  end
  return tbl
end

---Add a response to the chat buffer regarding a tool's execution
---@param llm_message string
---@param user_message? string
function Orchestrator:_send_response_to_chat(llm_message, user_message)
  self.tools.chat:add_tool_output(self.tool, llm_message, user_message)
end

---Execute a shell command with platform-specific handling
---@param cmd table
---@param callback function
function Orchestrator:_execute_shell_command(cmd, callback)
  if vim.fn.has("win32") == 1 then
    -- See PR #2186
    local shell_cmd = table.concat(cmd, " ") .. "\r\nEXIT %ERRORLEVEL%\r\n"
    self.tool_handle = vim.system({ "cmd.exe", "/Q", "/K" }, {
      stdin = shell_cmd,
      env = { PROMPT = "\r\n" },
    }, callback)
  else
    self.tool_handle = vim.system(os_utils.build_shell_command(cmd), {}, callback)
  end
end

---Converts a cmd-based tool to a function-based tool.
---@param tool CodeCompanion.Tools.Tool
---@return CodeCompanion.Tools.Tool
function Orchestrator:_cmd_to_func_tool(tool)
  tool.cmds = vim
    .iter(tool.cmds)
    :map(function(cmd)
      if type(cmd) == "function" then
        return cmd
      end

      local flag = cmd.flag
      cmd = cmd.cmd or cmd
      if type(cmd) == "string" then
        cmd = vim.split(cmd, " ", { trimempty = true })
      end

      ---@param tools CodeCompanion.Tools
      return function(tools, _, _, cb)
        cb = vim.schedule_wrap(cb)
        self:_execute_shell_command(cmd, function(out)
          if flag then
            tools.chat.tool_registry.flags = tools.chat.tool_registry.flags or {}
            tools.chat.tool_registry.flags[flag] = (out.code == 0)
          end

          local eol_pattern = vim.fn.has("win32") == 1 and "\r?\n" or "\n"

          if out.code == 0 then
            cb({
              status = "success",
              data = strip_ansi(vim.split(out.stdout, eol_pattern, { trimempty = true })),
            })
          else
            local combined = {}
            if out.stderr and out.stderr ~= "" then
              vim.list_extend(combined, strip_ansi(vim.split(out.stderr, eol_pattern, { trimempty = true })))
            end
            if out.stdout and out.stdout ~= "" then
              vim.list_extend(combined, strip_ansi(vim.split(out.stdout, eol_pattern, { trimempty = true })))
            end
            cb({ status = "error", data = combined })
          end
        end)
      end
    end)
    :totable()

  return tool
end

---@param tools CodeCompanion.Tools
---@param id number
function Orchestrator.new(tools, id)
  local self = setmetatable({
    id = id,
    queue = Queue.new(),
    tools = tools,
  }, { __index = Orchestrator })

  return self
end

---Add the tool's handlers to the executor
---@return nil
function Orchestrator:_setup_handlers()
  self.handlers = {
    setup = function()
      if not self.tool then
        return
      end
      if self.tool.handlers and self.tool.handlers.setup then
        return self.tool.handlers.setup(self.tool, self.tools)
      end
    end,
    prompt_condition = function()
      if not self.tool then
        return
      end

      if self.tool.handlers and self.tool.handlers.prompt_condition then
        return self.tool.handlers.prompt_condition(self.tool, self.tools, self.tools.tools_config)
      end
      return true
    end,
    on_exit = function()
      if not self.tool then
        return
      end

      if self.tool.handlers and self.tool.handlers.on_exit then
        return self.tool.handlers.on_exit(self.tool, self.tools)
      end
    end,
  }

  self.output = {
    cmd_string = function()
      if not self.tool then
        return
      end
      if self.tool.output and self.tool.output.cmd_string then
        return self.tool.output.cmd_string(self.tool, { tools = self.tools })
      end
      return nil
    end,

    prompt = function()
      if not self.tool then
        return
      end
      if self.tool.output and self.tool.output.prompt then
        return self.tool.output.prompt(self.tool, self.tools)
      end
    end,
    rejected = function(cmd, opts)
      if not self.tool then
        return
      end

      opts = opts or {}

      if self.tool.output and self.tool.output.rejected then
        self.tool.output.rejected(self.tool, self.tools, cmd, opts)
      else
        local rejection = fmt("\nThe user rejected the execution of the %s tool", self.tool.name)
        if opts.reason then
          rejection = rejection .. fmt(': "%s"', opts.reason)
        end
        -- If no handler is set then return a default message
        self:_send_response_to_chat(rejection)
      end
    end,
    error = function(cmd)
      if not self.tool then
        return
      end

      if self.tool.output and self.tool.output.error then
        self.tool.output.error(self.tool, self.tools, cmd, self.tools.stderr)
      else
        self:_send_response_to_chat(fmt("Error calling `%s`", self.tool.name))
      end
    end,
    cancelled = function(cmd)
      if not self.tool then
        return
      end

      if self.tool.output and self.tool.output.cancelled then
        self.tool.output.cancelled(self.tool, self.tools, cmd)
      else
        self:_send_response_to_chat(
          fmt("The user cancelled the execution of the %s tool", self.tool.name),
          fmt("Cancelled `%s`", self.tool.name)
        )
      end
    end,
    success = function(cmd)
      if not self.tool then
        return
      end

      if self.tool.output and self.tool.output.success then
        self.tool.output.success(self.tool, self.tools, cmd, self.tool_output or {})
      else
        self:_send_response_to_chat(fmt("Executed `%s`", self.tool.name))
      end
    end,
  }
end

---When the tools coordinator is finished, finalize it via an autocmd
---@return nil
function Orchestrator:_finalize_tools()
  self.tools.tool = nil
  self.tools.chat.tool_orchestrator = nil

  return utils.fire("ToolsFinished", { id = self.id, bufnr = self.tools.bufnr })
end

---Setup the tool to be executed
---@param input? any
---@return nil
function Orchestrator:setup_next_tool(input)
  if self.queue:is_empty() then
    return self:_finalize_tools()
  end

  -- Get the next tool to run
  self.tool = self.queue:pop()
  self.tool_output = {}

  self:_setup_handlers()
  self.handlers.setup() -- Call this early as cmd_runner needs to setup its cmds dynamically

  -- Transform cmd-based tools to func-based
  self.tool = self:_cmd_to_func_tool(self.tool)

  -- Get the first command to run
  local cmd = self.tool.cmds[1]
  log:debug("[Orchestrator::setup_next_tool] `%s` tool", self.tool.name)

  -- Check if the tool requires approval
  if
    self.tool.opts
    and not Approvals:is_approved(self.tools.bufnr, { cmd = self.output.cmd_string(), tool_name = self.tool.name })
  then
    local require_approval_before = self.tool.opts.require_approval_before

    if require_approval_before and type(require_approval_before) == "function" then
      require_approval_before = require_approval_before(self.tool, self.tools)
    end
    if require_approval_before and type(require_approval_before) ~= "boolean" then
      require_approval_before = self.handlers.prompt_condition()
    end

    if require_approval_before then
      log:debug("[Orchestrator::setup_next_tool] Asking for approval")

      local prompt = self.output.prompt()
      if prompt == nil or prompt == "" then
        prompt = ("Run the %q tool?"):format(self.tool.name)
      end

      -- Schedule the confirmation dialog to ensure UI is ready
      vim.schedule(function()
        local choice = ui_utils.confirm(prompt, { "1 Allow always", "2 Allow once", "3 Reject", "4 Cancel" })
        log:debug("[Orchestrator::setup_next_tool] User choice: %s", choice)

        -- Handle invalid/failed dialog (returns 0 or nil)
        if not choice or choice == 0 then
          self.output.rejected(cmd, { reason = "Confirmation dialog failed" })
          self:finalize_tool()
          return self:setup_next_tool()
        end

        if choice == 1 or choice == 2 then
          if choice == 1 then
            Approvals:always(self.tools.bufnr, { cmd = self.output.cmd_string(), tool_name = self.tool.name })
          end
          return self:execute_tool({ cmd = cmd, input = input })
        elseif choice == 3 then
          ui_utils.input({ prompt = fmt("Reason for rejecting `%s`", self.tool.name) }, function(i)
            self.output.rejected(cmd, { reason = i })
            return self:setup_next_tool()
          end)
        else
          -- NOTE: Cancel current tool, then cancel all queued tools
          self.output.cancelled(cmd)
          self:finalize_tool()
          self:cancel_pending_tools()
          return self:_finalize_tools()
        end
      end)
    else
      return self:execute_tool({ cmd = cmd, input = input })
    end
  else
    log:debug("[Orchestrator::setup_next_tool] No tool approval required")
    return self:execute_tool({ cmd = cmd, input = input })
  end
end

---Cancel all pending tools in the queue
---@return nil
function Orchestrator:cancel_pending_tools()
  while not self.queue:is_empty() do
    local pending_tool = self.queue:pop()
    self.tool = pending_tool

    -- Prepare handlers/output first
    self:_setup_handlers()
    local first_cmd = pending_tool.cmds and pending_tool.cmds[1] or nil

    local ok, err = pcall(function()
      self.output.cancelled(first_cmd)
    end)
    if not ok then
      return log:error("Failed to run cancelled handler for tool %s: %s", tostring(pending_tool.name), err)
    end
  end
end

---Execute the tool command
---@param args { cmd: function, input?: any }
---@return nil
function Orchestrator:execute_tool(args)
  utils.fire("ToolStarted", { id = self.id, tool = self.tool.name, bufnr = self.tools.bufnr })

  pcall(function()
    local edit_tracker = require("codecompanion.interactions.chat.edit_tracker")
    self.execution_id = edit_tracker.start_tool_monitoring(self.tool.name, self.tools.chat, self.tool.args)
  end)

  return Runner.new({ index = 1, orchestrator = self, cmd = args.cmd }):setup(args.input)
end

---Handle an error from a tool
---@param args { action: table, error?: any }
---@return nil
function Orchestrator:error(args)
  self.tools.status = self.tools.constants.STATUS_ERROR
  if args.error then
    table.insert(self.tools.stderr, args.error)
  end

  if self.tool and self.tool.name then
    -- Wrap this in error handling to avoid breaking the tool execution
    pcall(function()
      local edit_tracker = require("codecompanion.interactions.chat.edit_tracker")
      edit_tracker.finish_tool_monitoring(self.tool.name, self.tools.chat, false, self.execution_id)
    end)
  end

  local ok, err = pcall(function()
    self.output.error(args.action)
  end)
  if not ok then
    if self.tool and self.tool.function_call then
      self.tools.chat:add_tool_output(
        self.tool,
        string.format("Internal error with `%s` tool: %s", self.tool.name, err)
      )
    end
  end

  self:finalize_tool()
  self:setup_next_tool()
end

---Handle a successful completion of a tool
---@param args { action: table, output?: any }
---@return nil
function Orchestrator:success(args)
  self.tools.status = self.tools.constants.STATUS_SUCCESS

  if self.tool and self.tool.name then
    pcall(function()
      local edit_tracker = require("codecompanion.interactions.chat.edit_tracker")
      edit_tracker.finish_tool_monitoring(self.tool.name, self.tools.chat, true, self.execution_id)
    end)
  end

  if args.output then
    table.insert(self.tools.stdout, args.output)
    if not self.tool_output then
      self.tool_output = {}
    end
    table.insert(self.tool_output, args.output)
  end
  local ok, err = pcall(function()
    self.output.success(args.action)
  end)

  if not ok then
    log:error("Internal error with the %s success handler: %s", self.tool.name, err)
    if self.tool and self.tool.function_call then
      self.tools.chat:add_tool_output(self.tool, string.format("Internal error with `%s` tool", self.tool.name))
    end
  end
end

---Finalize the execution of the tool
---@return nil
function Orchestrator:finalize_tool()
  if self.tool then
    pcall(function()
      self.handlers.on_exit()
    end)
    utils.fire("ToolFinished", { id = self.id, name = self.tool.name, bufnr = self.tools.bufnr })
    self.tool = nil
  end
end

---Cancel the currently running tool execution
---@return nil
function Orchestrator:cancel()
  if self.cancelled then
    return
  end

  log:debug("Orchestrator:cancel")
  self.cancelled = true

  -- Kill the running system command if one exists
  if self.tool_handle then
    if vim.fn.has("win32") == 1 then
      -- /F flag forces process to end
      -- /T ends child process (required since we are wrapping the child
      --         process in a parent cmd.exe process)
      vim.system({ "taskkill", "/F", "/T", "/PID", tostring(self.tool_handle.pid) })
    else
      self.tool_handle:kill("sigkill")
    end

    self.tool_handle = nil
  end

  -- Output the cancellation message to the chat buffer.
  if self.tool and self.output then
    self.output.cancelled(self.tool.cmds[1])
  end

  -- Clean up the cancelled tool.
  self:finalize_tool()

  -- Cancel any pending tools in the queue.
  self:cancel_pending_tools()

  -- Close the current tool.
  self:_finalize_tools()
end

return Orchestrator