fix popup focus issue

6 files changed, 31 insertions, 45 deletions

diff --git a/pkg/gui/confirmation_panel.go b/pkg/gui/confirmation_panel.go
index 19f8745a8..4177a689b 100644
--- a/pkg/gui/confirmation_panel.go
+++ b/pkg/gui/confirmation_panel.go
@@ -16,9 +16,9 @@ import (
 // This file is for the rendering of confirmation panels along with setting and handling associated
 // keybindings.
 
-func (gui *Gui) wrappedConfirmationFunction(handlersManageFocus bool, function func() error) func() error {
+func (gui *Gui) wrappedConfirmationFunction(function func() error) func() error {
 	return func() error {
-		if err := gui.closeConfirmationPrompt(handlersManageFocus); err != nil {
+		if err := gui.c.PopContext(); err != nil {
 			return err
 		}
 
@@ -32,9 +32,9 @@ func (gui *Gui) wrappedConfirmationFunction(handlersManageFocus bool, function f
 	}
 }
 
-func (gui *Gui) wrappedPromptConfirmationFunction(handlersManageFocus bool, function func(string) error, getResponse func() string) func() error {
+func (gui *Gui) wrappedPromptConfirmationFunction(function func(string) error, getResponse func() string) func() error {
 	return func() error {
-		if err := gui.closeConfirmationPrompt(handlersManageFocus); err != nil {
+		if err := gui.c.PopContext(); err != nil {
 			return err
 		}
 
@@ -48,27 +48,15 @@ func (gui *Gui) wrappedPromptConfirmationFunction(handlersManageFocus bool, func
 	}
 }
 
-func (gui *Gui) closeConfirmationPrompt(handlersManageFocus bool) error {
+func (gui *Gui) deactivateConfirmationPrompt() {
 	gui.Mutexes.PopupMutex.Lock()
 	gui.State.CurrentPopupOpts = nil
 	gui.Mutexes.PopupMutex.Unlock()
 
-	// we've already closed it so we can just return
-	if !gui.Views.Confirmation.Visible {
-		return nil
-	}
-
-	if !handlersManageFocus {
-		if err := gui.c.PopContext(); err != nil {
-			return err
-		}
-	}
-
-	gui.clearConfirmationViewKeyBindings()
 	gui.Views.Confirmation.Visible = false
 	gui.Views.Suggestions.Visible = false
 
-	return nil
+	gui.clearConfirmationViewKeyBindings()
 }
 
 func (gui *Gui) getMessageHeight(wrap bool, message string, width int) int {
@@ -219,14 +207,13 @@ func (gui *Gui) setKeyBindings(opts types.CreatePopupPanelOpts) error {
 	_ = gui.renderString(gui.Views.Options, actions)
 	var onConfirm func() error
 	if opts.HandleConfirmPrompt != nil {
-		onConfirm = gui.wrappedPromptConfirmationFunction(opts.HandlersManageFocus, opts.HandleConfirmPrompt, func() string { return gui.Views.Confirmation.TextArea.GetContent() })
+		onConfirm = gui.wrappedPromptConfirmationFunction(opts.HandleConfirmPrompt, func() string { return gui.Views.Confirmation.TextArea.GetContent() })
 	} else {
-		onConfirm = gui.wrappedConfirmationFunction(opts.HandlersManageFocus, opts.HandleConfirm)
+		onConfirm = gui.wrappedConfirmationFunction(opts.HandleConfirm)
 	}
 
 	keybindingConfig := gui.c.UserConfig.Keybinding
 	onSuggestionConfirm := gui.wrappedPromptConfirmationFunction(
-		opts.HandlersManageFocus,
 		opts.HandleConfirmPrompt,
 		gui.getSelectedSuggestionValue,
 	)
@@ -248,7 +235,7 @@ func (gui *Gui) setKeyBindings(opts types.CreatePopupPanelOpts) error {
 			ViewName: "confirmation",
 			Contexts: []string{string(context.CONFIRMATION_CONTEXT_KEY)},
 			Key:      gui.getKey(keybindingConfig.Universal.Return),
-			Handler:  gui.wrappedConfirmationFunction(opts.HandlersManageFocus, opts.HandleClose),
+			Handler:  gui.wrappedConfirmationFunction(opts.HandleClose),
 		},
 		{
 			ViewName: "confirmation",
@@ -277,7 +264,7 @@ func (gui *Gui) setKeyBindings(opts types.CreatePopupPanelOpts) error {
 			ViewName: "suggestions",
 			Contexts: []string{string(context.SUGGESTIONS_CONTEXT_KEY)},
 			Key:      gui.getKey(keybindingConfig.Universal.Return),
-			Handler:  gui.wrappedConfirmationFunction(opts.HandlersManageFocus, opts.HandleClose),
+			Handler:  gui.wrappedConfirmationFunction(opts.HandleClose),
 		},
 		{
 			ViewName: "suggestions",
diff --git a/pkg/gui/context_config.go b/pkg/gui/context_config.go
index bc9df8ba4..a983ffdd5 100644
--- a/pkg/gui/context_config.go
+++ b/pkg/gui/context_config.go
@@ -124,6 +124,10 @@ func (gui *Gui) contextTree() *context.ContextTree {
 			}),
 			context.ContextCallbackOpts{
 				OnFocus: OnFocusWrapper(gui.handleAskFocused),
+				OnFocusLost: func() error {
+					gui.deactivateConfirmationPrompt()
+					return nil
+				},
 			},
 		),
 		CommitMessage: context.NewSimpleContext(
diff --git a/pkg/gui/controllers/helpers/merge_and_rebase_helper.go b/pkg/gui/controllers/helpers/merge_and_rebase_helper.go
index ef9dd76a2..6d2c680e3 100644
--- a/pkg/gui/controllers/helpers/merge_and_rebase_helper.go
+++ b/pkg/gui/controllers/helpers/merge_and_rebase_helper.go
@@ -149,17 +149,12 @@ func (self *MergeAndRebaseHelper) CheckMergeOrRebase(result error) error {
 		return nil
 	} else if isMergeConflictErr(result.Error()) {
 		return self.c.Confirm(types.ConfirmOpts{
-			Title:               self.c.Tr.FoundConflictsTitle,
-			Prompt:              self.c.Tr.FoundConflicts,
-			HandlersManageFocus: true,
+			Title:  self.c.Tr.FoundConflictsTitle,
+			Prompt: self.c.Tr.FoundConflicts,
 			HandleConfirm: func() error {
 				return self.c.PushContext(self.contexts.Files)
 			},
 			HandleClose: func() error {
-				if err := self.c.PopContext(); err != nil {
-					return err
-				}
-
 				return self.genericMergeCommand(REBASE_OPTION_ABORT)
 			},
 		})
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 04b53e074..e7f1e3499 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -424,7 +424,8 @@ func NewGui(
 		cmn,
 		gui.createPopupPanel,
 		func() error { return gui.c.Refresh(types.RefreshOptions{Mode: types.ASYNC}) },
-		func() error { return gui.closeConfirmationPrompt(false) },
+		gui.popContext,
+		gui.currentContext,
 		gui.createMenu,
 		gui.withWaitingStatus,
 		gui.toast,
diff --git a/pkg/gui/popup/popup_handler.go b/pkg/gui/popup/popup_handler.go
index 53c52c74e..de892437f 100644
--- a/pkg/gui/popup/popup_handler.go
+++ b/pkg/gui/popup/popup_handler.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/common"
+	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
@@ -17,7 +18,8 @@ type PopupHandler struct {
 	sync.Mutex
 	createPopupPanelFn  func(types.CreatePopupPanelOpts) error
 	onErrorFn           func() error
-	closePopupFn        func() error
+	popContextFn        func() error
+	currentContextFn    func() types.Context
 	createMenuFn        func(types.CreateMenuOptions) error
 	withWaitingStatusFn func(message string, f func() error) error
 	toastFn             func(message string)
@@ -30,7 +32,8 @@ func NewPopupHandler(
 	common *common.Common,
 	createPopupPanelFn func(types.CreatePopupPanelOpts) error,
 	onErrorFn func() error,
-	closePopupFn func() error,
+	popContextFn func() error,
+	currentContextFn func() types.Context,
 	createMenuFn func(types.CreateMenuOptions) error,
 	withWaitingStatusFn func(message string, f func() error) error,
 	toastFn func(message string),
@@ -41,7 +44,8 @@ func NewPopupHandler(
 		index:               0,
 		createPopupPanelFn:  createPopupPanelFn,
 		onErrorFn:           onErrorFn,
-		closePopupFn:        closePopupFn,
+		popContextFn:        popContextFn,
+		currentContextFn:    currentContextFn,
 		createMenuFn:        createMenuFn,
 		withWaitingStatusFn: withWaitingStatusFn,
 		toastFn:             toastFn,
@@ -93,11 +97,10 @@ func (self *PopupHandler) Confirm(opts types.ConfirmOpts) error {
 	self.Unlock()
 
 	return self.createPopupPanelFn(types.CreatePopupPanelOpts{
-		Title:               opts.Title,
-		Prompt:              opts.Prompt,
-		HandleConfirm:       opts.HandleConfirm,
-		HandleClose:         opts.HandleClose,
-		HandlersManageFocus: opts.HandlersManageFocus,
+		Title:         opts.Title,
+		Prompt:        opts.Prompt,
+		HandleConfirm: opts.HandleConfirm,
+		HandleClose:   opts.HandleClose,
 	})
 }
 
@@ -139,8 +142,8 @@ func (self *PopupHandler) WithLoaderPanel(message string, f func() error) error
 		}
 
 		self.Lock()
-		if index == self.index {
-			_ = self.closePopupFn()
+		if index == self.index && self.currentContextFn().GetKey() == context.CONFIRMATION_CONTEXT_KEY {
+			_ = self.popContextFn()
 		}
 		self.Unlock()
 	})
diff --git a/pkg/gui/types/common.go b/pkg/gui/types/common.go
index 21808705a..5bc0e5856 100644
--- a/pkg/gui/types/common.go
+++ b/pkg/gui/types/common.go
@@ -85,9 +85,6 @@ type CreatePopupPanelOpts struct {
 	HandleConfirmPrompt func(string) error
 	HandleClose         func() error
 
-	// when HandlersManageFocus is true, do not return from the confirmation context automatically. It's expected that the handlers will manage focus, whether that means switching to another context, or manually returning the context.
-	HandlersManageFocus bool
-
 	FindSuggestionsFunc func(string) []*Suggestion
 	Mask                bool
 }
@@ -97,7 +94,6 @@ type ConfirmOpts struct {
 	Prompt              string
 	HandleConfirm       func() error
 	HandleClose         func() error
-	HandlersManageFocus bool
 	HasLoader           bool
 	FindSuggestionsFunc func(string) []*Suggestion
 	Editable            bool