feat(ssh): pre-check the Remote SSH extension before starting the tunnel

packages/databricks-vscode/src/ssh/SshCommands.test.ts

@@ -6,15 +6,17 @@ import {SshCommands} from "./SshCommands";
 import {HostUtils} from "../utils";
 
 /**
- * Exercises the advisory host-CLI PATH warning in isolation. The prompt is
- * fired-and-forgotten from startTunnelCommand, so tests reach the private
- * warnIfHostCliMissing directly and flush the detached prompt chain it kicks
- * off before asserting on the message and any follow-up command.
+ * Exercises the two tunnel pre-checks in isolation, reaching the private methods
+ * directly. The host-CLI PATH warning is fired-and-forgotten from
+ * startTunnelCommand, so those tests flush the detached prompt chain it kicks off
+ * before asserting; the Remote SSH extension check is awaited and needs no flush.
  */
 describe(__filename, () => {
     let originalIsHostCliOnPath: typeof HostUtils.isHostCliOnPath;
     let originalGetHostCliCommand: typeof HostUtils.getHostCliCommand;
     let originalIsCursor: typeof HostUtils.isCursor;
+    let originalGetSshExtensionStatus: typeof HostUtils.getSshExtensionStatus;
+    let originalGetHostSshExtension: typeof HostUtils.getHostSshExtension;
     let originalShowWarningMessage: typeof window.showWarningMessage;
     let originalExecuteCommand: typeof commands.executeCommand;
     let originalPlatform: PropertyDescriptor | undefined;
@@ -52,6 +54,8 @@ describe(__filename, () => {
         originalIsHostCliOnPath = HostUtils.isHostCliOnPath;
         originalGetHostCliCommand = HostUtils.getHostCliCommand;
         originalIsCursor = HostUtils.isCursor;
+        originalGetSshExtensionStatus = HostUtils.getSshExtensionStatus;
+        originalGetHostSshExtension = HostUtils.getHostSshExtension;
         originalShowWarningMessage = window.showWarningMessage;
         originalExecuteCommand = commands.executeCommand;
         originalPlatform = Object.getOwnPropertyDescriptor(process, "platform");
@@ -76,6 +80,9 @@ describe(__filename, () => {
         (HostUtils as any).isHostCliOnPath = originalIsHostCliOnPath;
         (HostUtils as any).getHostCliCommand = originalGetHostCliCommand;
         (HostUtils as any).isCursor = originalIsCursor;
+        (HostUtils as any).getSshExtensionStatus =
+            originalGetSshExtensionStatus;
+        (HostUtils as any).getHostSshExtension = originalGetHostSshExtension;
         (window as any).showWarningMessage = originalShowWarningMessage;
         (commands as any).executeCommand = originalExecuteCommand;
         if (originalPlatform) {
@@ -170,4 +177,82 @@ describe(__filename, () => {
         assert.strictEqual(shownWarnings.length, 1);
         assert.strictEqual(executedCommands.length, 0);
     });
+
+    describe("offerToInstallSshExtension", () => {
+        function stubExtension(status: HostUtils.SshExtensionStatus) {
+            (HostUtils as any).getSshExtensionStatus = () => status;
+            (HostUtils as any).getHostSshExtension = () => ({
+                id: "ms-vscode-remote.remote-ssh",
+                name: "Remote - SSH",
+                minVersion: "0.120.0",
+            });
+        }
+
+        function offer(sshCommands: SshCommands) {
+            return (sshCommands as any).offerToInstallSshExtension();
+        }
+
+        it("says nothing when the extension is usable", async () => {
+            stubExtension({kind: "ok"});
+
+            await offer(newSshCommands());
+
+            assert.strictEqual(shownWarnings.length, 0);
+            assert.strictEqual(executedCommands.length, 0);
+        });
+
+        it("offers to install a missing extension", async () => {
+            stubExtension({kind: "missing"});
+            warningResponse = "Install";
+
+            await offer(newSshCommands());
+
+            assert.strictEqual(shownWarnings.length, 1);
+            assert.ok(shownWarnings[0].message.includes('"Remote - SSH"'));
+            assert.deepStrictEqual(shownWarnings[0].items, ["Install"]);
+            assert.deepStrictEqual(executedCommands, [
+                {
+                    command: "workbench.extensions.installExtension",
+                    args: ["ms-vscode-remote.remote-ssh"],
+                },
+            ]);
+        });
+
+        it("offers to update an outdated extension, naming the version", async () => {
+            stubExtension({kind: "outdated", installed: "0.100.0"});
+            warningResponse = "Update";
+
+            await offer(newSshCommands());
+
+            assert.strictEqual(shownWarnings.length, 1);
+            assert.ok(shownWarnings[0].message.includes("0.100.0"));
+            assert.deepStrictEqual(shownWarnings[0].items, ["Update"]);
+            assert.strictEqual(
+                executedCommands[0].command,
+                "workbench.extensions.installExtension"
+            );
+        });
+
+        it("installs nothing when the prompt is dismissed", async () => {
+            stubExtension({kind: "missing"});
+            warningResponse = undefined;
+
+            await offer(newSshCommands());
+
+            assert.strictEqual(shownWarnings.length, 1);
+            assert.strictEqual(executedCommands.length, 0);
+        });
+
+        // The pre-check must never gate the tunnel: the CLI repeats the check and
+        // reports the real error, so a broken install here has to fall through.
+        it("does not throw when the install fails", async () => {
+            stubExtension({kind: "missing"});
+            warningResponse = "Install";
+            (commands as any).executeCommand = () =>
+                Promise.reject(new Error("marketplace unreachable"));
+
+            await offer(newSshCommands());
+        });
+    });
 });
+

packages/databricks-vscode/src/ssh/SshCommands.ts

@@ -182,6 +182,7 @@ export class SshCommands implements Disposable {
         // terminal report the real error rather than blocking a tunnel that
         // would have worked.
         await this.warnIfHostCliMissing();
+        await this.offerToInstallSshExtension();
         const context = await this.resolveTunnelContext();
         if (context === undefined) {
             return;
@@ -482,6 +483,52 @@ export class SshCommands implements Disposable {
         await commands.executeCommand("vscode.open", Uri.parse(url));
     }
 
+    /**
+     * Offers to install the host's Remote SSH extension when it is missing or
+     * too old.
+     *
+     * `databricks ssh connect` checks this too, but from inside the terminal and
+     * by shelling out to `<command> --list-extensions`; that check is the
+     * largest single source of failed IDE-mode tunnels. Here the editor's own
+     * registry answers directly and its marketplace client does the install.
+     *
+     * Unlike warnIfHostCliMissing this awaits the choice rather than firing and
+     * forgetting: installing takes a moment, and the point is to let the attempt
+     * the user just started succeed instead of failing and relying on a retry.
+     * It still never blocks the tunnel — a dismissed or failed install falls
+     * through to the CLI's own check, which reports the real error.
+     */
+    private async offerToInstallSshExtension(): Promise<void> {
+        const status = HostUtils.getSshExtensionStatus();
+        if (status.kind === "ok") {
+            return;
+        }
+        const {id, name} = HostUtils.getHostSshExtension();
+        const detail =
+            status.kind === "missing"
+                ? "is not installed"
+                : `is version ${status.installed}, which is too old`;
+        const action = status.kind === "missing" ? "Install" : "Update";
+        const message =
+            `The "${name}" extension ${detail}. The Databricks SSH tunnel ` +
+            `needs it to open the remote window.`;
+
+        if ((await window.showWarningMessage(message, action)) !== action) {
+            return;
+        }
+        try {
+            await commands.executeCommand(
+                "workbench.extensions.installExtension",
+                id
+            );
+        } catch (e) {
+            logging.NamedLogger.getOrCreate(Loggers.Extension).error(
+                `Failed to install the "${name}" extension`,
+                e
+            );
+        }
+    }
+
     private async launchSshTunnel(
         authProvider: AuthProvider,
         compute: Compute

packages/databricks-vscode/src/utils/hostUtils.test.ts

@@ -1,6 +1,12 @@
-import {env} from "vscode";
+import {env, extensions} from "vscode";
 import assert from "assert";
-import {getHostCliCommand, isCursor, isHostCliOnPath} from "./hostUtils";
+import {
+    getHostCliCommand,
+    getHostSshExtension,
+    getSshExtensionStatus,
+    isCursor,
+    isHostCliOnPath,
+} from "./hostUtils";
 import {cancellableExecFile} from "../cli/CliWrapper";
 
 describe(__filename, () => {
@@ -118,4 +124,77 @@ describe(__filename, () => {
             );
         });
     });
+
+    describe("getSshExtensionStatus", () => {
+        let originalGetExtension: typeof extensions.getExtension;
+
+        // Stands in for the one field getSshExtensionStatus reads off an extension.
+        function stubInstalled(version: string | undefined) {
+            (extensions as any).getExtension = () =>
+                version === undefined
+                    ? undefined
+                    : {packageJSON: {version}};
+        }
+
+        beforeEach(() => {
+            originalGetExtension = extensions.getExtension;
+            stubUriScheme("vscode");
+        });
+
+        afterEach(() => {
+            (extensions as any).getExtension = originalGetExtension;
+        });
+
+        it("resolves the extension id per host", () => {
+            stubUriScheme("cursor");
+            assert.strictEqual(
+                getHostSshExtension().id,
+                "anysphere.remote-ssh"
+            );
+            stubUriScheme("vscode");
+            assert.strictEqual(
+                getHostSshExtension().id,
+                "ms-vscode-remote.remote-ssh"
+            );
+        });
+
+        it("is missing when the extension is not in the registry", () => {
+            stubInstalled(undefined);
+            assert.deepStrictEqual(getSshExtensionStatus(), {kind: "missing"});
+        });
+
+        it("is ok at and above the minimum version", () => {
+            stubInstalled("0.120.0");
+            assert.deepStrictEqual(getSshExtensionStatus(), {kind: "ok"});
+            stubInstalled("0.130.1");
+            assert.deepStrictEqual(getSshExtensionStatus(), {kind: "ok"});
+        });
+
+        it("is outdated below the minimum version, and reports it", () => {
+            stubInstalled("0.100.0");
+            assert.deepStrictEqual(getSshExtensionStatus(), {
+                kind: "outdated",
+                installed: "0.100.0",
+            });
+        });
+
+        it("treats an unparseable version as outdated, like the CLI does", () => {
+            stubInstalled("not-a-version");
+            assert.deepStrictEqual(getSshExtensionStatus(), {
+                kind: "outdated",
+                installed: "not-a-version",
+            });
+        });
+
+        it("applies the Cursor floor in Cursor", () => {
+            stubUriScheme("cursor");
+            // Below Cursor's 1.0.32 floor but far above VS Code's 0.120.0 one.
+            stubInstalled("1.0.10");
+            assert.deepStrictEqual(getSshExtensionStatus(), {
+                kind: "outdated",
+                installed: "1.0.10",
+            });
+        });
+    });
 });
+

packages/databricks-vscode/src/utils/hostUtils.ts

@@ -1,5 +1,6 @@
-import {env} from "vscode";
+import {env, extensions} from "vscode";
 import {ExecUtils, logging} from "@databricks/sdk-experimental";
+import * as semver from "semver";
 import {cancellableExecFile} from "../cli/CliWrapper";
 import {Loggers} from "../logger";
 
@@ -85,3 +86,58 @@ export async function isHostCliOnPath(
         return true;
     }
 }
+
+/**
+ * The Remote SSH extension `databricks ssh connect` opens the remote window
+ * through, and the lowest version it accepts. Mirrors the CLI's IDE descriptors
+ * the way getHostCliCommand mirrors its `--ide` handling: the CLI enforces this
+ * id and floor itself, so a value that drifts from it passes here and then
+ * fails in the terminal.
+ */
+export function getHostSshExtension(): {
+    id: string;
+    name: string;
+    minVersion: string;
+} {
+    if (isCursor()) {
+        return {
+            id: "anysphere.remote-ssh",
+            name: "Remote - SSH",
+            minVersion: "1.0.32",
+        };
+    }
+    return {
+        id: "ms-vscode-remote.remote-ssh",
+        name: "Remote - SSH",
+        minVersion: "0.120.0",
+    };
+}
+
+export type SshExtensionStatus =
+    | {kind: "ok"}
+    | {kind: "missing"}
+    | {kind: "outdated"; installed: string};
+
+/**
+ * Whether the host's Remote SSH extension is installed at a version the tunnel
+ * can use.
+ *
+ * The CLI answers the same question by shelling out to
+ * `<command> --list-extensions`, which can fail for reasons that say nothing
+ * about the extension. Reading the running editor's registry instead cannot.
+ * Extensions disabled in this window are absent from it and so read as missing,
+ * which is the right answer here: a disabled Remote SSH cannot open a window.
+ */
+export function getSshExtensionStatus(): SshExtensionStatus {
+    const {id, minVersion} = getHostSshExtension();
+    const installed = extensions.getExtension(id);
+    if (installed === undefined) {
+        return {kind: "missing"};
+    }
+    // Matches the CLI, which also treats an unparseable version as too old.
+    const version = String(installed.packageJSON.version);
+    if (semver.valid(version) !== null && semver.gte(version, minVersion)) {
+        return {kind: "ok"};
+    }
+    return {kind: "outdated", installed: version};
+}