Hold input back until the selection has moved on

Two space presses in quick succession only staged one hunk. Input is
withheld until the refresh has landed, which is enough where the diff is
rebuilt on the spot, but the main view re-renders asynchronously: the
selection only moves to the next change once that render is on screen, so
the second press acted on lines that were no longer in the diff, and
staged nothing.

So the wait is now for the selection to be where the work carries on from,
rather than for the model. A restore therefore has to say when it is done —
which it can be either way, since a view given a message rather than a
re-render now gives up the restore it was holding instead of leaving it to
claim some later render.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

pkg/gui/controllers/helpers/diff_line_restore.go

@@ -261,6 +261,7 @@ func (self *DiffLineHelper) restoreDiffLinePositionOnRerender(
 			return bufferLine, true
 		},
 		func(viewLine int) { place(found, viewLine) },
+		nil,
 	)
 }
 
@@ -294,7 +295,13 @@ func (self *DiffLineHelper) ChangeLineOrdinal(view *gocui.View, viewLine int) (i
 // user, this is where what they were doing carries on. When the new diff has fewer
 // changes than that, because the ones acted on were its last, it lands on the last
 // change left.
-func (self *DiffLineHelper) RevealChangeLineAtOrdinal(view *gocui.View, ordinal int, place func(viewLine int)) {
+//
+// done is called once the selection is where it belongs, or once it turns out that no
+// render is coming to put it there, for a caller that must not let the user act again
+// in between.
+func (self *DiffLineHelper) RevealChangeLineAtOrdinal(
+	view *gocui.View, ordinal int, place func(viewLine int), done func(),
+) {
 	// How many change lines the incremental search has passed, so that it can carry on
 	// counting where it left off.
 	seen := 0
@@ -326,6 +333,7 @@ func (self *DiffLineHelper) RevealChangeLineAtOrdinal(view *gocui.View, ordinal
 			return last, last != -1
 		},
 		place,
+		done,
 	)
 }
 
@@ -345,12 +353,16 @@ func (self *DiffLineHelper) installDiffLineRestore(
 	findEarly func(rows []gocui.DiffLineContent, offset int) (int, bool),
 	findComplete func(contents []gocui.DiffLineContent) (int, bool),
 	place func(viewLine int),
+	done func(),
 ) {
 	// Get-or-create, because the pane may not have rendered anything yet: a file whose
 	// diff has only just become split has a second pane whose first render is the one
 	// this restore is for.
 	manager := self.c.GetOrCreateViewBufferManagerForView(view)
 	if manager == nil {
+		if done != nil {
+			done()
+		}
 		return
 	}
 
@@ -396,6 +408,7 @@ func (self *DiffLineHelper) installDiffLineRestore(
 				place(viewLine)
 			}
 		},
+		Done: done,
 	})
 }
 

pkg/gui/controllers/main_view_controller.go

@@ -344,11 +344,17 @@ func (self *MainViewController) primaryAction() error {
 // target pane inherits that select mode, this being the same piece of work continuing
 // in another pane — and shows no selection until the restore places one, so that what
 // it was left showing the last time it was used doesn't appear for a frame.
+//
+// done is called once the selection is where it belongs, or once it turns out that no
+// render is coming to put it there, for a caller that must not let the user act again
+// in between.
 func revealSelectionAfterAction(
-	c *ControllerCommon, source types.DiffPaneContext, target types.DiffPaneContext, firstLineIdx int,
+	c *ControllerCommon, source types.DiffPaneContext, target types.DiffPaneContext,
+	firstLineIdx int, done func(),
 ) {
 	ordinal, ok := c.Helpers().DiffLine.ChangeLineOrdinal(source.GetView(), firstLineIdx)
 	if !ok {
+		done()
 		return
 	}
 
@@ -373,7 +379,7 @@ func revealSelectionAfterAction(
 		}
 		targetView.CancelRangeSelect()
 		c.Helpers().DiffLine.ShowSelectionAtLine(targetView, viewLine, true)
-	})
+	}, done)
 }
 
 // discardSelection takes the selected diff lines back out of what they are part of,

