UBERF-7118: Fix upgrade/refresh on reconnect (#5704)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
hcengineering · May 30, 2024 · 42d8c45 · 42d8c45
1 parent 3e913c9
commit 42d8c45
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/server/ws/src/server_http.ts b/server/ws/src/server_http.ts
@@ -276,10 +276,10 @@ export function startHttpServer (
 
     if (webSocketData.session instanceof Promise) {
       void webSocketData.session.then((s) => {
-        if ('upgrade' in s || 'error' in s) {
-          if ('error' in s) {
-            ctx.error('error', { error: s.error?.message, stack: s.error?.stack })
-          }
+        if ('error' in s) {
+          ctx.error('error', { error: s.error?.message, stack: s.error?.stack })
+        }
+        if ('upgrade' in s) {
           void cs
             .send(ctx, { id: -1, result: { state: 'upgrading', stats: (s as any).upgradeInfo } }, false, false)
             .then(() => {

diff --git a/server/ws/src/server_u.ts b/server/ws/src/server_u.ts
@@ -126,7 +126,8 @@ export function startUWebsocketServer (
         mode: data.payload.extra?.mode,
         model: data.payload.extra?.model
       }
-      data.connectionSocket = createWebSocketClientSocket(wrData, ws, data)
+      const cs = createWebSocketClientSocket(wrData, ws, data)
+      data.connectionSocket = cs
 
       data.session = sessions.addSession(
         ctx,
@@ -138,6 +139,21 @@ export function startUWebsocketServer (
         undefined,
         accountsUrl
       )
+
+      if (data.session instanceof Promise) {
+        void data.session.then((s) => {
+          if ('error' in s) {
+            ctx.error('error', { error: s.error?.message, stack: s.error?.stack })
+          }
+          if ('upgrade' in s) {
+            void cs
+              .send(ctx, { id: -1, result: { state: 'upgrading', stats: (s as any).upgradeInfo } }, false, false)
+              .then(() => {
+                cs.close()
+              })
+          }
+        })
+      }
     },
     message: (ws, message, isBinary) => {
       const data = ws.getUserData()