fix: #698 custom variables (#711)

Co-authored-by: Oli Morris <[email protected]>

Author: olimorris

lua/codecompanion/strategies/chat/variables/init.lua

@@ -26,26 +26,36 @@ end
 ---@param params? string
 ---@return table
 local function resolve(chat, var_config, params)
-  local splits = vim.split(var_config.callback, ".", { plain = true })
-  local path = table.concat(splits, ".", 1, #splits - 1)
-  local variable = splits[#splits]
-
-  local ok, module = pcall(require, "codecompanion." .. path .. "." .. variable)
-
-  local init = {
-    Chat = chat,
-    config = var_config,
-    params = params,
-  }
+  if type(var_config.callback) == "string" then
+    local splits = vim.split(var_config.callback, ".", { plain = true })
+    local path = table.concat(splits, ".", 1, #splits - 1)
+    local variable = splits[#splits]
+
+    local ok, module = pcall(require, "codecompanion." .. path .. "." .. variable)
+
+    local init = {
+      Chat = chat,
+      config = var_config,
+      params = params,
+    }
+
+    -- User is using a custom callback
+    if not ok then
+      log:trace("Calling variable: %s", path .. "." .. variable)
+      return require(path .. "." .. variable).new(init):output()
+    end
 
-  -- User is using a custom callback
-  if not ok then
     log:trace("Calling variable: %s", path .. "." .. variable)
-    return require(path .. "." .. variable).new(init):output()
+    return module.new(init):output()
   end
 
-  log:trace("Calling variable: %s", path .. "." .. variable)
-  return module.new(init):output()
+  return require("codecompanion.strategies.chat.variables.user")
+    .new({
+      Chat = chat,
+      config = var_config,
+      params = params,
+    })
+    :output()
 end
 
 ---@class CodeCompanion.Variables
@@ -92,6 +102,8 @@ function Variables:parse(chat, message)
       local var_config = self.vars[var]
       log:debug("Variable found: %s", var)
 
+      var_config["name"] = var
+
       if (var_config.opts and var_config.opts.contains_code) and not config.can_send_code() then
         log:warn("Sending of code has been disabled")
         goto continue

lua/codecompanion/strategies/chat/variables/user.lua

@@ -0,0 +1,35 @@
+local buf_utils = require("codecompanion.utils.buffers")
+local config = require("codecompanion.config")
+
+---@class CodeCompanion.Variable.User: CodeCompanion.Variable
+local Variable = {}
+
+---@param args CodeCompanion.VariableArgs
+function Variable.new(args)
+  local self = setmetatable({
+    Chat = args.Chat,
+    config = args.config,
+    params = args.params,
+  }, { __index = Variable })
+
+  return self
+end
+
+---Return the user's custom variable
+---@return nil
+function Variable:output()
+  local id = "<var>" .. self.config.name .. "</var>"
+
+  self.Chat:add_message({
+    role = config.constants.USER_ROLE,
+    content = self.config.callback(),
+  }, { reference = id, tag = "variable", visible = false })
+
+  self.Chat.References:add({
+    bufnr = self.Chat.bufnr,
+    id = id,
+    source = "codecompanion.strategies.chat.variables.user",
+  })
+end
+
+return Variable