Checklist tool for CodeCompanion

checklist.lua

-- checklist_tools.lua

local data_path = vim.fn.stdpath('data') .. '/codecompanion'
local workspace_root = vim.fn.getcwd()
local workspace_id = vim.fn.substitute(workspace_root, '[/\\:*?"<>|]', '_', 'g')
local checklist_file = data_path .. '/checklists_v3_' .. workspace_id .. '.json'

vim.fn.mkdir(data_path, 'p')

---@class ChecklistTask
---@field id integer
---@field text string
---@field done boolean
---@field created_at integer
---@field completed_at? integer

---@class ChecklistHistoryEntry
---@field action string
---@field commit_message string
---@field context string
---@field timestamp integer
---@field file_paths? string[]
---@field completed_task_ids? integer[]  -- IDs of tasks completed in this commit

---@class Checklist
---@field id integer
---@field goal string
---@field created_at integer
---@field tasks table<integer, ChecklistTask>
---@field next_task_id integer
---@field history ChecklistHistoryEntry[]

---@class SerializableChecklist: Checklist
---@field tasks_array ChecklistTask[]

---@class ChecklistStorage
---@field load fun(): (table<integer, Checklist>, integer)
---@field save fun(checklists: table<integer, Checklist>, next_id: integer)
local checklist_storage = {
  load = function()
    local ok, data = pcall(vim.fn.readfile, checklist_file)
    if not ok or not data or #data == 0 then
      return {}, 1
    end
    local json_ok, parsed = pcall(vim.json.decode, table.concat(data, '\n'))
    if not json_ok or not parsed then
      return {}, 1
    end
    local checklists = {}
    for checklist_id, checklist in pairs(parsed.checklists or {}) do
      local restored_checklist = vim.deepcopy(checklist)
      if checklist.tasks_array then
        restored_checklist.tasks = {}
        for _, task in ipairs(checklist.tasks_array) do
          restored_checklist.tasks[task.id] = task
        end
        restored_checklist.tasks_array = nil
      else
        restored_checklist.tasks = restored_checklist.tasks or {}
      end
      checklists[checklist_id] = restored_checklist
    end
    return checklists, parsed.next_id or 1
  end,

  save = function(checklists, next_id)
    local serializable_checklists = {}
    for checklist_id, checklist in pairs(checklists) do
      local serializable_checklist = vim.deepcopy(checklist)
      local tasks_array = {}
      for task_id, task in pairs(checklist.tasks or {}) do
        table.insert(tasks_array, task)
      end
      serializable_checklist.tasks_array = tasks_array
      serializable_checklist.tasks = nil
      serializable_checklists[checklist_id] = serializable_checklist
    end
    local data = {
      workspace = workspace_root,
      checklists = serializable_checklists,
      next_id = next_id,
      last_updated = os.time()
    }
    pcall(vim.fn.writefile, { vim.json.encode(data) }, checklist_file)
  end
}

local checklists, next_checklist_id = checklist_storage.load()

---@param checklist Checklist
---@return ChecklistTask[]
local function get_sorted_tasks(checklist)
  local sorted_tasks = {}
  for _, task in pairs(checklist.tasks) do
    table.insert(sorted_tasks, task)
  end
  table.sort(sorted_tasks, function(a, b) return a.id < b.id end)
  return sorted_tasks
end

