diff --git a/libs/react-client/README.md b/libs/react-client/README.md index 9cbce8f8d1..4f2834d1b5 100644 --- a/libs/react-client/README.md +++ b/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 () => { diff --git a/libs/react-client/src/useChatSession.ts b/libs/react-client/src/useChatSession.ts index fc1de3fbd5..94f8d74eb5 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, + requireWebSocket = false }: { userEnv: Record; 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);