Refactor UI formatters to use centralized spacing management

Co-authored-by: olimorris <[email protected]>

Author: Copilot

lua/codecompanion/strategies/chat/ui/builder.lua

@@ -3,6 +3,7 @@ local config = require("codecompanion.config")
 local Reasoning = require("codecompanion.strategies.chat.ui.formatters.reasoning")
 local Standard = require("codecompanion.strategies.chat.ui.formatters.standard")
 local Tools = require("codecompanion.strategies.chat.ui.formatters.tools")
+local Spacing = require("codecompanion.strategies.chat.ui.spacing")
 
 local api = vim.api
 
@@ -143,17 +144,22 @@ 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
+---Add appropriate spacing before header using centralized spacing manager
 ---@param lines table
 ---@param state table
 ---@return nil
 function Builder:_add_header_spacing(lines, state)
-  if state.last_type == self.chat.MESSAGE_TYPES.TOOL_MESSAGE then
-    table.insert(lines, "")
-  else
-    table.insert(lines, "")
-    table.insert(lines, "")
-  end
+  local spacing_context = {
+    is_new_role = true,
+    is_new_section = false,
+    previous_type = state.last_type,
+    current_type = nil, -- Headers are for roles, not specific message types
+    has_reasoning_transition = false,
+    is_reasoning_start = false,
+  }
+
+  local spacing_lines = Spacing.get_header_spacing(spacing_context, self.chat.MESSAGE_TYPES)
+  vim.list_extend(lines, spacing_lines)
 end
 
 ---Get the appropriate formatter

lua/codecompanion/strategies/chat/ui/formatters/reasoning.lua

@@ -1,4 +1,5 @@
 local BaseFormatter = require("codecompanion.strategies.chat.ui.formatters.base")
+local Spacing = require("codecompanion.strategies.chat.ui.spacing")
 
 ---@class CodeCompanion.Chat.UI.Formatters.Reasoning : CodeCompanion.Chat.UI.Formatters.Base
 local Reasoning = setmetatable({}, { __index = BaseFormatter })
@@ -15,14 +16,24 @@ end
 function Reasoning:format(message, opts, state)
   local lines = {}
 
-  -- Use rich state methods
+  -- Build spacing context
+  local spacing_context = {
+    is_new_role = false, -- Reasoning is always under LLM role
+    is_new_section = state.is_new_section,
+    previous_type = state.last_type,
+    current_type = self:get_type(),
+    has_reasoning_transition = false, -- We're starting reasoning, not transitioning from it
+    is_reasoning_start = not state.has_reasoning_output,
+  }
+
+  -- Add reasoning header if this is the start of reasoning output
   if not state.has_reasoning_output then
     table.insert(lines, "### Reasoning")
     table.insert(lines, "")
     state:mark_reasoning_started()
   end
 
-  -- Add reasoning content
+  -- Add reasoning content directly (no additional spacing needed)
   for _, line in ipairs(vim.split(message.content, "\n", { plain = true, trimempty = false })) do
     table.insert(lines, line)
   end
@@ -31,3 +42,6 @@ function Reasoning:format(message, opts, state)
 end
 
 return Reasoning
+end
+
+return Reasoning

lua/codecompanion/strategies/chat/ui/formatters/standard.lua

@@ -1,4 +1,5 @@
 local BaseFormatter = require("codecompanion.strategies.chat.ui.formatters.base")
+local Spacing = require("codecompanion.strategies.chat.ui.spacing")
 
 ---@class CodeCompanion.Chat.UI.Formatters.Standard : CodeCompanion.Chat.UI.Formatters.Base
 local Standard = setmetatable({}, { __index = BaseFormatter })
@@ -16,16 +17,26 @@ end
 function Standard:format(message, opts, state)
   local lines = {}
 
-  -- Handle transition from reasoning to response using rich state
+  -- Build spacing context
+  local spacing_context = {
+    is_new_role = false, -- Standard messages are typically under LLM role
+    is_new_section = state.is_new_section,
+    previous_type = state.last_type,
+    current_type = self:get_type(opts),
+    has_reasoning_transition = state.has_reasoning_output,
+    is_reasoning_start = false,
+  }
+
+  -- Handle transition from reasoning to response
   if state.has_reasoning_output then
     state:mark_reasoning_complete()
-    table.insert(lines, "")
-    table.insert(lines, "")
+    local transition_spacing = Spacing.get_pre_content_spacing(spacing_context, self.chat.MESSAGE_TYPES)
+    vim.list_extend(lines, transition_spacing)
     table.insert(lines, "### Response")
     table.insert(lines, "")
   end
 
-  -- Add content
+  -- Add content directly (spacing is handled at message boundary level)
   for _, line in ipairs(vim.split(message.content, "\n", { plain = true, trimempty = false })) do
     table.insert(lines, line)
   end
@@ -34,3 +45,6 @@ function Standard:format(message, opts, state)
 end
 
 return Standard
+end
+
+return Standard

lua/codecompanion/strategies/chat/ui/formatters/tools.lua

@@ -1,6 +1,7 @@
 local config = require("codecompanion.config")
 
 local BaseFormatter = require("codecompanion.strategies.chat.ui.formatters.base")
+local Spacing = require("codecompanion.strategies.chat.ui.spacing")
 
 ---@class CodeCompanion.Chat.UI.Formatters.Tools : CodeCompanion.Chat.UI.Formatters.Base
 local Tools = setmetatable({}, { __index = BaseFormatter })
@@ -17,37 +18,51 @@ end
 function Tools:format(message, opts, state)
   local lines = {}
 