---@param checklist Checklist
---@return string
local function format_checklist_output(checklist)
  if not checklist then
    return "No checklist data"
  end

  local output = string.format([[📋 **CHECKLIST %d**
 🎯 **GOAL**: %s
 📅 **CREATED**: %s
 📝 **TASKS**:]],
    checklist.id,
    checklist.goal or "No goal",
    os.date("%Y-%m-%d %H:%M", checklist.created_at)
  )

  if vim.tbl_isempty(checklist.tasks) then
    output = output .. "\n   No tasks defined"
  else
    local sorted_tasks = get_sorted_tasks(checklist)
    for _, task in ipairs(sorted_tasks) do
      local status_icon = task.done and "✅" or "❌"
      local completed_info = task.done and
          string.format(" (completed %s)", os.date("%m/%d %H:%M", task.completed_at)) or ""
      output = output .. string.format("\n%d. %s %s%s",
        task.id, status_icon, task.text, completed_info)
    end
  end

  -- Show history if available
  if checklist.history and #checklist.history > 0 then
    output = output .. "\n\n📜 **HISTORY**:"
    -- Sort history by most recent first
    local sorted_history = {}
    for _, entry in ipairs(checklist.history) do
      table.insert(sorted_history, entry)
    end
    table.sort(sorted_history, function(a, b)
      return a.timestamp > b.timestamp
    end)
    for _, entry in ipairs(sorted_history) do
      local completed_tasks_str = ""
      if entry.completed_task_ids and #entry.completed_task_ids > 0 then
        completed_tasks_str = string.format("\n  • Tasks completed: %s", table.concat(entry.completed_task_ids, ", "))
      end
      output = output .. string.format(
        "\n- [%s] %s\n  • %s\n  • Context: %s%s",
        os.date("%Y-%m-%d %H:%M", entry.timestamp),
        entry.action,
        entry.commit_message,
        entry.context,
        completed_tasks_str
      )
    end
  end

  -- Completion summary
  local total = vim.tbl_count(checklist.tasks)
  local completed = 0
  for _, task in pairs(checklist.tasks) do
    if task.done then completed = completed + 1 end
  end
  output = output .. string.format("\n\n📊 **PROGRESS**: %d/%d tasks complete (%.0f%%)",
    completed, total, total > 0 and (completed / total) * 100 or 0)

  return output
end

---@param id string|integer|nil
---@return Checklist|nil, string|nil
local function get_checklist(id)
  if not id then
    return nil, "No checklist ID provided"
  end
  local checklist_id = tonumber(id)
  if not checklist_id then
    return nil, "Invalid checklist ID format"
  end
  local checklist = checklists[checklist_id]
  if not checklist then
    return nil, "Checklist not found"
  end
  return checklist, nil
end

---@param goal string
---@param tasks string[]
---@param commit_message string
---@param context string
---@return Checklist
local function create_checklist(goal, tasks, commit_message, context)
  local id = next_checklist_id
  next_checklist_id = next_checklist_id + 1

  local checklist = {
    id = id,
    goal = goal,
    created_at = os.time(),
    tasks = {},
    next_task_id = 1,
    history = {
      {
        action = "create",
        commit_message = commit_message,
        context = context,
        timestamp = os.time()
      }
    }
  }

  for _, task_text in ipairs(tasks or {}) do
    if task_text and task_text:match("%S") then
      checklist.tasks[checklist.next_task_id] = {
        id = checklist.next_task_id,
        text = task_text:gsub("^%s*[-*+]?%s*", ""),
        done = false,
        created_at = os.time()
      }
      checklist.next_task_id = checklist.next_task_id + 1
    end
  end

  checklists[id] = checklist
  checklist_storage.save(checklists, next_checklist_id)
  return checklist
end

---@param agent CodeCompanion.Agent
---@param checklist Checklist
---@param complete_task_ids string[]|number[]
---@param commit_message string
---@param context string
---@return boolean, string
local function complete_tasks(agent, checklist, complete_task_ids, commit_message, context)
  local changed = false
  local actually_completed = {}
  local function extract_file_paths_from_refs(refs)
    local paths = {}
    local seen = {}
    if refs then
      for _, ref in pairs(refs) do
        local path = nil
        if ref.path then
          path = ref.path
        elseif ref.bufnr then
          path = vim.api.nvim_buf_get_name(ref.bufnr)
        end
        if path and not seen[path] then
          table.insert(paths, path)
          seen[path] = true
        end
      end
    end
    return paths
  end

  for _, task_id in ipairs(complete_task_ids or {}) do
    local task_id_num = tonumber(task_id)
    if task_id_num and checklist.tasks[task_id_num] and not checklist.tasks[task_id_num].done then
      checklist.tasks[task_id_num].done = true
      checklist.tasks[task_id_num].completed_at = os.time()
      changed = true
      table.insert(actually_completed, task_id_num)
    end
  end

  local file_paths = extract_file_paths_from_refs(agent.chat.refs)

  if changed then
    table.insert(checklist.history, {
      action = "complete_tasks",
      commit_message = commit_message,
      context = context,
      timestamp = os.time(),
      file_paths = file_paths,
      completed_task_ids = actually_completed,
    })
    checklist_storage.save(checklists, next_checklist_id)
    return true, "Tasks marked complete"
  else
    return false, "No valid incomplete tasks specified"
  end
