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 { - 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,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 { + return new Promise((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)[] = [];