diff --git a/vscode-extension/package.json b/vscode-extension/package.json index d080d1723..fe4a1154b 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -57,6 +57,11 @@ "title": "Open Custom Model Registry File", "category": "AIConfig" }, + { + "command": "vscode-aiconfig.restartActiveEditorServer", + "title": "Restart Active Editor Server", + "category": "AIConfig" + }, { "command": "vscode-aiconfig.setApiKeys", "title": "Set API Keys", @@ -181,4 +186,4 @@ "extensionDependencies": [ "ms-python.python" ] -} +} \ No newline at end of file diff --git a/vscode-extension/src/aiConfigEditor.ts b/vscode-extension/src/aiConfigEditor.ts index 709155463..924324da7 100644 --- a/vscode-extension/src/aiConfigEditor.ts +++ b/vscode-extension/src/aiConfigEditor.ts @@ -15,7 +15,7 @@ import { AIConfigEditorManager, AIConfigEditorState, } from "./aiConfigEditorManager"; -import { EditorServer } from "./editorServer"; +import { EditorServer, EditorServerState } from "./editorServer"; /** * Provider for AIConfig editors. @@ -57,7 +57,9 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { webviewPanel: vscode.WebviewPanel, _token: vscode.CancellationToken ): Promise { - let editorServer: EditorServer | null = null; + const editorServer: EditorServer = new EditorServer( + getCurrentWorkingDirectory(document) + ); let isWebviewDisposed = false; // TODO: saqadri - clean up console log @@ -95,44 +97,82 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { // Update webview immediately so we unblock the render; server init will happen in the background. updateWebview(); - // Do not start the server until we ensure the Python setup is ready - initializePythonFlow(this.context, this.extensionOutputChannel).then( - async () => { - // Start the AIConfig editor server process. Don't await at the top level here since that blocks the - // webview render (which happens only when resolveCustomTextEditor returns) - this.startEditorServer(document).then(async (startedServer) => { - editorServer = startedServer; - - this.aiconfigEditorManager.addEditor( - new AIConfigEditorState( - document, - webviewPanel, - startedServer, - this.aiconfigEditorManager - ) - ); + const setupServerState = async (server: EditorServer) => { + // Wait for server ready + await waitUntilServerReady(server.url); - // Wait for server ready - await waitUntilServerReady(startedServer.url); + // Now set up the server with the latest document content + await this.initializeServerStateWithRetry( + server.url, + document, + webviewPanel + ); - // Now set up the server with the latest document content - await this.startServerWithRetry( - startedServer.url, + // Inform the webview of the server URL + if (!isWebviewDisposed) { + webviewPanel.webview.postMessage({ + type: "set_server_url", + url: server.url, + }); + } + }; + + // Do not start the server until we ensure the Python setup is ready + // Don't await at the top level here since that blocks the webview render (which happens + // only when resolveCustomTextEditor returns) + initializePythonFlow(this.context, this.extensionOutputChannel).then(() => + this.startEditorServer(editorServer, document).then( + async (startedServer) => { + const editor = new AIConfigEditorState( document, - webviewPanel + webviewPanel, + startedServer, + this.aiconfigEditorManager ); - // Inform the webview of the server URL - if (!isWebviewDisposed) { - webviewPanel.webview.postMessage({ - type: "set_server_url", - url: startedServer.url, - }); - } - }); + this.aiconfigEditorManager.addEditor(editor); + await setupServerState(startedServer); + } + ) + ); + + const serverStateChangeSubscription = editorServer.onDidChangeState( + async (state) => { + switch (state) { + case EditorServerState.Stopped: + // Webview should be readonly until the server state is ready + if (!isWebviewDisposed) { + webviewPanel.webview.postMessage({ + type: "set_readonly_state", + isReadOnly: true, + }); + } + break; + case EditorServerState.Starting: + // Show notification with server starting progress if webview is focused + if (!isWebviewDisposed && webviewPanel.active) { + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: "Starting editor server...", + cancellable: false, + }, + async (progress) => { + progress.report({ + increment: 50, + }); + } + ); + } + break; + } } ); + const serverRestartSubscription = editorServer.onRestart(async (server) => { + await setupServerState(server); + }); + // Hook up event handlers so that we can synchronize the webview with the text document. // // The text document acts as our model, so we have to sync change in the document to our @@ -232,7 +272,7 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { updateWebview(); // Notify server of updated document - if (editorServer) { + if (editorServer.url) { await updateServerWithRetry(editorServer.url, e.document); } } @@ -277,12 +317,11 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { console.log(`${document.fileName}: Webview disposed`); changeDocumentSubscription.dispose(); + serverRestartSubscription.dispose(); + serverStateChangeSubscription.dispose(); willSaveDocumentSubscription.dispose(); - if (editorServer) { - editorServer.stop(); - editorServer = null; - } + editorServer.stop(); }); // Receive message from the webview. @@ -415,13 +454,19 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { if (e.webviewPanel.active) { if (!isWebviewDisposed) { updateWebviewEditorThemeMode(webviewPanel.webview); + + // Inform the webview if editor server updated in the background + webviewPanel.webview.postMessage({ + type: "set_server_url", + url: editorServer.url, + }); } } }); } } - private startServerWithRetry( + private initializeServerStateWithRetry( serverUrl: string, document: vscode.TextDocument, webviewPanel: vscode.WebviewPanel @@ -454,7 +499,11 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { } if (selection === "Retry") { - this.startServerWithRetry(serverUrl, document, webviewPanel); + this.initializeServerStateWithRetry( + serverUrl, + document, + webviewPanel + ); } }); }); @@ -465,27 +514,27 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { } private async startEditorServer( + editorServer: EditorServer, document: vscode.TextDocument ): Promise { this.extensionOutputChannel.info( this.prependMessage("Starting editor server", document) ); - const editorServer = new EditorServer(getCurrentWorkingDirectory(document)); await editorServer.start(); - editorServer.serverProc.stdout.on("data", (data) => { + editorServer.onStdout((data) => { this.extensionOutputChannel.info(this.prependMessage(data, document)); console.log(`server stdout: ${data}`); }); // TODO: saqadri - stderr is very noisy for some reason (duplicating INFO logs). Figure out why before enabling this. - editorServer.serverProc.stderr.on("data", (data) => { + editorServer.onStderr((data) => { this.extensionOutputChannel.error(this.prependMessage(data, document)); console.error(`server stderr: ${data}`); }); - editorServer.serverProc.on("spawn", () => { + editorServer.onSpawn(() => { this.extensionOutputChannel.info( this.prependMessage( `Started server at port=${editorServer.port}, pid=${editorServer.pid}`, @@ -495,7 +544,7 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { console.log(`server spawned: ${editorServer.pid}`); }); - editorServer.serverProc.on("close", (code) => { + editorServer.onClose((code) => { if (code !== 0) { this.extensionOutputChannel.error( this.prependMessage( @@ -503,7 +552,9 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { document ) ); - console.error(`server terminated unexpectedly: exit code=${code}`); + console.error( + `Server at port=${editorServer.port}, pid=${editorServer.pid} terminated unexpectedly: exit code=${code}` + ); } else { this.extensionOutputChannel.info( this.prependMessage( @@ -515,7 +566,7 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider { } }); - editorServer.serverProc.on("error", (err) => { + editorServer.onError((err) => { this.extensionOutputChannel.error( this.prependMessage(JSON.stringify(err), document) ); diff --git a/vscode-extension/src/editorServer.ts b/vscode-extension/src/editorServer.ts index 94e8d8fdc..d618041c2 100644 --- a/vscode-extension/src/editorServer.ts +++ b/vscode-extension/src/editorServer.ts @@ -4,6 +4,12 @@ import { EXTENSION_NAME } from "./util"; import { getPythonPath } from "./utilities/pythonSetupUtils"; import { ChildProcessWithoutNullStreams, spawn } from "child_process"; +export enum EditorServerState { + Starting = "Starting", + Running = "Running", + Stopped = "Stopped", +} + /** * Provider for AIConfig editors. * @@ -12,66 +18,135 @@ import { ChildProcessWithoutNullStreams, spawn } from "child_process"; */ export class EditorServer { private cwd: string; + private _onDidChangeState = new vscode.EventEmitter(); + private _onRestart = new vscode.EventEmitter(); + + // Readable and process listeners. Maintain at the class level so that subscribers don't + // need to care about underlying server process + private _onStdout = new vscode.EventEmitter(); + private _onStderr = new vscode.EventEmitter(); + private _onSpawn = new vscode.EventEmitter(); + private _onClose = new vscode.EventEmitter(); + private _onError = new vscode.EventEmitter(); public pid: number | null = null; public port: number | null = null; + // TODO: Should make this private and expose subscriptions to .stderr, .stdout, and .on public serverProc: ChildProcessWithoutNullStreams | null = null; + public serverState: EditorServerState = EditorServerState.Stopped; public url: string | null = null; + public readonly onDidChangeState = this._onDidChangeState.event; + public readonly onRestart = this._onRestart.event; + + public readonly onStdout = this._onStdout.event; + public readonly onStderr = this._onStderr.event; + public readonly onSpawn = this._onSpawn.event; + public readonly onClose = this._onClose.event; + public readonly onError = this._onError.event; + constructor(workingDirectory: string) { this.cwd = workingDirectory; } - public async start(): Promise { + private updateServerState(state: EditorServerState) { + this.serverState = state; + this._onDidChangeState.fire(state); + } + + public async start(): Promise { if (this.serverProc) { console.log( `Server process ${this.pid} already started, port ${this.port}` ); - return this; + return; } - // If there is a custom model registry path, pass it to the server - const config = vscode.workspace.getConfiguration(EXTENSION_NAME); - const modelRegistryPath = config.get("modelRegistryPath"); - const modelRegistryPathArgs = modelRegistryPath - ? ["--parsers-module-path", modelRegistryPath] - : []; - - this.port = await getPortPromise(); - - const pythonPath = await getPythonPath(); - - // TODO: saqadri - specify parsers_module_path - // `aiconfig` command not useable here because it relies on python. Instead invoke the module directly. - const startServer = spawn( - pythonPath, - [ - "-m", - "aiconfig.scripts.aiconfig_cli", - "start", - "--server-port", - this.port.toString(), - ...modelRegistryPathArgs, - ], - { - cwd: this.cwd, - } - ); + console.log("Starting editor server process"); + this.updateServerState(EditorServerState.Starting); + + try { + // If there is a custom model registry path, pass it to the server + const config = vscode.workspace.getConfiguration(EXTENSION_NAME); + const modelRegistryPath = config.get("modelRegistryPath"); + const modelRegistryPathArgs = modelRegistryPath + ? ["--parsers-module-path", modelRegistryPath] + : []; + + this.port = await getPortPromise(); + + const pythonPath = await getPythonPath(); + + // TODO: saqadri - specify parsers_module_path + // `aiconfig` command not useable here because it relies on python. Instead invoke the module directly. + const startServer = spawn( + pythonPath, + [ + "-m", + "aiconfig.scripts.aiconfig_cli", + "start", + "--server-port", + this.port.toString(), + ...modelRegistryPathArgs, + ], + { + cwd: this.cwd, + } + ); + + this.pid = startServer.pid; + this.serverProc = startServer; + this.url = `http://localhost:${this.port}`; + + startServer.stdout.on("data", (data) => { + this._onStdout.fire(data); + }); + + startServer.stderr.on("data", (data) => { + this._onStderr.fire(data); + }); + + startServer.on("spawn", () => { + this._onSpawn.fire(); + }); + + startServer.on("close", (code) => { + this._onClose.fire(code); + }); + + startServer.on("error", (err) => { + this._onError.fire(err); + }); - this.pid = startServer.pid; - this.serverProc = startServer; - this.url = `http://localhost:${this.port}`; + this.updateServerState(EditorServerState.Running); - return this; + console.log( + `Started editor server process ${this.pid}, port ${this.port}` + ); + } catch (e) { + console.error("Error starting editor server process", e); + this.updateServerState(EditorServerState.Stopped); + throw e; + } } public stop() { - console.log(`Killing editor server process ${this.pid}`); - this.serverProc.kill(); + console.log(`Killing editor server process ${this.pid}, port ${this.port}`); + this.serverProc?.kill(); this.serverProc = null; this.pid = null; this.port = null; this.url = null; + this.updateServerState(EditorServerState.Stopped); + } + + public async restart(): Promise { + console.log( + `Restarting editor server process ${this.pid}, port ${this.port}` + ); + this.stop(); + await this.start(); + this._onRestart.fire(this); } } diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index ae64c9eb8..86dfede68 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -129,6 +129,15 @@ export async function activate(context: vscode.ExtensionContext) { ); context.subscriptions.push(openModelParserCommand); + const restartActiveEditorCommand = vscode.commands.registerCommand( + COMMANDS.RESTART_ACTIVE_EDITOR_SERVER, + async () => { + const activeEditor = aiconfigEditorManager.getActiveEditor(); + activeEditor?.editorServer?.restart(); + } + ); + context.subscriptions.push(restartActiveEditorCommand); + // Register our custom editor providers const aiconfigEditorManager: AIConfigEditorManager = new AIConfigEditorManager(); diff --git a/vscode-extension/src/util.ts b/vscode-extension/src/util.ts index fc58ec816..7b8fc4174 100644 --- a/vscode-extension/src/util.ts +++ b/vscode-extension/src/util.ts @@ -17,6 +17,7 @@ export const COMMANDS = { CREATE_CUSTOM_MODEL_REGISTRY: `${EXTENSION_NAME}.createCustomModelRegistry`, OPEN_CONFIG_FILE: `${EXTENSION_NAME}.openConfigFile`, OPEN_MODEL_REGISTRY: `${EXTENSION_NAME}.openModelRegistry`, + RESTART_ACTIVE_EDITOR_SERVER: `${EXTENSION_NAME}.restartActiveEditorServer`, SET_API_KEYS: `${EXTENSION_NAME}.setApiKeys`, SHARE: `${EXTENSION_NAME}.share`, SHOW_WELCOME: `${EXTENSION_NAME}.showWelcome`,