+  -- Build spacing context for consistent spacing management
+  local spacing_context = {
+    is_new_role = false, -- Tools are always under LLM role
+    is_new_section = state.is_new_section,
+    previous_type = state.last_type,
+    current_type = self:get_type(),
+    has_reasoning_transition = state.has_reasoning_output,
+    is_reasoning_start = false,
+  }
+
+  -- Handle reasoning to response transition
   if state.has_reasoning_output then
     state:mark_reasoning_complete()
-    table.insert(lines, "")
-    table.insert(lines, "")
+    local transition_spacing = Spacing.get_pre_content_spacing(spacing_context, self.chat.MESSAGE_TYPES)
+    vim.list_extend(lines, transition_spacing)
     table.insert(lines, "### Response")
   end
 
-  if state.is_new_section then
-    table.insert(lines, "")
-  end
-  if state.last_type == self.chat.MESSAGE_TYPES.LLM_MESSAGE then
-    table.insert(lines, "")
+  -- Get pre-content spacing (but not for reasoning transitions as they're handled above)
+  if not state.has_reasoning_output then
+    local pre_spacing = Spacing.get_pre_content_spacing(spacing_context, self.chat.MESSAGE_TYPES)
+    vim.list_extend(lines, pre_spacing)
   end
 
   local content_start = #lines + 1
   local content = message.content or ""
   for _, line in ipairs(vim.split(content, "\n", { plain = true, trimempty = false })) do
     table.insert(lines, line)
   end
-  table.insert(lines, "")
+
+  -- Get post-content spacing
+  local post_spacing = Spacing.get_post_content_spacing(spacing_context, self.chat.MESSAGE_TYPES)
+  vim.list_extend(lines, post_spacing)
 
   if not config.strategies.chat.tools.opts.folds.enabled then
     return lines, nil
   end
 
-  -- Folds can only work up to the penultimate line in the buffer, so an extra
-  -- line has been added as a result. But, we don't want to fold this line
-  local content_end = #lines - 1
+  -- Folds can only work up to the penultimate line in the buffer, so we need
+  -- to account for the trailing spacing in fold calculations
+  local content_lines = vim.split(content, "\n", { plain = true, trimempty = false })
+  local content_end = content_start + #content_lines - 1
 
   local fold_info = nil
-  if content_end >= content_start then
+  if content_end >= content_start and #content_lines > 1 then
     fold_info = {
       start_offset = content_start - 1, -- 0-based index
       end_offset = content_end - 1, -- 0-based index
@@ -59,3 +74,6 @@ function Tools:format(message, opts, state)
 end
 
 return Tools
+end
+
+return Tools

lua/codecompanion/strategies/chat/ui/spacing.lua

@@ -0,0 +1,88 @@
+---Centralized spacing management for chat UI
+---Provides consistent spacing rules across all formatters
+
+---@class CodeCompanion.Chat.UI.Spacing
+local Spacing = {}
+
+---@class CodeCompanion.Chat.UI.SpacingContext
+---@field is_new_role boolean Whether this is a new role (triggers header)
+---@field is_new_section boolean Whether this is a new section under same role
+---@field previous_type string|nil The type of the previous message
+---@field current_type string The type of the current message
+---@field has_reasoning_transition boolean Whether transitioning from reasoning to response
+---@field is_reasoning_start boolean Whether starting reasoning output
+
+---Determine spacing needed before content
+---@param context CodeCompanion.Chat.UI.SpacingContext
+---@param message_types table Available message type constants
+---@return table lines Array of spacing lines (empty strings)
+function Spacing.get_pre_content_spacing(context, message_types)
+  local lines = {}
+
+  -- Handle reasoning transitions first (highest priority)
+  if context.has_reasoning_transition then
+    table.insert(lines, "")
+    table.insert(lines, "")
+    return lines
+  end
+
+  -- Handle new sections (when type changes but role stays the same)
+  if context.is_new_section then
+    -- Less spacing for section transitions within same role
+    if context.previous_type == message_types.LLM_MESSAGE 
+       and context.current_type == message_types.TOOL_MESSAGE then
+      table.insert(lines, "")
+    elseif context.previous_type == message_types.TOOL_MESSAGE 
+           and context.current_type == message_types.LLM_MESSAGE then
+      table.insert(lines, "")
+    end
+    return lines
+  end
+
+  -- For reasoning start, no extra spacing needed as header handles it
+  if context.is_reasoning_start then
+    return lines
+  end
+
+  return lines
+end
+
+---Determine spacing needed after content
+---@param context CodeCompanion.Chat.UI.SpacingContext
+---@param message_types table Available message type constants
+---@return table lines Array of spacing lines (empty strings)
+function Spacing.get_post_content_spacing(context, message_types)
+  local lines = {}
+
+  -- Most content doesn't need trailing spacing
+  -- The builder will handle spacing between messages
+  
+  -- Only tool output might need a trailing line for folding purposes
+  if context.current_type == message_types.TOOL_MESSAGE then
+    table.insert(lines, "")
+  end
+
+  return lines
+end
+
+---Determine spacing needed before headers
+---@param context CodeCompanion.Chat.UI.SpacingContext
+---@param message_types table Available message type constants
+---@return table lines Array of spacing lines (empty strings)
+function Spacing.get_header_spacing(context, message_types)
+  local lines = {}
+
+  -- Less aggressive spacing before headers
+  if context.previous_type == message_types.TOOL_MESSAGE then
+    -- Single line after tool messages
+    table.insert(lines, "")
+  elseif context.previous_type then
+    -- Double line for other transitions (standard)
+    table.insert(lines, "")
+    table.insert(lines, "")
+  end
+
+  return lines
+end
+
+return Spacing