From 6e6fe4d4bf15f15a86c2e2c5ff6ef1a8ecf9ef76 Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Mon, 3 Apr 2023 22:15:43 -0700
Subject: [PATCH] auto-install and reconnect flow working

---
 pkg/remote/remote.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go
index c827da494..007ab4bb5 100644
--- a/pkg/remote/remote.go
+++ b/pkg/remote/remote.go
@@ -501,6 +501,16 @@ func (msh *MShellProc) IsSudo() bool {
 	return msh.Remote.RemoteSudo
 }
 
+func (msh *MShellProc) tryAutoInstall() {
+	msh.Lock.Lock()
+	defer msh.Lock.Unlock()
+	if !msh.Remote.AutoInstall || !msh.NeedsMShellUpgrade || msh.InstallErr != nil {
+		return
+	}
+	msh.writeToPtyBuffer_nolock("trying auto-install\n")
+	go msh.RunInstall()
+}
+
 func (msh *MShellProc) GetRemoteRuntimeState() RemoteRuntimeState {
 	msh.Lock.Lock()
 	defer msh.Lock.Unlock()
@@ -1022,15 +1032,21 @@ func (msh *MShellProc) RunInstall() {
 		msh.setInstallErrorStatus(statusErr)
 		return
 	}
+	var connectMode string
 	msh.WithLock(func() {
 		msh.InstallStatus = StatusDisconnected
 		msh.InstallCancelFn = nil
 		msh.NeedsMShellUpgrade = false
 		msh.Status = StatusDisconnected
 		msh.Err = nil
+		connectMode = msh.Remote.ConnectMode
 	})
 	msh.WriteToPtyBuffer("successfully installed mshell %s to ~/.mshell\n", scbase.MShellVersion)
 	go msh.NotifyRemoteUpdate()
+	if connectMode == sstore.ConnectModeStartup || connectMode == sstore.ConnectModeAuto {
+		// the install was successful, and we don't have a manual connect mode, try to connect
+		go msh.Launch(true)
+	}
 	return
 }
 
@@ -1254,6 +1270,7 @@ func (msh *MShellProc) Launch(interactive bool) {
 	if err != nil {
 		msh.setErrorStatus(err)
 		msh.WriteToPtyBuffer("*error connecting to remote: %v\n", err)
+		go msh.tryAutoInstall()
 		return
 	}
 	msh.updateRemoteStateVars(context.Background(), msh.RemoteId, initPk)