feat(ui): don't stack extmarks in ACP chats (#2902)

Author: olimorris

lua/codecompanion/interactions/chat/acp/handler.lua

@@ -4,17 +4,14 @@ local log = require("codecompanion.utils.log")
 local utils = require("codecompanion.utils")
 local watch = require("codecompanion.interactions.shared.watch")
 
--- Keep a record of UI changes in the chat buffer
-
 ---@class CodeCompanion.Chat.ACPHandler
 ---@field chat CodeCompanion.Chat
 ---@field output table Standard output message from the Agent
 ---@field reasoning table Reasoning output from the Agent
 ---@field tools table<string, table> Cache of tool calls by their ID
+---@field ui_state table<string, table> Cache of tool call UI states (line_number, icon_id) by tool call ID
 local ACPHandler = {}
 
-local ACPHandlerUI = {} -- Cache of tool call UI states by chat buffer
-
 ---@param chat CodeCompanion.Chat
 ---@return CodeCompanion.Chat.ACPHandler
 function ACPHandler.new(chat)
@@ -23,19 +20,12 @@ function ACPHandler.new(chat)
     output = {},
     reasoning = {},
     tools = {},
+    ui_state = {},
   }, { __index = ACPHandler })
 
-  ACPHandlerUI[chat.bufnr] = {}
-
   return self --[[@type CodeCompanion.Chat.ACPHandler]]
 end
 
----Return the ACP client
----@return CodeCompanion.ACP.Connection
-local get_client = function()
-  return require("codecompanion.acp")
-end
-
 ---Merge an incoming tool call/update into the cache
 ---@param existing table|nil
 ---@param incoming table|nil
@@ -66,7 +56,7 @@ end
 ---@return boolean success
 function ACPHandler:ensure_connection()
   if not self.chat.acp_connection then
