From 0fbc411bc8b93e3acbdf28c84112c63695f6b4dd Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 2 Jul 2026 12:45:36 +0400 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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";