From 346cfbd4734fa4defecad8a6c65afa6dd135df8f Mon Sep 17 00:00:00 2001
From: Piyush Jain <piyushjain@duck.com>
Date: Thu, 8 Jun 2023 16:41:43 -0700
Subject: [PATCH] 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 <93281816+JasonWeill@users.noreply.github.com>

---------

Co-authored-by: Jason Weill <93281816+JasonWeill@users.noreply.github.com>
---
 packages/jupyter-ai/src/chat_handler.ts | 68 +++++++++++++++----------
 1 file changed, 41 insertions(+), 27 deletions(-)

diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts
index 2c4e2b87e..8c324e0ab 100644
--- 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)[] = [];