Follow the acted-on side to the right pane after staging from the main view

After staging or unstaging from the focused main view, focus stayed on the
pane the user acted in even when the side they were acting on moved to the
other pane. Two cases got it wrong:

- Unstaging the first hunk of an only-staged file splits the diff: the main
  half flips to show the just-unstaged change, and the staged remainder jumps
  to the secondary half — but focus stayed on the main half, away from the
  staged content the user was working through.
- Unstaging the last staged hunk from the secondary half empties the staged
  side and collapses the split, hiding the secondary half — leaving focus
  stranded on a hidden pane.

The rule is the same in both: focus the staged side while it survives. The
stage handler now reports which focused-main pane should hold focus —
the secondary half when unstaging leaves the file split, the main half
otherwise — and the controller re-selects the revealed change in that pane
and focuses it. The split is read from the model, which Refresh has already
updated synchronously by the time the handler returns; the re-render it
queues is what the reveal rides.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

pkg/gui/context/base_context.go

@@ -38,7 +38,7 @@ type (
 	onFocusFn                = func(types.OnFocusOpts)
 	onFocusLostFn            = func(types.OnFocusLostOpts)
 	onClickFocusedMainViewFn = func(mainViewName string, clickedLineIdx int) error
-	onStageFocusedMainViewFn = func(mainViewName string, firstLineIdx int, lastLineIdx int) error
+	onStageFocusedMainViewFn = func(mainViewName string, firstLineIdx int, lastLineIdx int) (focusViewName string, err error)
 )
 
 var _ types.IBaseContext = &BaseContext{}

pkg/gui/controllers/base_controller.go

@@ -23,7 +23,7 @@ func (self *baseController) GetOnClickFocusedMainView() func(mainViewName string
 	return nil
 }
 
-func (self *baseController) GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) error {
+func (self *baseController) GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) (focusViewName string, err error) {
 	return nil
 }
 

pkg/gui/controllers/files_controller.go

@@ -457,21 +457,21 @@ func (self *FilesController) diffSplitState(node *filetree.FileNode) (split bool
 	return split, mainShowsStaged
 }
 
-func (self *FilesController) GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) error {
-	return func(mainViewName string, firstLineIdx int, lastLineIdx int) error {
+func (self *FilesController) GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) (string, error) {
+	return func(mainViewName string, firstLineIdx int, lastLineIdx int) (string, error) {
 		if self.c.UserConfig().Git.DiffContextSize == 0 {
-			return fmt.Errorf(self.c.Tr.Actions.NotEnoughContextToStage,
+			return "", fmt.Errorf(self.c.Tr.Actions.NotEnoughContextToStage,
 				self.c.UserConfig().Keybinding.Universal.IncreaseContextInDiffView)
 		}
 
 		node := self.context().GetSelected()
 		if node == nil {
-			return nil
+			return "", nil
 		}
 
 		infos := self.c.Helpers().Staging.ChangeLinesInViewRange(mainViewName, firstLineIdx, lastLineIdx)
 		if len(infos) == 0 {
-			return nil
+			return "", nil
 		}
 
 		// The whole diff shown in the main view is on one side — the staged diff in
@@ -492,12 +492,28 @@ func (self *FilesController) GetOnStageFocusedMainView() func(mainViewName strin
 				continue
 			}
 			if err := self.stageDiffLines(file, fileInfos, reverse); err != nil {
-				return err
+				return "", err
 			}
 		}
 
 		self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES, types.STAGING}})
-		return nil
+
+		// Focus follows the side that was acted on. Staging keeps it in the main half
+		// (which always holds the unstaged side, or the staged side once the file has
+		// only staged changes). Unstaging keeps it on the staged side, which lives in
+		// the secondary half once the file is split into staged + unstaged, and moves
+		// back to the main half when the staged side empties and the split collapses.
+		// The model is up to date now (Refresh above is synchronous), so the post-op
+		// split is read from the freshly selected node.
+		focusViewName := self.c.Contexts().Normal.GetViewName()
+		if reverse {
+			if node := self.context().GetSelected(); node != nil {
+				if split, _ := self.diffSplitState(node); split {
+					focusViewName = self.c.Contexts().NormalSecondary.GetViewName()
+				}
+			}
+		}
+		return focusViewName, nil
 	}
 }
 

pkg/gui/controllers/main_view_controller.go

@@ -293,19 +293,42 @@ func (self *MainViewController) stageSelectedLine() error {
 		sel.RangeIsSticky = false
 	}
 