pkg/gui/controllers/working_tree_diff_actions.go

@@ -282,7 +282,13 @@ func (self *WorkingTreeDiffActions) revealSelectionInPaneItLandsIn(
 		target = self.otherPane(pane)
 	}
 
-	revealSelectionAfterAction(self.c, pane, target, firstLineIdx)
+	// Hold input back until the selection is on the change the work carries on from. The
+	// refresh holds it until the model is up to date, but the diff is re-rendered after
+	// that, and until it has been the selection is still on lines that aren't there any
+	// more — so a key pressed meanwhile would act on nothing.
+	self.c.GocuiGui().BeginBlockingEvents()
+	revealSelectionAfterAction(self.c, pane, target, firstLineIdx,
+		self.c.GocuiGui().EndBlockingEvents)
 }
 
 // otherPane returns the main pane that isn't the given one.

pkg/gui/main_panels.go

@@ -289,6 +289,10 @@ func (gui *Gui) clampDiffSelectionToContent(view *gocui.View) {
 // An emptied pane is showing nothing, so it also goes back to the top and stops
 // claiming the render it was showing: whatever it is given next is content the user
 // hasn't seen there, and is shown from the top like any other.
+//
+// A position waiting to be put back goes too: this pane is getting no render for it
+// to ride, and whoever is waiting for the view to be back where it belongs has to
+// hear that it never will be.
 func (gui *Gui) clearMainView(mainContext types.Context) {
 	view := mainContext.GetView()
 	view.Clear()
@@ -300,6 +304,7 @@ func (gui *Gui) clearMainView(mainContext types.Context) {
 	gui.State.ContextMgr.UpdateSelectionHighlights()
 	if manager := gui.getViewBufferManagerForView(view); manager != nil {
 		manager.ForgetRenderedContent()
+		manager.DropRestoreForNextTask()
 	}
 }
 

pkg/gui/tasks_adapter.go

@@ -87,6 +87,9 @@ func (gui *Gui) newStringTask(view *gocui.View, str string) error {
 
 func (gui *Gui) newStringTaskWithoutScroll(view *gocui.View, str string) error {
 	manager := gui.getManager(view)
+	// Whatever the view was going to be put back to belonged to a re-render of its
+	// content; this is a message instead, so there is nothing to put back.
+	manager.DropRestoreForNextTask()
 
 	f := func(tasks.TaskOpts) error {
 		return gui.g.OnUIThreadAndWaitBackground(func() {
@@ -105,6 +108,9 @@ func (gui *Gui) newStringTaskWithoutScroll(view *gocui.View, str string) error {
 
 func (gui *Gui) newStringTaskWithScroll(view *gocui.View, str string, originX int, originY int) error {
 	manager := gui.getManager(view)
+	// Whatever the view was going to be put back to belonged to a re-render of its
+	// content; this is a message instead, so there is nothing to put back.
+	manager.DropRestoreForNextTask()
 
 	f := func(tasks.TaskOpts) error {
 		return gui.g.OnUIThreadAndWaitBackground(func() {
@@ -124,6 +130,9 @@ func (gui *Gui) newStringTaskWithScroll(view *gocui.View, str string, originX in
 
 func (gui *Gui) newStringTaskWithKey(view *gocui.View, str string, key string) error {
 	manager := gui.getManager(view)
+	// Whatever the view was going to be put back to belonged to a re-render of its
+	// content; this is a message instead, so there is nothing to put back.
+	manager.DropRestoreForNextTask()
 
 	f := func(tasks.TaskOpts) error {
 		return gui.g.OnUIThreadAndWaitBackground(func() {

pkg/integration/tests/main_view/stage_hunks_with_rapid_keypresses.go

@@ -0,0 +1,54 @@
+package main_view
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+// The second space is pressed before the diff the first one changed has re-rendered.
+// The selection only moves to the next hunk once that render arrives, so until then it
+// still covers lines that are no longer in the diff. The second press has to wait for
+// the render rather than act on those lines.
+var StageHunksWithRapidKeypresses = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stage two hunks from the focused main view with two space presses in rapid succession",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(cfg *config.AppConfig) {
+		cfg.GetUserConfig().Gui.UseHunkModeInStagingView = true
+	},
+	SetupRepo: func(shell *Shell) {
+		// Seven context lines between the two change blocks, so that git makes them two
+		// hunks.
+		shell.CreateFileAndAdd("file1", "1\n2\na\nb\nc\nd\ne\nf\ng\n3\n4\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "1b\n2b\na\nb\nc\nd\ne\nf\ng\n3b\n4b\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("file1").IsSelected(),
+			).
+			Press(keys.Universal.FocusMainView)
+
+		t.Views().Main().
+			IsFocused().
+			PressRapidly(keys.Universal.Select, keys.Universal.Select)
+
+		// Both presses were acted on, so the whole file is staged.
+		t.Views().Files().Lines(
+			Contains("M  file1"),
+		)
+		t.Views().Secondary().
+			Title(Equals("Staged changes")).
+			ContainsLines(
+				Contains("+1b"),
+				Contains("+2b"),
+			).
+			ContainsLines(
+				Contains("+3b"),
+				Contains("+4b"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -425,6 +425,7 @@ var tests = []*components.IntegrationTest{
 	main_view.StageDeletedFile,
 	main_view.StageDiffLines,
 	main_view.StageDiffLinesOfAPathWithASpace,
+	main_view.StageHunksWithRapidKeypresses,
 	main_view.StageRangeSpanningFiles,
 	main_view.StageUnderConformingDiffRenderer,
 	main_view.StageUnderUnsupportedDiffRenderer,

pkg/tasks/tasks.go

@@ -195,6 +195,24 @@ type RenderRestore struct {
 	// case the view keeps the position the paint gave it: the offset it had, or the
 	// top for content the view hasn't seen.
 	Apply func(swapIn func())
+
+	// Done is called once the restore has had its render — after Apply, or when it
+	// is given up because the view is being shown something other than a re-render
+	// of what it was remembered from. It is how a caller that has to wait for the
+	// view to be back where it belongs knows that it either is, or never will be.
+	// Optional, and called on the UI thread, as Apply is.
+	Done func()
+}
+
+// resolved reports that this restore's render has happened, or that there will not be
+// one. Called on the UI thread, from wherever the restore ends: once, whichever way it
+// ended.
+func (self *RenderRestore) resolved() {
+	if self.Done != nil {
+		done := self.Done
+		self.Done = nil
+		done()
+	}
 }
 
 // SetRestoreForNextTask arranges for the next command task to put the view back
@@ -254,6 +272,22 @@ func (self *ViewBufferManager) clearRestore(restore *RenderRestore) {
 	}
 }
 
+// DropRestoreForNextTask gives up a restore that has no render to ride, because the
+// view is being given something other than a re-render of the content it was
+// remembered from — a message where a diff was. Without this the restore would sit
+// there and claim some later render of that view, putting the user somewhere they
+// haven't been for a while.
+func (self *ViewBufferManager) DropRestoreForNextTask() {
+	self.taskIDMutex.Lock()
+	restore := self.restoreForNextTask
+	self.restoreForNextTask = nil
+	self.taskIDMutex.Unlock()
+
+	if restore != nil {
+		restore.resolved()
+	}
+}
+
 func (self *ViewBufferManager) GetTaskKey() string {
 	self.taskIDMutex.Lock()
 	defer self.taskIDMutex.Unlock()
@@ -564,6 +598,7 @@ func (self *ViewBufferManager) NewCmdTask(start func() (Cmd, io.Reader), prefix
 					// user was in the new content before it is revealed.
 					restore.Apply(self.swapInRender)
 					self.clearRestore(restore)
+					restore.resolved()
 					return
 				}
 				self.swapInRender()