end

---@class ChecklistListArgs
---@field page string|nil
---@field per_page string|nil

---@class ChecklistListResult
---@field status string
---@field data string

---@type CodeCompanion.Agent.Tool
local ChecklistListTool = {
  name = "checklist_list",

  cmds = {
    ---@param agent CodeCompanion.Agent
    ---@param args ChecklistListArgs
    ---@param input any
    ---@param cb fun(result: ChecklistListResult)
    function(agent, args, input, cb)
      local page = tonumber(args.page) or 1
      local per_page = tonumber(args.per_page) or 10

      if page < 1 then
        return cb({ status = "error", data = "page must be >= 1" })
      end
      if per_page < 1 or per_page > 100 then
        return cb({ status = "error", data = "per_page must be between 1 and 100" })
      end

      if vim.tbl_isempty(checklists) then
        return cb({
          status = "success",
          data = "📋 **NO CHECKLISTS FOUND**\n\nUse the create tool to make a new checklist."
        })
      end

      local summaries = {}
      for id, checklist in pairs(checklists) do
        local total = vim.tbl_count(checklist.tasks)
        local completed = 0
        for _, task in pairs(checklist.tasks) do
          if task.done then completed = completed + 1 end
        end
        table.insert(summaries, {
          id = id,
          goal = checklist.goal or "No goal",
          created_at = checklist.created_at,
          progress = { completed = completed, total = total }
        })
      end
      table.sort(summaries, function(a, b)
        return a.created_at > b.created_at
      end)

      local total_count = #summaries
      local total_pages = math.max(1, math.ceil(total_count / per_page))
      local start_index = (page - 1) * per_page + 1
      local end_index = math.min(page * per_page, total_count)

      local output = string.format("📋 **CHECKLISTS** (Page %d of %d, showing %d-%d of %d):\n\n",
        page, total_pages, start_index, end_index, total_count)

      for i = start_index, end_index do
        local summary = summaries[i]
        output = output .. string.format("ID: %d | %s (%d/%d complete)\n",
          summary.id, summary.goal, summary.progress.completed, summary.progress.total)
        output = output .. string.format("     Created: %s\n\n",
          os.date("%m/%d %H:%M", summary.created_at))
      end

      if total_pages > 1 then
        output = output .. "\n📄 **PAGINATION**:\n"
        if page > 1 then
          output = output ..
              string.format("- Previous: `@checklist_list page=\"%d\" per_page=\"%d\"`\n", page - 1, per_page)
        end
        if page < total_pages then
          output = output ..
              string.format("- Next: `@checklist_list page=\"%d\" per_page=\"%d\"`\n", page + 1, per_page)
        end
        output = output .. string.format("- Go to page: `@checklist_list page=\"X\" per_page=\"%d\"`\n", per_page)
      end

      output = output .. "\n💡 **NEXT STEPS**:\n"
      output = output .. "- View checklist: `@checklist_status checklist_id=\"X\"`\n"
      output = output .. "- Create new: use the create tool"

      return cb({ status = "success", data = output })
    end,
  },

  function_call = {},

  schema = {
    type = "function",
    ["function"] = {
      name = "checklist_list",
      description = "List all checklists in the workspace.",
      parameters = {
        type = "object",
        properties = {
          page = {
            type = "string",
            description = "Page number to retrieve (default: 1)"
          },
          per_page = {
            type = "string",
            description = "Number of checklists per page (default: 10, max: 100)"
          }
        },
        additionalProperties = false
      },
      strict = true
    }
  },

  system_prompt = [[## Checklist List Tool

 - Use this tool to list all checklists in the workspace.
 - No modifications are allowed.
 - Use the create tool to add a new checklist.

### Usage Example

   @checklist_list page="1" per_page="5"
   @checklist_list page="2" per_page="10"
]],

  opts = {},

  env = nil,

  handlers = {
    setup = function(tool, agent)
      vim.notify(string.format("%s setup", tool.name), vim.log.levels.DEBUG)
    end,

    prompt_condition = function(tool, agent, config)
      vim.notify(string.format("%s prompt_condition", tool.name), vim.log.levels.DEBUG)
    end,

    on_exit = function(tool, agent)
      vim.notify(string.format("%s on_exit", tool.name), vim.log.levels.DEBUG)
    end,
  },

  output = {
    success = function(tool, agent, cmd, stdout)
      local chat = agent.chat
      chat:add_tool_output(tool, stdout[1])
    end,
    error = function(tool, agent, cmd, stderr)
      local chat = agent.chat
      local error_msg = stderr[1] or "Unknown error"
      chat:add_tool_output(tool, string.format("**Checklist List Tool Error**: %s", error_msg))
    end,
    rejected = function(tool, agent, cmd)
      local chat = agent.chat
      chat:add_tool_output(tool, "**Checklist List Tool**: User declined to execute the operation")
    end,
  },

  ["output.prompt"] = function(tool, agent)
    return ""
  end,

  args = {},

  tool = {},
}

