internal: Properly handle commands in the VSCode client when the server is stopped

editors/code/src/ctx.ts

@@ -18,6 +18,11 @@ export type Workspace =
           files: vscode.TextDocument[];
       };
 
+export type CommandFactory = {
+    enabled: (ctx: Ctx) => Cmd;
+    disabled?: (ctx: Ctx) => Cmd;
+};
+
 export class Ctx {
     readonly statusBar: vscode.StatusBarItem;
     readonly config: Config;
@@ -26,31 +31,40 @@ export class Ctx {
     private _serverPath: string | undefined;
     private traceOutputChannel: vscode.OutputChannel | undefined;
     private outputChannel: vscode.OutputChannel | undefined;
+    private clientSubscriptions: Disposable[];
     private state: PersistentState;
+    private commandFactories: Record<string, CommandFactory>;
+    private commandDisposables: Disposable[];
 
     workspace: Workspace;
 
-    constructor(readonly extCtx: vscode.ExtensionContext, workspace: Workspace) {
+    constructor(
-        this.statusBar = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left);
+        readonly extCtx: vscode.ExtensionContext,
-        extCtx.subscriptions.push(this.statusBar);
+        workspace: Workspace,
-        extCtx.subscriptions.push({
+        commandFactories: Record<string, CommandFactory>
-            dispose() {
+    ) {
-                this.dispose();
-            },
-        });
         extCtx.subscriptions.push(this);
+        this.statusBar = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left);
         this.statusBar.text = "rust-analyzer";
         this.statusBar.tooltip = "ready";
         this.statusBar.command = "rust-analyzer.analyzerStatus";
         this.statusBar.show();
         this.workspace = workspace;
+        this.clientSubscriptions = [];
+        this.commandDisposables = [];
+        this.commandFactories = commandFactories;
 
         this.state = new PersistentState(extCtx.globalState);
         this.config = new Config(extCtx);
+
+        this.updateCommands();
     }
 
     dispose() {
         this.config.dispose();
+        this.statusBar.dispose();
+        void this.disposeClient();
+        this.commandDisposables.forEach((disposable) => disposable.dispose());
     }
 
     clientFetcher() {
@@ -63,7 +77,6 @@ export class Ctx {
     }
 
     async getClient() {
-        // if server path changes -> dispose
         if (!this.traceOutputChannel) {
             this.traceOutputChannel = vscode.window.createOutputChannel(
                 "Rust Analyzer Language Server Trace"
@@ -118,7 +131,11 @@ export class Ctx {
                 initializationOptions,
                 serverOptions
             );
-            this.client.onNotification(ra.serverStatus, (params) => this.setServerStatus(params));
+            this.pushClientCleanup(
+                this.client.onNotification(ra.serverStatus, (params) =>
+                    this.setServerStatus(params)
+                )
+            );
         }
         return this.client;
     }
@@ -127,16 +144,25 @@ export class Ctx {
         log.info("Activating language client");
         const client = await this.getClient();
         await client.start();
+        this.updateCommands();
         return client;
     }
 
     async deactivate() {
         log.info("Deactivating language client");
         await this.client?.stop();
+        this.updateCommands();
     }
 
-    async disposeClient() {
+    async stop() {
-        log.info("Deactivating language client");
+        log.info("Stopping language client");
+        await this.disposeClient();
+        this.updateCommands();
+    }
+
+    private async disposeClient() {
+        this.clientSubscriptions?.forEach((disposable) => disposable.dispose());
+        this.clientSubscriptions = [];
         await this.client?.dispose();
         this._serverPath = undefined;
         this.client = undefined;
@@ -159,6 +185,25 @@ export class Ctx {
         return this._serverPath;
     }
 
+    private updateCommands() {
+        this.commandDisposables.forEach((disposable) => disposable.dispose());
+        this.commandDisposables = [];
+        const fetchFactory = (factory: CommandFactory, fullName: string) => {
+            return this.client && this.client.isRunning()
+                ? factory.enabled
+                : factory.disabled ||
+                      ((_) => () =>
+                          vscode.window.showErrorMessage(
+                              `command ${fullName} failed: rust-analyzer server is not running`
+                          ));
+        };
+        for (const [name, factory] of Object.entries(this.commandFactories)) {
+            const fullName = `rust-analyzer.${name}`;
+            const callback = fetchFactory(factory, fullName)(this);
+            this.commandDisposables.push(vscode.commands.registerCommand(fullName, callback));
+        }
+    }
+
     setServerStatus(status: ServerStatusParams) {
         let icon = "";
         const statusBar = this.statusBar;
@@ -194,16 +239,13 @@ export class Ctx {
         statusBar.text = `${icon}rust-analyzer`;
     }
 
-    registerCommand(name: string, factory: (ctx: Ctx) => Cmd) {
-        const fullName = `rust-analyzer.${name}`;
-        const cmd = factory(this);
-        const d = vscode.commands.registerCommand(fullName, cmd);
-        this.pushExtCleanup(d);
-    }
-
     pushExtCleanup(d: Disposable) {
         this.extCtx.subscriptions.push(d);
     }
+
+    private pushClientCleanup(d: Disposable) {
+        this.clientSubscriptions.push(d);
+    }
 }
 
 export interface Disposable {

editors/code/src/main.ts

@@ -2,7 +2,7 @@ import * as vscode from "vscode";
 import * as lc from "vscode-languageclient/node";
 
 import * as commands from "./commands";
-import { Ctx, Workspace } from "./ctx";
+import { CommandFactory, Ctx, Workspace } from "./ctx";
 import { isRustDocument } from "./util";
 import { activateTaskProvider } from "./tasks";
 import { setContextValue } from "./util";
@@ -57,7 +57,7 @@ export async function activate(
               }
             : { kind: "Workspace Folder" };
 
-    const ctx = new Ctx(context, workspace);
+    const ctx = new Ctx(context, workspace, createCommands());
     // VS Code doesn't show a notification when an extension fails to activate
     // so we do it ourselves.
     const api = await activateServer(ctx).catch((err) => {
@@ -75,8 +75,6 @@ async function activateServer(ctx: Ctx): Promise<RustAnalyzerExtensionApi> {
         ctx.pushExtCleanup(activateTaskProvider(ctx.config));
     }
 
-    await initCommonContext(ctx);
-
     vscode.workspace.onDidChangeConfiguration(
         async (_) => {
             await ctx
@@ -91,85 +89,78 @@ async function activateServer(ctx: Ctx): Promise<RustAnalyzerExtensionApi> {
     return ctx.clientFetcher();
 }
 
-async function initCommonContext(ctx: Ctx) {
+function createCommands(): Record<string, CommandFactory> {
-    // Register a "dumb" onEnter command for the case where server fails to
+    return {
-    // start.
+        onEnter: {
-    //
+            enabled: commands.onEnter,
-    // FIXME: refactor command registration code such that commands are
+            disabled: (_) => () => vscode.commands.executeCommand("default:type", { text: "\n" }),
-    // **always** registered, even if the server does not start. Use API like
+        },
-    // this perhaps?
+        reload: {
-    //
+            enabled: (ctx) => async () => {
-    // ```TypeScript
+                void vscode.window.showInformationMessage("Reloading rust-analyzer...");
-    // registerCommand(
+                // FIXME: We should re-use the client, that is ctx.deactivate() if none of the configs have changed
-    //    factory: (Ctx) => ((Ctx) => any),
+                await ctx.stop();
-    //    fallback: () => any = () => vscode.window.showErrorMessage(
+                await ctx.activate();
-    //        "rust-analyzer is not available"
+            },
-    //    ),
+            disabled: (ctx) => async () => {
-    // )
+                void vscode.window.showInformationMessage("Reloading rust-analyzer...");
-    const defaultOnEnter = vscode.commands.registerCommand("rust-analyzer.onEnter", () =>
+                await ctx.activate();
-        vscode.commands.executeCommand("default:type", { text: "\n" })
+            },
-    );
+        },
-    ctx.pushExtCleanup(defaultOnEnter);
+        startServer: {
+            enabled: (ctx) => async () => {
+                await ctx.activate();
+            },
+            disabled: (ctx) => async () => {
+                await ctx.activate();
+            },
+        },
+        stopServer: {
+            enabled: (ctx) => async () => {
+                // FIXME: We should re-use the client, that is ctx.deactivate() if none of the configs have changed
+                await ctx.stop();
+                ctx.setServerStatus({
+                    health: "ok",
+                    quiescent: true,
+                    message: "server is not running",
+                });
+            },
+        },
 
-    // Commands which invokes manually via command palette, shortcut, etc.
+        analyzerStatus: { enabled: commands.analyzerStatus },
-    ctx.registerCommand("reload", (_) => async () => {
+        memoryUsage: { enabled: commands.memoryUsage },
-        void vscode.window.showInformationMessage("Reloading rust-analyzer...");
+        shuffleCrateGraph: { enabled: commands.shuffleCrateGraph },
-        // FIXME: We should re-use the client, that is ctx.deactivate() if none of the configs have changed
+        reloadWorkspace: { enabled: commands.reloadWorkspace },
-        await ctx.disposeClient();
+        matchingBrace: { enabled: commands.matchingBrace },
-        await ctx.activate();
+        joinLines: { enabled: commands.joinLines },
-    });
+        parentModule: { enabled: commands.parentModule },
-
+        syntaxTree: { enabled: commands.syntaxTree },
-    ctx.registerCommand("startServer", (_) => async () => {
+        viewHir: { enabled: commands.viewHir },
-        await ctx.activate();
+        viewFileText: { enabled: commands.viewFileText },
-    });
+        viewItemTree: { enabled: commands.viewItemTree },
-    ctx.registerCommand("stopServer", (_) => async () => {
+        viewCrateGraph: { enabled: commands.viewCrateGraph },
-        // FIXME: We should re-use the client, that is ctx.deactivate() if none of the configs have changed
+        viewFullCrateGraph: { enabled: commands.viewFullCrateGraph },
-        await ctx.disposeClient();
+        expandMacro: { enabled: commands.expandMacro },
-        ctx.setServerStatus({
+        run: { enabled: commands.run },
-            health: "ok",
+        copyRunCommandLine: { enabled: commands.copyRunCommandLine },
-            quiescent: true,
+        debug: { enabled: commands.debug },
-            message: "server is not running",
+        newDebugConfig: { enabled: commands.newDebugConfig },
-        });
+        openDocs: { enabled: commands.openDocs },
-    });
+        openCargoToml: { enabled: commands.openCargoToml },
-    ctx.registerCommand("analyzerStatus", commands.analyzerStatus);
+        peekTests: { enabled: commands.peekTests },
-    ctx.registerCommand("memoryUsage", commands.memoryUsage);
+        moveItemUp: { enabled: commands.moveItemUp },
-    ctx.registerCommand("shuffleCrateGraph", commands.shuffleCrateGraph);
+        moveItemDown: { enabled: commands.moveItemDown },
-    ctx.registerCommand("reloadWorkspace", commands.reloadWorkspace);
+        cancelFlycheck: { enabled: commands.cancelFlycheck },
-    ctx.registerCommand("matchingBrace", commands.matchingBrace);
+        ssr: { enabled: commands.ssr },
-    ctx.registerCommand("joinLines", commands.joinLines);
+        serverVersion: { enabled: commands.serverVersion },
-    ctx.registerCommand("parentModule", commands.parentModule);
+        // Internal commands which are invoked by the server.
-    ctx.registerCommand("syntaxTree", commands.syntaxTree);
+        applyActionGroup: { enabled: commands.applyActionGroup },
-    ctx.registerCommand("viewHir", commands.viewHir);
+        applySnippetWorkspaceEdit: { enabled: commands.applySnippetWorkspaceEditCommand },
-    ctx.registerCommand("viewFileText", commands.viewFileText);
+        debugSingle: { enabled: commands.debugSingle },
-    ctx.registerCommand("viewItemTree", commands.viewItemTree);
+        gotoLocation: { enabled: commands.gotoLocation },
-    ctx.registerCommand("viewCrateGraph", commands.viewCrateGraph);
+        linkToCommand: { enabled: commands.linkToCommand },
-    ctx.registerCommand("viewFullCrateGraph", commands.viewFullCrateGraph);
+        resolveCodeAction: { enabled: commands.resolveCodeAction },
-    ctx.registerCommand("expandMacro", commands.expandMacro);
+        runSingle: { enabled: commands.runSingle },
-    ctx.registerCommand("run", commands.run);
+        showReferences: { enabled: commands.showReferences },
-    ctx.registerCommand("copyRunCommandLine", commands.copyRunCommandLine);
+    };
-    ctx.registerCommand("debug", commands.debug);
-    ctx.registerCommand("newDebugConfig", commands.newDebugConfig);
-    ctx.registerCommand("openDocs", commands.openDocs);
-    ctx.registerCommand("openCargoToml", commands.openCargoToml);
-    ctx.registerCommand("peekTests", commands.peekTests);
-    ctx.registerCommand("moveItemUp", commands.moveItemUp);
-    ctx.registerCommand("moveItemDown", commands.moveItemDown);
-    ctx.registerCommand("cancelFlycheck", commands.cancelFlycheck);
-
-    ctx.registerCommand("ssr", commands.ssr);
-    ctx.registerCommand("serverVersion", commands.serverVersion);
-
-    // Internal commands which are invoked by the server.
-    ctx.registerCommand("runSingle", commands.runSingle);
-    ctx.registerCommand("debugSingle", commands.debugSingle);
-    ctx.registerCommand("showReferences", commands.showReferences);
-    ctx.registerCommand("applySnippetWorkspaceEdit", commands.applySnippetWorkspaceEditCommand);
-    ctx.registerCommand("resolveCodeAction", commands.resolveCodeAction);
-    ctx.registerCommand("applyActionGroup", commands.applyActionGroup);
-    ctx.registerCommand("gotoLocation", commands.gotoLocation);
-
-    ctx.registerCommand("linkToCommand", commands.linkToCommand);
-
-    defaultOnEnter.dispose();
-    ctx.registerCommand("onEnter", commands.onEnter);
 }