refactor(ui): cleaner methods

Author: olimorris

lua/codecompanion/interactions/chat/ui/builder.lua

@@ -136,8 +136,8 @@ local function create_state(out, base_state)
 end
 
 ---Add message using centralized state
----@param data table
----@param opts table
+---@param data { content?: string, role?: string, reasoning?: { content: string } }
+---@param opts? { type?: string, force_role?: boolean, insert_at?: number, status?: string, _icon_info?: table, virt_text_pos?: string, fold_info?: table, state?: table }
 ---@return number,number|nil
 function Builder:add_message(data, opts)
   opts = opts or {}
@@ -197,7 +197,9 @@ function Builder:add_message(data, opts)
 
   local insert_line, icon_id
   if not vim.tbl_isempty(lines) then
-    insert_line, icon_id = self:_write_to_buffer(lines, opts, fold_info, state)
+    opts.fold_info = fold_info
+    opts.state = state
+    insert_line, icon_id = self:_write_to_buffer(lines, opts)
   end
 
   if current_type then
@@ -225,24 +227,24 @@ function Builder:add_message(data, opts)
 end
 
 ---Determine if we should start a new block under the header
----@param opts table
+---@param opts { type?: string }
 ---@param state table
 ---@return boolean
 function Builder:_should_start_new_block(opts, state)
   return opts.type ~= nil and opts.type ~= state.last_type
 end
 
 ---Check if we need to add a header to the chat buffer
----@param data table
----@param opts table
+---@param data { role?: string }
+---@param opts { force_role?: boolean }
 ---@param state table
 ---@return boolean
 function Builder:_should_add_header(data, opts, state)
   return (data.role and data.role ~= state.last_role) or (opts and opts.force_role)
 end
 
 ---Add appropriate spacing before header
----@param lines table
+---@param lines string[]
 ---@return nil
 function Builder:_add_header_spacing(lines)
   table.insert(lines, "")
@@ -263,12 +265,13 @@ function Builder:_get_formatter(data, opts)
 end
 
 ---Write lines to buffer with all the buffer management
----@param lines table
----@param opts table
----@param fold_info table|nil
----@param state table
----@return number The line number where the content was written and the extmark id of any icon applied
-function Builder:_write_to_buffer(lines, opts, fold_info, state)
+---@param lines string[]
+---@param opts { insert_at?: number, _icon_info?: table, virt_text_pos?: string, fold_info?: table, state?: table }
+---@return number, number|nil
+function Builder:_write_to_buffer(lines, opts)
+  local state = opts.state
+  local fold_info = opts.fold_info
+
   self.chat.ui:unlock_buf()
   local last_line, last_column, line_count = self.chat.ui:last()
 
@@ -345,7 +348,7 @@ end
 ---Update a specific line in the chat buffer
 ---@param line_number number The line number to update (1-based)
 ---@param content string The new content for the line
----@param opts? table Optional parameters
+---@param opts? { status?: string, icon_id?: number, priority?: number, virt_text_pos?: string }
 ---@return boolean success Whether the update was successful
 ---@return number|nil icon_id The new icon extmark ID, if an icon was placed
 function Builder:update_line(line_number, content, opts)

lua/codecompanion/interactions/chat/ui/formatters/base.lua

@@ -24,25 +24,25 @@ end
 
 ---Check if this formatter can handle the given data/opts
 ---@param message table
----@param opts table
----@param tags table
+---@param opts { type?: string, status?: string }
+---@param tags table<string, string>
 ---@return boolean
 function BaseFormatter:can_handle(message, opts, tags)
   error("Must implement can_handle method")
 end
 
 ---Get the message type for this formatter
----@param opts table
+---@param opts { type?: string }
 ---@return string
 function BaseFormatter:get_type(opts)
   error("Must implement get_type method")
 end
 
 ---Format the content into lines
 ---@param message table
----@param opts table
+---@param opts { type?: string, status?: string, _icon_info?: table }
 ---@param state table
----@return table lines, table? fold_info
+---@return string[] lines, table? fold_info
 function BaseFormatter:format(message, opts, state)
   error("Must implement format method")
 end

lua/codecompanion/interactions/chat/ui/icons.lua

@@ -12,7 +12,7 @@ local CONSTANTS = {
 ---@param bufnr number
 ---@param line number 0-based line number
 ---@param status string The tool status (in_progress, completed, failed)
----@param opts table|nil Additional options
+---@param opts? { priority?: number, virt_text_pos?: string }
 ---@return number|nil The extmark id or nil if status is invalid
 function Icons.apply(bufnr, line, status, opts)
   opts = vim.tbl_deep_extend("force", {