-    self.chat.acp_connection = get_client().new({
+    self.chat.acp_connection = require("codecompanion.acp").new({
       adapter = self.chat.adapter, --[[@type CodeCompanion.ACPAdapter]]
     })
 
@@ -155,16 +145,16 @@ function ACPHandler:create_and_send_prompt(payload)
       self:handle_thought_chunk(content)
     end)
     :on_tool_call(function(tool_call)
-      self:handle_tool_call(tool_call)
+      self:process_tool_call(tool_call)
     end)
-    :on_tool_update(function(tool_update)
-      self:handle_tool_update(tool_update)
+    :on_tool_update(function(tool_call)
+      self:process_tool_call(tool_call)
     end)
     :on_permission_request(function(request)
       self:handle_permission_request(request)
     end)
-    :on_complete(function(stop_reason)
-      self:handle_completion(stop_reason)
+    :on_complete(function()
+      self:handle_completion()
     end)
     :on_error(function(error)
       self:handle_error(error)
@@ -199,12 +189,10 @@ end
 ---@param tool_call table
 ---@return nil
 function ACPHandler:process_tool_call(tool_call)
-  -- Cache the tool call to handle processing later on, such as a later permission request
   local id = tool_call.toolCallId
 
-  local prev = self.tools[id]
-  local merged = merge_tool_call(prev, tool_call)
-  tool_call = merged or tool_call
+  local merged = merge_tool_call(self.tools[id], tool_call)
+  tool_call = merged
 
   local ok, content = pcall(formatter.tool_message, tool_call, self.chat.adapter)
   if not ok then
@@ -218,31 +206,31 @@ function ACPHandler:process_tool_call(tool_call)
     self.tools[id] = merged
   end
 
-  -- If the tool call has already written output to the chat buffer, then we can
-  -- update it rather than adding a new line. We do this by keeping track in
-  -- a global cache, segmented by chat buffer and tool call IDs
-  if ACPHandlerUI[self.chat.bufnr][id] then
-    local match = ACPHandlerUI[self.chat.bufnr][id]
-    -- Whilst I've tried to account for all types of ACP tool output, I'm taking
-    -- a cautious approach and wrapping line updates. Any failures and we'll
-    -- just write the tool output onto a new line in the chat buffer
-    ok, _ = pcall(function()
-      self.chat:update_buf_line(
-        match.line_number,
-        content,
-        { status = tool_call.status, icon_id = match.icon_id, priority = 120, virt_text_pos = "inline" }
-      )
-    end)
+  -- If the tool call has already written output to the chat buffer, update the
+  -- existing line rather than adding a new one
+  local cached = self.ui_state[id]
+  if cached then
+    local update_ok, _, new_icon_id = pcall(
+      self.chat.update_buf_line,
+      self.chat,
+      cached.line_number,
+      content,
+      { status = tool_call.status, icon_id = cached.icon_id, priority = 120, virt_text_pos = "inline" }
+    )
 
-    -- Cleanup the cache
-    if tool_call.status == "completed" then
-      ACPHandlerUI[self.chat.bufnr][id] = nil
+    if update_ok then
+      if tool_call.status == "completed" then
+        self.ui_state[id] = nil
+      elseif new_icon_id then
+        cached.icon_id = new_icon_id
+      end
+      return
     end
 
-    if ok then
-      return
+    if tool_call.status == "completed" then
+      self.ui_state[id] = nil
     end
-    log:debug("[ACP::Handler] Failed to update tool call line for toolCallId %s", tool_call.toolCallId)
+    log:debug("[ACP::Handler] Failed to update tool call line for toolCallId %s", id)
   end
 
   table.insert(self.output, content)
@@ -257,19 +245,7 @@ function ACPHandler:process_tool_call(tool_call)
     type = self.chat.MESSAGE_TYPES.TOOL_MESSAGE,
   })
 
-  ACPHandlerUI[self.chat.bufnr][id] = { line_number = line_number, icon_id = icon_id }
-end
-
----Handle tool call notifications
----@param tool_call table
-function ACPHandler:handle_tool_call(tool_call)
-  return self:process_tool_call(tool_call)
-end
-
----Handle tool call updates and their respective status
----@param tool_call table
-function ACPHandler:handle_tool_update(tool_call)
-  return self:process_tool_call(tool_call)
+  self.ui_state[id] = { line_number = line_number, icon_id = icon_id }
 end
 
 ---Handle permission requests from the agent
@@ -302,8 +278,7 @@ function ACPHandler:handle_permission_request(request)
 end
 
 ---Handle completion
----@param stop_reason string|nil
-function ACPHandler:handle_completion(stop_reason)
+function ACPHandler:handle_completion()
   if not self.chat.status or self.chat.status == "" then
     self.chat.status = "success"
   end

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

@@ -150,7 +150,7 @@ function Builder:add_message(data, opts)
 
   if needs_header then
     state:update_role(data.role)
-    self:_add_header_spacing(lines, state)
+    self:_add_header_spacing(lines)
     self.chat.ui:set_header(lines, config.interactions.chat.roles[data.role])
 
     -- Section started: reset block trackers
@@ -243,9 +243,8 @@ end
 
 ---Add appropriate spacing before header
 ---@param lines table
----@param state table
 ---@return nil
-function Builder:_add_header_spacing(lines, state)
+function Builder:_add_header_spacing(lines)
   table.insert(lines, "")
   table.insert(lines, "")
 end
@@ -348,6 +347,7 @@ end
 ---@param content string The new content for the line
 ---@param opts? table Optional parameters
 ---@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)
   opts = opts or {}
 
@@ -370,19 +370,23 @@ function Builder:update_line(line_number, content, opts)
   local start_line = zero_based_line
   local end_line = zero_based_line + 1
 
+  local new_icon_id
   local ok, _ = pcall(api.nvim_buf_set_lines, self.chat.bufnr, start_line, end_line, false, { content })
   if ok and opts.status then
-    vim.schedule(function()
-      Icons.clear_line(self.chat.bufnr, start_line)
-      Icons.apply(self.chat.bufnr, start_line, opts.status, opts)
-    end)
+    -- Clear by extmark ID first (handles extmarks that may have moved to a different line)
+    if opts.icon_id then
+      pcall(api.nvim_buf_del_extmark, self.chat.bufnr, Icons.ns(), opts.icon_id)
+    end
+    -- Also clear by line range as a safety net
+    Icons.clear_line(self.chat.bufnr, start_line)
+    new_icon_id = Icons.apply(self.chat.bufnr, start_line, opts.status, opts)
   end
 
   if self.state.last_role ~= config.constants.USER_ROLE then
     self.chat.ui:lock_buf()
   end
 
-  return true
+  return true, new_icon_id
 end
 
 return Builder

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

@@ -44,6 +44,9 @@ function Icons.apply(bufnr, line, status, opts)
     return
   end
 
+  -- Clear any existing tool icons on this line to prevent duplicates
+  api.nvim_buf_clear_namespace(bufnr, CONSTANTS.NS_TOOL_ICONS, line, line + 1)
+
   return api.nvim_buf_set_extmark(bufnr, CONSTANTS.NS_TOOL_ICONS, line, 0, {
     virt_text = { { config_entry.icon, config_entry.hl_group } },
     virt_text_pos = opts.virt_text_pos,
@@ -55,7 +58,7 @@ end
 ---@param bufnr number
 ---@param extmark_id number
 ---@return nil
-function Icons:clear_icon(bufnr, extmark_id)
+function Icons.clear_icon(bufnr, extmark_id)
   if not extmark_id then
     return
   end
@@ -75,4 +78,10 @@ function Icons.clear_icons(bufnr)
   api.nvim_buf_clear_namespace(bufnr, CONSTANTS.NS_TOOL_ICONS, 0, -1)
 end
 
+---Return the tool icons namespace ID
+---@return number
+function Icons.ns()
+  return CONSTANTS.NS_TOOL_ICONS
+end
+
 return Icons