From f831051a45050187001a191b5de163e3ad55b3f8 Mon Sep 17 00:00:00 2001
From: Yomi Ikuru <captyomjnr@gmail.com>
Date: Sun, 17 Mar 2024 12:43:22 +0000
Subject: [PATCH] feat: support for shell cmds in place of env vars (#24)

---
 README.md                        |  2 ++
 doc/codecompanion.txt            |  2 ++
 lua/codecompanion/adapter.lua    | 37 ++++++++++++++++++++++++++------
 lua/codecompanion/utils/util.lua |  7 ++++++
 4 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 73320b05..f40e0b88 100644
--- a/README.md
+++ b/README.md
@@ -185,6 +185,8 @@ require("codecompanion").setup({
     chat = require("codecompanion.adapters").use("openai", {
       env = {
         api_key = "DIFFERENT_OPENAI_KEY",
+        -- To execute a shell command, prefix it with "cmd:"
+        -- api_key = "cmd:gpg --decrypt ~/.openai-api-key.gpg 2>/dev/null",
       },
     }),
   },
diff --git a/doc/codecompanion.txt b/doc/codecompanion.txt
index 5768a03c..a47b75f6 100644
--- a/doc/codecompanion.txt
+++ b/doc/codecompanion.txt
@@ -164,6 +164,8 @@ passing in a table to the `use` method:
         chat = require("codecompanion.adapters").use("openai", {
           env = {
             api_key = "DIFFERENT_OPENAI_KEY",
+            -- To execute a shell command, prefix it with "cmd:"
+            -- api_key = "cmd:gpg --decrypt ~/.openai-api-key.gpg 2>/dev/null",
           },
         }),
       },
diff --git a/lua/codecompanion/adapter.lua b/lua/codecompanion/adapter.lua
index 143b4971..35897309 100644
--- a/lua/codecompanion/adapter.lua
+++ b/lua/codecompanion/adapter.lua
@@ -1,4 +1,5 @@
 local log = require("codecompanion.utils.log")
+local utils = require("codecompanion.utils.util")
 
 ---@class CodeCompanion.Adapter
 ---@field name string
@@ -84,19 +85,41 @@ function Adapter:replace_header_vars()
   if self.headers then
     for k, v in pairs(self.headers) do
       self.headers[k] = v:gsub("${(.-)}", function(var)
-        local env_var = self.env[var]
+        local env_var_or_cmd = self.env[var]
+
+        if not env_var_or_cmd then
+          log:error("Error: Could not find env var or command: %s", self.env[var])
+          return vim.notify(
+            string.format("[CodeCompanion.nvim]\nCould not find env var or command: %s", self.env[var]),
+            vim.log.levels.ERROR
+          )
+        end
 
-        if env_var then
-          env_var = os.getenv(env_var)
-          if not env_var then
-            log:error("Error: Could not find env var: %s", self.env[var])
+        if utils.is_cmd_var(env_var_or_cmd) then
+          local command = env_var_or_cmd:sub(5)
+          local handle = io.popen(command, "r")
+          if handle then
+            local result = handle:read("*a")
+            handle:close()
+            return result:gsub("%s+$", "")
+          else
+            log:error("Error: Could not execute command: %s", command)
             return vim.notify(
-              string.format("[CodeCompanion.nvim]\nCould not find env var: %s", self.env[var]),
+              string.format("[CodeCompanion.nvim]\nCould not execute command: %s", command),
               vim.log.levels.ERROR
             )
           end
-          return env_var
         end
+
+        local env_var = os.getenv(env_var_or_cmd)
+        if not env_var then
+          log:error("Error: Could not find env var: %s", self.env[var])
+          return vim.notify(
+            string.format("[CodeCompanion.nvim]\nCould not find env var: %s", self.env[var]),
+            vim.log.levels.ERROR
+          )
+        end
+        return env_var
       end)
     end
   end
diff --git a/lua/codecompanion/utils/util.lua b/lua/codecompanion/utils/util.lua
index 64c90c3b..7940cfd9 100644
--- a/lua/codecompanion/utils/util.lua
+++ b/lua/codecompanion/utils/util.lua
@@ -148,4 +148,11 @@ function M.replace_vars(msg, vars, mapping)
   return string.format(msg, unpack(replacements))
 end
 
+---Check if value starts with "cmd:"
+---@param value string
+---@return boolean
+function M.is_cmd_var(value)
+  return value:match("^cmd:")
+end
+
 return M