wip

Author: olimorris

lua/codecompanion/diff/keymaps.lua

@@ -1,5 +1,6 @@
 local config = require("codecompanion.config")
 local log = require("codecompanion.utils.log")
+local utils = require("codecompanion.utils")
 
 local api = vim.api
 
@@ -18,28 +19,19 @@ local function clear_map(keymaps, bufnr)
   end
 end
 
----Fire diff decision event
----@param diff_id number
----@param accepted boolean
----@param timeout? boolean
----@return nil
-local function notify(diff_id, accepted, timeout)
-  local event_name = accepted and "CodeCompanionDiffAccepted" or "CodeCompanionDiffRejected"
-  api.nvim_exec_autocmds("User", {
-    pattern = event_name,
-    data = {
-      diff_id = diff_id,
-      accepted = accepted,
-      timeout = timeout or false,
-    },
-  })
-end
-
 M.always_accept = {
   desc = "Always accept changes from this chat buffer",
   callback = function(diff_ui)
+    -- There might be an edge case where a user has made a decision on the diff
+    -- and the window has then closed, causing `WinClosed` to fire. So we put
+    -- a guard here to check for this and repeat across all the keymaps.
+    if diff_ui.resolved then
+      return
+    end
+    diff_ui.resolved = true
+
     log:trace("[Diff] Accepting diff for id=%s", diff_ui.diff_id)
-    notify(diff_ui.diff_id, true)
+    utils.fire("DiffAccepted", { id = diff_ui.diff_id })
 
     local approvals = require("codecompanion.interactions.chat.tools.approvals")
     approvals:always(diff_ui.chat_bufnr, diff_ui.tool_name)
@@ -55,8 +47,13 @@ M.always_accept = {
 M.accept_change = {
   desc = "Accept all changes",
   callback = function(diff_ui)
+    if diff_ui.resolved then
+      return
+    end
+    diff_ui.resolved = true
+
     log:trace("[Diff] Accepting diff for id=%s", diff_ui.diff_id)
-    notify(diff_ui.diff_id, true)
+    utils.fire("DiffAccepted", { id = diff_ui.diff_id })
 
     local Diff = require("codecompanion.diff")
     Diff.clear(diff_ui.diff)
@@ -69,8 +66,13 @@ M.accept_change = {
 M.reject_change = {
   desc = "Reject all changes",
   callback = function(diff_ui)
+    if diff_ui.resolved then
+      return
+    end
+    diff_ui.resolved = true
+
     log:trace("[Diff] Rejecting diff for id=%s", diff_ui.diff_id)
-    notify(diff_ui.diff_id, false)
+    utils.fire("DiffRejected", { id = diff_ui.diff_id })
 
     local Diff = require("codecompanion.diff")
     Diff.clear(diff_ui.diff)
@@ -83,8 +85,13 @@ M.reject_change = {
 M.close_window = {
   desc = "Close window and reject",
   callback = function(diff_ui)
+    if diff_ui.resolved then
+      return
+    end
+    diff_ui.resolved = true
+
     log:trace("[Diff] Closing diff window for id=%s", diff_ui.diff_id)
-    notify(diff_ui.diff_id, false, true)
+    utils.fire("DiffRejected", { id = diff_ui.diff_id })
 
     local Diff = require("codecompanion.diff")
     Diff.clear(diff_ui.diff)

lua/codecompanion/diff/ui.lua

@@ -7,14 +7,15 @@ local api = vim.api
 local M = {}
 
 ---@class CodeCompanion.DiffUI
+---@field bufnr number
+---@field chat_bufnr? number If the diff has an associated chat buffer, pass in the chat buffer number
+---@field current_hunk number The current hunk index (1-based)
 ---@field diff CC.Diff
 ---@field diff_id number
 ---@field hunks number The total number of hunks in the diff
----@field current_hunk number The current hunk index (1-based)
----@field bufnr number
----@field winnr number
----@field chat_bufnr? number If the diff has an associated chat buffer, pass in the chat buffer number
+---@field resolved boolean Whether the diff has been resolved (accepted/rejected)
 ---@field tool_name? string If the diff is associated with a tool, pass in the tool name
+---@field winnr number
 local DiffUI = {}
 DiffUI.__index = DiffUI
 
@@ -48,7 +49,7 @@ local function show_keymaps(bufnr, opts)
 
   ui_utils.show_buffer_notification(bufnr, {
     text = string.format(
-      "(%d/%d)  [%s] Always Accept | [%s] Accept | [%s] Reject | [%s]/[%s] Next/Prev hunks | [q] Close",
+      "[%d/%d]  %s Always Accept | %s Accept | %s Reject | %s/%s Next/Prev hunks | q Close",
       opts.current_hunk or 1,
       opts.hunks or 1,
       always_accept,
@@ -66,6 +67,10 @@ end
 ---Navigate to next hunk
 ---@param line number
 function DiffUI:next_hunk(line)
+  if self.hunks == 1 then
+    return
+  end
+
   for index, hunk in ipairs(self.diff.hunks) do
     local hunk_line = hunk.pos[1] + 1
     if hunk_line > line then
@@ -100,6 +105,10 @@ end
 ---@param line number
 ---@return nil
 function DiffUI:previous_hunk(line)
+  if self.hunks == 1 then
+    return
+  end
+
   for i = #self.diff.hunks, 1, -1 do
     local hunk = self.diff.hunks[i]
     local hunk_line = hunk.pos[1] + 1
@@ -179,6 +188,8 @@ function M.show(diff, opts)
 
   local cfg = vim.tbl_deep_extend("force", config.display.chat.floating_window, config.display.chat.diff_window or {})
 
+  local title = opts.title or get_buf_name(diff.bufnr)
+
   local bufnr, winnr = ui_utils.create_float(diff.to.lines, {
     window = {
       width = cfg.width,
@@ -189,19 +200,20 @@ function M.show(diff, opts)
     col = cfg.col,
     filetype = diff.ft or "text",
     ignore_keymaps = true,
-    title = opts.title or get_buf_name(diff.bufnr),
+    title = " " .. title .. " ",
   })
 
   ---@type CodeCompanion.DiffUI
   local diff_ui = setmetatable({
+    bufnr = bufnr,
+    chat_bufnr = opts.chat_bufnr,
+    current_hunk = 1,
     diff = diff,
     diff_id = opts.diff_id or math.random(10000000),
     hunks = #diff.hunks,
-    current_hunk = 1,
-    bufnr = bufnr,
-    winnr = winnr,
-    chat_bufnr = opts.chat_bufnr,
     tool_name = opts.tool_name,
+    resolved = false,
+    winnr = winnr,
   }, DiffUI)
 
   -- Apply diff extmarks

lua/codecompanion/helpers.lua

@@ -39,8 +39,9 @@ function M.show_diff(args)
 
   local diff_ui = require("codecompanion.diff.ui")
   return diff_ui.show(diff_obj, {
-    diff_id = args.diff_id,
     chat_bufnr = args.chat_bufnr,
+    diff_id = args.diff_id,
+    title = args.title,
     tool_name = args.tool_name,
   })
 end

lua/codecompanion/interactions/chat/tools/builtin/insert_edit_into_file/init.lua

@@ -48,7 +48,6 @@ local approvals = require("codecompanion.interactions.chat.tools.approvals")
 local codecompanion = require("codecompanion")
 local config = require("codecompanion.config")
 local constants = require("codecompanion.interactions.chat.tools.builtin.insert_edit_into_file.constants")
-local diff = require("codecompanion.interactions.chat.helpers.diff")
 local helpers = require("codecompanion.interactions.chat.helpers")
 local match_selector = require("codecompanion.interactions.chat.tools.builtin.insert_edit_into_file.match_selector")
 local strategies = require("codecompanion.interactions.chat.tools.builtin.insert_edit_into_file.strategies")
@@ -151,24 +150,19 @@ local function restore_file_content(path, content, file_info)
 end
 
 ---Handle user decision for edits with diff approval
----@param opts table Options containing: diff_id, chat_bufnr, display_name, should_diff, success_response, output_handler, bufnr (optional), on_reject (optional for files)
+---@param opts table Options containing: diff_id, chat_bufnr, display_name, diff_ui, success_response, output_handler, bufnr (optional), on_reject (optional for files)
 ---@return any Result of output_handler call
 local function handle_user_decision(opts)
-  local accept = config.interactions.inline.keymaps.accept_change.modes.n
-  local reject = config.interactions.inline.keymaps.reject_change.modes.n
-  local is_buffer = opts.bufnr ~= nil
-
   local wait_opts = {
     chat_bufnr = opts.chat_bufnr,
-    notify = config.display.icons.warning
-      .. (is_buffer and " Waiting for diff approval ..." or " Waiting for decision ..."),
-    sub_text = fmt("`%s` - Accept edits / `%s` - Reject edits", accept, reject),
+    notify = config.display.icons.warning .. " Waiting for diff approval ...",
+    sub_text = "Review changes in the diff window",
   }
 
   return wait.for_decision(opts.diff_id, { "CodeCompanionDiffAccepted", "CodeCompanionDiffRejected" }, function(result)
     local response
     if result.accepted then
-      if is_buffer then
+      if opts.bufnr then
         pcall(function()
           api.nvim_buf_call(opts.bufnr, function()
             vim.cmd("silent! w")
@@ -180,10 +174,7 @@ local function handle_user_decision(opts)
       return opts.output_handler(response)
     else
       get_rejection_reason(function(reason)
-        if is_buffer then
-          if result.timeout and opts.should_diff and opts.should_diff.reject then
-            opts.should_diff:reject()
-          end
+        if opts.bufnr then
           response = error_response(
             result.timeout and "User failed to accept the edits in time"
               or fmt('User rejected the edits for `%s`, with the reason "%s"', opts.display_name, reason)
@@ -201,9 +192,6 @@ local function handle_user_decision(opts)
               )
             )
           else
-            if result.timeout and opts.should_diff and opts.should_diff.reject then
-              opts.should_diff:reject()
-            end
             response = error_response(
               result.timeout and "User failed to accept the edits in time"
                 or fmt('User rejected the edits for `%s`, with the reason "%s"', opts.display_name, reason)
@@ -793,20 +781,31 @@ local function edit_file(action, chat_bufnr, output_handler, opts)
   end
 
   local diff_id = math.random(10000000)
-  local should_diff = diff.create(path, diff_id, {
+  local from_lines = vim.split(current_content, "\n", { plain = true })
+  local to_lines = vim.split(dry_run_result.final_content, "\n", { plain = true })
+
+  -- Detect filetype from path
+  local ft = vim.filetype.match({ filename = path }) or "text"
+
+  local diff_helpers = require("codecompanion.helpers")
+  local diff_ui = diff_helpers.show_diff({
+    from_lines = from_lines,
+    to_lines = to_lines,
+    ft = ft,
+    title = action.filepath,
+    diff_id = diff_id,
     chat_bufnr = chat_bufnr,
-    original_content = vim.split(current_content, "\n", { plain = true }),
     tool_name = "insert_edit_into_file",
   })
 
   local final_success = success_response(fmt("Edited `%s` file%s", action.filepath, extract_explanation(action)))
 
-  if should_diff and opts.require_confirmation_after then
+  if opts.require_confirmation_after then
     return handle_user_decision({
       diff_id = diff_id,
       chat_bufnr = chat_bufnr,
       display_name = action.filepath,
-      should_diff = should_diff,
+      diff_ui = diff_ui,
       success_response = final_success,
       on_reject = function()
         return restore_file_content(path, current_content, file_info)
@@ -900,28 +899,26 @@ local function edit_buffer(bufnr, chat_bufnr, action, output_handler, opts)
     return output_handler(success)
   end
 
-  local should_diff = diff.create(bufnr, diff_id, {
+  local ft = vim.bo[bufnr].filetype or "text"
+
+  local diff_helpers = require("codecompanion.helpers")
+  local diff_ui = diff_helpers.show_diff({
     chat_bufnr = chat_bufnr,
-    original_content = original_content,
+    from_lines = original_content,
+    to_lines = final_lines,
+    ft = ft,
+    title = display_name,
+    diff_id = diff_id,
     tool_name = "insert_edit_into_file",
   })
 
-  local start_line = nil
-  if dry_run_result.edit_results[1] and dry_run_result.edit_results[1].start_line then
-    start_line = dry_run_result.edit_results[1].start_line
-  end
-
-  if start_line then
-    ui_utils.scroll_to_line(bufnr, start_line)
-  end
-
-  if should_diff and opts.require_confirmation_after then
+  if opts.require_confirmation_after then
     return handle_user_decision({
       diff_id = diff_id,
       chat_bufnr = chat_bufnr,
       display_name = display_name,
       bufnr = bufnr,
-      should_diff = should_diff,
+      diff_ui = diff_ui,
       success_response = success,
       output_handler = output_handler,
     })

lua/codecompanion/utils/diff_test.lua

@@ -86,27 +86,17 @@ def process_data(data, max_items=100):
     name = "Line additions",
     filetype = "lua",
     before = [[
-local M = {}
-
-function M.setup()
-  print("setup")
-end
-
-return M
+return {
+  "CodeCompanion is amazing - Oli Morris"
+}
 ]],
     after = [[
-local M = {}
-
-function M.setup()
-  print("setup")
-  vim.notify("initialized")
-end
-
-function M.teardown()
-  print("teardown")
-end
-
-return M
+return {
+  "CodeCompanion is amazing - Oli Morris"
+  "Lua and Neovim are amazing too - Oli Morris"
+  "Happy coding!"
+  "Hello world"
+}
 ]],
   },
   {