init.lua

---@class CodeCompanion.Tools
---@field adapter CodeCompanion.HTTPAdapter The adapter in use for the chat
---@field aug number The augroup for the tool
---@field bufnr number The buffer of the chat buffer
---@field constants table<string, string> The constants for the tool
---@field chat CodeCompanion.Chat The chat buffer that initiated the tool
---@field extracted table The extracted tools from the LLM's response
---@field messages table The messages in the chat buffer
---@field status string The status of the tool
---@field stdout table The stdout of the tool
---@field stderr table The stderr of the tool
---@field tool CodeCompanion.Tools.Tool The current tool that's being run
---@field tools_config table The available tools for the tool system
---@field tools_ns number The namespace for the virtual text that appears in the header

local EditTracker = require("codecompanion.interactions.chat.edit_tracker")
local Orchestrator = require("codecompanion.interactions.chat.tools.orchestrator")
local approvals = require("codecompanion.interactions.chat.tools.approvals")
local tool_filter = require("codecompanion.interactions.chat.tools.filter")

local config = require("codecompanion.config")
local log = require("codecompanion.utils.log")
local regex = require("codecompanion.utils.regex")
local ui_utils = require("codecompanion.utils.ui")
local utils = require("codecompanion.utils")

local api = vim.api

local show_tools_processing = config.display.chat.show_tools_processing

local CONSTANTS = {
  PREFIX = "@",

  NS_TOOLS = "CodeCompanion-tools",
  AUTOCMD_GROUP = "codecompanion.tools",

  STATUS_ERROR = "error",
  STATUS_SUCCESS = "success",

  PROCESSING_MSG = (config.display.chat.icons.tools_in_progress or "⚡") .. " Tools processing ...",
}

---@class CodeCompanion.Tools
local Tools = {}

-- Private helper methods

---Creates a regex pattern to match a tool name in a message
---@param tool string The tool name to create a pattern for
---@return string The compiled regex pattern
function Tools:_pattern(tool)
  return CONSTANTS.PREFIX .. "{" .. tool .. "}"
end

---Handle missing or invalid tool errors
---@param tool table The tool that failed
---@param error_message string The error message
---@return nil
function Tools:_handle_tool_error(tool, error_message)
  local name = tool["function"].name
  local tool_call = vim.deepcopy(tool)
  tool_call.name = name
  tool_call.function_call = tool_call

  log:error(error_message)

  local available_tools_msg = ""
  if self.chat and self.chat.tool_registry and self.chat.tool_registry.in_use then
    local available_tools = vim.tbl_keys(self.chat.tool_registry.in_use)
    if next(available_tools) then
      available_tools_msg = "The available tools are: "
        .. table.concat(
          vim.tbl_map(function(t)
            return "`" .. t .. "`"
          end, available_tools),
          ", "
        )
    else
      available_tools_msg = "No tools available"
    end
  else
    available_tools_msg = "No tools available"
  end

  self.chat:add_tool_output(tool_call, string.format("Tool `%s` not found. %s", name, available_tools_msg), "")
  return utils.fire("ToolsFinished", { bufnr = self.bufnr })
end

