1 files changed, 94 insertions, 64 deletions
diff --git a/lua/99/request-context.lua b/lua/99/request-context.lua
index 2d950d5..4765772 100644
--- a/lua/99/request-context.lua
+++ b/lua/99/request-context.lua
@@ -22,95 +22,125 @@ RequestContext.__index = RequestContext
 --- @param xid number
 --- @return _99.RequestContext
 function RequestContext.from_current_buffer(_99, xid)
-    local buffer = vim.api.nvim_get_current_buf()
-    local full_path = vim.api.nvim_buf_get_name(buffer)
-    local file_type = vim.bo[buffer].ft
+  local buffer = vim.api.nvim_get_current_buf()
+  local full_path = vim.api.nvim_buf_get_name(buffer)
+  local file_type = vim.bo[buffer].ft
 
-    if file_type == "typescriptreact" then
-        file_type = "typescript"
-    end
+  if file_type == "typescriptreact" then
+    file_type = "typescript"
+  end
 
-    local mds = {}
-    for _, md in ipairs(_99.md_files) do
-        table.insert(mds, md)
-    end
+  local mds = {}
+  for _, md in ipairs(_99.md_files) do
+    table.insert(mds, md)
+  end
 
-    return setmetatable({
-        _99 = _99,
-        md_file_names = mds,
-        ai_context = {},
-        tmp_file = random_file(),
-        buffer = buffer,
-        full_path = full_path,
-        file_type = file_type,
-        logger = Logger:set_id(xid),
-        xid = xid,
-        model = _99.model,
-        marks = {},
-    }, RequestContext)
+  return setmetatable({
+    _99 = _99,
+    md_file_names = mds,
+    ai_context = {},
+    tmp_file = random_file(),
+    buffer = buffer,
+    full_path = full_path,
+    file_type = file_type,
+    logger = Logger:set_id(xid),
+    xid = xid,
+    model = _99.model,
+    marks = {},
+  }, RequestContext)
 end
 
 --- @param md_file_name string
 --- @return self
 function RequestContext:add_md_file_name(md_file_name)
-    table.insert(self.md_file_names, md_file_name)
-    return self
+  table.insert(self.md_file_names, md_file_name)
+  return self
 end
 
-function RequestContext:_read_md_files()
-    local cwd = vim.uv.cwd()
-    local dir = vim.fn.fnamemodify(self.full_path, ":h")
+--- TODO: Dedupe any rules that have already been added
+--- @param rules (_99.Agents.Rule | string)[]
+function RequestContext:add_agent_rules(rules)
+  for _, rule in ipairs(rules) do
+    -- Handle both string paths and rule objects
+    self.logger:debug("adding custom rule to agent", "rule", rule)
+    local ok, file = pcall(io.open, rule.path, "r")
+    if ok and file then
+      local content = file:read("*a")
+      file:close()
+      self.logger:info(
+        "Context#adding agent file to the context",
+        "agent_path",
+        rule.path
+      )
+      table.insert(
+        self.ai_context,
+        string.format(
+          [[
+<%s>
+%s
+</%s>]],
+          rule.name,
+          content,
+          rule.name
+        )
+      )
+    else
+      self.logger:debug("unable to read agent rule", "rule", rule)
+    end
+  end
+end
 
-    while dir:find(cwd, 1, true) == 1 do
-        for _, md_file_name in ipairs(self.md_file_names) do
-            local md_path = dir .. "/" .. md_file_name
-            local file = io.open(md_path, "r")
-            if file then
-                local content = file:read("*a")
-                file:close()
-                self.logger:info(
-                    "Context#adding md file to the context",
-                    "md_path",
-                    md_path
-                )
-                table.insert(self.ai_context, content)
-            end
-        end
+function RequestContext:_read_md_files()
+  local cwd = vim.uv.cwd()
+  local dir = vim.fn.fnamemodify(self.full_path, ":h")
 
-        if dir == cwd then
-            break
-        end
+  while dir:find(cwd, 1, true) == 1 do
+    for _, md_file_name in ipairs(self.md_file_names) do
+      local md_path = dir .. "/" .. md_file_name
+      local file = io.open(md_path, "r")
+      if file then
+        local content = file:read("*a")
+        file:close()
+        self.logger:info(
+          "Context#adding md file to the context",
+          "md_path",
+          md_path
+        )
+        table.insert(self.ai_context, content)
+      end
+    end
 
-        dir = vim.fn.fnamemodify(dir, ":h")
+    if dir == cwd then
+      break
     end
+
+    dir = vim.fn.fnamemodify(dir, ":h")
+  end
 end
 
 --- @return string[]
 function RequestContext:content()
-    return self.ai_context
+  return self.ai_context
 end
 
 --- @return self
 function RequestContext:finalize()
-    self:_read_md_files()
-    if self.range then
-        table.insert(self.ai_context, self._99.prompts.get_file_location(self))
-        table.insert(
-            self.ai_context,
-            self._99.prompts.get_range_text(self.range)
-        )
-    end
-    table.insert(
-        self.ai_context,
-        self._99.prompts.tmp_file_location(self.tmp_file)
-    )
-    return self
+  self:_read_md_files()
+  if self.range then
+    table.insert(self.ai_context, self._99.prompts.get_file_location(self))
+    table.insert(self.ai_context, self._99.prompts.get_range_text(self.range))
+  end
+  table.insert(
+    self.ai_context,
+    self._99.prompts.tmp_file_location(self.tmp_file)
+  )
+  return self
 end
 
 function RequestContext:clear_marks()
-    for _, mark in pairs(self.marks) do
-        mark:delete()
-    end
+  for _, mark in pairs(self.marks) do
+    mark:delete()
+  end
 end
 
 return RequestContext