Fine-tune spacing logic and add comprehensive tests

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

Author: Copilot

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

@@ -28,22 +28,14 @@ function Spacing.get_pre_content_spacing(context, message_types)
 
   -- 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
+    -- Only add spacing for LLM to Tool transitions
     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
 

tests/strategies/chat/test_ui_spacing.lua

@@ -25,7 +25,7 @@ T = new_set({
 
 T["UI Spacing Refactor"] = new_set()
 
-T["UI Spacing Refactor"]["basic tool output has reduced spacing"] = function()
+T["UI Spacing Refactor"]["reduced spacing between LLM and tool messages"] = function()
   local result = child.lua([[
     -- Add a user message first
     _G.chat.builder:add_message({
@@ -47,28 +47,39 @@ T["UI Spacing Refactor"]["basic tool output has reduced spacing"] = function()
 
     local lines = h.get_buf_lines(_G.chat.bufnr)
     
-    -- Count consecutive empty lines after tool output
-    local consecutive_empty = 0
-    local max_consecutive_empty = 0
+    -- Find the position of LLM message and tool message
+    local llm_line_idx = nil
+    local tool_line_idx = nil
     for i, line in ipairs(lines) do
-      if line == "" then
-        consecutive_empty = consecutive_empty + 1
-        max_consecutive_empty = math.max(max_consecutive_empty, consecutive_empty)
-      else
-        consecutive_empty = 0
+      if line:match("I'll help you with that") then
+        llm_line_idx = i
+      elseif line:match("Tool executed successfully") then
+        tool_line_idx = i
+      end
+    end
+
+    -- Count empty lines between LLM and tool message
+    local empty_lines_between = 0
+    if llm_line_idx and tool_line_idx then
+      for i = llm_line_idx + 1, tool_line_idx - 1 do
+        if lines[i] == "" then
+          empty_lines_between = empty_lines_between + 1
+        end
       end
     end
 
     return {
       buffer_lines = lines,
-      max_consecutive_empty = max_consecutive_empty,
-      total_lines = #lines
+      llm_line_idx = llm_line_idx,
+      tool_line_idx = tool_line_idx,
+      empty_lines_between = empty_lines_between
     }
   ]])
 
-  -- Should have significantly fewer consecutive empty lines
-  h.leq(result.max_consecutive_empty, 2) -- At most 2 consecutive empty lines
-  h.neq(result.total_lines, 0) -- Should have content
+  -- Should have at most 1 empty line between LLM message and tool output (improved from before)
+  h.leq(result.empty_lines_between, 1)
+  h.neq(result.llm_line_idx, nil) -- Should find LLM message
+  h.neq(result.tool_line_idx, nil) -- Should find tool message
 end
 
 T["UI Spacing Refactor"]["reasoning to response transition works"] = function()
@@ -93,26 +104,66 @@ T["UI Spacing Refactor"]["reasoning to response transition works"] = function()
 
     local lines = h.get_buf_lines(_G.chat.bufnr)
     
-    -- Check for reasoning header
+    -- Check for reasoning header and response header
     local has_reasoning_header = false
     local has_response_header = false
-    for _, line in ipairs(lines) do
+    local response_header_line = nil
+    for i, line in ipairs(lines) do
       if line:match("### Reasoning") then
         has_reasoning_header = true
       elseif line:match("### Response") then
         has_response_header = true
+        response_header_line = i
       end
     end
 
     return {
       buffer_lines = lines,
       has_reasoning_header = has_reasoning_header,
-      has_response_header = has_response_header
+      has_response_header = has_response_header,
+      response_header_line = response_header_line
     }
   ]])
 
   h.eq(result.has_reasoning_header, true)
   h.eq(result.has_response_header, true)
+  h.neq(result.response_header_line, nil)
+end
+
+T["UI Spacing Refactor"]["tool output has proper trailing spacing for folding"] = function()
+  local result = child.lua([[
+    -- Add tool output
+    _G.chat.builder:add_message({
+      role = "llm",
+      content = "Tool executed successfully\\nWith multiple lines\\nOf output"
+    }, { type = _G.chat.MESSAGE_TYPES.TOOL_MESSAGE })
+
+    local lines = h.get_buf_lines(_G.chat.bufnr)
+    
+    -- Find the last line of tool content
+    local last_content_line_idx = nil
+    for i = #lines, 1, -1 do
+      if lines[i]:match("Of output") then
+        last_content_line_idx = i
+        break
+      end
+    end
+
+    -- Check if there's a trailing empty line after tool content
+    local has_trailing_empty = false
+    if last_content_line_idx and last_content_line_idx < #lines then
+      has_trailing_empty = lines[last_content_line_idx + 1] == ""
+    end
+
+    return {
+      buffer_lines = lines,
+      last_content_line_idx = last_content_line_idx,
+      has_trailing_empty = has_trailing_empty
+    }
+  ]])
+
+  h.neq(result.last_content_line_idx, nil) -- Should find tool content
+  h.eq(result.has_trailing_empty, true) -- Should have trailing empty line for folding
 end
 
 return T