---@type CodeCompanion.Agent.Tool
---@class ChecklistStatusArgs
---@field checklist_id string

---@class ChecklistStatusResult
---@field status string
---@field data string

---@type CodeCompanion.Agent.Tool
local ChecklistStatusTool = {
  name = "checklist_status",

  cmds = {
    function(tool, args, input, cb)
      local checklist_id = args.checklist_id
      if not checklist_id then
        return cb({ status = "error", data = "checklist_id is required" })
      end
      local checklist, err = get_checklist(checklist_id)
      if not checklist then
        return cb({ status = "error", data = err })
      end
      return cb({
        status = "success",
        data = format_checklist_output(checklist)
      })
    end,
  },

  function_call = {},

  schema = {
    type = "function",
    ["function"] = {
      name = "checklist_status",
      description = "Show the status of a specific checklist.",
      parameters = {
        type = "object",
        properties = {
          checklist_id = {
            type = "string",
            description = "Checklist ID to show status for"
          }
        },
        required = { "checklist_id" },
        additionalProperties = false
      },
      strict = true
    }
  },

  system_prompt = [[## Checklist Status Tool

- Use this tool to show the status of a specific checklist.
- No modifications are allowed.

### Usage Example

- Show the status of a checklist (replace X with the checklist ID):
  @checklist_status checklist_id="X"
]],

  opts = {},

  env = nil,

  handlers = {
    setup = function(tool, agent)
      vim.notify(string.format("%s setup", tool.name), vim.log.levels.DEBUG)
    end,

    prompt_condition = function(tool, agent, config)
      vim.notify(string.format("%s prompt_condition", tool.name), vim.log.levels.DEBUG)
    end,

    on_exit = function(tool, agent)
      vim.notify(string.format("%s on_exit", tool.name), vim.log.levels.DEBUG)
    end,
  },

  output = {
    success = function(tool, agent, cmd, stdout)
      local chat = agent.chat
      chat:add_tool_output(tool, stdout[1])
    end,
    error = function(tool, agent, cmd, stderr)
      local chat = agent.chat
      local error_msg = stderr[1] or "Unknown error"
      chat:add_tool_output(tool, string.format("**Checklist Status Tool Error**: %s", error_msg))
    end,
    rejected = function(tool, agent, cmd)
      local chat = agent.chat
      chat:add_tool_output(tool, "**Checklist Status Tool**: User declined to execute the operation")
    end,
  },

  ["output.prompt"] = function(tool, agent)
    return ""
  end,

  args = {},

  tool = {},
}

---@class ChecklistCreateArgs
---@field goal string
---@field tasks string[]
---@field commit_message string
---@field context string

---@class ChecklistCreateResult
---@field status string
---@field data string

