From 0fbc411bc8b93e3acbdf28c84112c63695f6b4dd Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 2 Jul 2026 12:45:36 +0400 Subject: [PATCH 01/23] Start SSH tunnel --- packages/databricks-vscode/package.json | 7 + .../src/cli/CliWrapper.test.ts | 42 ++++ .../databricks-vscode/src/cli/CliWrapper.ts | 37 ++++ packages/databricks-vscode/src/extension.ts | 35 +++ .../src/language/SshTunnelStatusBarButton.ts | 46 ++++ .../databricks-vscode/src/ssh/SshCommands.ts | 200 ++++++++++++++++++ 6 files changed, 367 insertions(+) create mode 100644 packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts create mode 100644 packages/databricks-vscode/src/ssh/SshCommands.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 2a18106ba..e1477fc73 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -52,6 +52,13 @@ "types": "out/extension.d.ts", "contributes": { "commands": [ + { + "command": "databricks.ssh.startTunnel", + "title": "Start SSH Tunnel", + "category": "Databricks", + "icon": "$(remote)", + "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode" + }, { "command": "databricks.connection.logout", "title": "Logout", diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index 810e177f2..e4c7eaf2c 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -128,6 +128,48 @@ describe(__filename, function () { assert.equal([command, ...args].join(" "), syncCommand); }); + it("should create ssh connect commands", () => { + const logFilePath = getTempLogFilePath(); + const cli = createCliWrapper(logFilePath); + const loggingArgs = [ + "--log-level", + "debug", + "--log-file", + logFilePath, + "--log-format", + "json", + ]; + + // Serverless: no --cluster / --auto-start-cluster. + let {args} = cli.getSshConnectCommand({compute: {type: "serverless"}}); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + ...loggingArgs, + ]); + + // Dedicated cluster: --cluster and --auto-start-cluster. + ({args} = cli.getSshConnectCommand({ + compute: {type: "cluster", clusterId: "1234-clusterid"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--cluster=1234-clusterid", + "--auto-start-cluster", + ...loggingArgs, + ]); + + // No logging args when logging is disabled. + const configsSpy = spy(workspaceConfigs); + mocks.push(configsSpy); + when(configsSpy.loggingEnabled).thenReturn(false); + ({args} = cli.getSshConnectCommand({compute: {type: "serverless"}})); + assert.deepStrictEqual(args, ["ssh", "connect", "--ide=vscode"]); + }); + it("should list profiles when no config file exists", async () => { const logFilePath = getTempLogFilePath(); const cli = createCliWrapper(logFilePath); diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 17ca9170b..62074eaae 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -574,6 +574,43 @@ export class CliWrapper { }); } + /** + * Env vars for interactive CLI commands run in a terminal (e.g. `ssh + * connect`). Auth is forwarded via env vars, matching the bundle init flow. + */ + getSshConnectEnvVars(authProvider: AuthProvider) { + return removeUndefinedKeys({ + ...EnvVarGenerators.getEnvVarsForCli( + this.extensionContext, + workspaceConfigs.databrickscfgLocation + ), + ...EnvVarGenerators.getProxyEnvVars(), + ...this.getLogginEnvVars(), + ...authProvider.toEnv(), + // eslint-disable-next-line @typescript-eslint/naming-convention + DATABRICKS_OUTPUT_FORMAT: "text", + }); + } + + /** + * Constructs the `databricks ssh connect` command args for opening a remote + * VS Code window. Serverless is the default when no cluster is given. + */ + getSshConnectCommand(opts: { + compute: {type: "serverless"} | {type: "cluster"; clusterId: string}; + }): {args: string[]} { + const args = ["ssh", "connect", "--ide=vscode"]; + if (opts.compute.type === "cluster") { + args.push(`--cluster=${opts.compute.clusterId}`); + // Start a stopped single-user cluster when connecting. Host-key + // trust and extension-install prompts stay interactive (we do not + // pass --auto-approve). + args.push("--auto-start-cluster"); + } + args.push(...this.getLoggingArguments()); + return {args}; + } + async bundleInit( templateDirPath: string, outputDirPath: string, diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 1cfacf89e..493e8d2c9 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -4,6 +4,7 @@ import { env, ExtensionContext, extensions, + Uri, window, workspace, } from "vscode"; @@ -49,6 +50,8 @@ import {Events, Metadata} from "./telemetry/constants"; import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller"; import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits"; import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton"; +import {SshTunnelStatusBarButton} from "./language/SshTunnelStatusBarButton"; +import {SshCommands} from "./ssh/SshCommands"; import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager"; import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue"; import { @@ -267,6 +270,23 @@ export async function activate( ); } + // Auto-open the file the user was editing locally when the tunnel was + // started, if it resolves in the remote workspace. Forwarded out of + // band via env var since the CLI has no file flag; best-effort only. + const remoteOpenFile = process.env["DATABRICKS_REMOTE_OPEN_FILE"]; + if (remoteOpenFile) { + try { + const uri = Uri.file(remoteOpenFile); + await workspace.fs.stat(uri); + await window.showTextDocument(uri); + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).debug( + "Skipping remote auto-open of file (not found remotely)", + {remoteOpenFile} + ); + } + } + customWhenContext.setActivated(true); return; } @@ -825,6 +845,21 @@ export async function activate( ) ); + // SSH tunnel (remote development) group + const sshCommands = new SshCommands(connectionManager, clusterModel, cli); + const sshTunnelStatusBarButton = new SshTunnelStatusBarButton( + connectionManager + ); + context.subscriptions.push( + sshCommands, + sshTunnelStatusBarButton, + telemetry.registerCommand( + "databricks.ssh.startTunnel", + sshCommands.startTunnelCommand, + sshCommands + ) + ); + // Cluster group const clusterTreeDataProvider = new ClusterListDataProvider(clusterModel); const clusterCommands = new ClusterCommands( diff --git a/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts b/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts new file mode 100644 index 000000000..4d3823fe1 --- /dev/null +++ b/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts @@ -0,0 +1,46 @@ +import { + Disposable, + StatusBarAlignment, + StatusBarItem, + window, +} from "vscode"; +import {ConnectionManager} from "../configuration/ConnectionManager"; + +export class SshTunnelStatusBarButton implements Disposable { + private disposables: Disposable[] = []; + statusBarButton: StatusBarItem; + + constructor(private readonly connectionManager: ConnectionManager) { + // Priority 999 places this just to the right of the "Databricks + // Connect" button (priority 1000) in the left status bar group. + this.statusBarButton = window.createStatusBarItem( + StatusBarAlignment.Left, + 999 + ); + this.statusBarButton.name = "Start SSH Tunnel"; + this.statusBarButton.text = "$(remote) Start SSH Tunnel"; + this.statusBarButton.tooltip = + "Start Databricks remote development via SSH"; + this.statusBarButton.command = { + title: "Start SSH Tunnel", + command: "databricks.ssh.startTunnel", + }; + this.disposables.push( + this.statusBarButton, + this.connectionManager.onDidChangeState(this.update, this) + ); + this.update(); + } + + private update() { + if (this.connectionManager.state === "CONNECTED") { + this.statusBarButton.show(); + } else { + this.statusBarButton.hide(); + } + } + + dispose() { + this.disposables.forEach((i) => i.dispose()); + } +} diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts new file mode 100644 index 000000000..5d93e6f40 --- /dev/null +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -0,0 +1,200 @@ +import { + Disposable, + QuickPick, + QuickPickItem, + QuickPickItemKind, + ThemeIcon, + window, +} from "vscode"; +import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; +import {ClusterModel} from "../cluster/ClusterModel"; +import {ConnectionManager} from "../configuration/ConnectionManager"; +import { + ClusterItem, + formatQuickPickClusterDetails, +} from "../configuration/ConnectionCommands"; +import {CliWrapper} from "../cli/CliWrapper"; +import {onError} from "../utils/onErrorDecorator"; + +const SERVERLESS_LABEL = "$(cloud) Serverless"; + +type Compute = {type: "serverless"} | {type: "cluster"; clusterId: string}; + +export class SshCommands implements Disposable { + private disposables: Disposable[] = []; + + constructor( + private readonly connectionManager: ConnectionManager, + private readonly clusterModel: ClusterModel, + private readonly cli: CliWrapper + ) {} + + /** + * Shows a compute picker (single-user clusters + serverless), then launches + * `databricks ssh connect --ide=vscode` in a terminal to open a remote + * VS Code window. + */ + @onError({popup: {prefix: "Error starting SSH tunnel."}}) + async startTunnelCommand() { + const workspaceClient = this.connectionManager.workspaceClient; + const me = this.connectionManager.databricksWorkspace?.userName; + if (!workspaceClient || !me) { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return; + } + + const compute = await this.pickCompute(); + if (compute === undefined) { + return; + } + await this.launchSshTunnel(compute); + } + + private pickCompute(): Promise { + return new Promise((resolve) => { + const quickPick = window.createQuickPick< + ClusterItem | QuickPickItem + >(); + quickPick.title = "Select compute for SSH tunnel"; + quickPick.keepScrollPosition = true; + quickPick.busy = true; + quickPick.canSelectMany = false; + + const staticItems: QuickPickItem[] = [ + { + label: SERVERLESS_LABEL, + detail: "Connect to serverless compute (no dedicated cluster)", + alwaysShow: true, + }, + { + label: "", + kind: QuickPickItemKind.Separator, + }, + ]; + quickPick.items = staticItems; + + this.clusterModel.refresh(); + const refreshItems = () => { + // Only dedicated single-user access-mode clusters can be used + // for an SSH tunnel. + const clusters = (this.clusterModel.roots ?? []).filter((c) => + c.isSingleUser() + ); + quickPick.items = staticItems.concat( + clusters.map((c) => { + const treeItem = + ClusterListDataProvider.clusterNodeToTreeItem(c); + return { + label: `$(${ + (treeItem.iconPath as ThemeIcon).id + }) ${c.name!} (${c.id})`, + detail: formatQuickPickClusterDetails(c), + cluster: c, + }; + }) + ); + this.preselect(quickPick); + }; + + const disposables = [ + this.clusterModel.onDidChange(refreshItems), + quickPick, + ]; + + refreshItems(); + quickPick.busy = false; + quickPick.show(); + + quickPick.onDidAccept(() => { + const selectedItem = quickPick.selectedItems[0]; + disposables.forEach((d) => d.dispose()); + if (selectedItem === undefined) { + resolve(undefined); + } else if ("cluster" in selectedItem) { + resolve({ + type: "cluster", + clusterId: selectedItem.cluster.id, + }); + } else { + resolve({type: "serverless"}); + } + }); + + quickPick.onDidHide(() => { + disposables.forEach((d) => d.dispose()); + // resolve(undefined); + }); + }); + } + + /** + * Pre-selects the compute the user already has configured locally: the + * attached single-user cluster if any, otherwise serverless. + */ + private preselect(quickPick: QuickPick) { + const currentCluster = this.connectionManager.cluster; + if (currentCluster?.isSingleUser()) { + const match = quickPick.items.find( + (i): i is ClusterItem => + "cluster" in i && i.cluster.id === currentCluster.id + ); + if (match) { + quickPick.activeItems = [match]; + return; + } + } + if (this.connectionManager.serverless) { + const serverlessItem = quickPick.items.find( + (i) => i.label === SERVERLESS_LABEL + ); + if (serverlessItem) { + quickPick.activeItems = [serverlessItem]; + } + } + } + + private async launchSshTunnel(compute: Compute) { + const authProvider = + this.connectionManager.databricksWorkspace?.authProvider; + const userName = this.connectionManager.databricksWorkspace?.userName; + if (!authProvider || !userName) { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return; + } + + const {args} = this.cli.getSshConnectCommand({compute}); + + const env: Record = { + ...this.cli.getSshConnectEnvVars(authProvider), + // The remote window opens at the user's home folder. Forward the + // file the user is currently editing so the remote extension can + // auto-open it (see the remote-mode branch in extension.ts). The + // CLI has no folder/file flag, so we pass it out of band. + /* eslint-disable @typescript-eslint/naming-convention */ + DATABRICKS_REMOTE_HOME_FOLDER: `/Users/${userName}`, + /* eslint-enable @typescript-eslint/naming-convention */ + }; + const activeFile = window.activeTextEditor?.document.uri.fsPath; + if (activeFile) { + env["DATABRICKS_REMOTE_OPEN_FILE"] = activeFile; + } + + const terminal = window.createTerminal({ + name: "Databricks SSH Tunnel", + isTransient: true, + env, + strictEnv: false, + }); + this.disposables.push(terminal); + terminal.show(); + terminal.sendText(`${this.cli.escapedCliPath} ${args.join(" ")}`); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } +} From 2f26576c925b8ba49abca1a4960df0eafb94fd3d Mon Sep 17 00:00:00 2001 From: misha-db Date: Fri, 3 Jul 2026 19:27:24 +0400 Subject: [PATCH 02/23] Add accelerator support --- .../src/cli/CliWrapper.test.ts | 34 ++++---- .../databricks-vscode/src/cli/CliWrapper.ts | 25 ++++-- .../src/language/SshTunnelStatusBarButton.ts | 7 +- .../databricks-vscode/src/ssh/SshCommands.ts | 84 +++++++++++++++---- 4 files changed, 105 insertions(+), 45 deletions(-) diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index e4c7eaf2c..3c30370b6 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -131,14 +131,9 @@ describe(__filename, function () { it("should create ssh connect commands", () => { const logFilePath = getTempLogFilePath(); const cli = createCliWrapper(logFilePath); - const loggingArgs = [ - "--log-level", - "debug", - "--log-file", - logFilePath, - "--log-format", - "json", - ]; + + // Logging is configured via env vars, not CLI flags, so no --log-* + // args appear on the ssh connect command line. // Serverless: no --cluster / --auto-start-cluster. let {args} = cli.getSshConnectCommand({compute: {type: "serverless"}}); @@ -146,7 +141,19 @@ describe(__filename, function () { "ssh", "connect", "--ide=vscode", - ...loggingArgs, + "--auto-approve", + ]); + + // Serverless GPU: --accelerator, no --cluster / --auto-start-cluster. + ({args} = cli.getSshConnectCommand({ + compute: {type: "serverless", accelerator: "GPU_1xA10"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + "--accelerator=GPU_1xA10", ]); // Dedicated cluster: --cluster and --auto-start-cluster. @@ -157,17 +164,10 @@ describe(__filename, function () { "ssh", "connect", "--ide=vscode", + "--auto-approve", "--cluster=1234-clusterid", "--auto-start-cluster", - ...loggingArgs, ]); - - // No logging args when logging is disabled. - const configsSpy = spy(workspaceConfigs); - mocks.push(configsSpy); - when(configsSpy.loggingEnabled).thenReturn(false); - ({args} = cli.getSshConnectCommand({compute: {type: "serverless"}})); - assert.deepStrictEqual(args, ["ssh", "connect", "--ide=vscode"]); }); it("should list profiles when no config file exists", async () => { diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 62074eaae..f7a3f89f8 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -10,6 +10,7 @@ import { Uri, commands, CancellationToken, + env, } from "vscode"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {promisify} from "node:util"; @@ -594,20 +595,30 @@ export class CliWrapper { /** * Constructs the `databricks ssh connect` command args for opening a remote - * VS Code window. Serverless is the default when no cluster is given. + * IDE window. Serverless is the default when no cluster is given. + * + * The --ide flag matches the host editor so the CLI opens the right remote + * window: Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. + * + * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see + * getSshConnectEnvVars), so we do not pass --log-* flags here. */ getSshConnectCommand(opts: { - compute: {type: "serverless"} | {type: "cluster"; clusterId: string}; + compute: + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; }): {args: string[]} { - const args = ["ssh", "connect", "--ide=vscode"]; + const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; + const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; if (opts.compute.type === "cluster") { + // Start a stopped single-user cluster when connecting. args.push(`--cluster=${opts.compute.clusterId}`); - // Start a stopped single-user cluster when connecting. Host-key - // trust and extension-install prompts stay interactive (we do not - // pass --auto-approve). args.push("--auto-start-cluster"); + } else if (opts.compute.accelerator) { + // Serverless GPU: request a specific accelerator type. + args.push(`--accelerator=${opts.compute.accelerator}`); } - args.push(...this.getLoggingArguments()); return {args}; } diff --git a/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts b/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts index 4d3823fe1..251550d02 100644 --- a/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts +++ b/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts @@ -11,11 +11,12 @@ export class SshTunnelStatusBarButton implements Disposable { statusBarButton: StatusBarItem; constructor(private readonly connectionManager: ConnectionManager) { - // Priority 999 places this just to the right of the "Databricks - // Connect" button (priority 1000) in the left status bar group. + // Priority 1001 places this just to the left of the "Databricks + // Connect" button (priority 1000) in the left status bar group + // (higher priority sits further left). this.statusBarButton = window.createStatusBarItem( StatusBarAlignment.Left, - 999 + 1001 ); this.statusBarButton.name = "Start SSH Tunnel"; this.statusBarButton.text = "$(remote) Start SSH Tunnel"; diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 5d93e6f40..c40a80719 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -18,7 +18,40 @@ import {onError} from "../utils/onErrorDecorator"; const SERVERLESS_LABEL = "$(cloud) Serverless"; -type Compute = {type: "serverless"} | {type: "cluster"; clusterId: string}; +type Compute = + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + +/** + * A serverless QuickPick item, optionally carrying a GPU accelerator type + * forwarded to `databricks ssh connect --accelerator`. Plain serverless has no + * accelerator. + */ +interface ServerlessItem extends QuickPickItem { + accelerator?: string; +} + +// Serverless compute options shown at the top of the picker: plain serverless +// plus the serverless GPU accelerator types supported by the CLI. +const SERVERLESS_ITEMS: ServerlessItem[] = [ + { + label: SERVERLESS_LABEL, + detail: "Connect to serverless compute (no dedicated cluster)", + alwaysShow: true, + }, + { + label: "$(cloud) Serverless GPU 1xA10", + detail: "Connect to serverless GPU compute (GPU_1xA10)", + alwaysShow: true, + accelerator: "GPU_1xA10", + }, + { + label: "$(cloud) Serverless GPU 8xH100", + detail: "Connect to serverless GPU compute (GPU_8xH100)", + alwaysShow: true, + accelerator: "GPU_8xH100", + }, +]; export class SshCommands implements Disposable { private disposables: Disposable[] = []; @@ -45,29 +78,25 @@ export class SshCommands implements Disposable { return; } - const compute = await this.pickCompute(); + const compute = await this.pickCompute(me); if (compute === undefined) { return; } await this.launchSshTunnel(compute); } - private pickCompute(): Promise { + private pickCompute(me: string): Promise { return new Promise((resolve) => { const quickPick = window.createQuickPick< - ClusterItem | QuickPickItem + ClusterItem | ServerlessItem >(); quickPick.title = "Select compute for SSH tunnel"; quickPick.keepScrollPosition = true; quickPick.busy = true; quickPick.canSelectMany = false; - const staticItems: QuickPickItem[] = [ - { - label: SERVERLESS_LABEL, - detail: "Connect to serverless compute (no dedicated cluster)", - alwaysShow: true, - }, + const staticItems: ServerlessItem[] = [ + ...SERVERLESS_ITEMS, { label: "", kind: QuickPickItemKind.Separator, @@ -75,12 +104,11 @@ export class SshCommands implements Disposable { ]; quickPick.items = staticItems; - this.clusterModel.refresh(); const refreshItems = () => { - // Only dedicated single-user access-mode clusters can be used - // for an SSH tunnel. + // Only dedicated single-user clusters owned by the current user + // can be used for an SSH tunnel. const clusters = (this.clusterModel.roots ?? []).filter((c) => - c.isSingleUser() + c.isValidSingleUser(me) ); quickPick.items = staticItems.concat( clusters.map((c) => { @@ -95,16 +123,33 @@ export class SshCommands implements Disposable { }; }) ); + // Clear the spinner only once clusters have actually loaded. + // On a cold first open the loader is still fetching, so we keep + // spinning and let onDidChange repaint when clusters arrive. + if (clusters.length > 0) { + quickPick.busy = false; + } this.preselect(quickPick); }; - const disposables = [ + // Fallback so the spinner can't hang forever for a user with no + // eligible clusters (onDidChange may never add any). + const spinnerTimeout = setTimeout(() => { + quickPick.busy = false; + }, 10_000); + + // Register the change listener before triggering refresh() so no + // onDidChange fired by the (re)started loader can be missed. + const disposables: Disposable[] = [ this.clusterModel.onDidChange(refreshItems), quickPick, + {dispose: () => clearTimeout(spinnerTimeout)}, ]; + // Paint whatever is already cached first (fast path on reopen), then + // trigger a reload; fresh results stream in via onDidChange. refreshItems(); - quickPick.busy = false; + this.clusterModel.refresh(); quickPick.show(); quickPick.onDidAccept(() => { @@ -118,7 +163,10 @@ export class SshCommands implements Disposable { clusterId: selectedItem.cluster.id, }); } else { - resolve({type: "serverless"}); + resolve({ + type: "serverless", + accelerator: selectedItem.accelerator, + }); } }); @@ -133,7 +181,7 @@ export class SshCommands implements Disposable { * Pre-selects the compute the user already has configured locally: the * attached single-user cluster if any, otherwise serverless. */ - private preselect(quickPick: QuickPick) { + private preselect(quickPick: QuickPick) { const currentCluster = this.connectionManager.cluster; if (currentCluster?.isSingleUser()) { const match = quickPick.items.find( From 132d3f32e6bf04eca143d4800d4874e6d520203c Mon Sep 17 00:00:00 2001 From: misha-db Date: Fri, 10 Jul 2026 17:41:05 +0400 Subject: [PATCH 03/23] Rename cluster to compute. Move SSH button to configuration section --- .../DATABRICKS.quickstart.md | 6 +-- packages/databricks-vscode/package.json | 28 +++++++---- .../src/configuration/ConnectionCommands.ts | 2 +- packages/databricks-vscode/src/extension.ts | 5 -- .../src/language/SshTunnelStatusBarButton.ts | 47 ------------------- .../databricks-vscode/src/ssh/SshCommands.ts | 7 +-- .../ui/configuration-view/ClusterComponent.ts | 4 +- .../ConfigurationDataProvider.ts | 2 + .../configuration-view/SshTunnelComponent.ts | 45 ++++++++++++++++++ 9 files changed, 73 insertions(+), 73 deletions(-) delete mode 100644 packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts create mode 100644 packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index 688416493..00ae9b7fd 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -48,12 +48,12 @@ The Databricks extension for Visual Studio Code enables you to connect to your r If your folder has multiple [Declarative Automation Bundles](#dabs), you can select which one to use by clicking "Open Existing Databricks project" button and selecting the desired project. -## Select a cluster +## Select a compute -The extension uses an interactive cluster to run code. To select an interactive cluster: +The extension uses an interactive compute to run code. To select an interactive compute: 1. Open the Databricks panel by clicking on the Databricks icon on the left -2. Click on the "Select Cluster" button. +2. Click on the "Select Compute" button. - If you wish to change the selected cluster, click on the "Configure Cluster" gear icon, next to the name of the selected cluster. ## Run Python code diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 63b85c6c9..4cb1c52e9 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -80,20 +80,20 @@ }, { "command": "databricks.connection.attachCluster", - "title": "Attach cluster", + "title": "Attach compute", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(plug)" }, { "command": "databricks.connection.attachClusterQuickPick", - "title": "Configure cluster", + "title": "Configure compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(gear)" }, { "command": "databricks.connection.detachCluster", - "title": "Detach cluster", + "title": "Detach compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(debug-disconnect)" @@ -158,14 +158,14 @@ }, { "command": "databricks.cluster.start", - "title": "Start Cluster", + "title": "Start Compute", "icon": "$(debug-start)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" }, { "command": "databricks.cluster.stop", - "title": "Stop Cluster", + "title": "Stop Compute", "icon": "$(stop-circle)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" @@ -549,7 +549,7 @@ { "id": "clusterView", "when": "databricks.feature.views.cluster && !databricks.context.remoteMode", - "name": "Clusters" + "name": "Computes" }, { "id": "dabsResourceExplorerView", @@ -1108,6 +1108,16 @@ "command": "databricks.sync.stop", "when": "view == configurationView && viewItem =~ /^databricks.*sync.*is-running.*$/ && databricks.context.bundle.isDevTarget", "group": "navigation@0" + }, + { + "command": "databricks.ssh.startTunnel", + "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", + "group": "inline@0" + }, + { + "command": "databricks.ssh.startTunnel", + "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", + "group": "navigation@0" } ], "editor/title": [ @@ -1240,7 +1250,7 @@ "submenus": [ { "id": "databricks.cluster.filter", - "label": "Filter clusters ...", + "label": "Filter compute ...", "icon": "$(filter)" }, { @@ -1446,7 +1456,7 @@ "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": false, - "description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)" + "description": "Enable/disable filtering for only accessible computes (computes on which the current user can run code)" }, "databricks.overrideDatabricksConfigFile": { "type": "string", @@ -1462,7 +1472,7 @@ "views.workspace" ], "enumDescriptions": [ - "Show cluster view in the explorer.", + "Show compute view in the explorer.", "Show workspace browser in the explorer." ], "type": "string" diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 1f4e63e1d..b9ac7005a 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -133,7 +133,7 @@ export class ConnectionCommands implements Disposable { ClusterItem | QuickPickItem >(); quickPick.title = - typeof title === "string" ? title : "Select Cluster"; + typeof title === "string" ? title : "Select Compute"; quickPick.keepScrollPosition = true; quickPick.busy = true; quickPick.canSelectMany = false; diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index abc957bca..981ea2d99 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -50,7 +50,6 @@ import {Events, Metadata} from "./telemetry/constants"; import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller"; import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits"; import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton"; -import {SshTunnelStatusBarButton} from "./language/SshTunnelStatusBarButton"; import {SshCommands} from "./ssh/SshCommands"; import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager"; import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue"; @@ -847,12 +846,8 @@ export async function activate( // SSH tunnel (remote development) group const sshCommands = new SshCommands(connectionManager, clusterModel, cli); - const sshTunnelStatusBarButton = new SshTunnelStatusBarButton( - connectionManager - ); context.subscriptions.push( sshCommands, - sshTunnelStatusBarButton, telemetry.registerCommand( "databricks.ssh.startTunnel", sshCommands.startTunnelCommand, diff --git a/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts b/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts deleted file mode 100644 index 251550d02..000000000 --- a/packages/databricks-vscode/src/language/SshTunnelStatusBarButton.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { - Disposable, - StatusBarAlignment, - StatusBarItem, - window, -} from "vscode"; -import {ConnectionManager} from "../configuration/ConnectionManager"; - -export class SshTunnelStatusBarButton implements Disposable { - private disposables: Disposable[] = []; - statusBarButton: StatusBarItem; - - constructor(private readonly connectionManager: ConnectionManager) { - // Priority 1001 places this just to the left of the "Databricks - // Connect" button (priority 1000) in the left status bar group - // (higher priority sits further left). - this.statusBarButton = window.createStatusBarItem( - StatusBarAlignment.Left, - 1001 - ); - this.statusBarButton.name = "Start SSH Tunnel"; - this.statusBarButton.text = "$(remote) Start SSH Tunnel"; - this.statusBarButton.tooltip = - "Start Databricks remote development via SSH"; - this.statusBarButton.command = { - title: "Start SSH Tunnel", - command: "databricks.ssh.startTunnel", - }; - this.disposables.push( - this.statusBarButton, - this.connectionManager.onDidChangeState(this.update, this) - ); - this.update(); - } - - private update() { - if (this.connectionManager.state === "CONNECTED") { - this.statusBarButton.show(); - } else { - this.statusBarButton.hide(); - } - } - - dispose() { - this.disposables.forEach((i) => i.dispose()); - } -} diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index c40a80719..28b22e379 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -62,11 +62,6 @@ export class SshCommands implements Disposable { private readonly cli: CliWrapper ) {} - /** - * Shows a compute picker (single-user clusters + serverless), then launches - * `databricks ssh connect --ide=vscode` in a terminal to open a remote - * VS Code window. - */ @onError({popup: {prefix: "Error starting SSH tunnel."}}) async startTunnelCommand() { const workspaceClient = this.connectionManager.workspaceClient; @@ -172,7 +167,7 @@ export class SshCommands implements Disposable { quickPick.onDidHide(() => { disposables.forEach((d) => d.dispose()); - // resolve(undefined); + // resolve(undefined); }); }); } diff --git a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts index 6a0814f74..6303e132c 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts @@ -114,7 +114,7 @@ export class ClusterComponent extends BaseComponent { // We are logged in -> Select cluster prompt return [ { - label: LabelUtils.highlightedLabel("Select a cluster"), + label: LabelUtils.highlightedLabel("Select a compute"), collapsibleState: TreeItemCollapsibleState.None, contextValue: getContextValue("none"), iconPath: new ThemeIcon( @@ -123,7 +123,7 @@ export class ClusterComponent extends BaseComponent { ), id: TREE_ICON_ID, command: { - title: "Select a cluster", + title: "Select a compute", command: "databricks.connection.attachClusterQuickPick", }, }, diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index 2fd63c23c..b977f496b 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -20,6 +20,7 @@ import {logging} from "@databricks/sdk-experimental"; import {Loggers} from "../../logger"; import {FeatureManager} from "../../feature-manager/FeatureManager"; import {EnvironmentComponent} from "./EnvironmentComponent"; +import {SshTunnelComponent} from "./SshTunnelComponent"; import {WorkspaceFolderComponent} from "./WorkspaceFolderComponent"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; import {CodeSynchronizer} from "../../sync"; @@ -53,6 +54,7 @@ export class ConfigurationDataProvider private readonly workspaceFolderManager: WorkspaceFolderManager ) { this.components = [ + new SshTunnelComponent(this.connectionManager), new WorkspaceFolderComponent(this.workspaceFolderManager), new BundleTargetComponent(this.configModel), new AuthTypeComponent( diff --git a/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts new file mode 100644 index 000000000..6a697a398 --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts @@ -0,0 +1,45 @@ +import {ThemeIcon, TreeItemCollapsibleState} from "vscode"; +import {BaseComponent} from "./BaseComponent"; +import {ConfigurationTreeItem} from "./types"; +import {ConnectionManager} from "../../configuration/ConnectionManager"; + +export class SshTunnelComponent extends BaseComponent { + constructor(private readonly connectionManager: ConnectionManager) { + super(); + this.disposables.push( + this.connectionManager.onDidChangeState(() => { + this.onDidChangeEmitter.fire(); + }) + ); + } + + private async getRoot(): Promise { + if (this.connectionManager.state !== "CONNECTED") { + return []; + } + + return [ + { + label: "Start SSH Tunnel", + iconPath: new ThemeIcon("remote"), + contextValue: "databricks.configuration.sshTunnel", + collapsibleState: TreeItemCollapsibleState.None, + tooltip: "Start Databricks remote development via SSH", + command: { + title: "Start SSH Tunnel", + command: "databricks.ssh.startTunnel", + }, + }, + ]; + } + + public async getChildren( + parent?: ConfigurationTreeItem + ): Promise { + if (parent === undefined) { + return this.getRoot(); + } + + return []; + } +} From fb27fef4264115c01a1de9c5fb1250256b0318bb Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 15 Jul 2026 19:18:04 +0400 Subject: [PATCH 04/23] Remove redundant start ssh icon --- packages/databricks-vscode/package.json | 5 ----- 1 file changed, 5 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 96a3cd2fe..47e2ad016 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -1109,11 +1109,6 @@ "when": "view == configurationView && viewItem =~ /^databricks.*sync.*is-running.*$/ && databricks.context.bundle.isDevTarget", "group": "navigation@0" }, - { - "command": "databricks.ssh.startTunnel", - "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", - "group": "inline@0" - }, { "command": "databricks.ssh.startTunnel", "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", From ff0a5f9dafa3e017888228e52203d6df5967e849 Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 15 Jul 2026 19:29:59 +0400 Subject: [PATCH 05/23] Remove subheading from serverless options --- packages/databricks-vscode/src/ssh/SshCommands.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 28b22e379..1cc57ef0f 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -36,18 +36,15 @@ interface ServerlessItem extends QuickPickItem { const SERVERLESS_ITEMS: ServerlessItem[] = [ { label: SERVERLESS_LABEL, - detail: "Connect to serverless compute (no dedicated cluster)", alwaysShow: true, }, { label: "$(cloud) Serverless GPU 1xA10", - detail: "Connect to serverless GPU compute (GPU_1xA10)", alwaysShow: true, accelerator: "GPU_1xA10", }, { label: "$(cloud) Serverless GPU 8xH100", - detail: "Connect to serverless GPU compute (GPU_8xH100)", alwaysShow: true, accelerator: "GPU_8xH100", }, From 44172dd5b67234d2b14f2c02c78e92e4f6e94761 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 10:16:14 +0400 Subject: [PATCH 06/23] Cleanup --- packages/databricks-vscode/src/extension.ts | 17 ----------------- .../databricks-vscode/src/ssh/SshCommands.ts | 4 ---- 2 files changed, 21 deletions(-) diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 981ea2d99..56ae1264a 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -269,23 +269,6 @@ export async function activate( ); } - // Auto-open the file the user was editing locally when the tunnel was - // started, if it resolves in the remote workspace. Forwarded out of - // band via env var since the CLI has no file flag; best-effort only. - const remoteOpenFile = process.env["DATABRICKS_REMOTE_OPEN_FILE"]; - if (remoteOpenFile) { - try { - const uri = Uri.file(remoteOpenFile); - await workspace.fs.stat(uri); - await window.showTextDocument(uri); - } catch (e) { - logging.NamedLogger.getOrCreate(Loggers.Extension).debug( - "Skipping remote auto-open of file (not found remotely)", - {remoteOpenFile} - ); - } - } - customWhenContext.setActivated(true); return; } diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 1cc57ef0f..5d91186af 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -218,10 +218,6 @@ export class SshCommands implements Disposable { DATABRICKS_REMOTE_HOME_FOLDER: `/Users/${userName}`, /* eslint-enable @typescript-eslint/naming-convention */ }; - const activeFile = window.activeTextEditor?.document.uri.fsPath; - if (activeFile) { - env["DATABRICKS_REMOTE_OPEN_FILE"] = activeFile; - } const terminal = window.createTerminal({ name: "Databricks SSH Tunnel", From 98d98ed00108796ecc046fdaab01822a072c1f21 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 10:20:14 +0400 Subject: [PATCH 07/23] Fix linter error --- packages/databricks-vscode/src/extension.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 56ae1264a..050455f83 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -4,7 +4,6 @@ import { env, ExtensionContext, extensions, - Uri, window, workspace, } from "vscode"; From 43c5dbe774162fe3b19bf0ca8a8fe57cf21a959c Mon Sep 17 00:00:00 2001 From: misha-db Date: Mon, 20 Jul 2026 19:48:40 +0400 Subject: [PATCH 08/23] SSH tunnel view --- packages/databricks-vscode/package.json | 16 +- packages/databricks-vscode/src/extension.ts | 14 +- .../databricks-vscode/src/ssh/SshCommands.ts | 184 +++++++++++++++--- .../ConfigurationDataProvider.ts | 2 - .../configuration-view/SshTunnelComponent.ts | 45 ----- 5 files changed, 176 insertions(+), 85 deletions(-) delete mode 100644 packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 27465bcd6..dd2400ff1 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -57,7 +57,7 @@ "title": "Start SSH Tunnel", "category": "Databricks", "icon": "$(remote)", - "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode" + "enablement": "!databricks.context.remoteMode" }, { "command": "databricks.connection.logout", @@ -546,6 +546,11 @@ "name": "Configuration", "when": "!databricks.context.remoteMode" }, + { + "id": "sshTunnelView", + "name": "SSH Tunnel", + "when": "!databricks.context.remoteMode" + }, { "id": "clusterView", "when": "databricks.feature.views.cluster && !databricks.context.remoteMode", @@ -714,6 +719,10 @@ { "view": "configurationView", "contents": "To learn more about how to use the Databricks extension for Visual Studio Code [read our docs](https://docs.databricks.com/dev-tools/vscode-ext.html) or [Quickstart guide](command:databricks.quickstart.open)" + }, + { + "view": "sshTunnelView", + "contents": "Connect your IDE to a Databricks SSH Tunnel so you can run workloads using Databricks compute.\n[Start SSH Tunnel](command:databricks.ssh.startTunnel)\nTo learn more about the SSH tunnel [read our docs](https://docs.databricks.com/aws/en/dev-tools/ssh-tunnel)." } ], "menus": { @@ -1108,11 +1117,6 @@ "command": "databricks.sync.stop", "when": "view == configurationView && viewItem =~ /^databricks.*sync.*is-running.*$/ && databricks.context.bundle.isDevTarget", "group": "navigation@0" - }, - { - "command": "databricks.ssh.startTunnel", - "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", - "group": "navigation@0" } ], "editor/title": [ diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 050455f83..606e0f645 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -166,6 +166,18 @@ export async function activate( } ) ); + // The SSH Tunnel panel is always visible, including on the start screen + // with no folder open. There is no ConnectionManager/ClusterModel here, + // so wire the command to the standalone (login-then-tunnel) flow. + const sshCommands = new SshCommands(cli); + context.subscriptions.push( + sshCommands, + telemetry.registerCommand( + "databricks.ssh.startTunnel", + sshCommands.startTunnelCommand, + sshCommands + ) + ); // We show a welcome view when there's no workspace folders, prompting users // to either open a new folder or to initialize a new databricks project. // In both cases we expect the workspace to be reloaded and the extension will @@ -827,7 +839,7 @@ export async function activate( ); // SSH tunnel (remote development) group - const sshCommands = new SshCommands(connectionManager, clusterModel, cli); + const sshCommands = new SshCommands(cli, connectionManager, clusterModel); context.subscriptions.push( sshCommands, telemetry.registerCommand( diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 5d91186af..d5544d6b4 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -1,11 +1,14 @@ import { Disposable, + Event, + EventEmitter, QuickPick, QuickPickItem, QuickPickItemKind, ThemeIcon, window, } from "vscode"; +import {WorkspaceClient} from "@databricks/sdk-experimental"; import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; import {ClusterModel} from "../cluster/ClusterModel"; import {ConnectionManager} from "../configuration/ConnectionManager"; @@ -14,6 +17,9 @@ import { formatQuickPickClusterDetails, } from "../configuration/ConnectionCommands"; import {CliWrapper} from "../cli/CliWrapper"; +import {AuthProvider} from "../configuration/auth/AuthProvider"; +import {LoginWizard} from "../configuration/LoginWizard"; +import {Cluster} from "../sdk-extensions"; import {onError} from "../utils/onErrorDecorator"; const SERVERLESS_LABEL = "$(cloud) Serverless"; @@ -50,34 +56,155 @@ const SERVERLESS_ITEMS: ServerlessItem[] = [ }, ]; +/** + * Minimal cluster feed the compute picker needs. `ClusterModel` (connected + * path) satisfies this directly; `StandaloneClusterSource` (start-screen path, + * no workspace folder) implements the same shape from a bare WorkspaceClient. + */ +interface ClusterSource extends Disposable { + readonly roots: Cluster[] | undefined; + readonly onDidChange: Event; + refresh(): void; +} + +/** + * Fetches eligible clusters directly from a WorkspaceClient for the standalone + * (no workspace folder) tunnel flow, where no `ClusterModel` exists. + */ +class StandaloneClusterSource implements ClusterSource { + private _clusters: Cluster[] | undefined; + private readonly onDidChangeEmitter = new EventEmitter(); + readonly onDidChange = this.onDidChangeEmitter.event; + + constructor(private readonly workspaceClient: WorkspaceClient) {} + + get roots(): Cluster[] | undefined { + return this._clusters; + } + + refresh() { + void this.load(); + } + + private async load() { + const clusters: Cluster[] = []; + for await (const cluster of Cluster.list( + this.workspaceClient.apiClient + )) { + clusters.push(cluster); + } + this._clusters = clusters; + this.onDidChangeEmitter.fire(); + } + + dispose() { + this.onDidChangeEmitter.dispose(); + } +} + +/** + * The auth + compute context needed to launch a tunnel, resolved either from an + * already-connected workspace or from a standalone login on the start screen. + */ +interface TunnelContext { + authProvider: AuthProvider; + userName: string; + clusterSource: ClusterSource; + // Cluster sources we create ourselves (standalone) must be disposed after + // the picker; the shared ClusterModel is owned by the extension and is not. + ownsClusterSource: boolean; +} + export class SshCommands implements Disposable { private disposables: Disposable[] = []; + /** + * `connectionManager`/`clusterModel` are only available once a workspace + * folder is open. When they are undefined (start screen) the command falls + * back to a standalone login flow so the tunnel can be started from the + * dedicated SSH Tunnel panel with no folder open. + */ constructor( - private readonly connectionManager: ConnectionManager, - private readonly clusterModel: ClusterModel, - private readonly cli: CliWrapper + private readonly cli: CliWrapper, + private readonly connectionManager?: ConnectionManager, + private readonly clusterModel?: ClusterModel ) {} @onError({popup: {prefix: "Error starting SSH tunnel."}}) async startTunnelCommand() { - const workspaceClient = this.connectionManager.workspaceClient; - const me = this.connectionManager.databricksWorkspace?.userName; - if (!workspaceClient || !me) { - window.showErrorMessage( - "Please connect to a Databricks workspace before starting an SSH tunnel." - ); + const context = await this.resolveTunnelContext(); + if (context === undefined) { return; } + try { + const compute = await this.pickCompute( + context.userName, + context.clusterSource + ); + if (compute === undefined) { + return; + } + await this.launchSshTunnel( + context.authProvider, + context.userName, + compute + ); + } finally { + if (context.ownsClusterSource) { + context.clusterSource.dispose(); + } + } + } - const compute = await this.pickCompute(me); - if (compute === undefined) { - return; + /** + * Resolves the auth provider, user and cluster feed for the tunnel. Uses the + * connected workspace when available (connecting first if needed), otherwise + * runs a standalone login wizard so the tunnel works with no folder open. + */ + private async resolveTunnelContext(): Promise { + if (this.connectionManager && this.clusterModel) { + if (this.connectionManager.state !== "CONNECTED") { + await this.connectionManager.login(true); + } + const workspace = this.connectionManager.databricksWorkspace; + if (!workspace || this.connectionManager.state !== "CONNECTED") { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return undefined; + } + return { + authProvider: workspace.authProvider, + userName: workspace.userName, + clusterSource: this.clusterModel, + ownsClusterSource: false, + }; + } + + const authProvider = await LoginWizard.run(this.cli); + if (authProvider === undefined || !(await authProvider.check())) { + return undefined; + } + const workspaceClient = await authProvider.getWorkspaceClient(); + const userName = (await workspaceClient.currentUser.me()).userName; + if (!userName) { + window.showErrorMessage( + "Could not determine the current user for the SSH tunnel." + ); + return undefined; } - await this.launchSshTunnel(compute); + return { + authProvider, + userName, + clusterSource: new StandaloneClusterSource(workspaceClient), + ownsClusterSource: true, + }; } - private pickCompute(me: string): Promise { + private pickCompute( + me: string, + clusterSource: ClusterSource + ): Promise { return new Promise((resolve) => { const quickPick = window.createQuickPick< ClusterItem | ServerlessItem @@ -99,7 +226,7 @@ export class SshCommands implements Disposable { const refreshItems = () => { // Only dedicated single-user clusters owned by the current user // can be used for an SSH tunnel. - const clusters = (this.clusterModel.roots ?? []).filter((c) => + const clusters = (clusterSource.roots ?? []).filter((c) => c.isValidSingleUser(me) ); quickPick.items = staticItems.concat( @@ -133,7 +260,7 @@ export class SshCommands implements Disposable { // Register the change listener before triggering refresh() so no // onDidChange fired by the (re)started loader can be missed. const disposables: Disposable[] = [ - this.clusterModel.onDidChange(refreshItems), + clusterSource.onDidChange(refreshItems), quickPick, {dispose: () => clearTimeout(spinnerTimeout)}, ]; @@ -141,7 +268,7 @@ export class SshCommands implements Disposable { // Paint whatever is already cached first (fast path on reopen), then // trigger a reload; fresh results stream in via onDidChange. refreshItems(); - this.clusterModel.refresh(); + clusterSource.refresh(); quickPick.show(); quickPick.onDidAccept(() => { @@ -171,10 +298,11 @@ export class SshCommands implements Disposable { /** * Pre-selects the compute the user already has configured locally: the - * attached single-user cluster if any, otherwise serverless. + * attached single-user cluster if any, otherwise serverless. Only the + * connected path has a configured cluster/serverless preference. */ private preselect(quickPick: QuickPick) { - const currentCluster = this.connectionManager.cluster; + const currentCluster = this.connectionManager?.cluster; if (currentCluster?.isSingleUser()) { const match = quickPick.items.find( (i): i is ClusterItem => @@ -185,7 +313,7 @@ export class SshCommands implements Disposable { return; } } - if (this.connectionManager.serverless) { + if (this.connectionManager?.serverless) { const serverlessItem = quickPick.items.find( (i) => i.label === SERVERLESS_LABEL ); @@ -195,17 +323,11 @@ export class SshCommands implements Disposable { } } - private async launchSshTunnel(compute: Compute) { - const authProvider = - this.connectionManager.databricksWorkspace?.authProvider; - const userName = this.connectionManager.databricksWorkspace?.userName; - if (!authProvider || !userName) { - window.showErrorMessage( - "Please connect to a Databricks workspace before starting an SSH tunnel." - ); - return; - } - + private async launchSshTunnel( + authProvider: AuthProvider, + userName: string, + compute: Compute + ) { const {args} = this.cli.getSshConnectCommand({compute}); const env: Record = { diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index b977f496b..2fd63c23c 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -20,7 +20,6 @@ import {logging} from "@databricks/sdk-experimental"; import {Loggers} from "../../logger"; import {FeatureManager} from "../../feature-manager/FeatureManager"; import {EnvironmentComponent} from "./EnvironmentComponent"; -import {SshTunnelComponent} from "./SshTunnelComponent"; import {WorkspaceFolderComponent} from "./WorkspaceFolderComponent"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; import {CodeSynchronizer} from "../../sync"; @@ -54,7 +53,6 @@ export class ConfigurationDataProvider private readonly workspaceFolderManager: WorkspaceFolderManager ) { this.components = [ - new SshTunnelComponent(this.connectionManager), new WorkspaceFolderComponent(this.workspaceFolderManager), new BundleTargetComponent(this.configModel), new AuthTypeComponent( diff --git a/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts deleted file mode 100644 index 6a697a398..000000000 --- a/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts +++ /dev/null @@ -1,45 +0,0 @@ -import {ThemeIcon, TreeItemCollapsibleState} from "vscode"; -import {BaseComponent} from "./BaseComponent"; -import {ConfigurationTreeItem} from "./types"; -import {ConnectionManager} from "../../configuration/ConnectionManager"; - -export class SshTunnelComponent extends BaseComponent { - constructor(private readonly connectionManager: ConnectionManager) { - super(); - this.disposables.push( - this.connectionManager.onDidChangeState(() => { - this.onDidChangeEmitter.fire(); - }) - ); - } - - private async getRoot(): Promise { - if (this.connectionManager.state !== "CONNECTED") { - return []; - } - - return [ - { - label: "Start SSH Tunnel", - iconPath: new ThemeIcon("remote"), - contextValue: "databricks.configuration.sshTunnel", - collapsibleState: TreeItemCollapsibleState.None, - tooltip: "Start Databricks remote development via SSH", - command: { - title: "Start SSH Tunnel", - command: "databricks.ssh.startTunnel", - }, - }, - ]; - } - - public async getChildren( - parent?: ConfigurationTreeItem - ): Promise { - if (parent === undefined) { - return this.getRoot(); - } - - return []; - } -} From 9cff51ff3a9e8b565018dcf7f10b730b38394577 Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 22 Jul 2026 12:28:31 +0400 Subject: [PATCH 09/23] Add loading cluster message --- .../src/configuration/ConnectionManager.ts | 14 +++++ .../databricks-vscode/src/ssh/SshCommands.ts | 56 ++++++++++++++----- 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 6da1e99cc..7acf4d0f4 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -66,6 +66,19 @@ export class ConnectionManager implements Disposable { this.onDidChangeSyncDestinationEmitter.event; private readonly initialization = new Barrier(); + // Set once init() has resolved the initialization barrier. Guards callers + // (e.g. the SSH tunnel flow) that must not `await login()` before init(), + // since that would block on the barrier forever when the workspace is not a + // Databricks project and init() is never called. + private _initialized = false; + + /** + * Whether the connection manager has been initialized (init() ran). + * Only true once the workspace has been set up as a Databricks project. + */ + get isInitialized(): boolean { + return this._initialized; + } get projectRoot() { return this.workspaceFolderManager.activeProjectUri; @@ -212,6 +225,7 @@ export class ConnectionManager implements Disposable { ) ) ); + this._initialized = true; this.initialization.resolve(); } } diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index d5544d6b4..2842f97be 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -162,7 +162,12 @@ export class SshCommands implements Disposable { * runs a standalone login wizard so the tunnel works with no folder open. */ private async resolveTunnelContext(): Promise { - if (this.connectionManager && this.clusterModel) { + // Only use the connected path when the connection manager has actually + // been initialized. When a folder is open but is not a Databricks + // project, init() is never called, so awaiting login() would block on + // its initialization barrier forever (the button would silently do + // nothing). In that case we fall through to the standalone login flow. + if (this.connectionManager?.isInitialized && this.clusterModel) { if (this.connectionManager.state !== "CONNECTED") { await this.connectionManager.login(true); } @@ -210,18 +215,39 @@ export class SshCommands implements Disposable { ClusterItem | ServerlessItem >(); quickPick.title = "Select compute for SSH tunnel"; + quickPick.placeholder = "Loading dedicated clusters…"; quickPick.keepScrollPosition = true; quickPick.busy = true; quickPick.canSelectMany = false; - const staticItems: ServerlessItem[] = [ + // Whether the dedicated-cluster list is still being fetched. Drives + // the separator label and placeholder so the user can tell the list + // is still loading rather than waiting on them to pick. + let loading = true; + + // Serverless items are always ready; the separator that follows them + // reflects the dedicated-cluster loading state. + const buildStaticItems = ( + clusterCount: number + ): (ServerlessItem | QuickPickItem)[] => [ ...SERVERLESS_ITEMS, { - label: "", + label: loading + ? "Loading clusters…" + : clusterCount > 0 + ? "Dedicated clusters" + : "No dedicated clusters found", kind: QuickPickItemKind.Separator, }, ]; - quickPick.items = staticItems; + + const stopLoading = () => { + loading = false; + quickPick.busy = false; + quickPick.placeholder = undefined; + }; + + quickPick.items = buildStaticItems(0); const refreshItems = () => { // Only dedicated single-user clusters owned by the current user @@ -229,7 +255,15 @@ export class SshCommands implements Disposable { const clusters = (clusterSource.roots ?? []).filter((c) => c.isValidSingleUser(me) ); - quickPick.items = staticItems.concat( + // Clear the loading state only once clusters have actually + // loaded, before building items so the separator label matches + // this repaint. On a cold first open the loader is still + // fetching, so we keep loading and let onDidChange repaint when + // clusters arrive. + if (clusters.length > 0) { + stopLoading(); + } + quickPick.items = buildStaticItems(clusters.length).concat( clusters.map((c) => { const treeItem = ClusterListDataProvider.clusterNodeToTreeItem(c); @@ -242,19 +276,15 @@ export class SshCommands implements Disposable { }; }) ); - // Clear the spinner only once clusters have actually loaded. - // On a cold first open the loader is still fetching, so we keep - // spinning and let onDidChange repaint when clusters arrive. - if (clusters.length > 0) { - quickPick.busy = false; - } this.preselect(quickPick); }; // Fallback so the spinner can't hang forever for a user with no - // eligible clusters (onDidChange may never add any). + // eligible clusters (onDidChange may never add any). Repaint so the + // separator flips to "No dedicated clusters found". const spinnerTimeout = setTimeout(() => { - quickPick.busy = false; + stopLoading(); + refreshItems(); }, 10_000); // Register the change listener before triggering refresh() so no From 8cc27998b14755705b17d8bbf3f2b232ca181e90 Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 22 Jul 2026 12:46:44 +0400 Subject: [PATCH 10/23] Display status of dedicated cluster. change icon --- .../databricks-vscode/src/ssh/SshCommands.ts | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 2842f97be..bb0f72bac 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -5,11 +5,9 @@ import { QuickPick, QuickPickItem, QuickPickItemKind, - ThemeIcon, window, } from "vscode"; import {WorkspaceClient} from "@databricks/sdk-experimental"; -import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; import {ClusterModel} from "../cluster/ClusterModel"; import {ConnectionManager} from "../configuration/ConnectionManager"; import { @@ -264,17 +262,16 @@ export class SshCommands implements Disposable { stopLoading(); } quickPick.items = buildStaticItems(clusters.length).concat( - clusters.map((c) => { - const treeItem = - ClusterListDataProvider.clusterNodeToTreeItem(c); - return { - label: `$(${ - (treeItem.iconPath as ThemeIcon).id - }) ${c.name!} (${c.id})`, - detail: formatQuickPickClusterDetails(c), - cluster: c, - }; - }) + clusters.map((c) => ({ + // Use the same compute icon as the "Select a compute" + // entry in the configuration section; the per-state icon + // is dropped in favour of showing the state as text. + label: `$(server) ${c.name!} (${c.id})`, + detail: `${c.state} | ${formatQuickPickClusterDetails( + c + )}`, + cluster: c, + })) ); this.preselect(quickPick); }; From fb1d74f7830a7b458be949e75e16ec08ea450ee3 Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 29 Jul 2026 13:14:17 +0400 Subject: [PATCH 11/23] Change SSH tunnel view welcome message --- packages/databricks-vscode/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index cd9acb9da..da71fc2c3 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -722,7 +722,7 @@ }, { "view": "sshTunnelView", - "contents": "Connect your IDE to a Databricks SSH Tunnel so you can run workloads using Databricks compute.\n[Start SSH Tunnel](command:databricks.ssh.startTunnel)\nTo learn more about the SSH tunnel [read our docs](https://docs.databricks.com/aws/en/dev-tools/ssh-tunnel)." + "contents": "Connect your IDE to a Databricks SSH Tunnel so you can run all Python and SQL workloads using Databricks compute.\n[Start SSH Tunnel](command:databricks.ssh.startTunnel)\nTo learn more about the SSH tunnel [read our docs](https://docs.databricks.com/aws/en/dev-tools/ssh-tunnel)." } ], "menus": { From 872b89942c894499129af2c99f18f0bfdfe6a6c3 Mon Sep 17 00:00:00 2001 From: misha-db Date: Wed, 29 Jul 2026 14:04:53 +0400 Subject: [PATCH 12/23] Format compute status --- .../src/configuration/ConnectionCommands.test.ts | 15 ++++++++++++++- .../src/configuration/ConnectionCommands.ts | 15 +++++++++++++++ packages/databricks-vscode/src/ssh/SshCommands.ts | 7 ++++--- 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index b0facc8ba..626557093 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -3,7 +3,10 @@ import {ApiClient} from "@databricks/sdk-experimental"; import {Cluster} from "../sdk-extensions"; import assert from "assert"; import {mock} from "ts-mockito"; -import {formatQuickPickClusterDetails} from "./ConnectionCommands"; +import { + formatClusterState, + formatQuickPickClusterDetails, +} from "./ConnectionCommands"; describe(__filename, () => { it("attach cluster quickpick: correctly format cluster details", () => { @@ -21,4 +24,14 @@ describe(__filename, () => { assert.equal(clusterDetails, `2 GB | 4 Cores | spark-version | user-2`); }); + + it("formatClusterState: maps RUNNING/TERMINATED to Active/Inactive and title-cases the rest", () => { + assert.equal(formatClusterState("RUNNING"), "Active"); + assert.equal(formatClusterState("TERMINATED"), "Inactive"); + assert.equal(formatClusterState("PENDING"), "Pending"); + assert.equal(formatClusterState("RESTARTING"), "Restarting"); + assert.equal(formatClusterState("TERMINATING"), "Terminating"); + assert.equal(formatClusterState("ERROR"), "Error"); + assert.equal(formatClusterState("UNKNOWN"), "Unknown"); + }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index b9ac7005a..56deb99db 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -1,4 +1,5 @@ import {Cluster} from "../sdk-extensions"; +import {compute} from "@databricks/sdk-experimental"; import { Disposable, QuickPickItem, @@ -31,6 +32,20 @@ function formatQuickPickClusterSize(sizeInMB: number): string { return `${sizeInMB} MB`; } } +// Formats a compute state for display in the picker. RUNNING/TERMINATED map to +// the softer "Active"/"Inactive" (matching the Workspace UI); other states are +// title-cased ("PENDING" -> "Pending") so they read less harshly. +export function formatClusterState(state: compute.State): string { + switch (state) { + case "RUNNING": + return "Active"; + case "TERMINATED": + return "Inactive"; + default: + return state.charAt(0) + state.slice(1).toLowerCase(); + } +} + export function formatQuickPickClusterDetails(cluster: Cluster) { const details = []; if (cluster.memoryMb) { diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index bb0f72bac..1986e85b7 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -12,6 +12,7 @@ import {ClusterModel} from "../cluster/ClusterModel"; import {ConnectionManager} from "../configuration/ConnectionManager"; import { ClusterItem, + formatClusterState, formatQuickPickClusterDetails, } from "../configuration/ConnectionCommands"; import {CliWrapper} from "../cli/CliWrapper"; @@ -267,9 +268,9 @@ export class SshCommands implements Disposable { // entry in the configuration section; the per-state icon // is dropped in favour of showing the state as text. label: `$(server) ${c.name!} (${c.id})`, - detail: `${c.state} | ${formatQuickPickClusterDetails( - c - )}`, + detail: `${formatClusterState( + c.state + )} | ${formatQuickPickClusterDetails(c)}`, cluster: c, })) ); From 3e3a58bd423dd90f4cf1ccb5a861b8038f42264f Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 30 Jul 2026 19:49:14 +0400 Subject: [PATCH 13/23] Address comments --- .../src/cluster/ClusterModel.ts | 10 ++ .../databricks-vscode/src/ssh/SshCommands.ts | 107 +++++++++++++----- .../src/test/e2e/run_dbconnect.ucws.e2e.ts | 2 +- 3 files changed, 91 insertions(+), 28 deletions(-) diff --git a/packages/databricks-vscode/src/cluster/ClusterModel.ts b/packages/databricks-vscode/src/cluster/ClusterModel.ts index b1ab85dd2..a16abcb03 100644 --- a/packages/databricks-vscode/src/cluster/ClusterModel.ts +++ b/packages/databricks-vscode/src/cluster/ClusterModel.ts @@ -67,6 +67,16 @@ export class ClusterModel implements Disposable { ); } + /** + * All loaded clusters, ignoring the explorer's active filter. Consumers that + * need the complete set (e.g. the SSH tunnel picker, which starts stopped + * clusters via --auto-start-cluster) must not inherit the explorer's + * `ALL`/`ME`/`RUNNING` filter. + */ + public get allRoots(): Cluster[] | undefined { + return sortClusters(Array.from(this.clusterLoader.clusters.values())); + } + private applyFilter(nodes: Cluster[] | undefined): Cluster[] | undefined { if (!nodes) { return nodes; diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 1986e85b7..48e5ce47d 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -20,6 +20,8 @@ import {AuthProvider} from "../configuration/auth/AuthProvider"; import {LoginWizard} from "../configuration/LoginWizard"; import {Cluster} from "../sdk-extensions"; import {onError} from "../utils/onErrorDecorator"; +import {logging} from "@databricks/sdk-experimental"; +import {Loggers} from "../logger"; const SERVERLESS_LABEL = "$(cloud) Serverless"; @@ -86,13 +88,24 @@ class StandaloneClusterSource implements ClusterSource { } private async load() { - const clusters: Cluster[] = []; - for await (const cluster of Cluster.list( - this.workspaceClient.apiClient - )) { - clusters.push(cluster); + try { + const clusters: Cluster[] = []; + for await (const cluster of Cluster.list( + this.workspaceClient.apiClient + )) { + clusters.push(cluster); + } + this._clusters = clusters; + } catch (e) { + // On a list failure (network, permissions, expired token) surface an + // empty set rather than leaving the picker spinning until the 10s + // fallback, and still fire onDidChange so the state is honest. + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to list clusters for SSH tunnel", + e + ); + this._clusters = []; } - this._clusters = clusters; this.onDidChangeEmitter.fire(); } @@ -101,6 +114,30 @@ class StandaloneClusterSource implements ClusterSource { } } +/** + * Adapts the shared `ClusterModel` (connected path) to `ClusterSource`, exposing + * its *unfiltered* cluster set so the picker doesn't inherit the explorer's + * `ALL`/`ME`/`RUNNING` filter. The model is owned by the extension, so `dispose` + * is a no-op here. + */ +class ClusterModelSource implements ClusterSource { + constructor(private readonly clusterModel: ClusterModel) {} + + get roots(): Cluster[] | undefined { + return this.clusterModel.allRoots; + } + + get onDidChange(): Event { + return this.clusterModel.onDidChange; + } + + refresh() { + this.clusterModel.refresh(); + } + + dispose() {} +} + /** * The auth + compute context needed to launch a tunnel, resolved either from an * already-connected workspace or from a standalone login on the start screen. @@ -143,11 +180,7 @@ export class SshCommands implements Disposable { if (compute === undefined) { return; } - await this.launchSshTunnel( - context.authProvider, - context.userName, - compute - ); + await this.launchSshTunnel(context.authProvider, compute); } finally { if (context.ownsClusterSource) { context.clusterSource.dispose(); @@ -180,7 +213,11 @@ export class SshCommands implements Disposable { return { authProvider: workspace.authProvider, userName: workspace.userName, - clusterSource: this.clusterModel, + // Read the unfiltered cluster set so the picker is independent of + // the explorer's active filter — a "Running" filter would + // otherwise hide stopped single-user clusters that this flow can + // still start via --auto-start-cluster. + clusterSource: new ClusterModelSource(this.clusterModel), ownsClusterSource: false, }; } @@ -224,6 +261,13 @@ export class SshCommands implements Disposable { // is still loading rather than waiting on them to pick. let loading = true; + // Preselect only once, on the first populated repaint. refreshItems + // runs on every onDidChange (the loader fires once per cluster + // permission check), and re-applying activeItems would keep yanking + // the highlight back to the preselected entry while the user + // navigates. + let hasPreselected = false; + // Serverless items are always ready; the separator that follows them // reflects the dedicated-cluster loading state. const buildStaticItems = ( @@ -274,7 +318,9 @@ export class SshCommands implements Disposable { cluster: c, })) ); - this.preselect(quickPick); + if (!hasPreselected) { + hasPreselected = this.preselect(quickPick); + } }; // Fallback so the spinner can't hang forever for a user with no @@ -301,7 +347,12 @@ export class SshCommands implements Disposable { quickPick.onDidAccept(() => { const selectedItem = quickPick.selectedItems[0]; - disposables.forEach((d) => d.dispose()); + // Resolve with the accepted value BEFORE disposing. Disposing + // the quickPick (it is in `disposables`) synchronously fires + // onDidHide, whose handler resolves undefined; since resolve is + // idempotent, settling the real value first makes that later + // undefined a harmless no-op. Disposing first would let + // undefined win and drop the user's selection. if (selectedItem === undefined) { resolve(undefined); } else if ("cluster" in selectedItem) { @@ -315,11 +366,12 @@ export class SshCommands implements Disposable { accelerator: selectedItem.accelerator, }); } + disposables.forEach((d) => d.dispose()); }); quickPick.onDidHide(() => { disposables.forEach((d) => d.dispose()); - // resolve(undefined); + resolve(undefined); }); }); } @@ -327,9 +379,14 @@ export class SshCommands implements Disposable { /** * Pre-selects the compute the user already has configured locally: the * attached single-user cluster if any, otherwise serverless. Only the - * connected path has a configured cluster/serverless preference. + * connected path has a configured cluster/serverless preference. Returns + * whether a selection was actually applied so the caller can stop + * re-preselecting once the target item exists (the attached cluster may not + * be in the list yet on a cold first repaint). */ - private preselect(quickPick: QuickPick) { + private preselect( + quickPick: QuickPick + ): boolean { const currentCluster = this.connectionManager?.cluster; if (currentCluster?.isSingleUser()) { const match = quickPick.items.find( @@ -338,7 +395,7 @@ export class SshCommands implements Disposable { ); if (match) { quickPick.activeItems = [match]; - return; + return true; } } if (this.connectionManager?.serverless) { @@ -347,35 +404,31 @@ export class SshCommands implements Disposable { ); if (serverlessItem) { quickPick.activeItems = [serverlessItem]; + return true; } } + return false; } private async launchSshTunnel( authProvider: AuthProvider, - userName: string, compute: Compute ) { const {args} = this.cli.getSshConnectCommand({compute}); const env: Record = { ...this.cli.getSshConnectEnvVars(authProvider), - // The remote window opens at the user's home folder. Forward the - // file the user is currently editing so the remote extension can - // auto-open it (see the remote-mode branch in extension.ts). The - // CLI has no folder/file flag, so we pass it out of band. - /* eslint-disable @typescript-eslint/naming-convention */ - DATABRICKS_REMOTE_HOME_FOLDER: `/Users/${userName}`, - /* eslint-enable @typescript-eslint/naming-convention */ }; + // The transient terminal is not retained on `this.disposables`: doing so + // would only release it on extension deactivate/reload, tearing down an + // active tunnel. It is `isTransient` and disposes itself when closed. const terminal = window.createTerminal({ name: "Databricks SSH Tunnel", isTransient: true, env, strictEnv: false, }); - this.disposables.push(terminal); terminal.show(); terminal.sendText(`${this.cli.escapedCliPath} ${args.join(" ")}`); } diff --git a/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts b/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts index 63379ba2f..ce3057cc9 100644 --- a/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts @@ -643,7 +643,7 @@ describe("Run files on serverless compute", async function () { }); it("should select serverless compute", async () => { - await executeCommandWhenAvailable("Databricks: Configure cluster"); + await executeCommandWhenAvailable("Databricks: Configure compute"); const computeInput = await waitForInput(); await computeInput.selectQuickPick("Serverless"); }); From f82d27fb82dcb22f019f091bbdcfa82a5cb6d3ce Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 21:59:21 +0400 Subject: [PATCH 14/23] Update packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts Co-authored-by: Russell Clarey --- .../src/ui/configuration-view/ClusterComponent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts index 6303e132c..86e39ff20 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts @@ -123,7 +123,7 @@ export class ClusterComponent extends BaseComponent { ), id: TREE_ICON_ID, command: { - title: "Select a compute", + title: "Select compute", command: "databricks.connection.attachClusterQuickPick", }, }, From 0c9d5081b01579f69886379c98e786b459a8957d Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 21:59:32 +0400 Subject: [PATCH 15/23] Update packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts Co-authored-by: Russell Clarey --- .../src/ui/configuration-view/ClusterComponent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts index 86e39ff20..4c216b529 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts @@ -114,7 +114,7 @@ export class ClusterComponent extends BaseComponent { // We are logged in -> Select cluster prompt return [ { - label: LabelUtils.highlightedLabel("Select a compute"), + label: LabelUtils.highlightedLabel("Select compute"), collapsibleState: TreeItemCollapsibleState.None, contextValue: getContextValue("none"), iconPath: new ThemeIcon( From e878665c6ce7a9e7dbcf55616acfda5782ecb283 Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 21:59:57 +0400 Subject: [PATCH 16/23] Update packages/databricks-vscode/DATABRICKS.quickstart.md Co-authored-by: Russell Clarey --- packages/databricks-vscode/DATABRICKS.quickstart.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index 00ae9b7fd..af476f6a2 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -50,7 +50,7 @@ If your folder has multiple [Declarative Automation Bundles](#dabs), you can sel ## Select a compute -The extension uses an interactive compute to run code. To select an interactive compute: +The extension uses interactive compute to run code. To select interactive compute: 1. Open the Databricks panel by clicking on the Databricks icon on the left 2. Click on the "Select Compute" button. From 45e9eb6549da4649ada68c5b61d48bc03c88685f Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 22:00:30 +0400 Subject: [PATCH 17/23] Update packages/databricks-vscode/package.json Co-authored-by: Russell Clarey --- packages/databricks-vscode/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index da71fc2c3..e14dd96ef 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -554,7 +554,7 @@ { "id": "clusterView", "when": "databricks.feature.views.cluster && !databricks.context.remoteMode", - "name": "Computes" + "name": "Compute" }, { "id": "dabsResourceExplorerView", From 18e92b49f5fed164b5096cae8e1c19a6be18d83d Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 22:00:40 +0400 Subject: [PATCH 18/23] Update packages/databricks-vscode/package.json Co-authored-by: Russell Clarey --- packages/databricks-vscode/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index e14dd96ef..a47621ac9 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -1455,7 +1455,7 @@ "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": false, - "description": "Enable/disable filtering for only accessible computes (computes on which the current user can run code)" + "description": "Enable/disable filtering for only accessible compute (compute on which the current user can run code)" }, "databricks.overrideDatabricksConfigFile": { "type": "string", From 5ef92a503ad0988192a9f9d9583dae99ff2ff281 Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 22:01:00 +0400 Subject: [PATCH 19/23] Update packages/databricks-vscode/DATABRICKS.quickstart.md Co-authored-by: Russell Clarey --- packages/databricks-vscode/DATABRICKS.quickstart.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index af476f6a2..9f22f2547 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -54,7 +54,7 @@ The extension uses interactive compute to run code. To select interactive comput 1. Open the Databricks panel by clicking on the Databricks icon on the left 2. Click on the "Select Compute" button. - - If you wish to change the selected cluster, click on the "Configure Cluster" gear icon, next to the name of the selected cluster. + - If you wish to change the selected compute, click on the "Configure compute" gear icon, next to the name of the selected compute. ## Run Python code From 0aa56cf2a03c6681911b2f2578ee3eac884d5ae1 Mon Sep 17 00:00:00 2001 From: Misha Kezherashvili Date: Thu, 30 Jul 2026 22:01:19 +0400 Subject: [PATCH 20/23] Update packages/databricks-vscode/DATABRICKS.quickstart.md Co-authored-by: Russell Clarey --- packages/databricks-vscode/DATABRICKS.quickstart.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index 9f22f2547..eac550bbe 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -48,7 +48,7 @@ The Databricks extension for Visual Studio Code enables you to connect to your r If your folder has multiple [Declarative Automation Bundles](#dabs), you can select which one to use by clicking "Open Existing Databricks project" button and selecting the desired project. -## Select a compute +## Select compute The extension uses interactive compute to run code. To select interactive compute: From b1f195b02e16a71ca14fab3c64dc4c0ac89a4a13 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 30 Jul 2026 23:26:33 +0400 Subject: [PATCH 21/23] Address comments --- .../src/cli/CliWrapper.test.ts | 11 ++-- .../databricks-vscode/src/cli/CliWrapper.ts | 60 ++++++++++--------- .../databricks-vscode/src/ssh/SshCommands.ts | 11 ++-- 3 files changed, 41 insertions(+), 41 deletions(-) diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index fc92c5486..226bd0b08 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -6,7 +6,7 @@ import {execFile as execFileCb} from "node:child_process"; import {withFile} from "tmp-promise"; import {writeFile, readFile} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; -import {CliWrapper, waitForProcess} from "./CliWrapper"; +import {CliWrapper, getSshConnectCommand, waitForProcess} from "./CliWrapper"; import path from "node:path"; import os from "node:os"; import crypto from "node:crypto"; @@ -129,14 +129,11 @@ describe(__filename, function () { }); it("should create ssh connect commands", () => { - const logFilePath = getTempLogFilePath(); - const cli = createCliWrapper(logFilePath); - // Logging is configured via env vars, not CLI flags, so no --log-* // args appear on the ssh connect command line. // Serverless: no --cluster / --auto-start-cluster. - let {args} = cli.getSshConnectCommand({compute: {type: "serverless"}}); + let {args} = getSshConnectCommand({compute: {type: "serverless"}}); assert.deepStrictEqual(args, [ "ssh", "connect", @@ -145,7 +142,7 @@ describe(__filename, function () { ]); // Serverless GPU: --accelerator, no --cluster / --auto-start-cluster. - ({args} = cli.getSshConnectCommand({ + ({args} = getSshConnectCommand({ compute: {type: "serverless", accelerator: "GPU_1xA10"}, })); assert.deepStrictEqual(args, [ @@ -157,7 +154,7 @@ describe(__filename, function () { ]); // Dedicated cluster: --cluster and --auto-start-cluster. - ({args} = cli.getSshConnectCommand({ + ({args} = getSshConnectCommand({ compute: {type: "cluster", clusterId: "1234-clusterid"}, })); assert.deepStrictEqual(args, [ diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index f7a3f89f8..77df0e7c9 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -105,6 +105,37 @@ export interface ConfigEntry { } export type SyncType = "full" | "incremental"; + +export type SshConnectCompute = + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + +/** + * Constructs the `databricks ssh connect` command args for opening a remote + * IDE window. Serverless is the default when no cluster is given. + * + * The --ide flag matches the host editor so the CLI opens the right remote + * window: Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. + * + * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see + * CliWrapper.getSshConnectEnvVars), so we do not pass --log-* flags here. + */ +export function getSshConnectCommand(opts: {compute: SshConnectCompute}): { + args: string[]; +} { + const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; + const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; + if (opts.compute.type === "cluster") { + // Start a stopped single-user cluster when connecting. + args.push(`--cluster=${opts.compute.clusterId}`); + args.push("--auto-start-cluster"); + } else if (opts.compute.accelerator) { + // Serverless GPU: request a specific accelerator type. + args.push(`--accelerator=${opts.compute.accelerator}`); + } + return {args}; +} export class ProcessError extends Error { constructor( message: string, @@ -593,35 +624,6 @@ export class CliWrapper { }); } - /** - * Constructs the `databricks ssh connect` command args for opening a remote - * IDE window. Serverless is the default when no cluster is given. - * - * The --ide flag matches the host editor so the CLI opens the right remote - * window: Cursor identifies itself via env.uriScheme === "cursor", - * everything else (VS Code, Insiders) uses vscode. - * - * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see - * getSshConnectEnvVars), so we do not pass --log-* flags here. - */ - getSshConnectCommand(opts: { - compute: - | {type: "serverless"; accelerator?: string} - | {type: "cluster"; clusterId: string}; - }): {args: string[]} { - const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; - const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; - if (opts.compute.type === "cluster") { - // Start a stopped single-user cluster when connecting. - args.push(`--cluster=${opts.compute.clusterId}`); - args.push("--auto-start-cluster"); - } else if (opts.compute.accelerator) { - // Serverless GPU: request a specific accelerator type. - args.push(`--accelerator=${opts.compute.accelerator}`); - } - return {args}; - } - async bundleInit( templateDirPath: string, outputDirPath: string, diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index 48e5ce47d..aa019ef7c 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -15,7 +15,11 @@ import { formatClusterState, formatQuickPickClusterDetails, } from "../configuration/ConnectionCommands"; -import {CliWrapper} from "../cli/CliWrapper"; +import { + CliWrapper, + getSshConnectCommand, + SshConnectCompute, +} from "../cli/CliWrapper"; import {AuthProvider} from "../configuration/auth/AuthProvider"; import {LoginWizard} from "../configuration/LoginWizard"; import {Cluster} from "../sdk-extensions"; @@ -25,9 +29,7 @@ import {Loggers} from "../logger"; const SERVERLESS_LABEL = "$(cloud) Serverless"; -type Compute = - | {type: "serverless"; accelerator?: string} - | {type: "cluster"; clusterId: string}; +type Compute = SshConnectCompute; /** * A serverless QuickPick item, optionally carrying a GPU accelerator type @@ -414,7 +416,6 @@ export class SshCommands implements Disposable { authProvider: AuthProvider, compute: Compute ) { - const {args} = this.cli.getSshConnectCommand({compute}); const env: Record = { ...this.cli.getSshConnectEnvVars(authProvider), From 190110ccb981e7a5ef5d5ad54cead2bdb38fd63d Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 30 Jul 2026 23:26:42 +0400 Subject: [PATCH 22/23] Update SshCommands.ts --- packages/databricks-vscode/src/ssh/SshCommands.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts index aa019ef7c..2ea734a62 100644 --- a/packages/databricks-vscode/src/ssh/SshCommands.ts +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -416,6 +416,7 @@ export class SshCommands implements Disposable { authProvider: AuthProvider, compute: Compute ) { + const {args} = getSshConnectCommand({compute}); const env: Record = { ...this.cli.getSshConnectEnvVars(authProvider), From 7c5bcb3d37b4b24183557327289b838cff9893d9 Mon Sep 17 00:00:00 2001 From: misha-db Date: Fri, 31 Jul 2026 11:31:49 +0400 Subject: [PATCH 23/23] Update commonUtils.ts --- packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts b/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts index c97b89d7c..71ac38e3b 100644 --- a/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts +++ b/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts @@ -12,6 +12,7 @@ import { const ViewSectionTypes = [ "CLUSTERS", "CONFIGURATION", + "SSH TUNNEL", "WORKSPACE FILE SYSTEM", "BUNDLE RESOURCE EXPLORER", "BUNDLE VARIABLES",