Restore the focused main view by patch identity on escape

Escaping a patch explorer (staging / patch building) back to the focused main
view it was entered from used to replay a numeric scroll position and selection
index captured on the way in. But the reason to escape after staging or dropping
a hunk is that the content changed, so a saved index points at the wrong line —
and the host auto-advances the explorer's selection to a still-valid line anyway,
which is the line the user actually cares about returning to.

Restore by *patch identity* instead. On escape, read the (file, type, source
line) the explorer currently has selected, then have the main view's re-render
land on the row that matches it: scan the incoming content as it loads (the
inverse of the diff-line primitive), and once the matching row plus a screenful
below it have loaded, swap the off-screen render in and scroll to / select that
row in one step. FocusPoint with scrollIntoView centres the row only if it's
off-screen, so the common unchanged-content escape — where the row is already
where it was — doesn't move at all. If the line is gone (the content really
changed), nothing is forced.

This generalizes the scroll restore from a fixed origin to a predicate
(RenderRestore: FirstPaintReady decides when the saved position is reachable,
Apply re-establishes it), folding the separate selection restore into the same
first paint — so it no longer rides a post-load callback that could fire early.

The restore also now survives task replacement, which the numeric version did
not: a periodic refresh can stop the escape's re-render before it first-paints.
The pending restore is held on the buffer manager and is *not* cleared when a
task starts, so the replacement task picks it up. It is not gated on the command
key — staging the last unstaged hunk re-renders `git diff` as `git diff --cached`,
a different command, yet the line to land on is still in the new content — but
validates itself: the scan finds the target line only when the content still
contains it, so applying it to a different item is a harmless no-op. A task
clears it once it has applied it (found or not), so it lives for exactly one
re-render. Because the restore is anchored on content identity and is idempotent,
"survive replacement" and "restore by identity" are one mechanism, not two.

With the identity in hand the snapshot no longer needs the captured scroll/index;
they're derived from the explorer's live selection.

pkg/gocui/view.go

@@ -1873,6 +1873,22 @@ func (v *View) OffscreenDiffLineContents() []DiffLineContent {
 	return diffLineContents(v.offscreen)
 }
 
