fix: duplicated replace blocks in diff (#2549)

lua/avante/llm_tools/replace_in_file.lua

@@ -3,7 +3,6 @@ local Helpers = require("avante.llm_tools.helpers")
 local Utils = require("avante.utils")
 local Highlights = require("avante.highlights")
 local Config = require("avante.config")
-local diff2search_replace = require("avante.utils.diff2search_replace")
 
 local PRIORITY = (vim.hl or vim.highlight).priorities.user
 local NAMESPACE = vim.api.nvim_create_namespace("avante-diff")
@@ -102,55 +101,6 @@ M.returns = {
   },
 }
 
---- Some models (e.g., gpt-4o) cannot correctly return diff content and often miss the SEARCH line, so this needs to be manually fixed in such cases.
----@param diff string
----@return string
-local function fix_diff(diff)
-  diff = diff2search_replace(diff)
-  -- Normalize block headers to the expected ones (fix for some LLMs output)
-  diff = diff:gsub("<<<<<<<%s*SEARCH", "------- SEARCH")
-  diff = diff:gsub(">>>>>>>%s*REPLACE", "+++++++ REPLACE")
-  diff = diff:gsub("-------%s*REPLACE", "+++++++ REPLACE")
-  diff = diff:gsub("-------  ", "------- SEARCH\n")
-  diff = diff:gsub("=======  ", "======= \n")
-
-  local fixed_diff_lines = {}
-  local lines = vim.split(diff, "\n")
-  local first_line = lines[1]
-  if first_line and first_line:match("^%s*```") then
-    table.insert(fixed_diff_lines, first_line)
-    table.insert(fixed_diff_lines, "------- SEARCH")
-    fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 2)
-  else
-    table.insert(fixed_diff_lines, "------- SEARCH")
-    if first_line:match("------- SEARCH") then
-      fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 2)
-    else
-      fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 1)
-    end
-  end
-  local the_final_diff_lines = {}
-  local has_split_line = false
-  local replace_block_closed = false
-  for _, line in ipairs(fixed_diff_lines) do
-    if line:match("^-------%s*SEARCH") then has_split_line = false end
-    if line:match("^=======") then has_split_line = true end
-    if line:match("^+++++++%s*REPLACE") then
-      if not has_split_line then
-        table.insert(the_final_diff_lines, "=======")
-        has_split_line = true
-        goto continue
-      else
-        replace_block_closed = true
-      end
-    end
-    table.insert(the_final_diff_lines, line)
-    ::continue::
-  end
-  if not replace_block_closed then table.insert(the_final_diff_lines, "+++++++ REPLACE") end
-  return table.concat(the_final_diff_lines, "\n")
-end
-
 --- IMPORTANT: Using "the_diff" instead of "diff" is to avoid LLM streaming generating function parameters in alphabetical order, which would result in generating "path" after "diff", making it impossible to achieve a streaming diff view.
 ---@type AvanteLLMToolFunc<{ path: string, the_diff?: string }>
 function M.func(input, opts)
@@ -186,7 +136,7 @@ function M.func(input, opts)
     session_ctx.streaming_diff_lines_count_history[opts.tool_use_id] = streaming_diff_lines_count
   end
 
-  local diff = fix_diff(input.the_diff)
+  local diff = Utils.fix_diff(input.the_diff)
 
   if on_log and diff ~= input.the_diff then on_log("diff fixed") end
 

lua/avante/utils/init.lua

@@ -3,6 +3,7 @@ local fn = vim.fn
 local lsp = vim.lsp
 
 local LRUCache = require("avante.utils.lru_cache")
+local diff2search_replace = require("avante.utils.diff2search_replace")
 
 ---@class avante.utils: LazyUtilCore
 ---@field tokens avante.utils.tokens
@@ -1621,4 +1622,61 @@ function M.call_once(func)
   end
 end
 
+--- Some models (e.g., gpt-4o) cannot correctly return diff content and often miss the SEARCH line, so this needs to be manually fixed in such cases.
+---@param diff string
+---@return string
+function M.fix_diff(diff)
+  diff = diff2search_replace(diff)
+  -- Normalize block headers to the expected ones (fix for some LLMs output)
+  diff = diff:gsub("<<<<<<<%s*SEARCH", "------- SEARCH")
+  diff = diff:gsub(">>>>>>>%s*REPLACE", "+++++++ REPLACE")
+  diff = diff:gsub("-------%s*REPLACE", "+++++++ REPLACE")
+  diff = diff:gsub("-------  ", "------- SEARCH\n")
+  diff = diff:gsub("=======  ", "=======\n")
+
+  local fixed_diff_lines = {}
+  local lines = vim.split(diff, "\n")
+  local first_line = lines[1]
+  if first_line and first_line:match("^%s*```") then
+    table.insert(fixed_diff_lines, first_line)
+    table.insert(fixed_diff_lines, "------- SEARCH")
+    fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 2)
+  else
+    table.insert(fixed_diff_lines, "------- SEARCH")
+    if first_line:match("------- SEARCH") then
+      fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 2)
+    else
+      fixed_diff_lines = vim.list_extend(fixed_diff_lines, lines, 1)
+    end
+  end
+  local the_final_diff_lines = {}
+  local has_split_line = false
+  local replace_block_closed = false
+  local should_delete_following_lines = false
+  for _, line in ipairs(fixed_diff_lines) do
+    if should_delete_following_lines then goto continue end
+    if line:match("^-------%s*SEARCH") then has_split_line = false end
+    if line:match("^=======") then
+      if has_split_line then
+        should_delete_following_lines = true
+        goto continue
+      end
+      has_split_line = true
+    end
+    if line:match("^+++++++%s*REPLACE") then
+      if not has_split_line then
+        table.insert(the_final_diff_lines, "=======")
+        has_split_line = true
+        goto continue
+      else
+        replace_block_closed = true
+      end
+    end
+    table.insert(the_final_diff_lines, line)
+    ::continue::
+  end
+  if not replace_block_closed then table.insert(the_final_diff_lines, "+++++++ REPLACE") end
+  return table.concat(the_final_diff_lines, "\n")
+end
+
 return M