From 42d8c4582ca0362511d4eaac28859081663bf66c Mon Sep 17 00:00:00 2001
From: Andrey Sobolev <haiodo@users.noreply.github.com>
Date: Fri, 31 May 2024 00:31:12 +0700
Subject: [PATCH] UBERF-7118: Fix upgrade/refresh on reconnect (#5704)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
---
 server/ws/src/server_http.ts |  8 ++++----
 server/ws/src/server_u.ts    | 18 +++++++++++++++++-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/server/ws/src/server_http.ts b/server/ws/src/server_http.ts
index b6186fa281..2ff771a822 100644
--- 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
index dbae1c5b51..2af66937fb 100644
--- 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()