-	// The diff re-renders after staging; install a restore (before the handler
-	// triggers that re-render) so the selection lands on the next change rather than
-	// at a now-meaningless position.
-	self.c.Helpers().Staging.RevealSelectionAfterStaging(v, v, first, last, func(viewLine int) {
-		if self.sel().Mode == context.DiffSelectModeHunk {
-			self.selectHunkAround(viewLine)
+	// Staging updates the model synchronously and queues the main-view re-render, and
+	// reports which pane should hold focus afterwards (staging/unstaging can move the
+	// acted-on side to the other pane). "" means nothing was staged.
+	focusViewName, err := handler(self.context.GetViewName(), first, last)
+	if err != nil {
+		return err
+	}
+	if focusViewName == "" {
+		return nil
+	}
+
+	// Re-select in whichever pane now holds the acted-on side, and focus it. The
+	// candidate change lines are read from the pane we acted in (its content is still
+	// the pre-staging diff until the queued re-render), and found again in the target
+	// pane's re-render. The target inherits our select mode (line/hunk).
+	targetContext := self.context
+	if focusViewName == self.otherContext.GetViewName() {
+		targetContext = self.otherContext
+	}
+	*targetContext.DiffSelectState() = *sel
+	targetView := targetContext.GetView()
+
+	self.c.Helpers().Staging.RevealSelectionAfterStaging(v, targetView, first, last, func(viewLine int) {
+		if sel.Mode == context.DiffSelectModeHunk {
+			selectDiffHunk(self.c, targetContext, viewLine)
 		} else {
-			v.CancelRangeSelect()
-			showSelectionAtLine(v, viewLine, true)
+			targetView.CancelRangeSelect()
+			showSelectionAtLine(targetView, viewLine, true)
 		}
 	})
 
-	return handler(self.context.GetViewName(), first, last)
+	if targetContext != self.context {
+		self.c.Context().Push(targetContext, types.OnFocusOpts{})
+	}
+
+	return nil
 }
 
 func (self *MainViewController) enter() error {

pkg/gui/types/context.go

@@ -102,7 +102,7 @@ type IBaseContext interface {
 	// (space), delegated to the side panel that owns the diff being shown. The
 	// inclusive view-line range is the current selection (a single line, a range, or
 	// a hunk).
-	AddOnStageFocusedMainViewFn(func(mainViewName string, firstLineIdx int, lastLineIdx int) error)
+	AddOnStageFocusedMainViewFn(func(mainViewName string, firstLineIdx int, lastLineIdx int) (focusViewName string, err error))
 	// Adding on to the above, this is so that a list-specific handler can register
 	// a hook for doing additional click handling
 	AddOnClickFn(func(opts gocui.ViewMouseBindingOpts) error)
@@ -337,8 +337,11 @@ type HasKeybindings interface {
 	// Implement this in a side-panel controller to stage/unstage (or, later, add to
 	// the custom patch) the selected diff line(s) when the user presses space in the
 	// focused main view. The inclusive view-line range is the current selection (a
-	// single line, a range, or a hunk). Return nil to do nothing.
-	GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) error
+	// single line, a range, or a hunk). It returns the name of the focused main view
+	// that should hold focus afterwards — staging/unstaging can move the acted-on
+	// side to the other pane — or "" when nothing was done. Return a nil func to do
+	// nothing.
+	GetOnStageFocusedMainView() func(mainViewName string, firstLineIdx int, lastLineIdx int) (focusViewName string, err error)
 }
 
 type IController interface {

pkg/integration/tests/staging/focus_follows_staged_side_to_secondary_after_unstaging.go

@@ -0,0 +1,54 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var FocusFollowsStagedSideToSecondaryAfterUnstaging = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Unstaging the first hunk of an only-staged file from the focused main view splits the diff; focus follows the staged remainder into the secondary half",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = true
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\neleven\n")
+		shell.Commit("one")
+
+		// Two staged hunks and no unstaged changes, so the main view shows the staged
+		// diff in the main half without a split.
+		shell.UpdateFileAndAdd("file1", "one\ntwo\nTHREE\nfour\nfive\nsix\nseven\neight\nNINE\nten\neleven\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		// The only-staged file shows the staged diff in the main half, first hunk selected.
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			).
+			// Unstage the first staged hunk, which splits the file into staged + unstaged.
+			PressPrimaryAction()
+
+		// Focus follows the staged remainder into the secondary half, landing on the
+		// next staged hunk rather than staying on the now-unstaged half.
+		t.Views().Secondary().
+			IsFocused().
+			SelectedLines(
+				Contains("-nine"),
+				Contains("+NINE"),
+			)
+
+		// The main half now shows the hunk we just unstaged.
+		t.Views().Main().
+			ContainsLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			)
+	},
+})

pkg/integration/tests/staging/focus_returns_to_main_after_unstaging_last_staged_hunk.go

@@ -0,0 +1,56 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var FocusReturnsToMainAfterUnstagingLastStagedHunk = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Unstaging the last staged hunk from the secondary half collapses the split; focus returns to the main half on the now-unstaged change",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = true
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\neleven\n")
+		shell.Commit("one")
+
+		// One staged hunk...
+		shell.UpdateFileAndAdd("file1", "one\ntwo\nTHREE\nfour\nfive\nsix\nseven\neight\nnine\nten\neleven\n")
+		// ...plus an unstaged change, so the main view splits into staged/unstaged.
+		shell.UpdateFile("file1", "one\ntwo\nTHREE\nfour\nfive\nSIX\nseven\neight\nnine\nten\neleven\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		// The unstaged half is focused first; switch to the staged half.
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-six"),
+				Contains("+SIX"),
+			).
+			Press(keys.Universal.TogglePanel)
+
+		t.Views().Secondary().
+			IsFocused().
+			SelectedLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			).
+			// Unstage the only staged hunk, emptying the staged side and collapsing the split.
+			PressPrimaryAction()
+
+		// Focus returns to the main half, which now shows the unstaged diff, landing
+		// on the hunk we just unstaged.
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -409,6 +409,8 @@ var tests = []*components.IntegrationTest{
 	staging.DiffChangeScreenMode,
 	staging.DiffContextChange,
 	staging.DiscardAllChanges,
+	staging.FocusFollowsStagedSideToSecondaryAfterUnstaging,
+	staging.FocusReturnsToMainAfterUnstagingLastStagedHunk,
 	staging.Search,
 	staging.SelectHunkOnFocusingMainView,
 	staging.SelectNextHunkAfterStagingFromMainView,