From c46d94b44c224cbb2c973208c51509356cc21faa Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Thu, 8 Jun 2023 00:07:20 -0700 Subject: [PATCH 1/3] Creates a new websocket connection in case of disconnect --- packages/jupyter-ai/src/chat_handler.ts | 65 +++++++++++++++---------- 1 file changed, 38 insertions(+), 27 deletions(-) diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts index 2c4e2b87e..fc4f2bf5a 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 { - return new Promise((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,41 @@ 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") + setTimeout(async () => await this._initialize(), 500); + } + + private _initialize(): Promise { + return new Promise((resolve, reject) => { + if (this.isDisposed) { + return; + } + console.log("Creating a new websocket connection..."); + 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)[] = []; From d0e873daa22658b3289e9057709a574b34edc5e9 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Thu, 8 Jun 2023 15:51:39 -0700 Subject: [PATCH 2/3] Added a console statement for re-connecting --- packages/jupyter-ai/src/chat_handler.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts index fc4f2bf5a..123fe9eca 100644 --- a/packages/jupyter-ai/src/chat_handler.ts +++ b/packages/jupyter-ai/src/chat_handler.ts @@ -143,7 +143,9 @@ export class ChatHandler implements IDisposable { private _onClose(reject: any) { reject(new Error("Chat UI websocket disconnected")) console.error("Chat UI websocket disconnected") - setTimeout(async () => await this._initialize(), 500); + const delaySeconds = 1 + console.info(`Will try to reconnect in ${delaySeconds} seconds.`) + setTimeout(async () => await this._initialize(), delaySeconds * 1000); } private _initialize(): Promise { @@ -151,7 +153,7 @@ export class ChatHandler implements IDisposable { if (this.isDisposed) { return; } - console.log("Creating a new websocket connection..."); + console.log("Creating a new websocket connection for chat..."); const { token, WebSocket, wsUrl } = this.serverSettings; const url = URLExt.join(wsUrl, CHAT_SERVICE_URL) + @@ -162,6 +164,7 @@ export class ChatHandler implements IDisposable { 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; From 6ab3fad93afd45952870e0b5252db620918fc099 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Thu, 8 Jun 2023 16:12:25 -0700 Subject: [PATCH 3/3] Update packages/jupyter-ai/src/chat_handler.ts Co-authored-by: Jason Weill <93281816+JasonWeill@users.noreply.github.com> --- packages/jupyter-ai/src/chat_handler.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts index 123fe9eca..8c324e0ab 100644 --- a/packages/jupyter-ai/src/chat_handler.ts +++ b/packages/jupyter-ai/src/chat_handler.ts @@ -144,7 +144,7 @@ export class ChatHandler implements IDisposable { reject(new Error("Chat UI websocket disconnected")) console.error("Chat UI websocket disconnected") const delaySeconds = 1 - console.info(`Will try to reconnect in ${delaySeconds} seconds.`) + console.info(`Will try to reconnect in ${delaySeconds} s.`) setTimeout(async () => await this._initialize(), delaySeconds * 1000); }