feat(session): retrieve session after connection is down

go.mod

@@ -16,6 +16,7 @@ require (
 require (
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
 	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
+	github.com/google/uuid v1.3.1
 	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/iancoleman/orderedmap v0.0.0-20190318233801-ac98e3ecb4b0 // indirect
 	github.com/labstack/echo/v4 v4.8.0 // indirect

go.sum

@@ -10,6 +10,8 @@ github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keL
 github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
+github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
+github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
 github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/iancoleman/orderedmap v0.0.0-20190318233801-ac98e3ecb4b0 h1:i462o439ZjprVSFSZLZxcsoAe592sZB1rci2Z8j4wdk=
@@ -52,10 +54,12 @@ golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211103235746-7861aae1554b h1:1VkfZQv42XQlA/jchYumAnv1UPo6RgF9rJFkTgZIxO4=
 golang.org/x/sys v0.0.0-20211103235746-7861aae1554b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 h1:Hir2P/De0WpUhtrKGGjvSb2YxUgyZ7EFOSLIcSSpiwE=
 golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/ui/events.go

@@ -55,3 +55,8 @@ type UpdateCheckedOptionsParams[T []string | string] struct {
 type RunParams struct {
 	DieWithConn any
 }
+
+type SessionProtocolParams struct {
+	ClientId        string `json:"clientId"`
+	ServerSignature string `json:"serverSignature"`
+}

pkg/ui/flat/events.go

@@ -5,6 +5,7 @@ import (
 	"CLI2UI/pkg/ui"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/labstack/gommon/log"
 	"github.com/yuyz0112/sunmao-ui-go-binding/pkg/runtime"
 )
@@ -76,6 +77,8 @@ func (u UI) registerEvents() {
 			return err
 		}
 
+		sess.CurrentConnId = &connId
+
 		go func() {
 			for {
 				select {
@@ -103,7 +106,6 @@ func (u UI) registerEvents() {
 			}
 
 			for sess.Exec.State.IsRunning {
-				// TODO(xinxi.guo): this can be extended to send more useful messages
 				err := u.Runtime.Ping(&connId, "Ping")
 
 				// this fails when a WebSocket connection drops **loudly**
@@ -148,8 +150,23 @@ func (u UI) registerEvents() {
 		return nil
 	})
 
-	u.Runtime.Handle("EstablishedConnection", func(m *runtime.Message, connId int) error {
-		ui.GetOrCreateSession(0, u.FormTemplates, connId)
+	u.Runtime.Handle("ConnectionEstablished", func(m *runtime.Message, connId int) error {
+		p := ui.ToStruct[ui.SessionProtocolParams](m.Params)
+		if p.ServerSignature != ui.ServerSignature {
+			clientId := uuid.NewString()
+			err := u.Runtime.Ping(&connId, ui.SessionProtocolParams{
+				ClientId:        clientId,
+				ServerSignature: ui.ServerSignature,
+			})
+			if err != nil {
+				return err
+			}
+			ui.UpdateConnIdToClientId(connId, clientId)
+		}
+		s := ui.GetOrCreateSession(0, u.FormTemplates, connId)
+		*s.CurrentConnId = connId
+		// TODO(xinxi.guo): UI won't be updated according to the form, thus clear the form to be consistent
+		s.Form.Clear()
 		return nil
 	})
 }

pkg/ui/naive/events.go

@@ -4,6 +4,7 @@ import (
 	"CLI2UI/pkg/ui"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/labstack/gommon/log"
 	"github.com/yuyz0112/sunmao-ui-go-binding/pkg/runtime"
 )
@@ -69,6 +70,8 @@ func (u UI) registerEvents() {
 			return err
 		}
 
+		sess.CurrentConnId = &connId
+
 		go func() {
 			for {
 				select {
@@ -126,8 +129,23 @@ func (u UI) registerEvents() {
 		return nil
 	})
 
-	u.Runtime.Handle("EstablishedConnection", func(m *runtime.Message, connId int) error {
-		ui.GetOrCreateSession(0, u.FormTemplates, connId)
+	u.Runtime.Handle("ConnectionEstablished", func(m *runtime.Message, connId int) error {
+		p := ui.ToStruct[ui.SessionProtocolParams](m.Params)
+		if p.ServerSignature != ui.ServerSignature {
+			clientId := uuid.NewString()
+			err := u.Runtime.Ping(&connId, ui.SessionProtocolParams{
+				ClientId:        clientId,
+				ServerSignature: ui.ServerSignature,
+			})
+			if err != nil {
+				return err
+			}
+			ui.UpdateConnIdToClientId(connId, clientId)
+		}
+		s := ui.GetOrCreateSession(0, u.FormTemplates, connId)
+		*s.CurrentConnId = connId
+		// TODO(xinxi.guo): UI won't be updated according to the form, thus clear the form to be consistent
+		s.Form.Clear()
 		return nil
 	})
 

pkg/ui/session.go

@@ -3,31 +3,42 @@ package ui
 import (
 	"CLI2UI/pkg/config"
 	"CLI2UI/pkg/executor"
+
+	"github.com/google/uuid"
 )
 
-var sessions = map[int]*session{}
+var ServerSignature = uuid.NewString()
+var sessions = map[string]*session{}
+var connIdToClientId = map[int]string{}
 
 type session struct {
-	Form        *config.Form
-	Exec        *executor.Executor
-	HeartbeatCh chan struct{}
-	cliIndex    int
+	Form          *config.Form
+	Exec          *executor.Executor
+	HeartbeatCh   chan struct{}
+	cliIndex      int
+	CurrentConnId *int
+}
+
+func UpdateConnIdToClientId(connId int, clientId string) {
+	connIdToClientId[connId] = clientId
 }
 
 func GetOrCreateSession(cliIndex int, templates []*config.Form, connId int) *session {
-	s, ok := sessions[connId]
+	sId := connIdToClientId[connId]
+	s, ok := sessions[sId]
 	if !ok {
 		f := templates[cliIndex].Clone()
 		hbCh := make(chan struct{})
 		exec := executor.NewExecutor()
 
 		s = &session{
-			Form:        f,
-			Exec:        &exec,
-			HeartbeatCh: hbCh,
-			cliIndex:    cliIndex,
+			Form:          f,
+			Exec:          &exec,
+			HeartbeatCh:   hbCh,
+			cliIndex:      cliIndex,
+			CurrentConnId: &connId,
 		}
-		sessions[connId] = s
+		sessions[sId] = s
 	}
 
 	if cliIndex != s.cliIndex {

ui/src/application/index.tsx

@@ -15,36 +15,45 @@ export function renderApp(options: MainOptions) {
     modulesPatch,
   } = options;
   const ws = new WebSocket(wsUrl);
-  ws.onopen = () => {
+
+  ws.addEventListener('open', () => {
     ws.send(JSON.stringify({
       type: 'Action',
-      handler: 'EstablishedConnection',
-      params: {},
+      handler: 'ConnectionEstablished',
+      params: {
+        clientId: localStorage.getItem('clientId'),
+        serverSignature: localStorage.getItem('serverSignature')
+      },
       store: null
     }))
-  };
-  ws.onclose = () => {
+  });
+
+  ws.addEventListener('close', () => {
     if (reloadWhenWsDisconnected) {
       setTimeout(() => {
         window.location.reload();
       }, 1500);
     }
-  };
+  });
 
   ws.addEventListener('message', e => {
     const data = JSON.parse(e.data)
 
-    if (data.handler !== 'Heartbeat') {
-      return
-    }
+    if (data.handler === 'Heartbeat') {
+      if (data.params.hasOwnProperty('serverSignature')) {
+        localStorage.setItem('clientId', data.params.clientId)
+        localStorage.setItem('serverSignature', data.params.serverSignature)
+        // FIXME(xinxi.guo): this is a workaround, right after clientId is set, websocket does not work as expected
+        window.location.reload();
+      }
 
-    ws.send(JSON.stringify({
-      type: 'Action',
-      handler: 'Heartbeat',
-      // TODO(xinxi.guo): this can be extended later
-      params: "Pong",
-      store: null
-    }))
+      ws.send(JSON.stringify({
+        type: 'Action',
+        handler: 'Heartbeat',
+        params: "Pong",
+        store: null
+      }))
+    }
   })
 
   ReactDOM.render(