---Resolve and prepare a tool for execution
---@param tool table The tool call from the LLM
---@param id number The execution ID for event firing
---@return table|nil The resolved tool or nil if failed
---@return string|nil Error message if resolution failed
---@return boolean|nil Whether this is a JSON parsing error that needs special handling
function Tools:_resolve_and_prepare_tool(tool, id)
  local name = tool["function"].name
  local tool_config = self.tools_config[name]

  -- Allow for hybrid tools that use an adapter's tool alongside a CodeCompanion tool
  if tool_config and tool_config._adapter_tool == true and tool_config._has_client_tool then
    tool_config = utils.resolve_nested_value(config, tool_config.opts.client_tool)
  end

  if not tool_config then
    return nil, string.format("Couldn't find the tool `%s`", name), false
  end

  local ok, resolved_tool = pcall(function()
    return Tools.resolve(tool_config)
  end)

  if not ok or not resolved_tool then
    log:debug("Tool resolution failed for `%s`: %s", name, resolved_tool)
    return nil, string.format("Couldn't resolve the tool `%s`", name), false
  end

  -- NOTE: We deepcopy here to avoid mutating the original tool definition which
  -- has disastrous side effects.
  local prepared_tool = vim.deepcopy(resolved_tool)
  prepared_tool.name = name
  prepared_tool.function_call = tool

  -- Parse and set arguments - handle JSON errors specially like the original code
  if tool["function"].arguments then
    local args = tool["function"].arguments
    -- For some adapter's that aren't streaming, the args are strings rather than tables
    if type(args) == "string" then
      if args == "" then
        args = "{}"
      end
      local ok, decoded = pcall(vim.json.decode, args)
      if not ok then
        log:error("Couldn't decode the tool arguments: %s", args)
        self.chat:add_tool_output(
          prepared_tool,
          string.format('You made an error in calling the %s tool: "%s"', name, decoded),
          ""
        )
        self.status = CONSTANTS.STATUS_ERROR
        utils.fire("ToolsFinished", { id = id, bufnr = self.bufnr })
        return nil, "JSON parsing failed", true -- Special flag to indicate this was handled
      end

      args = decoded
    end
    prepared_tool.args = args
  end

  -- Merge options
  prepared_tool.opts = vim.tbl_extend("force", prepared_tool.opts or {}, tool_config.opts or {})

  -- Handle environment variables
  if prepared_tool.env then
    local env = type(prepared_tool.env) == "function" and prepared_tool.env(vim.deepcopy(prepared_tool)) or {}
    utils.replace_placeholders(prepared_tool.cmds, env)
  end

  return prepared_tool, nil, false
end

---Start edit tracking for all tools
---@param tools table The tools to track
---@return nil
function Tools:_start_edit_tracking(tools)
  for _, tool in ipairs(tools) do
    local tool_name = tool["function"].name
    local tool_args = tool["function"].arguments

    -- Handle argument parsing more robustly, like the original code
    if type(tool_args) == "string" then
      local success, decoded = pcall(vim.json.decode, tool_args)
      if success then
        tool_args = decoded
      else
        tool_args = nil
      end
    end

    EditTracker.start_tool_monitoring(tool_name, self.chat, tool_args)
  end
end

-- Public interface methods

---@param args table
function Tools.new(args)
  local self = setmetatable({
    adapter = args.adapter,
    aug = api.nvim_create_augroup(CONSTANTS.AUTOCMD_GROUP .. ":" .. args.bufnr, { clear = true }),
    bufnr = args.bufnr,
    chat = {},
    constants = CONSTANTS,
    extracted = {},
    messages = args.messages,
    stdout = {},
    stderr = {},
    tool = {},
    tools_config = tool_filter.filter_enabled_tools(config.interactions.chat.tools, { adapter = args.adapter }),
    tools_ns = api.nvim_create_namespace(CONSTANTS.NS_TOOLS),
  }, { __index = Tools })

  -- Listen for any adapter and model changes on the chat buffer and update the available tools
  api.nvim_create_autocmd("User", {
    group = api.nvim_create_augroup(CONSTANTS.AUTOCMD_GROUP .. ".list:" .. args.bufnr, { clear = true }),
    pattern = "CodeCompanionChatModel",
    callback = function(autocmd_args)
      if autocmd_args.data.bufnr ~= self.bufnr then
        return
      end
      self.tools_config =
        tool_filter.filter_enabled_tools(config.interactions.chat.tools, { adapter = autocmd_args.data.adapter })
    end,
  })

  return self
end

