From 783e82f3efa2bd7f3e56ad76cd5d91e8368696a7 Mon Sep 17 00:00:00 2001 From: Ivan Gabriele Date: Thu, 20 Jul 2023 22:59:26 +0200 Subject: [PATCH] feat: handle multiple instances via WebSocket satellites --- .vscode/settings.json | 1 + package.json | 5 + src/commands/addOrRemoveCurrentDocument.ts | 2 +- .../evaluateAndSendCurrentDocumentOrStack.ts | 19 +- src/commands/sendCurrentDocumentOrStack.ts | 19 +- src/commands/startOrStopServer.ts | 9 + src/commands/welcome.ts | 50 +-- src/extension.ts | 28 +- src/helpers/fromWebSocketMessage.ts | 9 + src/helpers/toWebSocketMessage.ts | 5 + src/helpers/waitFor.ts | 5 + src/libs/Logger.ts | 13 + src/libs/Server.ts | 331 ++++++++++++++++++ src/libs/{stackManager.ts => StackManager.ts} | 0 src/libs/{stateManager.ts => StateManager.ts} | 44 ++- src/libs/server.ts | 98 ------ src/types/Communication.ts | 43 +++ src/types/index.ts | 12 +- yarn.lock | 8 + 19 files changed, 503 insertions(+), 198 deletions(-) create mode 100644 .vscode/settings.json create mode 100644 src/commands/startOrStopServer.ts create mode 100644 src/helpers/fromWebSocketMessage.ts create mode 100644 src/helpers/toWebSocketMessage.ts create mode 100644 src/helpers/waitFor.ts create mode 100644 src/libs/Server.ts rename src/libs/{stackManager.ts => StackManager.ts} (100%) rename src/libs/{stateManager.ts => StateManager.ts} (65%) delete mode 100644 src/libs/server.ts create mode 100644 src/types/Communication.ts diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..9e26dfe --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/package.json b/package.json index dac7672..bfe12f0 100644 --- a/package.json +++ b/package.json @@ -37,6 +37,7 @@ "@typescript-eslint/parser": "6.0.0", "@vscode/test-electron": "2.3.3", "@vscode/vsce": "2.19.0", + "cuid": "3.0.0", "eslint": "8.45.0", "eslint-config-airbnb": "19.0.4", "eslint-config-airbnb-typescript": "17.1.0", @@ -84,6 +85,10 @@ { "command": "openai-forge.sendCurrentDocument", "title": "OpenAI Forge: Send current document or stack" + }, + { + "command": "openai-forge.startOrStopServer", + "title": "OpenAI Forge: Start/Stop OAIF Server/Satellite" } ], "configuration": { diff --git a/src/commands/addOrRemoveCurrentDocument.ts b/src/commands/addOrRemoveCurrentDocument.ts index e581694..62cc744 100644 --- a/src/commands/addOrRemoveCurrentDocument.ts +++ b/src/commands/addOrRemoveCurrentDocument.ts @@ -1,7 +1,7 @@ import { window } from 'vscode' import { DocumentInfo } from '../libs/DocumentInfo' -import { stackManager } from '../libs/stackManager' +import { stackManager } from '../libs/StackManager' import { getCurrentDocumentPath } from '../utils/getCurrentDocumentPath' type AddOrRemoveCurrentDocumentArgs = { diff --git a/src/commands/evaluateAndSendCurrentDocumentOrStack.ts b/src/commands/evaluateAndSendCurrentDocumentOrStack.ts index 153637e..347f81b 100644 --- a/src/commands/evaluateAndSendCurrentDocumentOrStack.ts +++ b/src/commands/evaluateAndSendCurrentDocumentOrStack.ts @@ -1,17 +1,16 @@ import { ProgressLocation, window } from 'vscode' import { DocumentInfo } from '../libs/DocumentInfo' -import { stackManager } from '../libs/stackManager' -import { WebSocketDataAction, type WebSocketData } from '../types' +import { server } from '../libs/Server' +import { stackManager } from '../libs/StackManager' +import { Communication } from '../types' import { getChatGptPrompt } from '../utils/getChatGptPrompt' import { getUserSetting } from '../utils/getUserSetting' import { getUserWorkspaceEvaluators } from '../utils/getUserWorkspaceEvaluators' import { getUserWorkspaceInfo } from '../utils/getUserWorkspaceInfo' import { runEvaluator } from '../utils/runEvaluator' -import type { WebSocket } from 'ws' - -export async function evaluateAndSendCurrentDocumentOrStack(webSocket: WebSocket) { +export async function evaluateAndSendCurrentDocumentOrStack() { await window.withProgress( { location: ProgressLocation.Notification, @@ -38,16 +37,16 @@ export async function evaluateAndSendCurrentDocumentOrStack(webSocket: WebSocket progress.report({ message: 'OpenAI Forge: Sending source code & errors to ChatGPT...' }) const excludeProjectInfo = getUserSetting('prompt', 'excludeProjectInfo') || false - const message = await getChatGptPrompt(currentOrStackDocumentInfos, { + const messageMessage = await getChatGptPrompt(currentOrStackDocumentInfos, { errorOutput, workspaceInfo: !excludeProjectInfo ? workspaceInfo : undefined, }) - const webSocketData: WebSocketData = { - action: WebSocketDataAction.ASK, - message, + const message: Communication.Message = { + action: Communication.MessageAction.ASK, + message: messageMessage, } - webSocket.send(JSON.stringify(webSocketData)) + server.send(message) }, ) } diff --git a/src/commands/sendCurrentDocumentOrStack.ts b/src/commands/sendCurrentDocumentOrStack.ts index e9f9e80..370e5d8 100644 --- a/src/commands/sendCurrentDocumentOrStack.ts +++ b/src/commands/sendCurrentDocumentOrStack.ts @@ -1,15 +1,14 @@ import { ProgressLocation, window } from 'vscode' import { DocumentInfo } from '../libs/DocumentInfo' -import { stackManager } from '../libs/stackManager' -import { WebSocketDataAction, type WebSocketData } from '../types' +import { server } from '../libs/Server' +import { stackManager } from '../libs/StackManager' +import { Communication } from '../types' import { getChatGptPrompt } from '../utils/getChatGptPrompt' import { getUserSetting } from '../utils/getUserSetting' import { getUserWorkspaceInfo } from '../utils/getUserWorkspaceInfo' -import type { WebSocket } from 'ws' - -export async function sendCurrentDocument(webSocket: WebSocket) { +export async function sendCurrentDocument() { await window.withProgress( { location: ProgressLocation.Notification, @@ -33,16 +32,16 @@ export async function sendCurrentDocument(webSocket: WebSocket) { progress.report({ message: 'OpenAI Forge: Sending source code & errors to ChatGPT...' }) - const message = await getChatGptPrompt(currentOrStackDocumentInfos, { + const messageMessage = await getChatGptPrompt(currentOrStackDocumentInfos, { userMessage, workspaceInfo: currentWorkspaceInfo, }) - const webSocketData: WebSocketData = { - action: WebSocketDataAction.ASK, - message, + const message: Communication.Message = { + action: Communication.MessageAction.ASK, + message: messageMessage, } - webSocket.send(JSON.stringify(webSocketData)) + server.send(message) }, ) } diff --git a/src/commands/startOrStopServer.ts b/src/commands/startOrStopServer.ts new file mode 100644 index 0000000..0af67cc --- /dev/null +++ b/src/commands/startOrStopServer.ts @@ -0,0 +1,9 @@ +import { server } from '../libs/Server' + +export function startOrStopServer() { + if (server.isStarted) { + server.stop() + } else { + server.start() + } +} diff --git a/src/commands/welcome.ts b/src/commands/welcome.ts index b117bac..a7dd727 100644 --- a/src/commands/welcome.ts +++ b/src/commands/welcome.ts @@ -1,13 +1,9 @@ -import { existsSync } from 'fs' -import { join } from 'path' -import { Position, Range, Uri, ViewColumn, WorkspaceEdit, window, workspace } from 'vscode' +import { window } from 'vscode' import { DocumentationPath } from '../constants' import { handleMessageItems } from '../helpers/handleMessageItems' -import { isEmpty } from '../helpers/isEmpty' import { GlobalStateKey, getGlobalStateManager } from '../libs/GlobalStateManager' import { MessageItemType, type MessageButton } from '../types' -import { getUserWorkspaceRootPath } from '../utils/getUserWorkspaceRootPath' import { showDocumentation } from '../utils/showDocumentation' const HIDE_MESSAGE_BUTTON: MessageButton = { @@ -26,45 +22,7 @@ export async function welcome() { await showDocumentation(DocumentationPath.WELCOME) - const workspaceSettingsPath = join(getUserWorkspaceRootPath(), '.vscode', 'settings.json') - if (!existsSync(workspaceSettingsPath)) { - await workspace.fs.writeFile(Uri.file(workspaceSettingsPath), Buffer.from('{}')) - } - const workspaceSettingsDocument = await workspace.openTextDocument(workspaceSettingsPath) - const workspaceSettingsText = workspaceSettingsDocument.getText() - - if (!workspaceSettingsText.includes('openai-forge.customEvaluators')) { - const workspaceSettings = !isEmpty(workspaceSettingsText) ? JSON.parse(workspaceSettingsText) : {} - const suggestedWorkspaceSettings = JSON.stringify( - { - ...workspaceSettings, - 'openai-forge.customEvaluators': [ - { - command: 'npm', - commandArgs: ['run', 'build'], - extensions: ['.js', '.jsx', '.ts', '.tsx'], - }, - ], - }, - // eslint-disable-next-line no-null/no-null - null, - 2, - ) - const workspaceSettingsEdit = new WorkspaceEdit() - workspaceSettingsEdit.replace( - workspaceSettingsDocument.uri, - new Range( - new Position(0, 0), - workspaceSettingsDocument.lineAt(workspaceSettingsDocument.lineCount - 1).range.end, - ), - suggestedWorkspaceSettings, - ) - - await window.showTextDocument(workspaceSettingsDocument, ViewColumn.Two) - await workspace.applyEdit(workspaceSettingsEdit) - - window - .showInformationMessage("Once you've read the welcome page:", HIDE_MESSAGE_BUTTON.label) - .then(handleMessageItems([HIDE_MESSAGE_BUTTON])) - } + window + .showInformationMessage("Once you've read the welcome page:", HIDE_MESSAGE_BUTTON.label) + .then(handleMessageItems([HIDE_MESSAGE_BUTTON])) } diff --git a/src/extension.ts b/src/extension.ts index 1616bca..650ff65 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -3,12 +3,14 @@ import { type ExtensionContext, workspace, commands, ExtensionMode } from 'vscod import { addOrRemoveCurrentDocument } from './commands/addOrRemoveCurrentDocument' import { evaluateAndSendCurrentDocumentOrStack } from './commands/evaluateAndSendCurrentDocumentOrStack' import { sendCurrentDocument } from './commands/sendCurrentDocumentOrStack' +import { startOrStopServer } from './commands/startOrStopServer' import { welcome } from './commands/welcome' import { handleError } from './helpers/handleError' import { getGlobalStateManager, initializeGlobalStateManager } from './libs/GlobalStateManager' -import { server } from './libs/server' -import { stackManager } from './libs/stackManager' -import { stateManager } from './libs/stateManager' +import { logger } from './libs/Logger' +import { server } from './libs/Server' +import { stackManager } from './libs/StackManager' +import { stateManager } from './libs/StateManager' export async function activate(context: ExtensionContext) { try { @@ -24,6 +26,7 @@ export async function activate(context: ExtensionContext) { await initializeGlobalStateManager(context) if (context.extensionMode === ExtensionMode.Development) { + logger.isDevelopment = true await getGlobalStateManager().clear() } @@ -36,17 +39,18 @@ export async function activate(context: ExtensionContext) { ) const evaluateAndSendCurrentDocumentOrStackDisposable = commands.registerCommand( 'openai-forge.evaluateAndSendCurrentDocumentOrStack', - () => { - stateManager.clients.forEach(evaluateAndSendCurrentDocumentOrStack) - }, + evaluateAndSendCurrentDocumentOrStack, ) - const sendCurrentDocumentDisposable = commands.registerCommand('openai-forge.sendCurrentDocument', () => { - stateManager.clients.forEach(sendCurrentDocument) - }) + const sendCurrentDocumentDisposable = commands.registerCommand( + 'openai-forge.sendCurrentDocument', + sendCurrentDocument, + ) + const startOrStopServerDisposable = commands.registerCommand('openai-forge.startOrStopServer', startOrStopServer) context.subscriptions.push(addOrRemoveCurrentDocumentDisposable) context.subscriptions.push(evaluateAndSendCurrentDocumentOrStackDisposable) context.subscriptions.push(sendCurrentDocumentDisposable) + context.subscriptions.push(startOrStopServerDisposable) // ------------------------------------------------------------------------- // Status Bar Items @@ -57,7 +61,7 @@ export async function activate(context: ExtensionContext) { // ------------------------------------------------------------------------- // WebSocket Server - server.start(context) + server.start() // ------------------------------------------------------------------------- // Welcome Documentation @@ -71,6 +75,6 @@ export async function activate(context: ExtensionContext) { } } -export function deactivate() { - server.stop() +export async function deactivate() { + await server.stop() } diff --git a/src/helpers/fromWebSocketMessage.ts b/src/helpers/fromWebSocketMessage.ts new file mode 100644 index 0000000..1aa5d8a --- /dev/null +++ b/src/helpers/fromWebSocketMessage.ts @@ -0,0 +1,9 @@ +import { type MessageEvent, type RawData } from 'ws' + +import { type Communication } from '../types' + +export function fromWebSocketMessage(rawDataOrMessageEvent: RawData | MessageEvent): Communication.Message { + const message: Communication.Message = JSON.parse(rawDataOrMessageEvent.toString()) + + return message +} diff --git a/src/helpers/toWebSocketMessage.ts b/src/helpers/toWebSocketMessage.ts new file mode 100644 index 0000000..1735772 --- /dev/null +++ b/src/helpers/toWebSocketMessage.ts @@ -0,0 +1,5 @@ +import type { Communication } from '../types' + +export function toWebSocketMessage(message: Communication.Message): string { + return JSON.stringify(message) +} diff --git a/src/helpers/waitFor.ts b/src/helpers/waitFor.ts new file mode 100644 index 0000000..80426d6 --- /dev/null +++ b/src/helpers/waitFor.ts @@ -0,0 +1,5 @@ +export async function waitFor(inMs: number): Promise { + await new Promise(resolve => { + setTimeout(resolve, inMs) + }) +} diff --git a/src/libs/Logger.ts b/src/libs/Logger.ts index dc87163..f42f257 100644 --- a/src/libs/Logger.ts +++ b/src/libs/Logger.ts @@ -23,12 +23,25 @@ function parseLogArgs(dataOrLevel: any, maybeLevel: LogLevel | undefined): [LogL } class Logger { + #isDevelopment: boolean = false #outputChannel: OutputChannel constructor() { this.#outputChannel = window.createOutputChannel('OpenAI Forge') } + set isDevelopment(isDevelopment: boolean) { + this.#isDevelopment = isDevelopment + } + + debug(...args: any[]): void { + if (!this.#isDevelopment) { + return + } + + console.debug(...args) + } + log(message: string): void log(message: string, level: LogLevel): void log(message: string, data: any): void diff --git a/src/libs/Server.ts b/src/libs/Server.ts new file mode 100644 index 0000000..0e123ea --- /dev/null +++ b/src/libs/Server.ts @@ -0,0 +1,331 @@ +import cuid from 'cuid' +import { dissoc } from 'ramda' +import { WebSocket, WebSocketServer, type RawData } from 'ws' + +import { InternalError } from './InternalError' +import { logger } from './Logger' +import { stateManager } from './StateManager' +import { UserError } from './UserError' +import { fromWebSocketMessage } from '../helpers/fromWebSocketMessage' +import { toWebSocketMessage } from '../helpers/toWebSocketMessage' +import { waitFor } from '../helpers/waitFor' +import { Communication, State } from '../types' + +class Server { + #clients: Record = {} + // @ts-ignore + #id: string | undefined + #isStarted: boolean = false + #wasManualStop: boolean = false + #satellites: Record = {} + #webSocketServer: WebSocketServer | undefined = undefined + #webSocketServerAsClient: WebSocket | undefined = undefined + + get #clientsCount(): number { + return Object.entries(this.#clients).length + } + + get #satelitesCount(): number { + return Object.entries(this.#satellites).length + } + + get isStarted(): boolean { + return this.#isStarted + } + + send(message: { action: Communication.MessageAction.ASK; message: string }) { + logger.debug('Server.send()', { message }) + + if (!this.#isStarted) { + throw new UserError( + [ + 'OAIF Server/Satellite is stopped.', + 'You can start it again from the command palette: "OpenAI Forge: Start/Stop OAIF Server/Satellite".', + ].join(' '), + ) + } + + if (this.#webSocketServer) { + Object.values(this.#clients).forEach(client => client.send(toWebSocketMessage(message))) + } else if (this.#webSocketServerAsClient) { + this.#webSocketServerAsClient.send( + toWebSocketMessage({ + $id: this.#id, + action: Communication.MessageAction.CONVEY, + value: message, + }), + ) + } else { + throw new InternalError('Both `this.#webSocketServer` and `this.#webSocketServerAsClient` are undefined.') + } + } + + start(isClient: boolean = false) { + logger.debug('Server.start()', { isClient }) + + if (stateManager.state !== State.RESTARTING) { + stateManager.state = State.STARTING + } + + try { + if (!isClient) { + this.#webSocketServer = new WebSocketServer({ port: 4242 }) + + this.#webSocketServer.on('connection', this.#handleServerConnection.bind(this)) + this.#webSocketServer.on('error', this.#handleServerError.bind(this)) + + stateManager.isSatellite = false + } else { + this.#webSocketServerAsClient = new WebSocket('ws://localhost:4242') + + this.#webSocketServerAsClient.on('close', this.#handleSatelliteClose.bind(this)) + this.#webSocketServerAsClient.on('error', this.#handleSatelliteError.bind(this)) + this.#webSocketServerAsClient.on('message', this.#handleSatelliteMessage.bind(this)) + this.#webSocketServerAsClient.on('open', this.#handleSatelliteOpen.bind(this)) + + stateManager.isSatellite = true + } + + this.#isStarted = true + + stateManager.state = State.RUNNING + } catch (err) { + throw new InternalError('WebSocket server failed to start.', err) + } + } + + async stop() { + logger.debug('Server.stop()') + + try { + if (this.#webSocketServer) { + stateManager.state = State.STOPPING + + Object.values(this.#clients).forEach(client => client.close()) + Object.values(this.#satellites).forEach(client => client.close()) + + await this.#waitForClientsAndSatellitesToCloseAndStop() + } else if (this.#webSocketServerAsClient) { + this.#wasManualStop = true + + this.#webSocketServerAsClient.close() + } else { + throw new InternalError('Both `this.#webSocketServer` and `this.#webSocketServerAsClient` are undefined.') + } + } catch (err) { + throw new InternalError('WebSocket server failed to stop.', err) + } + } + + #handleSatelliteClose(error: Error) { + logger.debug('Server.#handleSatelliteError()', error) + logger.log('Soft WebSocket Server-as-Client error.', error) + + this.#reset() + + if (this.#wasManualStop) { + this.#wasManualStop = false + + stateManager.state = State.STOPPED + + return + } + + stateManager.state = State.WILL_RESTART + + setTimeout(() => { + stateManager.state = State.RESTARTING + + this.start() + }, 5000) + } + + // eslint-disable-next-line class-methods-use-this + #handleSatelliteError(error: Error) { + logger.debug('Server.#handleSatelliteError()', { error }) + + this.#reset() + + stateManager.state = State.FAILED + + throw new InternalError(`WebSocket Client received an unknow error: \`${(error as any).code}\`.`, error) + } + + #handleSatelliteMessage(rawData: RawData) { + logger.debug('Server.#handleSatelliteMessage()', { rawData }) + + const message = fromWebSocketMessage(rawData) + + if (message.action === Communication.MessageAction.ORDER) { + if (message.message === Communication.MessageMessage.SET_CLIENTS_COUNT) { + this.#updateClientsCount(message.value) + + return + } + + if (message.message === Communication.MessageMessage.SET_ID) { + this.#id = message.value + + return + } + } + + throw new InternalError(`WebSocket server-as-client received an unknow message ("${JSON.stringify(message)}").`) + } + + #handleSatelliteOpen() { + logger.debug('Server.#handleSatelliteOpen()') + + if (!this.#webSocketServerAsClient) { + throw new InternalError('WebSocket server-as-client failed to start.') + } + + this.#webSocketServerAsClient.send( + toWebSocketMessage({ + action: Communication.MessageAction.DECLARE, + message: Communication.MessageMessage.AS_SATELLITE, + }), + ) + } + + #handleServerConnection(webSocket: WebSocket) { + logger.debug('Server.#handleServerConnection()') + + const clientId = cuid() + + this.#clients = { + ...this.#clients, + [clientId]: webSocket, + } + + this.#updateClientsCount() + + webSocket.on('message', rawData => { + const message = fromWebSocketMessage(rawData) + + if (message.action === Communication.MessageAction.CONVEY) { + this.send(message.value) + + return + } + + if (message.action === Communication.MessageAction.DECLARE) { + if (message.message === Communication.MessageMessage.AS_SATELLITE) { + if (this.#clients[clientId]) { + this.#clients = dissoc(clientId, this.#clients) + } else { + throw new InternalError(`WebSocket client with id \`${clientId}\` doesn't exist.`) + } + + this.#clients = dissoc(clientId, this.#clients) + this.#satellites = { + ...this.#satellites, + [clientId]: webSocket, + } + + this.#updateClientsCount() + + return + } + } + + throw new InternalError(`WebSocket server received an unknow message ("${JSON.stringify(message)}").`) + }) + + webSocket.on('close', () => { + if (this.#clients[clientId]) { + this.#clients = dissoc(clientId, this.#clients) + } else if (this.#satellites[clientId]) { + this.#satellites = dissoc(clientId, this.#satellites) + } else { + throw new InternalError(`WebSocket client or satellite with id \`${clientId}\` doesn't exist.`) + } + + this.#updateClientsCount() + }) + } + + #handleServerError(error: Error) { + logger.debug('Server.#handleServerError()', { error }) + + this.#reset() + + if ((error as any).code === 'EADDRINUSE') { + this.start(true) + + return + } + + stateManager.state = State.FAILED + + throw new InternalError(`WebSocket Server received an unknow error: \`${(error as any).code}\`.`, error) + } + + #reset() { + logger.debug('Server.#reset()') + + this.#clients = {} + this.#id = undefined + this.#isStarted = false + this.#clients = {} + this.#satellites = {} + this.#webSocketServer = undefined + this.#webSocketServerAsClient = undefined + + this.#updateClientsCount(0) + } + + #updateClientsCount(clientsCount?: number) { + logger.debug('Server.#updateClientsCount()', { clientsCount }) + + if (!this.#webSocketServer && clientsCount === undefined) { + throw new InternalError('Both `this.#webSocketServer` and `clientsCount` are undefined.') + } + + const controlledClientsCount = clientsCount || this.#clientsCount + + if (this.#webSocketServer) { + this.#webSocketServer.clients.forEach(client => { + client.send( + toWebSocketMessage({ + action: Communication.MessageAction.ORDER, + message: Communication.MessageMessage.SET_CLIENTS_COUNT, + value: controlledClientsCount, + }), + ) + }) + } + + stateManager.clientCount = controlledClientsCount + } + + async #waitForClientsAndSatellitesToCloseAndStop() { + logger.debug('Server.#updateClientsCount()') + + if (!this.#webSocketServer) { + throw new InternalError('`this.#webSocketServer` is undefined.') + } + + if (this.#clientsCount || this.#satelitesCount) { + await waitFor(500) + + this.#waitForClientsAndSatellitesToCloseAndStop() + + return + } + + this.#webSocketServer.close(error => { + logger.debug('Server.#webSocketServer.close()', { error }) + + if (error) { + throw new InternalError('WebSocket server failed to stop.', error) + } + + this.#reset() + + stateManager.state = State.STOPPED + }) + } +} + +export const server = new Server() diff --git a/src/libs/stackManager.ts b/src/libs/StackManager.ts similarity index 100% rename from src/libs/stackManager.ts rename to src/libs/StackManager.ts diff --git a/src/libs/stateManager.ts b/src/libs/StateManager.ts similarity index 65% rename from src/libs/stateManager.ts rename to src/libs/StateManager.ts index 2bdaee5..10fc816 100644 --- a/src/libs/stateManager.ts +++ b/src/libs/StateManager.ts @@ -3,26 +3,29 @@ import { window, type StatusBarItem, StatusBarAlignment } from 'vscode' import { InternalError } from './InternalError' import { State } from '../types' -import type { WebSocket } from 'ws' - export const STATE_ICON: Record = { [State.FAILED]: 'error', - [State.RUNNING]: 'radio-tower', + [State.RESTARTING]: 'gear-spin', + [State.RUNNING]: '', [State.STARTING]: 'gear-spin', [State.STOPPED]: 'circle-slash', [State.STOPPING]: 'gear-spin', + [State.WILL_RESTART]: 'debug-disconnect', } export const STATE_LABEL: Record = { [State.FAILED]: 'Failed to start', + [State.RESTARTING]: 'Restarting...', [State.RUNNING]: '', [State.STARTING]: 'Starting...', [State.STOPPED]: 'Stopped', [State.STOPPING]: 'Stopping...', + [State.WILL_RESTART]: 'Restarting in 5s...', } class StateManager { - clients: Set = new Set() + #clientsCount: number = 0 + #isSatellite: boolean = false #state: State = State.STARTING #errorMessage: string | undefined = undefined #statusBarItem: StatusBarItem @@ -32,6 +35,12 @@ class StateManager { this.#statusBarItem.show() } + set clientCount(clientsCount: number) { + this.#clientsCount = clientsCount + + this.updateStatusBarItem() + } + set errorMessage(errorMessage: string) { this.#errorMessage = errorMessage this.#state = State.FAILED @@ -39,6 +48,12 @@ class StateManager { this.updateStatusBarItem() } + set isSatellite(isSatellite: boolean) { + this.#isSatellite = isSatellite + + this.updateStatusBarItem() + } + get state(): State { return this.#state } @@ -51,25 +66,30 @@ class StateManager { } updateStatusBarItem(): void { + const label = this.#isSatellite ? 'OAIF Satellite' : 'OAIF Server' + switch (stateManager.state) { case State.FAILED: - this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) OAIF Server: ${ + this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) ${label}: ${ this.#errorMessage || STATE_LABEL.FAILED }` break case State.RUNNING: - this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) OAIF Server: ${ - stateManager.clients.size || 'No' - } client${stateManager.clients.size > 1 ? 's' : ''}` + this.#statusBarItem.text = `$(${this.#isSatellite ? 'rss' : 'radio-tower'}) ${label}: ${ + this.#clientsCount || 'No' + } client${this.#clientsCount > 1 ? 's' : ''}` break case State.STARTING: - case State.STOPPED: case State.STOPPING: - this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) OAIF Server: ${ - STATE_LABEL[stateManager.state] - }` + this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) ${label}: ${STATE_LABEL[stateManager.state]}` + break + + case State.RESTARTING: + case State.STOPPED: + case State.WILL_RESTART: + this.#statusBarItem.text = `$(${STATE_ICON[stateManager.state]}) OAIF: ${STATE_LABEL[stateManager.state]}` break default: diff --git a/src/libs/server.ts b/src/libs/server.ts deleted file mode 100644 index a522842..0000000 --- a/src/libs/server.ts +++ /dev/null @@ -1,98 +0,0 @@ -import { window, type ExtensionContext } from 'vscode' -import { WebSocketServer } from 'ws' - -import { GlobalStateKey, getGlobalStateManager } from './GlobalStateManager' -import { InternalError } from './InternalError' -import { stackManager } from './stackManager' -import { stateManager } from './stateManager' -import { NotificationAction } from '../constants' -import { State } from '../types' - -class Server { - #webSocketServer: WebSocketServer | undefined - - start(context: ExtensionContext) { - try { - if (this.#webSocketServer) { - return - } - - this.#webSocketServer = new WebSocketServer({ port: 4242 }) - - this.#webSocketServer.on('connection', webSocket => { - stateManager.clients.add(webSocket) - - // webSocket.on('message', message => { - // console.debug(`New message received: ${message}.`) - // }) - - webSocket.on('close', () => { - stateManager.clients.delete(webSocket) - }) - }) - - this.#webSocketServer.on('error', async error => { - stateManager.state = State.FAILED - - if ((error as any).code === 'EADDRINUSE') { - const isNotificationHidden = await getGlobalStateManager().get( - GlobalStateKey.NOTIFICATION__HIDE_SERVER_PORT_IN_USE_ERROR, - ) - - stackManager.hideStatusBarItem() - stateManager.errorMessage = 'Port 4242 is already in use' - - context.subscriptions.forEach(disposable => disposable.dispose()) - - if (!isNotificationHidden) { - const answer = await window.showErrorMessage( - [ - 'OpenAI Forge: Handling multiple Visual Studio Code instances is not yet supported.', - 'Only the first Visual Studio Code instance will be able to communicate with ChatGPT.', - 'You can close other instances and reload this instance from the command palette:', - '"Developer: Reload Window".', - // 'You can close other instances and restart it from the command palette:', - // '"OpenAI Forge: Restart WebSocket Server".', - ].join(' '), - // NotificationAction.SHOW_HELP, - NotificationAction.NEVER_SHOW_AGAIN, - ) - - if (answer === NotificationAction.NEVER_SHOW_AGAIN) { - await getGlobalStateManager().set(GlobalStateKey.NOTIFICATION__HIDE_SERVER_PORT_IN_USE_ERROR, true) - } - } - - return - } - - throw new InternalError(`WebSocket started with an unknow error: \`${(error as any).code}\`.`, error) - }) - - stateManager.state = State.RUNNING - } catch (err) { - throw new InternalError('WebSocket server failed to start.', err) - } - } - - stop() { - try { - if (!this.#webSocketServer) { - return - } - - this.#webSocketServer.close(err => { - if (err) { - throw new InternalError('WebSocket server failed to stop.', err) - } - - this.#webSocketServer = undefined - stateManager.state = State.STOPPED - }) - } catch (err) { - throw new InternalError('WebSocket server failed to stop.', err) - } - } -} - -export const server = new Server() diff --git a/src/types/Communication.ts b/src/types/Communication.ts new file mode 100644 index 0000000..f0c592b --- /dev/null +++ b/src/types/Communication.ts @@ -0,0 +1,43 @@ +export namespace Communication { + export enum MessageAction { + ASK = 'ASK', + CONVEY = 'CONVEY', + DECLARE = 'DECLARE', + ORDER = 'ORDER', + } + + export enum MessageMessage { + AS_SATELLITE = 'AS_SATELLITE', + PASS_MESSAGE = 'AS_SATELLITE', + SET_CLIENTS_COUNT = 'SET_CLIENTS_COUNT', + SET_ID = 'SET_ID', + } + + export type Message = + | { + action: MessageAction.ASK + message: string + } + | { + $id: string | undefined + action: MessageAction.CONVEY + value: { + action: MessageAction.ASK + message: string + } + } + | { + action: MessageAction.DECLARE + message: MessageMessage.AS_SATELLITE + } + | { + action: MessageAction.ORDER + message: MessageMessage.SET_CLIENTS_COUNT + value: number + } + | { + action: MessageAction.ORDER + message: MessageMessage.SET_ID + value: string + } +} diff --git a/src/types/index.ts b/src/types/index.ts index 098924c..4715377 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -19,20 +19,14 @@ export type MessageLink = { export enum State { 'FAILED' = 'FAILED', + 'RESTARTING' = 'RESTARTING', 'RUNNING' = 'RUNNING', 'STARTING' = 'STARTING', 'STOPPED' = 'STOPPED', 'STOPPING' = 'STOPPING', + 'WILL_RESTART' = 'WILL_RESTART', } -export enum WebSocketDataAction { - ASK = 'ASK', -} - -export type WebSocketData = { - action: WebSocketDataAction - message: string -} - +export { Communication } from './Communication' export { type UserSettings } from './UserSettings' export { UserWorkspace } from './UserWorkspace' diff --git a/yarn.lock b/yarn.lock index 9b6fa94..d68f098 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3761,6 +3761,13 @@ __metadata: languageName: node linkType: hard +"cuid@npm:3.0.0": + version: 3.0.0 + resolution: "cuid@npm:3.0.0" + checksum: f9a344bd90e26b62f75a86b616ad717bad88b20b02df25c6e4008dd2c0948b081bd72ad8075b25ae644dce7fc367a43c622954e88aed972b72b1f132416faa4e + languageName: node + linkType: hard + "dateformat@npm:^3.0.0, dateformat@npm:^3.0.3": version: 3.0.3 resolution: "dateformat@npm:3.0.3" @@ -8552,6 +8559,7 @@ __metadata: "@typescript-eslint/parser": 6.0.0 "@vscode/test-electron": 2.3.3 "@vscode/vsce": 2.19.0 + cuid: 3.0.0 eslint: 8.45.0 eslint-config-airbnb: 19.0.4 eslint-config-airbnb-typescript: 17.1.0