Optionally require websocket connection in react-client (#1379)

Signed-off-by: San Nguyen <[email protected]>

libs/react-client/README.md

@@ -42,8 +42,8 @@ This hook is responsible for managing the chat session's connection to the WebSo
 
 #### Methods
 
-- `connect`: Establishes a connection to the WebSocket server.
-- `disconnect`: Disconnects from the WebSocket server.
+- `connect`: Establishes a connection to the SocketIO server.
+- `disconnect`: Disconnects from the SocketIO server.
 - `setChatProfile`: Sets the chat profile state.
 
 #### Example
@@ -60,7 +60,8 @@ const ChatComponent = () => {
       userEnv: {
         /* user environment variables */
       },
-      accessToken: 'Bearer YOUR_ACCESS_TOKEN' // Optional Chainlit auth token
+      accessToken: 'Bearer YOUR_ACCESS_TOKEN', // Optional Chainlit auth token
+      requireWebSocket: true // Optional, require WebSocket upgrade to be successful before user can interact with the chat bot. Will retry upgrade request every 500ms until successful. Please note if your server is behind a proxy, you will have to configure it to accept websocket upgrade request, otherwise users won't be able to interact with the app. You can check an example using nginx proxy here: https://nginx.org/en/docs/http/websocket.html. Default to false.
     });
 
     return () => {

libs/react-client/src/useChatSession.ts

@@ -86,10 +86,12 @@ const useChatSession = () => {
   const _connect = useCallback(
     ({
       userEnv,
-      accessToken
+      accessToken,
+      requireWebSocket = false
     }: {
       userEnv: Record<string, string>;
       accessToken?: string;
+      requireWebSocket?: boolean;
     }) => {
       const { protocol, host, pathname } = new URL(client.httpEndpoint);
       const uri = `${protocol}//${host}`;
@@ -119,14 +121,49 @@ const useChatSession = () => {
         };
       });
 
-      socket.on('connect', () => {
+      const onConnect = () => {
         socket.emit('connection_successful');
         setSession((s) => ({ ...s!, error: false }));
-      });
+      };
 
-      socket.on('connect_error', (_) => {
+      const onConnectError = () => {
         setSession((s) => ({ ...s!, error: true }));
-      });
+      };
+
+      // https://socket.io/docs/v4/how-it-works/#upgrade-mechanism
+      // Require WebSocket when connecting to backend
+      if (requireWebSocket) {
+        // https://socket.io/docs/v4/client-socket-instance/#socketio
+        // 'connect' event is emitted when the underlying connection is established with polling transport
+        // 'upgrade' event is emitted when the underlying connection is upgraded to WebSocket and polling request is stopped.
+        const engine = socket.io.engine;
+        // https://github.com/socketio/socket.io/tree/main/packages/engine.io-client#events
+        engine.once('upgrade', () => {
+          // Set session on connect event, otherwise user can not interact with text input UI.
+          // Upgrade event is required to make sure user won't interact with the session before websocket upgrade success
+          socket.on('connect', onConnect);
+        });
+        // Socket.io will not retry upgrade request.
+        // Retry upgrade to websocket when error can only be done via reconnect.
+        // This will not be an issue for users if they are using persistent sticky session.
+        // In case they are using soft session affinity like Istio, then sometimes upgrade request will fail
+        engine.once('upgradeError', () => {
+          onConnectError();
+          setTimeout(() => {
+            socket.removeAllListeners();
+            socket.close();
+            _connect({
+              userEnv,
+              accessToken,
+              requireWebSocket
+            });
+          }, 500);
+        });
+      } else {
+        socket.on('connect', onConnect);
+      }
+
+      socket.on('connect_error', onConnectError);
 
       socket.on('task_start', () => {
         setLoading(true);