Creates a new websocket connection in case of disconnect (#219)

* Creates a new websocket connection in case of disconnect * Added a console statement for re-connecting * Update packages/jupyter-ai/src/chat_handler.ts Co-authored-by: Jason Weill <[email protected]> --------- Co-authored-by: Jason Weill <[email protected]>
jupyterlab · Jun 8, 2023 · 3e0238c · 3e0238c
1 parent 5be0a21
commit 3e0238c
Showing 1 changed file with 41 additions and 27 deletions.
diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts
@@ -29,34 +29,11 @@ export class ChatHandler implements IDisposable {
    * resolved when server acknowledges connection and sends the client ID. This
    * must be awaited before calling any other method.
    */
-  public initialize(): Promise<void> {
-    return new Promise<void>((resolve, reject) => {
-      if (this.isDisposed) {
-        return;
-      }
-      const { token, WebSocket, wsUrl } = this.serverSettings;
-      const url =
-        URLExt.join(wsUrl, CHAT_SERVICE_URL) +
-        (token ? `?token=${encodeURIComponent(token)}` : '');
-
-      const socket = (this._socket = new WebSocket(url));
-      socket.onerror = (e) => reject(e);
-      socket.onmessage = msg =>
-        msg.data && this._onMessage(JSON.parse(msg.data));
-
-      const listenForConnection = (message: AiService.Message) => {
-        if (message.type !== 'connection') {
-          return;
-        }
-        this.id = message.client_id;
-        resolve();
-        this.removeListener(listenForConnection);
-      };
-
-      this.addListener(listenForConnection);
-    });
+  public async initialize() {
+    await this._initialize();
   }
 
+
   /**
    * Sends a message across the WebSocket. Promise resolves to the message ID
    * when the server sends the same message back, acknowledging receipt.
@@ -116,7 +93,6 @@ export class ChatHandler implements IDisposable {
       return;
     }
     this._isDisposed = true;
-
     this._listeners = [];
 
     // Clean up socket.
@@ -164,6 +140,44 @@ export class ChatHandler implements IDisposable {
     (value: AiService.AgentChatMessage) => void
   > = {};
 
+  private _onClose(reject: any) {
+    reject(new Error("Chat UI websocket disconnected"))
+    console.error("Chat UI websocket disconnected")
+    const delaySeconds = 1
+    console.info(`Will try to reconnect in ${delaySeconds} s.`)
+    setTimeout(async () => await this._initialize(), delaySeconds * 1000);
+  }
+
+  private _initialize(): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+      if (this.isDisposed) {
+        return;
+      }
+      console.log("Creating a new websocket connection for chat...");
+      const { token, WebSocket, wsUrl } = this.serverSettings;
+      const url =
+        URLExt.join(wsUrl, CHAT_SERVICE_URL) +
+        (token ? `?token=${encodeURIComponent(token)}` : '');
+
+        const socket = (this._socket = new WebSocket(url));
+        socket.onclose = () => this._onClose(reject);
+        socket.onerror = (e) => reject(e);
+        socket.onmessage = msg =>
+          msg.data && this._onMessage(JSON.parse(msg.data));
+
+        const listenForConnection = (message: AiService.Message) => {
+          if (message.type !== 'connection') {
+            return;
+          }
+          this.id = message.client_id;
+          resolve();
+          this.removeListener(listenForConnection);
+        };
+
+        this.addListener(listenForConnection);
+    });
+  }
+
   private _isDisposed = false;
   private _socket: WebSocket | null = null;
   private _listeners: ((msg: any) => void)[] = [];