---@type CodeCompanion.Agent.Tool
local ChecklistCreateTool = {
  name = "checklist_create",

  cmds = {
    function(tool, args, input, cb)
      local goal = args.goal
      local tasks = args.tasks or {}
      local commit_message = args.commit_message
      local context = args.context

      if not goal or goal == "" then
        return cb({ status = "error", data = "Goal is required" })
      end
      if not tasks or #tasks == 0 then
        return cb({ status = "error", data = "At least one task is required" })
      end
      if not commit_message or commit_message == "" then
        return cb({ status = "error", data = "commit_message is required" })
      end
      if not context or context == "" then
        return cb({ status = "error", data = "context is required" })
      end
      local checklist = create_checklist(goal, tasks, commit_message, context)
      return cb({
        status = "success",
        data = format_checklist_output(checklist)
      })
    end,
  },

  function_call = {},

  schema = {
    type = "function",
    ["function"] = {
      name = "checklist_create",
      description = "Create a new checklist. Requires goal, tasks, commit_message, and context.",
      parameters = {
        type = "object",
        properties = {
          goal = {
            type = "string",
            description = "Goal of the checklist"
          },
          tasks = {
            type = "array",
            items = { type = "string" },
            description = "Initial tasks as array of strings"
          },
          commit_message = {
            type = "string",
            description = "Justification for the change"
          },
          context = {
            type = "string",
            description = "All information used to make the decision"
          }
        },
        required = { "goal", "tasks", "commit_message", "context" },
        additionalProperties = false
      },
      strict = true
    }
  },

  system_prompt = [[## Checklist Create Tool

- Use this tool to create a new checklist.
- All fields are required.
- All changes require:
  - commit_message: a detailed justification for the change.
  - context: all information used to make the decision (user prompt, file, code, etc).

### Usage Example

- Create a new checklist:
  @checklist_create goal="Refactor module X" tasks=["Update API", "Write tests"] commit_message="Refactoring for maintainability" context={...}
]],

  opts = {
    requires_approval = true,
  },

  env = nil,

  handlers = {
    setup = function(tool, agent)
      vim.notify(string.format("%s setup", tool.name), vim.log.levels.DEBUG)
    end,

    prompt_condition = function(tool, agent, config)
      vim.notify(string.format("%s prompt_condition", tool.name), vim.log.levels.DEBUG)
    end,

    on_exit = function(tool, agent)
      vim.notify(string.format("%s on_exit", tool.name), vim.log.levels.DEBUG)
    end,
  },

  output = {
    success = function(tool, agent, cmd, stdout)
      local chat = agent.chat
      chat:add_tool_output(tool, stdout[1])
    end,
    error = function(tool, agent, cmd, stderr)
      local chat = agent.chat
      local error_msg = stderr[1] or "Unknown error"
      chat:add_tool_output(tool, string.format("**Checklist Create Tool Error**: %s", error_msg))
    end,
    rejected = function(tool, agent, cmd)
      local chat = agent.chat
      chat:add_tool_output(tool, "**Checklist Create Tool**: User declined to execute the operation")
    end,
  },

  ["output.prompt"] = function(tool, agent)
    return string.format(
      "Create checklist with goal: '%s' and %d tasks?\n\nCommit message: %s\nContext: %s",
      tool.args.goal or "",
      tool.args.tasks and #tool.args.tasks or 0,
      tool.args.commit_message or "",
      type(tool.args.context) == "string" and tool.args.context or vim.inspect(tool.args.context)
    )
  end,

  args = {},

  tool = {},
}

---@class ChecklistCompleteTasksArgs
---@field checklist_id string
---@field complete_task_ids string[]
---@field commit_message string
---@field context string

---@class ChecklistCompleteTasksResult
---@field status string
---@field data string

