From 446109795539b3bb9cb6a0b260c60a3ec6d49a39 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:09:30 +0000 Subject: [PATCH] feat(ssh): pre-check the Remote SSH extension before starting the tunnel `databricks ssh connect` needs the host editor's Remote SSH extension to open the remote window, and checks for it by shelling out to ` --list-extensions` from inside the terminal. That check is the largest single source of failed IDE-mode tunnels, and its failures are sticky: most users who hit it never get a working tunnel, even across several attempts. Check the running editor's own extension registry instead, before the terminal exists, and offer a button that installs through the editor's marketplace client. Reading the registry cannot fail the way spawning a process can, so this removes the list-failure path rather than working around it. Unlike the host-CLI PATH warning next to it, this awaits the user's choice: installing takes a moment, and the point is to let the attempt they just started succeed rather than fail and depend on a retry. It still never gates the tunnel -- a dismissed or failed install falls through to the CLI's own check, which reports the real error. Adds HostUtils.getHostSshExtension/getSshExtensionStatus, mirroring the CLI's IDE descriptors the way getHostCliCommand already mirrors its --ide handling. Co-authored-by: Isaac --- .../src/ssh/SshCommands.test.ts | 93 ++++++++++++++++++- .../databricks-vscode/src/ssh/SshCommands.ts | 47 ++++++++++ .../src/utils/hostUtils.test.ts | 83 ++++++++++++++++- .../databricks-vscode/src/utils/hostUtils.ts | 58 +++++++++++- 4 files changed, 274 insertions(+), 7 deletions(-) diff --git a/packages/databricks-vscode/src/ssh/SshCommands.test.ts b/packages/databricks-vscode/src/ssh/SshCommands.test.ts index 015d57a54..fc69bd10b 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.test.ts +++ b/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()); + }); + }); }); + diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 1514fafd8..ea3bc77e0 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/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 ` --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 { + 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 diff --git a/packages/databricks-vscode/src/utils/hostUtils.test.ts b/packages/databricks-vscode/src/utils/hostUtils.test.ts index fddb45da7..d2f78a4af 100644 --- a/packages/databricks-vscode/src/utils/hostUtils.test.ts +++ b/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", + }); + }); + }); }); + diff --git a/packages/databricks-vscode/src/utils/hostUtils.ts b/packages/databricks-vscode/src/utils/hostUtils.ts index 0fc0d9c95..f8b4476ed 100644 --- a/packages/databricks-vscode/src/utils/hostUtils.ts +++ b/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 + * ` --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}; +}