fix(diff): user-initiated accept/reject should clean up windows and buffers

The user commands :ClaudeCodeDiffAccept and :ClaudeCodeDiffDeny call
_resolve_diff_as_saved/_rejected, which only mark the diff status. Real
cleanup (close windows, force-delete buffers, remove from active_diffs)
is performed by _cleanup_diff_state, currently triggered exclusively by
the close_tab MCP tool.

But close_tab has schema = nil (it is intentionally internal-only, as
documented in the source), so the Claude CLI never invokes it. The result
is that user-initiated accept/reject leaves the diff windows and buffers
behind, requiring manual cleanup.

The fix: have the user commands call _cleanup_diff_state themselves,
since they are not waiting for an MCP response from the CLI. Also
restore terminal focus when diff_opts.keep_terminal_focus is set
(previously this only applied at diff-open time).

Adds unit tests covering both commands and the no-marker no-op case.

Closes #205, partially addresses #238.

Author: mightreya

lua/claudecode/diff.lua

@@ -1442,8 +1442,22 @@ function M.reload_file_buffers_manual(file_path, original_cursor_pos)
   return reload_file_buffers(file_path, original_cursor_pos)
 end
 
+-- Restore focus to the Claude terminal after a user-initiated resolve, when configured.
+local function restore_terminal_focus_after_user_resolve()
+  if not (config and config.diff_opts and config.diff_opts.keep_terminal_focus) then
+    return
+  end
+  local terminal_win = find_claudecode_terminal_window()
+  if terminal_win and vim.api.nvim_win_is_valid(terminal_win) then
+    pcall(vim.api.nvim_set_current_win, terminal_win)
+    pcall(vim.cmd, "startinsert")
+  end
+end
+
 ---Accept the current diff (user command version)
----This function reads the diff context from buffer variables
+---This function reads the diff context from buffer variables.
+---Unlike Claude-initiated resolve (which waits for the close_tab MCP tool call),
+---user commands clean up the diff UI immediately since close_tab is internal-only.
 function M.accept_current_diff()
   local current_buffer = vim.api.nvim_get_current_buf()
   local tab_name = vim.b[current_buffer].claudecode_diff_tab_name
@@ -1454,10 +1468,14 @@ function M.accept_current_diff()
   end
 
   M._resolve_diff_as_saved(tab_name, current_buffer)
+  M._cleanup_diff_state(tab_name, "accepted via :ClaudeCodeDiffAccept")
+  restore_terminal_focus_after_user_resolve()
 end
 
 ---Deny/reject the current diff (user command version)
----This function reads the diff context from buffer variables
+---This function reads the diff context from buffer variables.
+---Unlike Claude-initiated resolve (which waits for the close_tab MCP tool call),
+---user commands clean up the diff UI immediately since close_tab is internal-only.
 function M.deny_current_diff()
   local current_buffer = vim.api.nvim_get_current_buf()
   local tab_name = vim.b[current_buffer].claudecode_diff_tab_name
@@ -1467,8 +1485,9 @@ function M.deny_current_diff()
     return
   end
 
-  -- Do not close windows/tabs here; just mark as rejected.
   M._resolve_diff_as_rejected(tab_name)
+  M._cleanup_diff_state(tab_name, "rejected via :ClaudeCodeDiffDeny")
+  restore_terminal_focus_after_user_resolve()
 end
 
 return M

tests/unit/diff_user_command_cleanup_spec.lua

@@ -0,0 +1,118 @@
+require("tests.busted_setup")
+
+local function reset_vim_state()
+  assert(vim and vim._mock and vim._mock.reset, "Expected vim mock with _mock.reset()")
+
+  vim._mock.reset()
+
+  vim._tabs = { [1] = true }
+  vim._current_tabpage = 1
+  vim._current_window = 1000
+  vim._next_winid = 1001
+
+  vim._mock.add_buffer(1, "/home/user/project/test.lua", "local test = {}\nreturn test", { modified = false })
+  vim._mock.add_window(1000, 1, { 1, 0 })
+  vim._win_tab[1000] = 1
+  vim._tab_windows[1] = { 1000 }
+end
+
+describe("Diff user command cleanup (accept/deny)", function()
+  local diff
+  local test_old_file = "/tmp/test_user_command_cleanup_old.txt"
+  local tab_name = "test_user_command_cleanup_tab"
+
+  before_each(function()
+    reset_vim_state()
+
+    local f = assert(io.open(test_old_file, "w"))
+    f:write("line1\nline2\n")
+    f:close()
+
+    package.loaded["claudecode.logger"] = {
+      debug = function() end,
+      error = function() end,
+      info = function() end,
+      warn = function() end,
+    }
+
+    package.loaded["claudecode.diff"] = nil
+    diff = require("claudecode.diff")
+
+    diff.setup({
+      diff_opts = {
+        layout = "vertical",
+        open_in_new_tab = false,
+        keep_terminal_focus = false,
+      },
+      terminal = {},
+    })
+  end)
+
+  after_each(function()
+    os.remove(test_old_file)
+    if diff and diff._cleanup_all_active_diffs then
+      diff._cleanup_all_active_diffs("test teardown")
+    end
+    package.loaded["claudecode.diff"] = nil
+  end)
+
+  it("accept_current_diff cleans up windows and removes diff state", function()
+    local params = {
+      old_file_path = test_old_file,
+      new_file_path = test_old_file,
+      new_file_contents = "new1\nnew2\n",
+      tab_name = tab_name,
+    }
+
+    diff._setup_blocking_diff(params, function() end)
+
+    local state = diff._get_active_diffs()[tab_name]
+    assert.is_table(state)
+
+    local new_win = state.new_window
+    local new_buffer = state.new_buffer
+
+    -- Make the proposed buffer current so accept_current_diff can read the
+    -- buffer-local tab_name marker.
+    vim.api.nvim_set_current_buf(new_buffer)
+
+    diff.accept_current_diff()
+
+    -- After accept, windows must be closed and state must be cleared.
+    assert.is_false(vim.api.nvim_win_is_valid(new_win))
+    assert.is_nil(diff._get_active_diffs()[tab_name])
+  end)
+
+  it("deny_current_diff cleans up windows and removes diff state", function()
+    local params = {
+      old_file_path = test_old_file,
+      new_file_path = test_old_file,
+      new_file_contents = "new1\nnew2\n",
+      tab_name = tab_name,
+    }
+
+    diff._setup_blocking_diff(params, function() end)
+
+    local state = diff._get_active_diffs()[tab_name]
+    assert.is_table(state)
+
+    local new_win = state.new_window
+    local new_buffer = state.new_buffer
+
+    vim.api.nvim_set_current_buf(new_buffer)
+
+    diff.deny_current_diff()
+
+    assert.is_false(vim.api.nvim_win_is_valid(new_win))
+    assert.is_nil(diff._get_active_diffs()[tab_name])
+  end)
+
+  it("accept_current_diff is a no-op when buffer has no diff marker", function()
+    -- No diff registered; current buffer is a plain non-diff buffer.
+    diff.accept_current_diff()
+    -- Should not raise; no diffs should be present.
+    local diffs = diff._get_active_diffs()
+    assert.is_table(diffs)
+    assert.is_nil(next(diffs))
+  end)
+end)