---Refresh the tools configuration to pick up any dynamically added tools
---@param opts? table Options for refreshing the tools
---@return CodeCompanion.Tools
function Tools:refresh(opts)
  opts = opts or {}
  self.tools_config = tool_filter.filter_enabled_tools(config.interactions.chat.tools, opts)
  return self
end

---Set the autocmds for the tool
---@return nil
function Tools:set_autocmds()
  api.nvim_create_autocmd("User", {
    desc = "Handle responses from the Tool system",
    group = self.aug,
    pattern = "CodeCompanionTools*",
    callback = function(request)
      if request.data.bufnr ~= self.bufnr then
        return
      end

      if request.match == "CodeCompanionToolsStarted" then
        log:info("[Tool System] Initiated")
        if show_tools_processing then
          local namespace = CONSTANTS.NS_TOOLS .. "_" .. tostring(self.bufnr)
          ui_utils.show_buffer_notification(self.bufnr, {
            namespace = namespace,
            text = CONSTANTS.PROCESSING_MSG,
            main_hl = "CodeCompanionChatInfo",
            spacer = true,
          })
        end
      elseif request.match == "CodeCompanionToolsFinished" then
        return vim.schedule(function()
          local auto_submit = function()
            return self.chat:submit({
              auto_submit = true,
              callback = function()
                self:reset({ auto_submit = true })
              end,
            })
          end

          if approvals:is_approved(self.bufnr) then
            return auto_submit()
          end
          if self.status == CONSTANTS.STATUS_ERROR and self.tools_config.opts.auto_submit_errors then
            return auto_submit()
          end
          if self.status == CONSTANTS.STATUS_SUCCESS and self.tools_config.opts.auto_submit_success then
            return auto_submit()
          end

          self:reset({ auto_submit = false })
        end)
      end
    end,
  })
end

---Execute the tool in the chat buffer based on the LLM's response
---@param chat CodeCompanion.Chat
---@param tools table The tools requested by the LLM
---@return nil
function Tools:execute(chat, tools)
  local id = math.random(10000000)
  self.chat = chat

  self:_start_edit_tracking(tools)

  local function safe_execute()
    -- NOTE: Set autocmds early so that errors can be handled properly
    self:set_autocmds()

    chat.tool_orchestrator = Orchestrator.new(self, id)

    for _, tool in ipairs(tools) do
      local resolved_tool, error_msg, is_json_error = self:_resolve_and_prepare_tool(tool, id)

      if not resolved_tool then
        if is_json_error then
          -- JSON error was already handled by _resolve_and_prepare_tool
          return
        else
          return self:_handle_tool_error(tool, error_msg or "Unknown Error occurred")
        end
      end

      self.tool = resolved_tool --[[@as CodeCompanion.Tools.Tool]]
      chat.tool_orchestrator.queue:push(resolved_tool)
    end

    utils.fire("ToolsStarted", { id = id, bufnr = self.bufnr })
    chat.tool_orchestrator:setup_next_tool()
  end

  local ok, err = pcall(safe_execute)
  if not ok then
    log:error("chat::tools::init::execute - Execution error %s", err)
    self.status = CONSTANTS.STATUS_ERROR
    vim.schedule(function()
      utils.fire("ToolsFinished", { id = id, bufnr = self.bufnr })
    end)
  end
end

---Look for tools in a given message
---@param message table
---@return table?, table?
function Tools:find(message)
  if not message.content then
    return nil, nil
  end

  local groups = {}
  local tools = {}

  ---@param tool string The tool name to search for
  ---@return number?,number? The start position of the match, or nil if not found
  local function is_found(tool)
    local pattern = self:_pattern(tool)
    return regex.find(message.content, pattern)
  end

  -- Process groups
  vim.iter(self.tools_config.groups):each(function(tool)
    if is_found(tool) then
      table.insert(groups, tool)
    end
  end)

  -- Process tools
  vim
    .iter(self.tools_config)
    :filter(function(name)
      return name ~= "opts" and name ~= "groups"
    end)
    :each(function(tool)
      if is_found(tool) and not vim.tbl_contains(tools, tool) then
        table.insert(tools, tool)
      end
    end)

  if #tools == 0 and #groups == 0 then
    return nil, nil
  end

  return tools, groups