+// OffscreenLineCount returns the number of unwrapped buffer lines read so far
+// into an in-progress off-screen re-render (see BeginOffscreenRender), or 0 if
+// none is underway. The escape restore uses it to tell, cheaply, once it has
+// found its target line, when a screenful below it has loaded too — so the swap
+// shows the target with context rather than at the very bottom of a part-filled
+// view.
+func (v *View) OffscreenLineCount() int {
+	v.writeMutex.Lock()
+	defer v.writeMutex.Unlock()
+
+	if v.offscreen == nil {
+		return 0
+	}
+	return len(v.offscreen.lines)
+}
+
 func diffLineContents(buf *viewBuffer) []DiffLineContent {
 	contents := make([]DiffLineContent, len(buf.lines))
 	for i, line := range buf.lines {

pkg/gui/controllers.go

@@ -53,8 +53,8 @@ func (gui *Gui) resetHelpersAndControllers() {
 	gpgHelper := helpers.NewGpgHelper(helperCommon)
 	viewHelper := helpers.NewViewHelper(helperCommon, gui.State.Contexts)
 	windowHelper := helpers.NewWindowHelper(helperCommon, viewHelper)
-	patchBuildingHelper := helpers.NewPatchBuildingHelper(helperCommon)
 	stagingHelper := helpers.NewStagingHelper(helperCommon, windowHelper)
+	patchBuildingHelper := helpers.NewPatchBuildingHelper(helperCommon, stagingHelper)
 	mergeConflictsHelper := helpers.NewMergeConflictsHelper(helperCommon)
 	searchHelper := helpers.NewSearchHelper(helperCommon)
 

pkg/gui/controllers/commits_files_controller.go

@@ -549,7 +549,7 @@ func (self *CommitFilesController) expandAll() error {
 func (self *CommitFilesController) GetOnClickFocusedMainView() func(mainViewName string, clickedLineIdx int) error {
 	return func(mainViewName string, clickedLineIdx int) error {
 		// Capture before any mutation below that might re-render the main view.
-		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context(), clickedLineIdx)
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context())
 
 		info, ok := self.c.Helpers().Staging.GetDiffLineInfo(mainViewName, clickedLineIdx)
 		line := -1

pkg/gui/controllers/files_controller.go

@@ -410,7 +410,7 @@ func (self *FilesController) GetOnDoubleClick() func() error {
 func (self *FilesController) GetOnClickFocusedMainView() func(mainViewName string, clickedLineIdx int) error {
 	return func(mainViewName string, clickedLineIdx int) error {
 		// Capture before any mutation below that might re-render the main view.
-		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context(), clickedLineIdx)
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context())
 
 		info, ok := self.c.Helpers().Staging.GetDiffLineInfo(mainViewName, clickedLineIdx)
 		line := -1

pkg/gui/controllers/helpers/patch_building_helper.go

@@ -9,14 +9,17 @@ import (
 )
 
 type PatchBuildingHelper struct {
-	c *HelperCommon
+	c             *HelperCommon
+	stagingHelper *StagingHelper
 }
 
 func NewPatchBuildingHelper(
 	c *HelperCommon,
+	stagingHelper *StagingHelper,
 ) *PatchBuildingHelper {
 	return &PatchBuildingHelper{
-		c: c,
+		c:             c,
+		stagingHelper: stagingHelper,
 	}
 }
 
@@ -35,15 +38,15 @@ func (self *PatchBuildingHelper) ShowHunkStagingHint() {
 // takes us from the patch building panel back to the commit files panel, or to
 // the focused main view if that's where we entered it from
 func (self *PatchBuildingHelper) Escape() {
-	EscapeFromPatchExplorer(self.c, self.c.Contexts().CustomPatchBuilder)
+	EscapeFromPatchExplorer(self.c, self.stagingHelper, self.c.Contexts().CustomPatchBuilder)
 }
 
 // EscapeFromPatchExplorer returns from a patch explorer context (staging or
 // patch building). If we entered it from a focused main view, we go back to
 // where we came from (re-rendering the side panel's content into the main view,
-// like the plain escape does), then focus the main view and restore its scroll
-// position and selection. Otherwise we just pop to the side panel.
-func EscapeFromPatchExplorer(c *HelperCommon, context types.IPatchExplorerContext) {
+// like the plain escape does), then focus the main view and land on the line the
+// explorer currently has selected. Otherwise we just pop to the side panel.
+func EscapeFromPatchExplorer(c *HelperCommon, stagingHelper *StagingHelper, context types.IPatchExplorerContext) {
 	snapshot := context.GetFocusedMainViewSnapshot()
 	if snapshot == nil {
 		c.Context().Pop()
@@ -59,47 +62,26 @@ func EscapeFromPatchExplorer(c *HelperCommon, context types.IPatchExplorerContex
 		listContext.GetList().SetSelectedLineIdx(snapshot.SidePanelSelectedLineIdx)
 	}
 
-	view := snapshot.MainView.GetView()
-
-	restore := func() {
-		view.FocusPoint(0, snapshot.SelectedLineIdx, false)
-		view.Highlight = true
-		view.HighlightInactive = false
-	}
-
-	// Ask the upcoming re-render to restore the scroll position and selection.
-	// Pushing the side panel re-renders its content into the main view via a
-	// cmd/pty task. Until that content is ready, the main view keeps showing the
-	// placeholder that CopyContent left in it (the view we're leaving) at its
-	// current scroll; the task then scrolls to the saved position as part of the
-	// first paint that shows the real content (rather than setting the origin up
-	// front, which would jump to the top or show a misplaced placeholder frame).
-	// The selection needs the content loaded down to the selected line, so it
-	// rides the same task and fires at the end of its initial read. Threading it
-	// through the task (rather than a ReadToEnd issued after the pushes) avoids a
-	// race: ReadToEnd fires synchronously when the freshly-created task's read
-	// channel isn't live yet, which would run FocusPoint before the content is
-	// loaded and silently drop the selection.
-	manager := c.GetViewBufferManagerForView(view)
-	if manager != nil {
-		manager.ScrollToOriginYForNextTask(snapshot.OriginY)
-		manager.ThenForNextTask(func() {
-			c.OnUIThread(func() error {
-				restore()
-				return nil
-			})
-		})
+	// Ask the upcoming re-render of the main view to land on the line the explorer
+	// currently has selected. Read that identity now, before the pushes: pushing
+	// the side panel re-renders its content into the main view, and the restore
+	// rides that re-render — finding the matching row as it loads and scrolling to
+	// and selecting it as the content first appears. See RestoreFocusedMainViewOnEscape.
+	//
+	// Anchor on the *first* line of the explorer's selection: in hunk or range mode
+	// the selection spans several lines and its cursor sits at the last one, but
+	// returning to the start of the hunk is what reads as "the same place".
+	selectedViewLine := context.GetView().SelectedLineIdx()
+	if state := context.GetState(); state != nil {
+		selectedViewLine, _ = state.SelectedViewRange()
 	}
+	stagingHelper.RestoreFocusedMainViewOnEscape(
+		context.GetView(), snapshot.MainView.GetView(), selectedViewLine)
 
 	// Land on the side panel first (this re-renders the original content into the
 	// main view), then focus the main view on top of it.
 	c.Context().Push(snapshot.SidePanel, types.OnFocusOpts{})
 	c.Context().Push(snapshot.MainView, types.OnFocusOpts{})
-
-	// Without a buffer manager there is no re-render task to ride, so restore now.
-	if manager == nil {
-		restore()
-	}
 }
 
 // kills the custom patch and returns us back to the commit files panel if needed

pkg/gui/controllers/helpers/staging_helper.go

@@ -8,6 +8,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gocui"
 	"github.com/jesseduffield/lazygit/pkg/gui/patch_exploring"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/tasks"
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
@@ -159,7 +160,14 @@ func (self *StagingHelper) GetDiffLineInfo(windowName string, viewLineIdx int) (
 	if v == nil {
 		return types.DiffLineInfo{}, false
 	}
+	return self.GetDiffLineInfoForView(v, viewLineIdx)
+}
 
+// GetDiffLineInfoForView is GetDiffLineInfo against a specific view rather than
+// one looked up by window. It is used to read the identity of the line the patch
+// explorer currently has selected when escaping back to the focused main view,
+// where we hold the explorer's view directly.
+func (self *StagingHelper) GetDiffLineInfoForView(v *gocui.View, viewLineIdx int) (types.DiffLineInfo, bool) {
 	// A click/cursor lands on a (wrapped) view line; resolve it to the unwrapped
 	// buffer line all three backends key off, then read that buffer line's content.
 	bufferLineIdx, ok := v.BufferLineForViewLine(viewLineIdx)
@@ -169,6 +177,103 @@ func (self *StagingHelper) GetDiffLineInfo(windowName string, viewLineIdx int) (
 	return self.diffLineInfoFromContents(v.DiffLineContents(), bufferLineIdx)
 }
 
+// FindDiffLine returns the index of the first buffer line, at or after from,
+// whose patch identity matches target (see types.DiffLineInfo.SamePatchLine), or
+// -1 if none does. It is the inverse direction of the diff-line primitive:
+// instead of resolving the line under a cursor, it scans a rendered diff for a
+// known identity. The escape restore uses it to locate, in the focused main view
+// it is returning to, the line the patch explorer had selected — scanning the
+// re-render's content as it loads off-screen. (from lets the caller resume the
+// scan as more content arrives without re-checking lines already scanned; the
+// backends still see all of contents, so a buffer-parse match can look back to
+// its file/hunk headers.)
+func (self *StagingHelper) FindDiffLine(contents []gocui.DiffLineContent, target types.DiffLineInfo, from int) int {
+	for i := from; i < len(contents); i++ {
+		if info, ok := self.diffLineInfoFromContents(contents, i); ok && info.SamePatchLine(target) {
+			return i
+		}
+	}
+	return -1
+}
+
+// RestoreFocusedMainViewOnEscape arranges, when escaping a patch explorer back to
+// the focused main view it was entered from, for that view to re-render and then
+// land on the line the explorer currently has selected. After staging or dropping
+// hunks the explorer's selection auto-advances, so the line the user ended up on
+// is more useful to return to than the one they entered on — and, since the diff
+// has changed, more reliable than replaying a numeric scroll/index.
+//
+// It reads that line's patch identity from the explorer now (explorerView,
+// explorerSelectedLineIdx), then installs a restore that, as the main view
+// re-renders, scans the incoming content for the row matching that identity and —
+// once it and a screenful below have loaded — swaps in, scrolls there and selects
+// it. If the identity can't be read, or the line never turns up in the re-render
+// (the content changed out from under it), no restore is applied and the view
+// just re-renders normally.
+func (self *StagingHelper) RestoreFocusedMainViewOnEscape(explorerView, mainView *gocui.View, explorerSelectedLineIdx int) {
+	target, ok := self.GetDiffLineInfoForView(explorerView, explorerSelectedLineIdx)
+	if !ok {
+		return
+	}
+
+	manager := self.c.GetViewBufferManagerForView(mainView)
+	if manager == nil {
+		return
+	}
+
+	// targetBufferLine is found once: either incrementally as the content loads
+	// (scanned tracks how far we've scanned, so each line is checked once), or — if
+	// the incremental scan can't resolve it — once more on the complete content
+	// when we swap in (see Apply).
+	targetBufferLine := -1
+	scanned := 0
+	manager.SetRestoreForNextTask(&tasks.RenderRestore{
+		FirstPaintReady: func() bool {
+			// Scan the incoming content for the target as it loads, so we can paint
+			// at the saved position as soon as it's reachable. This finds the line
+			// for backends that resolve a row on its own — OSC metadata, lazygit-edit
+			// hyperlinks. The buffer-parse backend can't resolve here: it parses whole
+			// hunks against their @@ lengths, and the trailing hunk is incomplete
+			// while loading, so the parse is rejected as not-well-formed until the
+			// diff is fully read. That case is handled in Apply instead.
+			if targetBufferLine == -1 {
+				contents := mainView.OffscreenDiffLineContents()
+				targetBufferLine = self.FindDiffLine(contents, target, scanned)
+				scanned = len(contents)
+				if targetBufferLine == -1 {
+					return false
+				}
+			}
+			return mainView.OffscreenLineCount() >= targetBufferLine+mainView.InnerHeight()
+		},
+		Apply: func() {
+			// The off-screen render has just been swapped in, so the displayed buffer
+			// now holds it. If the incremental scan never found the target — the
+			// common case for buffer-parse, which only becomes well-formed once the
+			// whole diff has loaded, at which point the swap happens at end of input —
+			// scan the now-complete content once more.
+			if targetBufferLine == -1 {
+				targetBufferLine = self.FindDiffLine(mainView.DiffLineContents(), target, 0)
+			}
+			// If the target line still isn't there (the content changed and the line
+			// is gone), leave the scroll and selection as they are rather than showing
+			// a selection on a line that no longer means what it did.
+			if targetBufferLine != -1 {
+				if viewLine, ok := mainView.ViewLineForBufferLine(targetBufferLine); ok {
+					// scrollIntoView centres the line if it's off-screen, and leaves the
+					// scroll untouched if it's already visible — so for the common
+					// unchanged-content escape (the placeholder is the same content at
+					// the same scroll) nothing moves.
+					mainView.FocusPoint(0, viewLine, true)
+					mainView.Highlight = true
+					mainView.HighlightInactive = false
+				}
+			}
+			manager.ClearRestoreForNextTask()
+		},
+	})
+}
+
 // diffLineInfoFromContents recovers the patch-space identity of the buffer line
 // at idx within a snapshot of a diff's per-line content (see gocui.DiffLineContent).
 // It is the single resolver behind both directions of the diff-line primitive —

pkg/gui/controllers/main_view_controller.go

@@ -201,14 +201,15 @@ func focusedMainViewContextForViewName(c *ControllerCommon, viewName string) typ
 	return c.Contexts().Normal
 }
 
-// focusedMainViewSnapshot captures where a focused main view is (scroll +
-// selected line) when diving into a patch explorer from it, so escaping can
-// return there with the main view focused. sidePanel is the panel to land on
-// first (which re-renders the content); for commits/stash it's the originating
-// panel, skipping the commit files panel we pass through. selectedLineIdx is the
-// view line that was selected in the focused main view. Call this before any
-// mutation that might re-render the main view.
-func focusedMainViewSnapshot(c *ControllerCommon, mainViewName string, sidePanel types.Context, selectedLineIdx int) *types.FocusedMainViewSnapshot {
+// focusedMainViewSnapshot records the focused main view to return to when diving
+// into a patch explorer from it, so escaping can come back with the main view
+// focused. sidePanel is the panel to land on first (which re-renders the
+// content); for commits/stash it's the originating panel, skipping the commit
+// files panel we pass through. Where to scroll to and select on return isn't
+// captured: escape lands on the line the explorer ended up on (see
+// EscapeFromPatchExplorer). Call this before any mutation that might re-render
+// the main view.
+func focusedMainViewSnapshot(c *ControllerCommon, mainViewName string, sidePanel types.Context) *types.FocusedMainViewSnapshot {
 	mainView := focusedMainViewContextForViewName(c, mainViewName)
 	sidePanelSelectedLineIdx := -1
 	if listContext, ok := sidePanel.(types.IListContext); ok {
@@ -218,8 +219,6 @@ func focusedMainViewSnapshot(c *ControllerCommon, mainViewName string, sidePanel
 		SidePanel:                sidePanel,
 		SidePanelSelectedLineIdx: sidePanelSelectedLineIdx,
 		MainView:                 mainView,
-		OriginY:                  mainView.GetView().OriginY(),
-		SelectedLineIdx:          selectedLineIdx,
 	}
 }
 

pkg/gui/controllers/staging_controller.go

@@ -176,7 +176,7 @@ func (self *StagingController) Escape() error {
 		return nil
 	}
 
-	helpers.EscapeFromPatchExplorer(self.c.HelperCommon, self.context)
+	helpers.EscapeFromPatchExplorer(self.c.HelperCommon, self.c.Helpers().Staging, self.context)
 	return nil
 }
 

pkg/gui/controllers/switch_to_diff_files_controller.go

@@ -63,7 +63,7 @@ func (self *SwitchToDiffFilesController) GetOnClickFocusedMainView() func(mainVi
 		// Capture before self.enter() pushes the commit files panel, which
 		// re-renders the main view. We escape "all the way out" to this side
 		// panel (skipping the commit files panel), then focus the main view.
-		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context, clickedLineIdx)
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context)
 
 		if err := self.enter(); err != nil {
 			return err

pkg/gui/pty.go

@@ -82,11 +82,6 @@ func (gui *Gui) newPtyTask(view *gocui.View, cmd *exec.Cmd, prefix string) error
 	// (see the matching call in newCmdTask).
 	view.FreezeScrollbarHeight()
 
-	// Read any requested scroll-restore now so we can size the initial read to it
-	// in afterLayout; the task itself clears the request and applies the scroll at
-	// its first paint.
-	targetOriginY := gui.getManager(view).GetScrollToOriginYForNextTask()
-
 	// Run the pty after layout so that it gets the correct size
 	gui.afterLayout(func() error {
 		// Need to get the width and the pager command again because the layout might have
@@ -158,11 +153,15 @@ func (gui *Gui) newPtyTask(view *gocui.View, cmd *exec.Cmd, prefix string) error
 			gui.Mutexes.PtyMutex.Unlock()
 		}
 
-		linesToRead := gui.linesToReadFromCmdTask(view, targetOriginY)
-		// As in newCmdTask: let the task run any requested after-load callback at
-		// the end of its initial read (e.g. restoring a focused main view's
-		// selection on escape). The task clears the request when it starts.
-		linesToRead.Then = manager.GetThenForNextTask()
+		linesToRead := gui.linesToReadFromCmdTask(view)
+		// As in newCmdTask: if a restore is pending for this content (returning to a
+		// focused main view on escape), let the task re-establish the scroll
+		// position and selection as it first paints, reading to end of input so a
+		// deep target line is found and the scrollbar ends up accurate.
+		if restore := manager.GetRestoreForNextTask(); restore != nil {
+			linesToRead.Restore = restore
+			linesToRead.Total = -1
+		}
 		return manager.NewTask(manager.NewCmdTask(start, prefix, linesToRead, onClose), cmdStr)
 	})
 

pkg/gui/tasks_adapter.go

@@ -35,11 +35,6 @@ func (gui *Gui) newCmdTask(view *gocui.View, cmd *exec.Cmd, prefix string) error
 	// still-running writes (see View.SetContentWidth).
 	contentWidth := view.InnerWidth()
 
-	// If a caller asked us to restore a scroll position for this render, size the
-	// initial read to it (below) and let the task scroll there at its first paint.
-	// The task clears the request and suppresses the origin reset when it starts.
-	targetOriginY := manager.GetScrollToOriginYForNextTask()
-
 	var r io.ReadCloser
 	start := func() (tasks.Cmd, io.Reader) {
 		view.SetContentWidth(contentWidth)
@@ -56,12 +51,15 @@ func (gui *Gui) newCmdTask(view *gocui.View, cmd *exec.Cmd, prefix string) error
 		}
 	}
 
-	linesToRead := gui.linesToReadFromCmdTask(view, targetOriginY)
-	// If a caller asked us to run something once this re-render has loaded (e.g.
-	// restoring a focused main view's selection on escape), let the task own it,
-	// firing at the end of its initial read. The task clears the request when it
-	// starts.
-	linesToRead.Then = manager.GetThenForNextTask()
+	linesToRead := gui.linesToReadFromCmdTask(view)
+	// If a restore is pending for this content (returning to a focused main view
+	// on escape), let the task re-establish the scroll position and selection as
+	// it first paints. It also reads to end of input so a deep target line is
+	// found and the scrollbar ends up accurate. See RenderRestore.
+	if restore := manager.GetRestoreForNextTask(); restore != nil {
+		linesToRead.Restore = restore
+		linesToRead.Total = -1
+	}
 	if err := manager.NewTask(manager.NewCmdTask(start, prefix, linesToRead, onClose), cmdStr); err != nil {
 		gui.c.Log.Error(err)
 	}

pkg/gui/types/context.go

@@ -228,11 +228,11 @@ type FocusedMainViewSnapshot struct {
 	// to a file in the files panel); restoring it makes the main view show the
 	// same content again. -1 if the side panel isn't a list.
 	SidePanelSelectedLineIdx int
-	// The focused main view context to focus afterwards.
+	// The focused main view context to focus afterwards. Where in it to scroll to
+	// and select is not captured here: on escape we land on the line the patch
+	// explorer ended up selecting, found by its patch identity in the re-rendered
+	// content, which survives the diff changing in a way a saved index wouldn't.
 	MainView Context
-	// The scroll position and selected line to restore in the main view.
-	OriginY         int
-	SelectedLineIdx int
 }
 
 type IViewTrait interface {

pkg/gui/types/diff_line_info.go

@@ -45,6 +45,25 @@ func (self DiffLineInfo) PatchSelectLine() (lineNumber int, isDeletion bool) {
 	return self.NewLine, false
 }
 
+// SamePatchLine reports whether two identities point at the same source line of
+// the same file. It is how the escape restore matches the line the patch explorer
+// had selected against the rows of the focused main view as it re-renders: the
+// comparison is in source-line-number space (PatchSelectLine), which survives the
+// diff being regenerated, rather than a fragile view-line index.
+//
+// A backend that can't determine the side (delta's lazygit-edit hyperlinks report
+// DiffLineOther) yields a non-deletion identity, so a deletion captured from a
+// full-fidelity backend won't match such a row — the restore then just doesn't
+// find its line, which is the acceptable degradation for that pager config.
+func (self DiffLineInfo) SamePatchLine(other DiffLineInfo) bool {
+	if self.Path != other.Path {
+		return false
+	}
+	selfLine, selfIsDeletion := self.PatchSelectLine()
+	otherLine, otherIsDeletion := other.PatchSelectLine()
+	return selfLine == otherLine && selfIsDeletion == otherIsDeletion
+}
+
 // PullRequestAnchor returns the side ("L"/"R") and line number to anchor a
 // GitHub PR deep-link at: the left/old side for a deletion, the right/new side
 // otherwise.

pkg/gui/view_helpers.go

@@ -19,22 +19,9 @@ func (gui *Gui) resetViewOrigin(v *gocui.View) {
 // Returns the number of lines that we should read initially from a cmd task so
 // that the scrollbar has the correct size, along with the number of lines after
 // which the view is filled and we can do a first refresh.
-//
-// If targetOriginY is non-nil, the read is sized to that scroll position rather
-// than the view's current one, and the returned LinesToRead carries an
-// ApplyInitialScroll that scrolls the view there at the first refresh. This is
-// used when re-rendering content the user was already scrolled into, so the
-// saved position is applied exactly when the content first paints.
-func (gui *Gui) linesToReadFromCmdTask(v *gocui.View, targetOriginY *int) tasks.LinesToRead {
+func (gui *Gui) linesToReadFromCmdTask(v *gocui.View) tasks.LinesToRead {
 	height := v.InnerHeight()
 	oy := v.OriginY()
-	var applyInitialScroll func()
-	if targetOriginY != nil {
-		oy = *targetOriginY
-		applyInitialScroll = func() {
-			v.SetOrigin(v.OriginX(), *targetOriginY)
-		}
-	}
 
 	linesForFirstRefresh := height + oy + 10
 
@@ -50,7 +37,6 @@ func (gui *Gui) linesToReadFromCmdTask(v *gocui.View, targetOriginY *int) tasks.
 	return tasks.LinesToRead{
 		Total:               linesToReadForAccurateScrollbar,
 		InitialRefreshAfter: linesForFirstRefresh,
-		ApplyInitialScroll:  applyInitialScroll,
 	}
 }
 

pkg/tasks/tasks.go

@@ -74,34 +74,20 @@ type ViewBufferManager struct {
 	taskKey   string
 	onNewKey  func()
 
-	// When non-nil, the next cmd/pty task restores this scroll position. Used
-	// when re-rendering content the user was already scrolled into (e.g. when
-	// returning to a focused main view on escape). It has two effects:
+	// When non-nil, the next cmd/pty task re-establishes the view's scroll
+	// position and selection once it has re-rendered the content (returning to a
+	// focused main view on escape). The task does not reset the view's origin to
+	// the top at start: instead it keeps the placeholder showing at its current
+	// scroll until the restore can show the saved position as part of the first
+	// paint, rather than flicking to the top. See RenderRestore.
 	//
-	//   - The task does not reset the view's origin to the top at start, even if
-	//     its command key differs from the previous task's. This keeps the
-	//     placeholder that CopyContent left in the view showing at its current
-	//     scroll position (i.e. "as if nothing changed") until the real content
-	//     is ready, rather than flicking it to the top.
-	//   - The task sizes its initial read to this position and scrolls there as
-	//     part of its first refresh, so the saved scroll is applied in the same
-	//     paint that shows the real content.
-	//
-	// Cleared once the next task has started.
-	scrollToOriginYForNextTask *int
-
-	// When non-nil, run once after the next cmd/pty task has read enough of its
-	// content (the end of its initial read). Used to restore state that depends
-	// on the re-rendered content being loaded — e.g. the selection in a focused
-	// main view we're returning to on escape — by riding the task's own
-	// lifecycle rather than a separate post-hoc ReadToEnd. A ReadToEnd issued
-	// right after triggering the re-render can fire synchronously, because the
-	// freshly-created task's read channel isn't live yet (it is created inside
-	// the task goroutine, after the previous task has been stopped); that would
-	// run the restore before any content is loaded and silently drop it.
-	//
-	// Cleared once the next task has started.
-	thenForNextTask func()
+	// It is set just before triggering the re-render. Unlike a per-task field, it
+	// is *not* cleared when a task starts: it must survive a task being stopped
+	// and replaced (a periodic refresh can stop the escape's re-render before it
+	// first-paints), so it is kept until a task successfully applies it, or until
+	// the content key changes (a different item was selected, so the saved
+	// position no longer applies). Guarded by taskIDMutex, like the task key.
+	restoreForNextTask *RenderRestore
 
 	// Whether a command task is currently reading content into the view. While
 	// this is true the content is still growing, so callers (e.g. the layout)
@@ -153,16 +139,38 @@ type LinesToRead struct {
 	// subsequent requests.
 	InitialRefreshAfter int
 
-	// When set, called once, just before the view is first refreshed, to scroll
-	// it to a saved position. Used so that content the user was scrolled into is
-	// painted at the saved scroll position the first time it appears, rather than
-	// at the top. Only set for the initial read request.
-	ApplyInitialScroll func()
+	// When set, re-establishes the view's scroll position and selection as the
+	// re-render first paints (returning to a focused main view on escape). When
+	// set, it — rather than InitialRefreshAfter — also decides when the first
+	// paint happens, since the saved position may be reachable only after more
+	// than a screenful has loaded. Only set for the initial read request. See
+	// RenderRestore.
+	Restore *RenderRestore
 
 	// Function to call after reading the lines is done
 	Then func()
 }
 
+// RenderRestore re-establishes a view's scroll position and selection after it
+// re-renders content the user was already looking at (returning to a focused main
+// view on escape). The render task reads the new content into an off-screen
+// buffer; RenderRestore decides, as that buffer fills, when the task has read far
+// enough to show the saved position (FirstPaintReady), and then — once the
+// off-screen buffer has been swapped in — scrolls there and restores the
+// selection (Apply). It is a predicate rather than a fixed scroll position so the
+// target can be a row matching a patch identity, located by scanning the loading
+// content, rather than a line number that the changed content may have moved.
+type RenderRestore struct {
+	// FirstPaintReady reports whether the task has now read enough of the new
+	// (off-screen) content to first-paint at the saved position. Evaluated after
+	// each line is read.
+	FirstPaintReady func() bool
+
+	// Apply runs once, just after the off-screen render is swapped in at the first
+	// paint, to scroll to the saved position and restore the selection.
+	Apply func()
+}
+
 func (self *ViewBufferManager) GetTaskKey() string {
 	return self.taskKey
 }
@@ -206,36 +214,40 @@ func (self *ViewBufferManager) ReadLines(totalLines int) {
 	}
 }
 
-// ScrollToOriginYForNextTask makes the next cmd/pty task restore the given
-// scroll position instead of rendering at the top. Call this right before
-// triggering a re-render of content the view is already scrolled into (e.g.
-// when returning to a focused main view on escape). See the field doc for the
-// two effects this has. It is cleared once the next task starts.
-func (self *ViewBufferManager) ScrollToOriginYForNextTask(originY int) {
-	self.scrollToOriginYForNextTask = &originY
+// SetRestoreForNextTask makes the next cmd/pty task re-establish the view's
+// scroll position and selection once it has re-rendered the content. Call this
+// right before triggering a re-render of content the view was already showing
+// (returning to a focused main view on escape). See the field doc and RenderRestore.
+func (self *ViewBufferManager) SetRestoreForNextTask(restore *RenderRestore) {
+	self.taskIDMutex.Lock()
+	defer self.taskIDMutex.Unlock()
+
+	self.restoreForNextTask = restore
 }
 
-// GetScrollToOriginYForNextTask returns the scroll position requested by a
-// preceding ScrollToOriginYForNextTask call, or nil if none. It does not clear
-// it; the task clears it when it starts.
-func (self *ViewBufferManager) GetScrollToOriginYForNextTask() *int {
-	return self.scrollToOriginYForNextTask
+// GetRestoreForNextTask returns the pending restore, or nil. It is not gated on
+// the command key: returning to a focused main view after staging changes the
+// command (e.g. the unstaged diff becomes the staged one once the last unstaged
+// hunk is gone), yet the line to land on is still in the new content. The restore
+// validates itself instead — its scan simply doesn't find the target line when
+// the content no longer contains it (a different item was selected), in which case
+// applying it is a no-op. So it is safe to hand to whatever renders next.
+func (self *ViewBufferManager) GetRestoreForNextTask() *RenderRestore {
+	self.taskIDMutex.Lock()
+	defer self.taskIDMutex.Unlock()
+
+	return self.restoreForNextTask
 }
 
-// ThenForNextTask makes the next cmd/pty task run the given function once it has
-// read enough of its content (the end of its initial read). Call this right
-// before triggering a re-render whose loaded content the function depends on.
-// See the field doc for why this is preferable to a separate ReadToEnd. It is
-// cleared once the next task starts.
-func (self *ViewBufferManager) ThenForNextTask(then func()) {
-	self.thenForNextTask = then
-}
+// ClearRestoreForNextTask drops the pending restore. A task clears it once it has
+// first-painted and applied it (whether or not it found its target line), so that
+// it lives for exactly one re-render — surviving a task being stopped and replaced
+// before it could paint, but not re-applying on every later render.
+func (self *ViewBufferManager) ClearRestoreForNextTask() {
+	self.taskIDMutex.Lock()
+	defer self.taskIDMutex.Unlock()
 
-// GetThenForNextTask returns the function requested by a preceding
-// ThenForNextTask call, or nil if none. It does not clear it; the task clears it
-// when it starts.
-func (self *ViewBufferManager) GetThenForNextTask() func() {
-	return self.thenForNextTask
+	self.restoreForNextTask = nil
 }
 
 // IsLoading reports whether a command task is currently reading content into the
@@ -413,15 +425,23 @@ func (self *ViewBufferManager) NewCmdTask(start func() (Cmd, io.Reader), prefix
 			// this to work out how many more lines, if any, we still need to read.
 			linesRead := 0
 
-			// The initial read request carries an optional scroll-to function (see
-			// LinesToRead.ApplyInitialScroll). We apply it exactly once, right before
-			// the view is first refreshed, so that content the user was scrolled into
-			// is painted at the saved position the first time it appears.
-			initialScroll := linesToRead.ApplyInitialScroll
-			var applyInitialScrollOnce sync.Once
-			applyInitialScroll := func() {
-				if initialScroll != nil {
-					applyInitialScrollOnce.Do(initialScroll)
+			// The first paint swaps the off-screen render in to reveal the new
+			// content. When a restore is pending (returning to a focused main view on
+			// escape, see RenderRestore), it scrolls to the saved position and
+			// restores the selection in the same step, so the real content first
+			// appears already at the right place rather than at the top. firstPaint
+			// happens once, either when we've read far enough (below) or at end of
+			// input for content shorter than that.
+			restore := linesToRead.Restore
+			painted := false
+			firstPaint := func() {
+				if painted {
+					return
+				}
+				painted = true
+				self.swapInRender()
+				if restore != nil {
+					restore.Apply()
 				}
 			}
 
@@ -494,17 +514,16 @@ func (self *ViewBufferManager) NewCmdTask(start func() (Cmd, io.Reader), prefix
 								break outer
 							default:
 							}
-							// Genuine end of input: swap in whatever we read (the content was
-							// shorter than a screenful, so we never hit the InitialRefreshAfter
-							// swap). Apply the saved scroll first (if any) so that onEndOfInput
-							// clamps it back into range when the content turned out shorter than
-							// expected.
+							// Genuine end of input: do the first paint now if it hasn't happened
+							// yet — the content was shorter than the first-paint point, or a
+							// restore's target line was never found. firstPaint swaps in whatever
+							// we read and, for a restore, scrolls to the saved position before
+							// onEndOfInput clamps the origin back into range for short content.
 							// onEndOfInput reads the view's dimensions (to decide
 							// whether to scroll) and sets the origin, both of which
 							// are UI-thread-only, so run it there.
 							_ = self.onUIThread(func() error {
-								self.swapInRender()
-								applyInitialScroll()
+								firstPaint()
 								self.onEndOfInput()
 								return nil
 							})
@@ -539,15 +558,25 @@ func (self *ViewBufferManager) NewCmdTask(start func() (Cmd, io.Reader), prefix
 							time.Sleep(slowRenderPerLine)
 						}
 
-						if linesRead == linesToRead.InitialRefreshAfter {
-							// We have read enough lines to fill the view, so swap the off-screen
-							// content in and do a first refresh to show it. Apply the saved
-							// scroll first (if any) so the first paint already lands at it.
-							// Continue reading and refresh again at the end to make sure the
-							// scrollbar has the right size.
-							self.swapInRender()
-							applyInitialScroll()
-							refreshViewIfStale()
+						// Do the first paint as soon as we've read far enough: for a restore,
+						// when it can show the saved position (RenderRestore.FirstPaintReady,
+						// e.g. its target line plus a screenful below it have loaded); otherwise
+						// when we've read enough lines to fill the view (InitialRefreshAfter).
+						// This swaps the off-screen content in and refreshes; we keep reading
+						// afterwards and refresh again at the end so the scrollbar ends up the
+						// right size.
+						if !painted {
+							var ready bool
+							if restore != nil {
+								ready = restore.FirstPaintReady()
+							} else {
+								ready = linesToRead.InitialRefreshAfter > 0 && linesRead >= linesToRead.InitialRefreshAfter
+							}
+							if ready {
+								// TODO: should probably use OnUIThread?
+								firstPaint()
+								refreshViewIfStale()
+							}
 						}
 					}
 					refreshViewIfStale()
@@ -665,13 +694,15 @@ func (self *ViewBufferManager) NewTask(f func(TaskOpts) error, key string) error
 			return
 		}
 
-		// Reset the origin to the top when the command changed, unless a caller
-		// asked us to restore a scroll position: in that case we keep the
-		// placeholder showing at its current scroll until the task scrolls to the
-		// saved position as part of its first paint (see scrollToOriginYForNextTask).
-		resetOrigin := self.GetTaskKey() != key && self.onNewKey != nil && self.scrollToOriginYForNextTask == nil
-		self.scrollToOriginYForNextTask = nil
-		self.thenForNextTask = nil
+		// Reset the origin to the top when the command changed, unless a restore is
+		// pending: a restore re-establishes the scroll position itself (and we keep
+		// the placeholder showing at its current scroll until it does), so resetting
+		// to the top would just flicker. The restore isn't cleared here: it must
+		// outlive a task being stopped and replaced before it could paint, and it
+		// validates itself against the content it lands in (see restoreForNextTask),
+		// so a stale one can't apply to the wrong place — it's cleared once a task
+		// has applied it.
+		resetOrigin := self.GetTaskKey() != key && self.onNewKey != nil && self.restoreForNextTask == nil
 		self.taskKey = key
 
 		self.taskIDMutex.Unlock()

pkg/tasks/tasks_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/jesseduffield/lazygit/pkg/gocui"
 	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/stretchr/testify/assert"
 )
 
 func getCounter() (func(), func() int) {
@@ -171,6 +172,71 @@ func TestNewCmdTask(t *testing.T) {
 
 // A dummy reader that simply yields as many blank lines as requested. The only
 // thing we want to do with the output is count the number of lines.
+// When a RenderRestore is set, the first paint is driven by its FirstPaintReady
+// predicate rather than the InitialRefreshAfter line count, and Apply runs exactly
+// once, right after the off-screen render is swapped in. This is the read-loop
+// half of the escape restore: scroll to and select the saved position as the new
+// content first appears. See RenderRestore.
+func TestNewCmdTaskRestore(t *testing.T) {
+	writer := bytes.NewBuffer(nil)
+	linesWritten := func() int { return strings.Count(writer.String(), "\n") }
+
+	swapped := false
+	applyCount := 0
+	applyAtLines := -1
+	applyAfterSwap := true
+
+	task := gocui.NewFakeTask()
+	manager := NewViewBufferManager(
+		utils.NewDummyLog(),
+		writer,
+		func() {},                 // beforeStart
+		func() {},                 // refreshView
+		func() {},                 // onEndOfInput
+		func() {},                 // onNewKey
+		func() {},                 // beginRender
+		func() { swapped = true }, // swapInRender
+		func() gocui.Task { return task },
+		// no UI thread in the test; run the view mutations inline
+		func(f func() error) error { return f() },
+	)
+
+	restore := &RenderRestore{
+		// Ready once five lines have loaded — well before InitialRefreshAfter (30).
+		FirstPaintReady: func() bool { return linesWritten() >= 5 },
+		Apply: func() {
+			applyCount++
+			applyAtLines = linesWritten()
+			if !swapped {
+				applyAfterSwap = false
+			}
+		},
+	}
+
+	stop := make(chan struct{})
+	reader := BlankLineReader{totalLinesToYield: 50}
+	start := func() (Cmd, io.Reader) {
+		cmd := exec.Command("blah")
+		return ExecCmd{Cmd: cmd}, &reader
+	}
+	fn := manager.NewCmdTask(start, "", LinesToRead{Total: 50, InitialRefreshAfter: 30, Restore: restore}, func() {})
+
+	wg := sync.WaitGroup{}
+	wg.Go(func() {
+		time.Sleep(100 * time.Millisecond)
+		close(stop)
+	})
+	_ = fn(TaskOpts{Stop: stop, InitialContentLoaded: func() { task.Done() }})
+	wg.Wait()
+
+	assert.Equal(t, 1, applyCount, "Apply should run exactly once")
+	assert.True(t, applyAfterSwap, "Apply should run after the off-screen render is swapped in")
+	// The first paint was driven by FirstPaintReady (>=5 lines), not by
+	// InitialRefreshAfter (30).
+	assert.GreaterOrEqual(t, applyAtLines, 5)
+	assert.Less(t, applyAtLines, 30)
+}
+
 type BlankLineReader struct {
 	totalLinesToYield int
 	linesYielded      int