From 8a0df2db7e52697b75c3044d632e2ca9946ebc8d Mon Sep 17 00:00:00 2001
From: San Nguyen <vinhsannguyen91@gmail.com>
Date: Sat, 12 Oct 2024 08:36:27 +0900
Subject: [PATCH] add parameter retryUpgradeWebSocket to useChatSession

Signed-off-by: San Nguyen <vinhsannguyen91@gmail.com>
---
 libs/react-client/src/useChatSession.ts | 27 +++++++++++++++----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/libs/react-client/src/useChatSession.ts b/libs/react-client/src/useChatSession.ts
index 81cdc565e8..f3a13c78fd 100644
--- a/libs/react-client/src/useChatSession.ts
+++ b/libs/react-client/src/useChatSession.ts
@@ -86,10 +86,12 @@ const useChatSession = () => {
   const _connect = useCallback(
     ({
       userEnv,
-      accessToken
+      accessToken,
+      retryUpgradeWebSocket = false
     }: {
       userEnv: Record<string, string>;
       accessToken?: string;
+      retryUpgradeWebSocket?: boolean;
     }) => {
       const { protocol, host, pathname } = new URL(client.httpEndpoint);
       const uri = `${protocol}//${host}`;
@@ -119,16 +121,19 @@ const useChatSession = () => {
         };
       });
 
-      // https://socket.io/docs/v4/how-it-works/#upgrade-mechanism
-      // Retry upgrading to websocket when error
-      const engine = socket.io.engine;
-      engine.on('upgradeError', () => {
-        setTimeout(() => {
-          socket.removeAllListeners();
-          socket.close();
-          _connect({ userEnv, accessToken });
-        }, 3000);
-      });
+      if (retryUpgradeWebSocket) {
+        // https://socket.io/docs/v4/how-it-works/#upgrade-mechanism
+        // Retry upgrading to websocket when error
+        // This happens sometimes when user is using soft session affinity (like Istio)
+        const engine = socket.io.engine;
+        engine.on('upgradeError', () => {
+          setTimeout(() => {
+            socket.removeAllListeners();
+            socket.close();
+            _connect({ userEnv, accessToken });
+          }, 3000);
+        });
+      }
 
       socket.on('connect', () => {
         socket.emit('connection_successful');