fix: diff cleanup refocus terminal

CLAUDE.md

@@ -298,7 +298,7 @@ require("claudecode").setup({
 The `diff_opts` configuration allows you to customize diff behavior:
 
 - `layout` ("vertical"|"horizontal", default: `"vertical"`) - Whether the diff panes open in a vertical or horizontal split.
-- `keep_terminal_focus` (boolean, default: `false`) - When enabled, keeps focus in the Claude Code terminal when a diff opens instead of moving focus to the diff buffer. This allows you to continue using terminal keybindings like `<CR>` for accepting/rejecting diffs without accidentally triggering other mappings.
+- `keep_terminal_focus` (boolean, default: `false`) - When enabled, keeps focus in the Claude Code terminal when a diff opens instead of moving focus to the diff buffer. This allows you to continue using terminal keybindings like `<CR>` for accepting/rejecting diffs without accidentally triggering other mappings. Also applies when a same-tab diff is cleaned up (accepted via `:w` or rejected): focus returns to the Claude terminal instead of wherever Neovim's default window selection would otherwise land.
 - `open_in_new_tab` (boolean, default: `false`) - Open diffs in a new tab instead of the current tab.
 - `hide_terminal_in_new_tab` (boolean, default: `false`) - When opening diffs in a new tab, do not show the Claude terminal split in that new tab. The terminal remains in the original tab, giving maximum screen estate for reviewing the diff.
 - `on_new_file_reject` ("keep_empty"|"close_window", default: `"keep_empty"`) - Behavior when rejecting a diff for a new file (where the old file did not exist).

README.md

@@ -382,7 +382,7 @@ For deep technical details, see [ARCHITECTURE.md](./ARCHITECTURE.md).
       --   Neovim >= 0.9.0. Highlight groups are customizable: ClaudeCodeInlineDiffAdd,
       --   ClaudeCodeInlineDiffDelete, ClaudeCodeInlineDiffAddSign, ClaudeCodeInlineDiffDeleteSign.
       open_in_new_tab = false,
-      keep_terminal_focus = false, -- If true, moves focus back to terminal after diff opens
+      keep_terminal_focus = false, -- If true, moves focus back to terminal after diff opens, and after same-tab accept/reject
       hide_terminal_in_new_tab = false,
       auto_resize_terminal = true, -- Let the plugin manage the terminal width across the diff lifecycle; set false to own it via the User autocmds below
       -- on_new_file_reject = "keep_empty", -- "keep_empty" or "close_window"

lua/claudecode/diff.lua

@@ -1245,7 +1245,19 @@ function M._cleanup_diff_state(tab_name, reason)
 
     -- After closing the diff in the same tab, restore the idle terminal width
     -- (unless the user opted out via auto_resize_terminal). Floating terminals are skipped.
-    resize_terminal_for_diff(find_claudecode_terminal_window(), "idle")
+    local terminal_win = find_claudecode_terminal_window()
+    resize_terminal_for_diff(terminal_win, "idle")
+
+    -- Closing the diff windows above leaves Neovim to pick whatever window it
+    -- considers "next" (often whatever else happens to be adjacent in the
+    -- split tree), not necessarily the Claude terminal. When keep_terminal_focus
+    -- is set, honor it here too (it otherwise only applies at diff-open time)
+    -- and explicitly refocus the terminal.
+    if config and config.diff_opts and config.diff_opts.keep_terminal_focus then
+      if terminal_win and vim.api.nvim_win_is_valid(terminal_win) then
+        pcall(vim.api.nvim_set_current_win, terminal_win)
+      end
+    end
   end
 
   -- ALWAYS clean up buffers regardless of tab mode (fixes buffer leak)

tests/unit/diff_cleanup_terminal_focus_spec.lua

@@ -0,0 +1,171 @@
+require("tests.busted_setup")
+
+-- Regression test: accepting a same-tab diff via `:w` (BufWriteCmd ->
+-- _resolve_diff_as_saved, then Claude's close_tab tool -> _cleanup_diff_state)
+-- should return focus to the Claude terminal window afterward when
+-- diff_opts.keep_terminal_focus is enabled, instead of leaving Neovim's
+-- default post-window-close selection to land anywhere.
+
+describe("Diff cleanup terminal focus (same-tab)", function()
+  local diff
+  local test_old_file = "/tmp/test_cleanup_terminal_focus_old.txt"
+  local tab_name = "test_cleanup_terminal_focus_tab"
+
+  local editor_win = 1000
+  local terminal_win = 1001
+  local terminal_buf
+
+  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 = editor_win
+    vim._next_winid = 1002
+
+    vim._mock.add_buffer(1, "/home/user/project/test.lua", "local test = {}\nreturn test", { modified = false })
+    vim._mock.add_window(editor_win, 1, { 1, 0 })
+    vim._win_tab[editor_win] = 1
+    vim._tab_windows[1] = { editor_win }
+  end
+
+  local function add_terminal_window()
+    terminal_buf = vim.api.nvim_create_buf(false, true)
+    vim.api.nvim_buf_set_option(terminal_buf, "buftype", "terminal")
+    vim._mock.add_window(terminal_win, terminal_buf, { 1, 0 })
+    vim._win_tab[terminal_win] = 1
+    table.insert(vim._tab_windows[1], terminal_win)
+  end
+
+  before_each(function()
+    reset_vim_state()
+    add_terminal_window()
+
+    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.terminal"] = {
+      get_active_terminal_bufnr = function()
+        return terminal_buf
+      end,
+      ensure_visible = function() end,
+    }
+
+    package.loaded["claudecode.diff"] = nil
+    diff = require("claudecode.diff")
+  end)
+
+  after_each(function()
+    os.remove(test_old_file)
+    package.loaded["claudecode.terminal"] = nil
+    if diff and diff._cleanup_all_active_diffs then
+      diff._cleanup_all_active_diffs("test teardown")
+    end
+    package.loaded["claudecode.diff"] = nil
+  end)
+
+  it("refocuses the claude terminal window when keep_terminal_focus is enabled", function()
+    diff.setup({
+      diff_opts = {
+        layout = "vertical",
+        open_in_new_tab = false,
+        keep_terminal_focus = true,
+      },
+      terminal = {},
+    })
+
+    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)
+
+    -- Simulate the user's cursor being in the diff buffer when they run :w.
+    vim.api.nvim_set_current_win(state.new_window)
+
+    -- BufWriteCmd -> _resolve_diff_as_saved, then Claude's close_tab tool call -> _cleanup_diff_state.
+    diff._resolve_diff_as_saved(tab_name, state.new_buffer)
+    diff._cleanup_diff_state(tab_name, "accepted via :w")
+
+    assert.is_false(vim.api.nvim_win_is_valid(state.new_window))
+    assert.equal(terminal_win, vim.api.nvim_get_current_win())
+  end)
+
+  it("leaves focus untouched when keep_terminal_focus is disabled (default)", function()
+    diff.setup({
+      diff_opts = {
+        layout = "vertical",
+        open_in_new_tab = false,
+        keep_terminal_focus = false,
+      },
+      terminal = {},
+    })
+
+    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]
+    vim.api.nvim_set_current_win(state.new_window)
+
+    local win_before_cleanup = vim.api.nvim_get_current_win()
+    diff._resolve_diff_as_saved(tab_name, state.new_buffer)
+    diff._cleanup_diff_state(tab_name, "accepted via :w")
+
+    -- No explicit refocus should have happened; the current window is whatever
+    -- it was left as (this plugin's existing, unchanged default behavior).
+    assert.equal(win_before_cleanup, vim.api.nvim_get_current_win())
+  end)
+
+  it("does not error when no claude terminal window is visible", function()
+    diff.setup({
+      diff_opts = {
+        layout = "vertical",
+        open_in_new_tab = false,
+        keep_terminal_focus = true,
+      },
+      terminal = {},
+    })
+
+    package.loaded["claudecode.terminal"].get_active_terminal_bufnr = function()
+      return nil
+    end
+
+    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]
+    vim.api.nvim_set_current_win(state.new_window)
+    diff._resolve_diff_as_saved(tab_name, state.new_buffer)
+
+    assert.has_no.errors(function()
+      diff._cleanup_diff_state(tab_name, "accepted via :w")
+    end)
+  end)
+end)