end

---Parse a user message looking for a tool
---@param chat CodeCompanion.Chat
---@param message table
---@return boolean
function Tools:parse(chat, message)
  local tools, groups = self:find(message)

  if tools or groups then
    if tools and not vim.tbl_isempty(tools) then
      for _, tool in ipairs(tools) do
        chat.tool_registry:add(tool, self.tools_config[tool])
      end
    end

    if groups and not vim.tbl_isempty(groups) then
      for _, group in ipairs(groups) do
        chat.tool_registry:add_group(group, self.tools_config)
      end
    end
    return true
  end

  return false
end

---Replace the tool tag in a given message
---@param message string
---@return string
function Tools:replace(message)
  for tool, _ in pairs(self.tools_config) do
    if tool ~= "opts" and tool ~= "groups" then
      local replacement = utils.replace_placeholders(self.tools_config.opts.tool_replacement_message, { tool = tool })
      message = vim.trim(regex.replace(message, self:_pattern(tool), replacement))
    end
  end

  for group, _ in pairs(self.tools_config.groups) do
    local replacement
    local group_config = self.tools_config.groups[group]
    local tools = table.concat(group_config.tools, ", ")

    if group_config.prompt then
      replacement = utils.replace_placeholders(group_config.prompt, { tools = tools .. " tools" })
    end

    message = vim.trim(regex.replace(message, self:_pattern(group), replacement or tools))
  end

  return message
end

---Reset the Tools class
---@param opts? table
---@return nil
function Tools:reset(opts)
  opts = opts or {}

  if show_tools_processing then
    ui_utils.clear_notification(self.bufnr, { namespace = CONSTANTS.NS_TOOLS .. "_" .. tostring(self.bufnr) })
  end

  api.nvim_clear_autocmds({ group = self.aug })

  self.extracted = {}
  self.status = CONSTANTS.STATUS_SUCCESS
  self.stderr = {}
  self.stdout = {}

  self.chat:tools_done(opts)
  log:info("[Tools] Completed")
end

---Add an error message to the chat buffer
---@param error string
---@return CodeCompanion.Tools
function Tools:add_error_to_chat(error)
  self.chat:add_message({
    role = config.constants.USER_ROLE,
    content = error,
  }, { visible = false })

  --- Alert the user that the error message has been shared
  self.chat:add_buf_message({
    role = config.constants.USER_ROLE,
    content = "Please correct for the error message I've shared",
  })

  if self.tools_config.opts and self.tools_config.opts.auto_submit_errors then
    self.chat:submit()
  end

  return self
end

---Resolve a tool from the config
---@param tool table The tool from the config
---@return CodeCompanion.Tools.Tool|nil
function Tools.resolve(tool)
  local callback = tool.callback

  if type(callback) == "table" then
    return callback --[[@as CodeCompanion.Tools.Tool]]
  end

  if type(callback) == "function" then
    return callback() --[[@as CodeCompanion.Tools.Tool]]
  end

  local ok, module = pcall(require, "codecompanion." .. callback)
  if ok then
    log:debug("[Tools] %s identified", callback)
    return module
  end

  -- Try loading the tool from the user's config using a module path
  ok, module = pcall(require, callback)
  if ok then
    log:debug("[Tools] %s identified", callback)
    return module
  end

  -- Try loading the tool from the user's config using a file path
  local err
  module, err = loadfile(vim.fs.normalize(callback))
  if err then
    return log:error("[Tools] Failed to load tool from %s: %s", callback, err)
  end

  if module then
    log:debug("[Tools] %s identified", callback)
    return module()
  end

  return nil
end

return Tools