Skip to content

Commit

Permalink
[2/n] Support Restarting Editor Server
Browse files Browse the repository at this point in the history 
# [2/n] Support Restarting Editor Server

Add an "AIConfig: Restart Active Editor Server" command so that users can restart the server for the active (in-view) editor. As part of this, implement the restart functionality for the EditorServer so that the server process is killed and a new one is started, ensuring the associated extension/webview is properly updated to associate with the new server process. As best effort, also implement an `onDidChangeState` so that we can listen to the editor server status changing and show a progress notification that the server is starting.

This logic will be leveraged in subsequent PR to handle restarting all editor servers when the workspace python interpreter changes.


https://github.com/lastmile-ai/aiconfig/assets/5060851/39739008-abc7-4538-9087-42972c1400aa


Note: The server heartbeat banner shows when restarting because of the existing server status request failing. This should be fixed by lastmile-ai#1325

## Testing:
- Restart server and ensure (via logging) that the old one is killed and new one is set up, ensuring prompts run in new one
- Close webview and ensure the server is killed
- Activate restart server command and toggle away then back to webview, ensure it's in readonly until it starts up and becomes editable again when toggle back
  • Loading branch information
Ryan Holinshead committed Feb 23, 2024
1 parent b395d22 commit e11e1a6
Show file tree
Hide file tree
Showing 5 changed files with 223 additions and 82 deletions.
7 changes: 6 additions & 1 deletion vscode-extension/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -181,4 +186,4 @@
"extensionDependencies": [
"ms-python.python"
]
}
}
143 changes: 97 additions & 46 deletions vscode-extension/src/aiConfigEditor.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ import {
AIConfigEditorManager,
AIConfigEditorState,
} from "./aiConfigEditorManager";
import { EditorServer } from "./editorServer";
import { EditorServer, EditorServerState } from "./editorServer";

/**
* Provider for AIConfig editors.
Expand Down Expand Up @@ -57,7 +57,9 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider {
webviewPanel: vscode.WebviewPanel,
_token: vscode.CancellationToken
): Promise<void> {
let editorServer: EditorServer | null = null;
const editorServer: EditorServer = new EditorServer(
getCurrentWorkingDirectory(document)
);
let isWebviewDisposed = false;

// TODO: saqadri - clean up console log
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -454,7 +499,11 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider {
}

if (selection === "Retry") {
this.startServerWithRetry(serverUrl, document, webviewPanel);
this.initializeServerStateWithRetry(
serverUrl,
document,
webviewPanel
);
}
});
});
Expand All @@ -465,27 +514,27 @@ export class AIConfigEditorProvider implements vscode.CustomTextEditorProvider {
}

private async startEditorServer(
editorServer: EditorServer,
document: vscode.TextDocument
): Promise<EditorServer> {
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}`,
Expand All @@ -495,15 +544,17 @@ 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(
`Server at port=${editorServer.port}, pid=${editorServer.pid} was terminated unexpectedly with exit code ${code}`,
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(
Expand All @@ -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)
);
Expand Down
Loading

0 comments on commit e11e1a6

Please sign in to comment.