Correctly request force-pushing in a triangular workflow (#3528)v0.42.0

- **PR Description**

Some people push to a different branch (or even remote) than they pull
from. One example is described in #3437. Our logic of when to request a
force push is not appropriate for these workflows: we check the
configured upstream branch for divergence, but that's the one you pull
from. We should instead check the push-to branch for divergence.

Fixes #3437.

18 files changed, 449 insertions, 140 deletions

diff --git a/docs/Custom_Command_Keybindings.md b/docs/Custom_Command_Keybindings.md
index 8604a5247..426d6f8f6 100644
--- a/docs/Custom_Command_Keybindings.md
+++ b/docs/Custom_Command_Keybindings.md
@@ -305,7 +305,7 @@ SelectedWorktree
 CheckedOutBranch
 ```
 
-To see what fields are available on e.g. the `SelectedFile`, see [here](https://github.com/jesseduffield/lazygit/blob/master/pkg/commands/models/file.go) (all the modelling lives in the same directory). Note that the custom commands feature does not guarantee backwards compatibility (until we hit Lazygit version 1.0 of course) which means a field you're accessing on an object may no longer be available from one release to the next. Typically however, all you'll need is `{{.SelectedFile.Name}}`, `{{.SelectedLocalCommit.Hash}}` and `{{.SelectedLocalBranch.Name}}`. In the future we will likely introduce a tighter interface that exposes a limited set of fields for each model.
+To see what fields are available on e.g. the `SelectedFile`, see [here](https://github.com/jesseduffield/lazygit/blob/master/pkg/gui/services/custom_commands/models.go) (all the modelling lives in the same file).
 
 ## Keybinding collisions
 
diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index b43c8c4e5..7e7d9354f 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -134,7 +134,7 @@ func NewGitCommandAux(
 	worktreeCommands := git_commands.NewWorktreeCommands(gitCommon)
 	blameCommands := git_commands.NewBlameCommands(gitCommon)
 
-	branchLoader := git_commands.NewBranchLoader(cmn, cmd, branchCommands.CurrentBranchInfo, configCommands)
+	branchLoader := git_commands.NewBranchLoader(cmn, gitCommon, cmd, branchCommands.CurrentBranchInfo, configCommands)
 	commitFileLoader := git_commands.NewCommitFileLoader(cmn, cmd)
 	commitLoader := git_commands.NewCommitLoader(cmn, cmd, statusCommands.RebaseMode, gitCommon)
 	reflogCommitLoader := git_commands.NewReflogCommitLoader(cmn, cmd)
diff --git a/pkg/commands/git_commands/branch_loader.go b/pkg/commands/git_commands/branch_loader.go
index 198368502..16777243a 100644
--- a/pkg/commands/git_commands/branch_loader.go
+++ b/pkg/commands/git_commands/branch_loader.go
@@ -40,6 +40,7 @@ type BranchInfo struct {
 // BranchLoader returns a list of Branch objects for the current repo
 type BranchLoader struct {
 	*common.Common
+	*GitCommon
 	cmd                  oscommands.ICmdObjBuilder
 	getCurrentBranchInfo func() (BranchInfo, error)
 	config               BranchLoaderConfigCommands
@@ -47,12 +48,14 @@ type BranchLoader struct {
 
 func NewBranchLoader(
 	cmn *common.Common,
+	gitCommon *GitCommon,
 	cmd oscommands.ICmdObjBuilder,
 	getCurrentBranchInfo func() (BranchInfo, error),
 	config BranchLoaderConfigCommands,
 ) *BranchLoader {
 	return &BranchLoader{
 		Common:               cmn,
+		GitCommon:            gitCommon,
 		cmd:                  cmd,
 		getCurrentBranchInfo: getCurrentBranchInfo,
 		config:               config,
@@ -61,7 +64,7 @@ func NewBranchLoader(
 
 // Load the list of branches for the current repo
 func (self *BranchLoader) Load(reflogCommits []*models.Commit) ([]*models.Branch, error) {
-	branches := self.obtainBranches()
+	branches := self.obtainBranches(self.version.IsAtLeast(2, 22, 0))
 
 	if self.AppState.LocalBranchSortOrder == "recency" {
 		reflogBranches := self.obtainReflogBranches(reflogCommits)
@@ -124,7 +127,7 @@ func (self *BranchLoader) Load(reflogCommits []*models.Commit) ([]*models.Branch
 	return branches, nil
 }
 
-func (self *BranchLoader) obtainBranches() []*models.Branch {
+func (self *BranchLoader) obtainBranches(canUsePushTrack bool) []*models.Branch {
 	output, err := self.getRawBranches()
 	if err != nil {
 		panic(err)
@@ -147,7 +150,7 @@ func (self *BranchLoader) obtainBranches() []*models.Branch {
 		}
 
 		storeCommitDateAsRecency := self.AppState.LocalBranchSortOrder != "recency"
-		return obtainBranch(split, storeCommitDateAsRecency), true
+		return obtainBranch(split, storeCommitDateAsRecency, canUsePushTrack), true
 	})
 }
 
@@ -183,23 +186,31 @@ var branchFields = []string{
 	"refname:short",
 	"upstream:short",
 	"upstream:track",
+	"push:track",
 	"subject",
 	"objectname",
 	"committerdate:unix",
 }
 
 // Obtain branch information from parsed line output of getRawBranches()
-func obtainBranch(split []string, storeCommitDateAsRecency bool) *models.Branch {
+func obtainBranch(split []string, storeCommitDateAsRecency bool, canUsePushTrack bool) *models.Branch {
 	headMarker := split[0]
 	fullName := split[1]
 	upstreamName := split[2]
 	track := split[3]
-	subject := split[4]
-	commitHash := split[5]
-	commitDate := split[6]
+	pushTrack := split[4]
+	subject := split[5]
+	commitHash := split[6]
+	commitDate := split[7]
 
 	name := strings.TrimPrefix(fullName, "heads/")
-	pushables, pullables, gone := parseUpstreamInfo(upstreamName, track)
+	aheadForPull, behindForPull, gone := parseUpstreamInfo(upstreamName, track)
+	var aheadForPush, behindForPush string
+	if canUsePushTrack {
+		aheadForPush, behindForPush, _ = parseUpstreamInfo(upstreamName, pushTrack)
+	} else {
+		aheadForPush, behindForPush = aheadForPull, behindForPull
+	}
 
 	recency := ""
 	if storeCommitDateAsRecency {
@@ -209,14 +220,16 @@ func obtainBranch(split []string, storeCommitDateAsRecency bool) *models.Branch
 	}
 
 	return &models.Branch{
-		Name:         name,
-		Recency:      recency,
-		Pushables:    pushables,
-		Pullables:    pullables,
-		UpstreamGone: gone,
-		Head:         headMarker == "*",
-		Subject:      subject,
-		CommitHash:   commitHash,
+		Name:          name,
+		Recency:       recency,
+		AheadForPull:  aheadForPull,
+		BehindForPull: behindForPull,
+		AheadForPush:  aheadForPush,
+		BehindForPush: behindForPush,
+		UpstreamGone:  gone,
+		Head:          headMarker == "*",
+		Subject:       subject,
+		CommitHash:    commitHash,
 	}
 }
 
@@ -232,10 +245,10 @@ func parseUpstreamInfo(upstreamName string, track string) (string, string, bool)
 		return "?", "?", true
 	}
 
-	pushables := parseDifference(track, `ahead (\d+)`)
-	pullables := parseDifference(track, `behind (\d+)`)
+	ahead := parseDifference(track, `ahead (\d+)`)
+	behind := parseDifference(track, `behind (\d+)`)
 
-	return pushables, pullables, false
+	return ahead, behind, false
 }
 
 func parseDifference(track string, regexStr string) string {
diff --git a/pkg/commands/git_commands/branch_loader_test.go b/pkg/commands/git_commands/branch_loader_test.go
index 9e56666fe..2236374e5 100644
--- a/pkg/commands/git_commands/branch_loader_test.go
+++ b/pkg/commands/git_commands/branch_loader_test.go
@@ -25,89 +25,101 @@ func TestObtainBranch(t *testing.T) {
 	scenarios := []scenario{
 		{
 			testName:                 "TrimHeads",
-			input:                    []string{"", "heads/a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "heads/a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
-				Name:       "a_branch",
-				Pushables:  "?",
-				Pullables:  "?",
-				Head:       false,
-				Subject:    "subject",
-				CommitHash: "123",
+				Name:          "a_branch",
+				AheadForPull:  "?",
+				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
+				Head:          false,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 		{
 			testName:                 "NoUpstream",
-			input:                    []string{"", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
-				Name:       "a_branch",
-				Pushables:  "?",
-				Pullables:  "?",
-				Head:       false,
-				Subject:    "subject",
-				CommitHash: "123",
+				Name:          "a_branch",
+				AheadForPull:  "?",
+				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
+				Head:          false,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 		{
 			testName:                 "IsHead",
-			input:                    []string{"*", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"*", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
-				Name:       "a_branch",
-				Pushables:  "?",
-				Pullables:  "?",
-				Head:       true,
-				Subject:    "subject",
-				CommitHash: "123",
+				Name:          "a_branch",
+				AheadForPull:  "?",
+				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
+				Head:          true,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 		{
 			testName:                 "IsBehindAndAhead",
-			input:                    []string{"", "a_branch", "a_remote/a_branch", "[behind 2, ahead 3]", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "a_remote/a_branch", "[behind 2, ahead 3]", "[behind 2, ahead 3]", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
-				Name:       "a_branch",
-				Pushables:  "3",
-				Pullables:  "2",
-				Head:       false,
-				Subject:    "subject",
-				CommitHash: "123",
+				Name:          "a_branch",
+				AheadForPull:  "3",
+				BehindForPull: "2",
+				AheadForPush:  "3",
+				BehindForPush: "2",
+				Head:          false,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 		{
 			testName:                 "RemoteBranchIsGone",
-			input:                    []string{"", "a_branch", "a_remote/a_branch", "[gone]", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "a_remote/a_branch", "[gone]", "[gone]", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
-				Name:         "a_branch",
-				UpstreamGone: true,
-				Pushables:    "?",
-				Pullables:    "?",
-				Head:         false,
-				Subject:      "subject",
-				CommitHash:   "123",
+				Name:          "a_branch",
+				UpstreamGone:  true,
+				AheadForPull:  "?",
+				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
+				Head:          false,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 		{
 			testName:                 "WithCommitDateAsRecency",
-			input:                    []string{"", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: true,
 			expectedBranch: &models.Branch{
-				Name:       "a_branch",
-				Recency:    "2h",
-				Pushables:  "?",
-				Pullables:  "?",
-				Head:       false,
-				Subject:    "subject",
-				CommitHash: "123",
+				Name:          "a_branch",
+				Recency:       "2h",
+				AheadForPull:  "?",
+				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
+				Head:          false,
+				Subject:       "subject",
+				CommitHash:    "123",
 			},
 		},
 	}
 
 	for _, s := range scenarios {
 		t.Run(s.testName, func(t *testing.T) {
-			branch := obtainBranch(s.input, s.storeCommitDateAsRecency)
+			branch := obtainBranch(s.input, s.storeCommitDateAsRecency, true)
 			assert.EqualValues(t, s.expectedBranch, branch)
 		})
 	}
diff --git a/pkg/commands/models/branch.go b/pkg/commands/models/branch.go
index c5fcfdaed..25d806fca 100644
--- a/pkg/commands/models/branch.go
+++ b/pkg/commands/models/branch.go
@@ -10,10 +10,14 @@ type Branch struct {
 	DisplayName string
 	// indicator of when the branch was last checked out e.g. '2d', '3m'
 	Recency string
-	// how many commits ahead we are from the remote branch (how many commits we can push)
-	Pushables string
+	// how many commits ahead we are from the remote branch (how many commits we can push, assuming we push to our tracked remote branch)
+	AheadForPull string
 	// how many commits behind we are from the remote branch (how many commits we can pull)
-	Pullables string
+	BehindForPull string
+	// how many commits ahead we are from the branch we're pushing to (which might not be the same as our upstream branch in a triangular workflow)
+	AheadForPush string
+	// how many commits behind we are from the branch we're pushing to (which might not be the same as our upstream branch in a triangular workflow)
+	BehindForPush string
 	// whether the remote branch is 'gone' i.e. we're tracking a remote branch that has been deleted
 	UpstreamGone bool
 	// whether this is the current branch. Exactly one branch should have this be true
@@ -80,26 +84,30 @@ func (b *Branch) IsTrackingRemote() bool {
 // we know that the remote branch is not stored locally based on our pushable/pullable
 // count being question marks.
 func (b *Branch) RemoteBranchStoredLocally() bool {
-	return b.IsTrackingRemote() && b.Pushables != "?" && b.Pullables != "?"
+	return b.IsTrackingRemote() && b.AheadForPull != "?" && b.BehindForPull != "?"
 }
 
 func (b *Branch) RemoteBranchNotStoredLocally() bool {
-	return b.IsTrackingRemote() && b.Pushables == "?" && b.Pullables == "?"
+	return b.IsTrackingRemote() && b.AheadForPull == "?" && b.BehindForPull == "?"
 }
 
 func (b *Branch) MatchesUpstream() bool {
-	return b.RemoteBranchStoredLocally() && b.Pushables == "0" && b.Pullables == "0"
+	return b.RemoteBranchStoredLocally() && b.AheadForPull == "0" && b.BehindForPull == "0"
 }
 
-func (b *Branch) HasCommitsToPush() bool {
-	return b.RemoteBranchStoredLocally() && b.Pushables != "0"
+func (b *Branch) IsAheadForPull() bool {
+	return b.RemoteBranchStoredLocally() && b.AheadForPull != "0"
 }
 
-func (b *Branch) HasCommitsToPull() bool {
-	return b.RemoteBranchStoredLocally() && b.Pullables != "0"
+func (b *Branch) IsBehindForPull() bool {
+	return b.RemoteBranchStoredLocally() && b.BehindForPull != "0"
+}
+
+func (b *Branch) IsBehindForPush() bool {
+	return b.BehindForPush != "" && b.BehindForPush != "0"
 }
 
 // for when we're in a detached head state
 func (b *Branch) IsRealBranch() bool {
-	return b.Pushables != "" && b.Pullables != ""
+	return b.AheadForPull != "" && b.BehindForPull != ""
 }
diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go
index b08ddd0cd..d7faa7811 100644
--- a/pkg/gui/controllers/branches_controller.go
+++ b/pkg/gui/controllers/branches_controller.go
@@ -620,7 +620,7 @@ func (self *BranchesController) fastForward(branch *models.Branch) error {
 	if !branch.RemoteBranchStoredLocally() {
 		return errors.New(self.c.Tr.FwdNoLocalUpstream)
 	}
-	if branch.HasCommitsToPush() {
+	if branch.IsAheadForPull() {
 		return errors.New(self.c.Tr.FwdCommitsToPush)
 	}
 
diff --git a/pkg/gui/controllers/sync_controller.go b/pkg/gui/controllers/sync_controller.go
index 403f31d94..7d7ca9eed 100644
--- a/pkg/gui/controllers/sync_controller.go
+++ b/pkg/gui/controllers/sync_controller.go
@@ -87,10 +87,10 @@ func (self *SyncController) branchCheckedOut(f func(*models.Branch) error) func(
 }
 
 func (self *SyncController) push(currentBranch *models.Branch) error {
-	// if we have pullables we'll ask if the user wants to force push
+	// if we are behind our upstream branch we'll ask if the user wants to force push
 	if currentBranch.IsTrackingRemote() {
 		opts := pushOpts{}
-		if currentBranch.HasCommitsToPull() {
+		if currentBranch.IsBehindForPush() {
 			return self.requestToForcePush(currentBranch, opts)
 		} else {
 			return self.pushAux(currentBranch, opts)
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index 0abf2d4cd..406a580d5 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -196,11 +196,11 @@ func BranchStatus(
 	}
 
 	result := ""
-	if branch.HasCommitsToPush() {
-		result = fmt.Sprintf("↑%s", branch.Pushables)
+	if branch.IsAheadForPull() {
+		result = fmt.Sprintf("↑%s", branch.AheadForPull)
 	}
-	if branch.HasCommitsToPull() {
-		result = fmt.Sprintf("%s↓%s", result, branch.Pullables)
+	if branch.IsBehindForPull() {
+		result = fmt.Sprintf("%s↓%s", result, branch.BehindForPull)
 	}
 
 	return result
diff --git a/pkg/gui/presentation/branches_test.go b/pkg/gui/presentation/branches_test.go
index 250b143e3..cf2f1d994 100644
--- a/pkg/gui/presentation/branches_test.go
+++ b/pkg/gui/presentation/branches_test.go
@@ -58,8 +58,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				Name:           "branch_name",
 				Recency:        "1m",
 				UpstreamRemote: "origin",
-				Pushables:      "0",
-				Pullables:      "0",
+				AheadForPull:   "0",
+				BehindForPull:  "0",
 			},
 			itemOperation:        types.ItemOperationNone,
 			fullDescription:      false,
@@ -73,8 +73,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				Name:           "branch_name",
 				Recency:        "1m",
 				UpstreamRemote: "origin",
-				Pushables:      "3",
-				Pullables:      "5",
+				AheadForPull:   "3",
+				BehindForPull:  "5",
 			},
 			itemOperation:        types.ItemOperationNone,
 			fullDescription:      false,
@@ -99,8 +99,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				CommitHash:     "1234567890",
 				UpstreamRemote: "origin",
 				UpstreamBranch: "branch_name",
-				Pushables:      "0",
-				Pullables:      "0",
+				AheadForPull:   "0",
+				BehindForPull:  "0",
 				Subject:        "commit title",
 			},
 			itemOperation:        types.ItemOperationNone,
@@ -144,8 +144,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				Name:           "branch_name",
 				Recency:        "1m",
 				UpstreamRemote: "origin",
-				Pushables:      "0",
-				Pullables:      "0",
+				AheadForPull:   "0",
+				BehindForPull:  "0",
 			},
 			itemOperation:        types.ItemOperationNone,
 			fullDescription:      false,
@@ -159,8 +159,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				Name:           "branch_name",
 				Recency:        "1m",
 				UpstreamRemote: "origin",
-				Pushables:      "3",
-				Pullables:      "5",
+				AheadForPull:   "3",
+				BehindForPull:  "5",
 			},
 			itemOperation:        types.ItemOperationNone,
 			fullDescription:      false,
@@ -212,8 +212,8 @@ func Test_getBranchDisplayStrings(t *testing.T) {
 				CommitHash:     "1234567890",
 				UpstreamRemote: "origin",
 				UpstreamBranch: "branch_name",
-				Pushables:      "0",
-				Pullables:      "0",
+				AheadForPull:   "0",
+				BehindForPull:  "0",
 				Subject:        "commit title",
 			},
 			itemOperation:        types.ItemOperationNone,
diff --git a/pkg/gui/services/custom_commands/models.go b/pkg/gui/services/custom_commands/models.go
new file mode 100644
index 000000000..261bace45
--- /dev/null
+++ b/pkg/gui/services/custom_commands/models.go
@@ -0,0 +1,100 @@
+package custom_commands
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/stefanhaller/git-todo-parser/todo"
+)
+
+// We create shims for all the model classes in order to get a more stable API
+// for custom commands. At the moment these are almost identical to the model
+// classes, but this allows us to add "private" fields to the model classes that
+// we don't want to expose to custom commands, or rename a model field to a
+// better name without breaking people's custom commands. In such a case we add
+// the new, better name to the shim but keep the old one for backwards
+// compatibility. We already did this for Commit.Sha, which was renamed to Hash.
+
+type Commit struct {
+	Hash          string // deprecated: use Sha
+	Sha           string
+	Name          string
+	Status        models.CommitStatus
+	Action        todo.TodoCommand
+	Tags          []string
+	ExtraInfo     string
+	AuthorName    string
+	AuthorEmail   string
+	UnixTimestamp int64
+	Divergence    models.Divergence
+	Parents       []string
+}
+
+type File struct {
+	Name                    string
+	PreviousName            string
+	HasStagedChanges        bool
+	HasUnstagedChanges      bool
+	Tracked                 bool
+	Added                   bool
+	Deleted                 bool
+	HasMergeConflicts       bool
+	HasInlineMergeConflicts bool
+	DisplayString           string
+	ShortStatus             string
+	IsWorktree              bool
+}
+
+type Branch struct {
+	Name           string
+	DisplayName    string
+	Recency        string
+	Pushables      string // deprecated: use AheadForPull
+	Pullables      string // deprecated: use BehindForPull
+	AheadForPull   string
+	BehindForPull  string
+	AheadForPush   string
+	BehindForPush  string
+	UpstreamGone   bool
+	Head           bool
+	DetachedHead   bool
+	UpstreamRemote string
+	UpstreamBranch string
+	Subject        string
+	CommitHash     string
+}
+
+type RemoteBranch struct {
+	Name       string
+	RemoteName string
+}
+
+type Remote struct {
+	Name     string
+	Urls     []string
+	Branches []*RemoteBranch
+}
+
+type Tag struct {
+	Name    string
+	Message string
+}
+
+type StashEntry struct {
+	Index   int
+	Recency string
+	Name    string
+}
+
+type CommitFile struct {
+	Name         string
+	ChangeStatus string
+}
+
+type Worktree struct {
+	IsMain        bool
+	IsCurrent     bool
+	Path          string
+	IsPathMissing bool
+	GitDir        string
+	Branch        string
+	Name          string
+}
diff --git a/pkg/gui/services/custom_commands/session_state_loader.go b/pkg/gui/services/custom_commands/session_state_loader.go
index 6a3068df8..6f39c5f8c 100644
--- a/pkg/gui/services/custom_commands/session_state_loader.go
+++ b/pkg/gui/services/custom_commands/session_state_loader.go
@@ -3,7 +3,7 @@ package custom_commands
 import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
-	"github.com/stefanhaller/git-todo-parser/todo"
+	"github.com/samber/lo"
 )
 
 // loads the session state at the time that a custom command is invoked, for use
@@ -20,22 +20,7 @@ func NewSessionStateLoader(c *helpers.HelperCommon, refsHelper *helpers.RefsHelp
 	}
 }
 
-type Commit struct {
-	Hash          string
-	Sha           string
-	Name          string
-	Status        models.CommitStatus
-	Action        todo.TodoCommand
-	Tags          []string
-	ExtraInfo     string
-	AuthorName    string
-	AuthorEmail   string
-	UnixTimestamp int64
-	Divergence    models.Divergence
-	Parents       []string
-}
-
-func commitWrapperFromModelCommit(commit *models.Commit) *Commit {
+func commitShimFromModelCommit(commit *models.Commit) *Commit {
 	if commit == nil {
 		return nil
 	}
@@ -56,39 +41,160 @@ func commitWrapperFromModelCommit(commit *models.Commit) *Commit {
 	}
 }
 
+func fileShimFromModelFile(file *models.File) *File {
+	if file == nil {
+		return nil
+	}
+
+	return &File{
+		Name:                    file.Name,
+		PreviousName:            file.PreviousName,
+		HasStagedChanges:        file.HasStagedChanges,
+		HasUnstagedChanges:      file.HasUnstagedChanges,
+		Tracked:                 file.Tracked,
+		Added:                   file.Added,
+		Deleted:                 file.Deleted,
+		HasMergeConflicts:       file.HasMergeConflicts,
+		HasInlineMergeConflicts: file.HasInlineMergeConflicts,
+		DisplayString:           file.DisplayString,
+		ShortStatus:             file.ShortStatus,
+		IsWorktree:              file.IsWorktree,
+	}
+}
+
+func branchShimFromModelBranch(branch *models.Branch) *Branch {
+	if branch == nil {
+		return nil
+	}
+
+	return &Branch{
+		Name:           branch.Name,
+		DisplayName:    branch.DisplayName,
+		Recency:        branch.Recency,
+		Pushables:      branch.AheadForPull,
+		Pullables:      branch.BehindForPull,
+		AheadForPull:   branch.AheadForPull,
+		BehindForPull:  branch.BehindForPull,
+		AheadForPush:   branch.AheadForPush,
+		BehindForPush:  branch.BehindForPush,
+		UpstreamGone:   branch.UpstreamGone,
+		Head:           branch.Head,
+		DetachedHead:   branch.DetachedHead,
+		UpstreamRemote: branch.UpstreamRemote,
+		UpstreamBranch: branch.UpstreamBranch,
+		Subject:        branch.Subject,
+		CommitHash:     branch.CommitHash,
+	}
+}
+
+func remoteBranchShimFromModelRemoteBranch(remoteBranch *models.RemoteBranch) *RemoteBranch {
+	if remoteBranch == nil {
+		return nil
+	}
+
+	return &RemoteBranch{
+		Name:       remoteBranch.Name,
+		RemoteName: remoteBranch.RemoteName,
+	}
+}
+
+func remoteShimFromModelRemote(remote *models.Remote) *Remote {
+	if remote == nil {
+		return nil
+	}
+
+	return &Remote{
+		Name: remote.Name,
+		Urls: remote.Urls,
+		Branches: lo.Map(remote.Branches, func(branch *models.RemoteBranch, _ int) *RemoteBranch {
+			return remoteBranchShimFromModelRemoteBranch(branch)
+		}),
+	}
+}
+
+func tagShimFromModelRemote(tag *models.Tag) *Tag {
+	if tag == nil {
+		return nil
+	}
+
+	return &Tag{
+		Name:    tag.Name,
+		Message: tag.Message,
+	}
+}
+
+func stashEntryShimFromModelRemote(stashEntry *models.StashEntry) *StashEntry {
+	if stashEntry == nil {
+		return nil
+	}
+
+	return &StashEntry{
+		Index:   stashEntry.Index,
+		Recency: stashEntry.Recency,
+		Name:    stashEntry.Name,
+	}
+}
+
+func commitFileShimFromModelRemote(commitFile *models.CommitFile) *CommitFile {
+	if commitFile == nil {
+		return nil
+	}
+
+	return &CommitFile{
+		Name:         commitFile.Name,
+		ChangeStatus: commitFile.ChangeStatus,
+	}
+}
+
+func worktreeShimFromModelRemote(worktree *models.Worktree) *Worktree {
+	if worktree == nil {
+		return nil
+	}
+
+	return &Worktree{
+		IsMain:        worktree.IsMain,
+		IsCurrent:     worktree.IsCurrent,
+		Path:          worktree.Path,
+		IsPathMissing: worktree.IsPathMissing,
+		GitDir:        worktree.GitDir,
+		Branch:        worktree.Branch,
+		Name:          worktree.Name,
+	}
+}
+
 // SessionState captures the current state of the application for use in custom commands
 type SessionState struct {
 	SelectedLocalCommit    *Commit
 	SelectedReflogCommit   *Commit
 	SelectedSubCommit      *Commit
-	SelectedFile           *models.File
+	SelectedFile           *File
 	SelectedPath           string
-	SelectedLocalBranch    *models.Branch
-	SelectedRemoteBranch   *models.RemoteBranch
-	SelectedRemote         *models.Remote
-	SelectedTag            *models.Tag
-	SelectedStashEntry     *models.StashEntry
-	SelectedCommitFile     *models.CommitFile
+	SelectedLocalBranch    *Branch
+	SelectedRemoteBranch   *RemoteBranch
+	SelectedRemote         *Remote
+	SelectedTag            *Tag
+	SelectedStashEntry     *StashEntry
+	SelectedCommitFile     *CommitFile
 	SelectedCommitFilePath string
-	SelectedWorktree       *models.Worktree
-	CheckedOutBranch       *models.Branch
+	SelectedWorktree       *Worktree
+	CheckedOutBranch       *Branch
 }
 
 func (self *SessionStateLoader) call() *SessionState {
 	return &SessionState{
-		SelectedFile:           self.c.Contexts().Files.GetSelectedFile(),
+		SelectedFile:           fileShimFromModelFile(self.c.Contexts().Files.GetSelectedFile()),
 		SelectedPath:           self.c.Contexts().Files.GetSelectedPath(),
-		SelectedLocalCommit:    commitWrapperFromModelCommit(self.c.Contexts().LocalCommits.GetSelected()),
-		SelectedReflogCommit:   commitWrapperFromModelCommit(self.c.Contexts().ReflogCommits.GetSelected()),
-		SelectedLocalBranch:    self.c.Contexts().Branches.GetSelected()