fix(tools): make insert_edit_into_file more robust for patch markers (olimorris#1653)

LLMs sometimes miss the BEGIN and END wrapping for the patches. This
fails the merge.

Making the algorithm lax to allow and handle these misses confidently.

Author: pratyushmittal

lua/codecompanion/strategies/chat/agents/tools/helpers/patch.lua

@@ -86,7 +86,7 @@ local function parse_changes_from_patch(patch)
     elseif line == "" and lines[i + 1] and lines[i + 1]:match("^@@") then
       -- empty lines can be part of pre/post context
       -- we treat empty lines as new change block and not as post context
-      -- only when the the next line uses @@ identifier
+      -- only when the next line uses @@ identifier
       table.insert(changes, change)
       change = get_new_change()
     elseif line:sub(1, 1) == "-" then
@@ -115,14 +115,20 @@ end
 
 ---Parse the full raw string from LLM for all patches, returning all Change objects parsed.
 ---@param raw string Raw text containing patch blocks
----@return Change[] All parsed Change objects
+---@return Change[], boolean All parsed Change objects, and whether the patch was properly parsed
 function M.parse_changes(raw)
   local patches = {}
   for patch in raw:gmatch("%*%*%* Begin Patch%s+(.-)%s+%*%*%* End Patch") do
     table.insert(patches, patch)
   end
+
+  local had_begin_end_markers = true
   if #patches == 0 then
-    error("Invalid patch format: missing Begin/End markers")
+    --- LLMs miss the begin / end markers sometimes
+    --- let's assume the raw content was correctly wrapped in these cases
+    --- setting a `markers_error` so that we can show this error in case the patch fails to apply
+    had_begin_end_markers = false
+    table.insert(patches, raw)
   end
 
   local all_changes = {}
@@ -132,7 +138,7 @@ function M.parse_changes(raw)
       table.insert(all_changes, change)
     end
   end
-  return all_changes
+  return all_changes, had_begin_end_markers
 end
 
 ---Score how many lines from needle match haystack lines.
@@ -190,7 +196,7 @@ end
 ---Determine best insertion spot for a Change and its match score.
 ---@param lines string[] File lines
 ---@param change Change Patch block
----@return integer,number location (1-based), Score (0-1)
+---@return integer, number location (1-based), Score (0-1)
 local function get_best_location(lines, change)
   -- try applying patch in flexible spaces mode
   -- there is no standardised way to of spaces in diffs

lua/codecompanion/strategies/chat/agents/tools/insert_edit_into_file.lua

@@ -36,7 +36,7 @@ local function edit_file(action)
 
   -- 1. extract list of changes from the code
   local raw = action.code or ""
-  local changes = patch.parse_changes(raw)
+  local changes, had_begin_end_markers = patch.parse_changes(raw)
 
   -- 2. read file into lines
   local content = p:read()
@@ -46,7 +46,11 @@ local function edit_file(action)
   for _, change in ipairs(changes) do
     local new_lines = patch.apply_change(lines, change)
     if new_lines == nil then
-      error(fmt("Bad/Incorrect diff:\n\n%s\n\nNo changes were applied", patch.get_change_string(change)))
+      if had_begin_end_markers then
+        error(fmt("Bad/Incorrect diff:\n\n%s\n\nNo changes were applied", patch.get_change_string(change)))
+      else
+        error("Invalid patch format: missing Begin/End markers")
+      end
     else
       lines = new_lines
     end
@@ -82,7 +86,7 @@ local function edit_buffer(bufnr, chat_bufnr, action, output_handler, opts)
 
   -- Parse and apply patches to buffer
   local raw = action.code or ""
-  local changes = patch.parse_changes(raw)
+  local changes, had_begin_end_markers = patch.parse_changes(raw)
 
   -- Get current buffer content as lines
   local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
@@ -92,7 +96,11 @@ local function edit_buffer(bufnr, chat_bufnr, action, output_handler, opts)
   for _, change in ipairs(changes) do
     local new_lines = patch.apply_change(lines, change)
     if new_lines == nil then
-      error(fmt("Bad/Incorrect diff:\n\n%s\n\nNo changes were applied", patch.get_change_string(change)))
+      if had_begin_end_markers then
+        error(fmt("Bad/Incorrect diff:\n\n%s\n\nNo changes were applied", patch.get_change_string(change)))
+      else
+        error("Invalid patch format: missing Begin/End markers")
+      end
     else
       if not start_line then
         start_line = patch.get_change_location(lines, change)

tests/fixtures/files-diff-1.5.patch

@@ -1,10 +1,12 @@
-*** Begin Patch
 @@ two
 three
 four
-five
 -five
 +5
 six
 seven
-*** End Patch
+
+@@
+seven
+eight
+-nine

tests/fixtures/files-output-1.5.html

@@ -6,5 +6,4 @@
 six
 seven
 eight
-nine
 ten

tests/strategies/chat/agents/tools/test_insert_edit_into_file.lua

@@ -214,6 +214,34 @@ T["File"]["insert_edit_into_file tool multiple patches"] = function()
   h.eq_info(output, expected, child.lua_get("chat.messages[#chat.messages].content"))
 end
 
+T["File"]["insert_edit_into_file tool multiple patches"] = function()
+  child.lua([[
+      -- read initial file from fixture
+      local initial = vim.fn.readfile("tests/fixtures/files-input-1.html")
+      local ok = vim.fn.writefile(initial, _G.TEST_TMPFILE_ABSOLUTE)
+      assert(ok == 0)
+
+      -- read contents for the tool from fixtures
+      local content = table.concat(vim.fn.readfile("tests/fixtures/files-diff-1.5.patch"), "\n")
+      local arguments = vim.json.encode({ filepath = _G.TEST_TMPFILE, explanation = "...", code = content })
+      local tool = {
+        {
+          ["function"] = {
+            name = "insert_edit_into_file",
+            arguments = arguments
+          },
+        },
+      }
+      agent:execute(chat, tool)
+      vim.wait(200)
+    ]])
+
+  -- Test that the file was updated as per the output fixture
+  local output = child.lua_get("vim.fn.readfile(_G.TEST_TMPFILE_ABSOLUTE)")
+  local expected = child.lua_get("vim.fn.readfile('tests/fixtures/files-output-1.5.html')")
+  h.eq_info(output, expected, child.lua_get("chat.messages[#chat.messages].content"))
+end
+
 T["File"]["insert_edit_into_file tool multiple continuation"] = function()
   child.lua([[
       -- read initial file from fixture