---@type CodeCompanion.Agent.Tool
local ChecklistCompleteTasksTool = {
  name = "checklist_complete_tasks",

  cmds = {
    function(agent, args, input, cb)
      local checklist_id = args.checklist_id
      local complete_task_ids = args.complete_task_ids
      local commit_message = args.commit_message
      local context = args.context

      if not checklist_id then
        return cb({ status = "error", data = "checklist_id is required" })
      end
      if not complete_task_ids or #complete_task_ids == 0 then
        return cb({ status = "error", data = "complete_task_ids is required" })
      end
      if not commit_message or commit_message == "" then
        return cb({ status = "error", data = "commit_message is required" })
      end
      if not context then
        return cb({ status = "error", data = "context is required" })
      end
      local checklist, err = get_checklist(checklist_id)
      if not checklist then
        return cb({ status = "error", data = err })
      end
      local success, msg = complete_tasks(agent, checklist, complete_task_ids, commit_message, context)
      if not success then
        return cb({ status = "error", data = msg })
      end
      return cb({
        status = "success",
        data = format_checklist_output(checklist)
      })
    end,
  },

  function_call = {},

  schema = {
    type = "function",
    ["function"] = {
      name = "checklist_complete_tasks",
      description =
      "Mark tasks complete in a checklist. Requires checklist_id, complete_task_ids, commit_message, and context.",
      parameters = {
        type = "object",
        properties = {
          checklist_id = {
            type = "string",
            description = "Checklist ID to update"
          },
          complete_task_ids = {
            type = "array",
            items = { type = "string" },
            description = "Task IDs to mark complete"
          },
          commit_message = {
            type = "string",
            description = "Justification for the change"
          },
          context = {
            type = "object",
            description = "All information used to make the decision"
          }
        },
        required = { "checklist_id", "complete_task_ids", "commit_message", "context" },
        additionalProperties = false
      },
      strict = true
    }
  },

  system_prompt = [[## Checklist Complete Tasks Tool

- Use this tool to mark tasks complete in a checklist.
- All fields are required.
- All changes require:
  - commit_message: a detailed justification for the change.
  - context: all information used to make the decision (user prompt, file, code, etc).

### Usage Example

- Mark tasks complete (replace X with checklist ID, and Y/Z with task IDs):
  @checklist_complete_tasks checklist_id="X" complete_task_ids=["Y", "Z"] commit_message="Tasks completed after review" context={...}
]],

  opts = {
    requires_approval = true,
  },

  env = nil,

  handlers = {
    setup = function(tool, agent)
      vim.notify(string.format("%s setup", tool.name), vim.log.levels.DEBUG)
    end,

    prompt_condition = function(tool, agent, config)
      vim.notify(string.format("%s prompt_condition", tool.name), vim.log.levels.DEBUG)
    end,

    on_exit = function(tool, agent)
      vim.notify(string.format("%s on_exit", tool.name), vim.log.levels.DEBUG)
    end,
  },

  output = {
    success = function(tool, agent, cmd, stdout)
      local chat = agent.chat
      chat:add_tool_output(tool, stdout[1])
    end,
    error = function(tool, agent, cmd, stderr)
      local chat = agent.chat
      local error_msg = stderr[1] or "Unknown error"
      chat:add_tool_output(tool, string.format("**Checklist Complete Tasks Tool Error**: %s", error_msg))
    end,
    rejected = function(tool, agent, cmd)
      local chat = agent.chat
      chat:add_tool_output(tool, "**Checklist Complete Tasks Tool**: User declined to execute the operation")
    end,
  },

  ["output.prompt"] = function(tool, agent)
    return string.format(
      "Mark tasks complete in checklist %s?\n\nCommit message: %s\nContext: %s",
      tool.args.checklist_id or "",
      tool.args.commit_message or "",
      type(tool.args.context) == "string" and tool.args.context or vim.inspect(tool.args.context)
    )
  end,

  args = {},

  tool = {},
}

local M = {
  checklist_list = {
    description = "List all checklists in the workspace.",
    callback = ChecklistListTool
  },
  checklist_status = {
    description = "Show the status of a specific checklist.",
    callback = ChecklistStatusTool
  },
  checklist_create = {
    description = "Create a new checklist with tasks.",
    callback = ChecklistCreateTool
  },
  checklist_complete_tasks = {
    description = "Mark tasks complete in a checklist.",
    callback = ChecklistCompleteTasksTool
  },
}

return M